Fix nil client context in AI carrier timing

Author: batuhan

pkg/connector/client.go

@@ -130,6 +130,17 @@ func (dc *DummyClient) Disconnect() {
 	dc.wg.Wait()
 }
 
+func (dc *DummyClient) clientContext() context.Context {
+	if dc != nil && dc.ctx != nil {
+		return dc.ctx
+	}
+	return context.Background()
+}
+
+func (dc *DummyClient) done() <-chan struct{} {
+	return dc.clientContext().Done()
+}
+
 func (dc *DummyClient) IsLoggedIn() bool {
 	return true
 }
@@ -280,7 +291,7 @@ func (dc *DummyClient) HandleMatrixReaction(ctx context.Context, msg *bridgev2.M
 	dc.wg.Add(1)
 	go func() {
 		defer dc.wg.Done()
-		dc.queueAIApprovalResponse(dc.ctx, portal, target, response)
+		dc.queueAIApprovalResponse(dc.clientContext(), portal, target, response)
 	}()
 
 	logger := log.Info().
@@ -470,7 +481,7 @@ func (dc *DummyClient) queueRemoteEcho(msg *bridgev2.MatrixMessage, transactionI
 		defer timer.Stop()
 
 		select {
-		case <-dc.ctx.Done():
+		case <-dc.done():
 			return
 		case <-timer.C:
 		}
@@ -540,7 +551,7 @@ func (dc *DummyClient) queueAIResponse(ctx context.Context, portal *bridgev2.Por
 				timer := time.NewTimer(delay)
 				select {
 				case <-timer.C:
-				case <-dc.ctx.Done():
+				case <-dc.done():
 					timer.Stop()
 					return
 				}
@@ -696,7 +707,7 @@ func (dc *DummyClient) sleepUntilCarrierTime(run aistream.Run, carrier aistream.
 	timer := time.NewTimer(delay)
 	select {
 	case <-timer.C:
-	case <-dc.ctx.Done():
+	case <-dc.done():
 		timer.Stop()
 	}
 }
@@ -769,10 +780,7 @@ func (dc *DummyClient) waitForMessageMXID(
 	if dc == nil || dc.UserLogin == nil || dc.UserLogin.Bridge == nil || dc.UserLogin.Bridge.DB == nil || portal == nil {
 		return ""
 	}
-	parent := dc.ctx
-	if parent == nil {
-		parent = context.Background()
-	}
+	parent := dc.clientContext()
 	ctx, cancel := context.WithTimeout(parent, timeout)
 	defer cancel()
 

pkg/connector/client_test.go

@@ -4,6 +4,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/beeper/ai-bridge/pkg/ag-ui"
+	"github.com/beeper/ai-bridge/pkg/ai-stream"
 	"maunium.net/go/mautrix/event"
 )
 
@@ -37,3 +39,35 @@ func TestGetRemoteEchoBehavior(t *testing.T) {
 		})
 	}
 }
+
+func TestSleepUntilCarrierTimeWithoutConnectedContext(t *testing.T) {
+	base := time.Now()
+	run := aistream.Run{
+		Events: []agui.Event{{
+			"type":      agui.EventRunStarted,
+			"timestamp": base.UnixMilli(),
+			"threadId":  "thread-1",
+		}},
+	}
+	carrier := aistream.Carrier{
+		Envelopes: []aistream.Envelope{{
+			Part: agui.Event{
+				"type":      agui.EventTextMessageContent,
+				"timestamp": base.Add(time.Millisecond).UnixMilli(),
+				"messageId": "message-1",
+			},
+		}},
+	}
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		(&DummyClient{}).sleepUntilCarrierTime(run, carrier, base)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("timed out waiting for carrier sleep without connected context")
+	}
+}