Propagate run context through agent streams

agent/agent.go

@@ -228,13 +228,22 @@ func (a *agentImpl) Ask(ctx context.Context, message string) (*Response, error)
 func (a *agentImpl) Stream(ctx context.Context, message string) (ai.Stream, error) {
 	a.mu.Lock()
 	defer a.mu.Unlock()
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
 	if a.model == nil {
 		a.setup()
 	}
 	toolList, err := a.discoverTools()
 	if err != nil {
 		return nil, fmt.Errorf("discover tools: %w", err)
 	}
+	runID := uuid.New().String()
+	ctx = ai.WithRunInfo(ctx, ai.RunInfo{
+		RunID:    runID,
+		ParentID: a.parentRunID,
+		Agent:    a.opts.Name,
+	})
 	messages := append([]ai.Message(nil), a.mem.Messages()...)
 	messages = append(messages, ai.Message{Role: "user", Content: message})
 	stream, err := a.model.Stream(ctx, &ai.Request{
@@ -246,6 +255,10 @@ func (a *agentImpl) Stream(ctx context.Context, message string) (ai.Stream, erro
 	if err != nil {
 		return nil, err
 	}
+	if err := ctx.Err(); err != nil {
+		_ = stream.Close()
+		return nil, err
+	}
 	a.mem.Add("user", message)
 	return &memoryRecordingStream{stream: stream, memory: a.mem}, nil
 }

agent/stream_test.go

@@ -118,8 +118,14 @@ func TestStreamAskHelperRejectsUnsupportedAgent(t *testing.T) {
 
 func TestAgentStreamUsesProviderStreamingAndRecordsAssistantMemory(t *testing.T) {
 	var sawRequest bool
+	var sawRunInfo bool
 	fakeStream = func(ctx context.Context, opts ai.Options, req *ai.Request) (ai.Stream, error) {
 		sawRequest = true
+		info, ok := ai.RunInfoFrom(ctx)
+		if !ok || info.RunID == "" || info.Agent != "provider-stream" {
+			t.Fatalf("RunInfo = %#v, %v; want provider stream run metadata", info, ok)
+		}
+		sawRunInfo = true
 		if req.Prompt != "stream the answer" {
 			t.Fatalf("Prompt = %q, want stream the answer", req.Prompt)
 		}
@@ -153,6 +159,9 @@ func TestAgentStreamUsesProviderStreamingAndRecordsAssistantMemory(t *testing.T)
 	if !sawRequest {
 		t.Fatal("provider Stream was not called")
 	}
+	if !sawRunInfo {
+		t.Fatal("provider Stream did not receive RunInfo")
+	}
 	if reply != "hello" {
 		t.Fatalf("reply = %q, want hello", reply)
 	}
@@ -162,6 +171,29 @@ func TestAgentStreamUsesProviderStreamingAndRecordsAssistantMemory(t *testing.T)
 	}
 }
 
+func TestAgentStreamCanceledContextSkipsProviderCallAndMemory(t *testing.T) {
+	calls := 0
+	fakeStream = func(ctx context.Context, opts ai.Options, req *ai.Request) (ai.Stream, error) {
+		calls++
+		return &sliceStream{chunks: []string{"late"}}, nil
+	}
+	defer func() { fakeStream = nil }()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	a := newTestAgent(Name("provider-stream-cancel"))
+	_, err := a.Stream(ctx, "do not start")
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("Stream error = %v, want context canceled", err)
+	}
+	if calls != 0 {
+		t.Fatalf("provider Stream calls = %d, want 0 after caller cancellation", calls)
+	}
+	if got := a.mem.Messages(); len(got) != 0 {
+		t.Fatalf("memory = %#v, want no recorded canceled stream turn", got)
+	}
+}
+
 func TestResumeStreamAskDoesNotReplayCompletedTool(t *testing.T) {
 	ctx := context.Background()
 	cp := flow.StoreCheckpoint(store.NewStore(), "stream-resume-agent")