Go: fix context (#18118)

Signed-off-by: Jin Hai <haijin.chn@gmail.com>

internal/agent/component/loop_test.go

@@ -27,7 +27,6 @@
 package component
 
 import (
-	"context"
 	"testing"
 )
 
@@ -84,7 +83,8 @@ func TestLoop_InvokeIsNoOp(t *testing.T) {
 			{"variable": "counter", "input_mode": "constant", "value": 7, "type": "number"},
 		},
 	})
-	out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1})
+	ctx := t.Context()
+	out, err := c.Invoke(ctx, nil, map[string]any{"in": 1})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -97,7 +97,8 @@ func TestLoop_InvokeIsNoOp(t *testing.T) {
 // empty-map chunk and closes.
 func TestLoop_StreamMirrorsInvoke(t *testing.T) {
 	c := NewLoopComponent(loopParam{})
-	ch, err := c.Stream(context.Background(), nil, nil)
+	ctx := t.Context()
+	ch, err := c.Stream(ctx, nil, nil)
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}

internal/agent/component/memory_save_test.go

@@ -28,7 +28,8 @@ import (
 func TestStubMemorySaver_DefaultReturnsError(t *testing.T) {
 	SetMemorySaver(nil)
 	saver := GetMemorySaver()
-	err := saver.Save(context.Background(), MemorySaveRequest{
+	ctx := t.Context()
+	err := saver.Save(ctx, MemorySaveRequest{
 		MemoryIDs: []string{"m1"},
 		AgentID:   "a1",
 	})
@@ -48,7 +49,8 @@ func TestSetMemorySaver_Roundtrip(t *testing.T) {
 	if got != custom {
 		t.Fatalf("saver not registered")
 	}
-	if err := got.Save(context.Background(), MemorySaveRequest{
+	ctx := t.Context()
+	if err := got.Save(ctx, MemorySaveRequest{
 		MemoryIDs:     []string{"m1"},
 		AgentResponse: "hi",
 	}); err != nil {

internal/agent/component/message_phase8b_test.go

@@ -49,7 +49,7 @@ func TestMessage_OutputFormatParam(t *testing.T) {
 		"output_format": "html",
 	})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"text": "hello", "stream": false})
 	if err != nil {
@@ -69,7 +69,7 @@ func TestMessage_OutputFormatInputOverride(t *testing.T) {
 		"output_format": "html",
 	})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"text":          "hi",
@@ -89,7 +89,7 @@ func TestMessage_OutputFormatInputOverride(t *testing.T) {
 func TestMessage_DownloadsExtraction(t *testing.T) {
 	c, _ := NewMessageComponent(map[string]any{"text": "see attachment"})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	dl := map[string]any{
 		"doc_id":    "d-1",
@@ -124,7 +124,7 @@ func TestMessage_DownloadsExtraction(t *testing.T) {
 func TestMessage_DownloadJSONStringSuppressesContent(t *testing.T) {
 	c, _ := NewMessageComponent(map[string]any{"text": "unused"})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	downloadJSON := `{"doc_id":"d-1","filename":"report.md","mime_type":"text/markdown","url":"/api/v1/agents/attachments/d-1/download","include_download_info_in_content":true}`
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
@@ -159,7 +159,7 @@ func TestMessage_AutoPlay_NoEngine(t *testing.T) {
 		"auto_play": true,
 	})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false})
 	if err != nil {
@@ -192,7 +192,7 @@ func TestMessage_AutoPlay_Success(t *testing.T) {
 		"lang":      "en",
 	})
 	state := canvas.NewCanvasState("r1", "t1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"text": "hi", "stream": false})
 	if err != nil {
@@ -224,7 +224,7 @@ func TestMessage_MemorySave_NoService(t *testing.T) {
 	c, _ := NewMessageComponent(map[string]any{"text": "hi"})
 	state := canvas.NewCanvasState("run-x", "task-x")
 	state.Sys["query"] = "what?"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"text":        "hi",
@@ -260,7 +260,7 @@ func TestMessage_MemorySave_Success(t *testing.T) {
 	state.Sys["canvas_id"] = "canvas-y"
 	state.Sys["session_id"] = "session-y"
 	state.Sys["agent_id"] = "agent-y"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	_, err := c.Invoke(ctx, nil, map[string]any{
 		"text":        "hi",
@@ -300,7 +300,7 @@ func TestMessage_MemorySave_FallbackIDs(t *testing.T) {
 	c, _ := NewMessageComponent(map[string]any{"text": "hi"})
 	state := canvas.NewCanvasState("run-fallback", "task-fallback")
 	state.Sys["query"] = "what?"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	_, err := c.Invoke(ctx, nil, map[string]any{
 		"text":        "hi",
@@ -338,7 +338,7 @@ func TestMessage_MemorySave_FromDSLParams(t *testing.T) {
 	})
 	state := canvas.NewCanvasState("run-dsl", "task-dsl")
 	state.Sys["query"] = "hello?"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	// Inputs simulate what the pipeline actually provides: only upstream
 	// outputs, NO memory_ids or memory_save keys.
@@ -377,7 +377,7 @@ func TestMessage_MemorySave_UserIDVariable(t *testing.T) {
 	state := canvas.NewCanvasState("run-uid", "task-uid")
 	state.Sys["query"] = "hello?"
 	state.SetVar("begin", "user_id", "resolved-user-123")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	_, err := c.Invoke(ctx, nil, map[string]any{
 		"text":        "hi",
@@ -410,7 +410,7 @@ func TestMessage_MemorySave_UserIDLiteral(t *testing.T) {
 	})
 	state := canvas.NewCanvasState("run-uid2", "task-uid2")
 	state.Sys["query"] = "hello?"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	_, err := c.Invoke(ctx, nil, map[string]any{
 		"text":        "hi",

internal/agent/component/message_test.go

@@ -32,7 +32,7 @@ func TestMessage_ResolveTemplate(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-1", "task-1")
 	state.Sys["query"] = "world"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"text":   "hello {{sys.query}}",
@@ -57,7 +57,7 @@ func TestMessage_ResolveListReferenceAsJSON(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-list", "task-list")
 	state.SetVar("list_0", "result", []any{"user: 1"})
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"text":   "{{list_0@result}}",
@@ -78,7 +78,7 @@ func TestMessage_Stream(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-2", "task-2")
 	state.Sys["query"] = "alice"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	ch, err := c.Stream(ctx, nil, map[string]any{
 		"text":   "hi",
@@ -107,7 +107,7 @@ func TestMessage_Stream(t *testing.T) {
 func TestMessage_NoTemplate(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-3", "task-3")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"text": "no refs here", "stream": false})
 	if err != nil {
@@ -121,7 +121,7 @@ func TestMessage_NoTemplate(t *testing.T) {
 func TestMessage_RuntimeContentInput(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-4", "task-4")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"content": "from upstream", "stream": false})
 	if err != nil {
@@ -135,7 +135,7 @@ func TestMessage_RuntimeContentInput(t *testing.T) {
 func TestMessage_EmitsAgentMessage(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-emit", "task-emit")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var emitted []string
 	ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) {
 		if contentDelta != "" {
@@ -164,7 +164,7 @@ func TestMessage_EmitsAgentMessage(t *testing.T) {
 func TestMessage_EmitsDirectCanvasMessage(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-direct", "task-direct")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var direct []string
 	var agent []string
 	ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) {
@@ -198,7 +198,7 @@ func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-normal-template", "task-normal-template")
 	state.Sys["query"] = "world"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var direct []string
 	ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) {
 		direct = append(direct, content)
@@ -224,7 +224,7 @@ func TestMessage_NormalTemplateEmitsOnlyRenderedMessage(t *testing.T) {
 func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-skip", "task-skip")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var emitted []string
 	ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) {
 		if contentDelta != "" {
@@ -261,7 +261,7 @@ func TestMessage_SkipsEmissionWhenAgentAlreadyStreamed(t *testing.T) {
 func TestMessage_EmitsContentDifferentFromAgentStream(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-distinct", "task-distinct")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var emitted []string
 	ctx = runtime.WithAgentMessageEmitter(ctx, func(contentDelta, thinkingDelta string) {
 		if contentDelta != "" {
@@ -301,7 +301,7 @@ func TestMessage_ConsumesDeferredAgentStream(t *testing.T) {
 			return map[string]any{"content": "hello world"}, nil
 		},
 	})
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var emitted []string
 	ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) {
 		if content != "" {
@@ -331,7 +331,7 @@ func TestMessage_DeferredStreamThinkingEvents(t *testing.T) {
 			return map[string]any{"content": "answer"}, nil
 		},
 	})
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var events []string
 	ctx = runtime.WithCanvasMessageEventEmitter(ctx, func(content string, startToThink, endToThink bool) {
 		switch {
@@ -361,7 +361,7 @@ func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) {
 			return map[string]any{"content": "grounded answer [ID:1]"}, nil
 		},
 	})
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 	var streamed []string
 	ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) {
 		streamed = append(streamed, content)
@@ -385,7 +385,7 @@ func TestMessage_DeferredStreamUsesCompletedContent(t *testing.T) {
 func TestMessage_FormalizedContentFallback(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-5", "task-5")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"formalized_content": "retrieved answer",
@@ -403,7 +403,7 @@ func TestMessage_FormalizedContentFallback(t *testing.T) {
 func TestMessage_SingleStringFallback(t *testing.T) {
 	c, _ := NewMessageComponent(nil)
 	state := canvas.NewCanvasState("run-6", "task-6")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"value":         "single upstream text",

internal/agent/component/parallel_test.go

@@ -34,7 +34,6 @@
 package component
 
 import (
-	"context"
 	"slices"
 	"testing"
 )
@@ -75,7 +74,9 @@ func TestParallel_InvokeIsNoOp(t *testing.T) {
 		ItemsRef:       "sys.arr",
 		MaxConcurrency: 3,
 	})
-	out, err := c.Invoke(context.Background(), nil, map[string]any{"in": 1})
+	ctx := t.Context()
+
+	out, err := c.Invoke(ctx, nil, map[string]any{"in": 1})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -88,7 +89,9 @@ func TestParallel_InvokeIsNoOp(t *testing.T) {
 // empty-map chunk and closes.
 func TestParallel_StreamMirrorsInvoke(t *testing.T) {
 	c := NewParallelComponent(ParallelParam{})
-	ch, err := c.Stream(context.Background(), nil, nil)
+	ctx := t.Context()
+
+	ch, err := c.Stream(ctx, nil, nil)
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}

internal/agent/component/sampler_params_test.go

@@ -35,7 +35,9 @@ func TestLLM_ForwardsTopP(t *testing.T) {
 		ModelID: "echo",
 		TopP:    &topP,
 	})
-	if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil {
+	ctx := t.Context()
+
+	if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 	if stub.calls != 1 {
@@ -57,9 +59,11 @@ func TestLLM_ForwardsTopP(t *testing.T) {
 func TestLLM_TopPFromInputs(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	if _, err := c.Invoke(context.Background(), nil, map[string]any{
+
+	if _, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt": "hi",
 		"top_p":       0.7,
 	}); err != nil {
@@ -78,9 +82,10 @@ func TestLLM_TopPFromInputs(t *testing.T) {
 func TestLLM_NoTopPByDefault(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil {
+	if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 	if stub.captured == nil {
@@ -128,6 +133,7 @@ func TestAgentParam_ForwardsTopP(t *testing.T) {
 		}
 		return &schema.Message{Content: "ok"}, nil
 	})
+	ctx := t.Context()
 
 	topP := 0.5
 	c := NewAgentComponent(AgentParam{
@@ -135,7 +141,7 @@ func TestAgentParam_ForwardsTopP(t *testing.T) {
 		TopP:      &topP,
 		MaxRounds: 1,
 	})
-	if _, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "hi"}); err != nil {
+	if _, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "hi"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 }
@@ -148,9 +154,10 @@ func TestAgent_TopPFromInputs(t *testing.T) {
 		}
 		return &schema.Message{Content: "ok"}, nil
 	})
+	ctx := t.Context()
 
 	c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1})
-	if _, err := c.Invoke(context.Background(), nil, map[string]any{
+	if _, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt": "hi",
 		"top_p":       0.42,
 	}); err != nil {

internal/agent/component/stagehand_runtime_integration_test.go

@@ -129,7 +129,7 @@ func TestStagehandRuntime_Extract(t *testing.T) {
 		Schema:      schema,
 	}
 
-	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Minute)
+	ctx, cancel := context.WithTimeout(t.Context(), 3*time.Minute)
 	defer cancel()
 
 	t.Logf("starting stagehand RunExtract (timeout 3m); spawns subprocess, calls LLM once with schema=%s",
@@ -240,7 +240,7 @@ func TestBrowser_E2E_Extract(t *testing.T) {
 		t.Fatalf("NewBrowserComponent: %v", err)
 	}
 
-	ctx := canvas.WithState(context.Background(), canvas.NewCanvasState("run-1", "task-1"))
+	ctx := canvas.WithState(t.Context(), canvas.NewCanvasState("run-1", "task-1"))
 	state, _, _ := runtime.GetStateFromContext[*runtime.CanvasState](ctx)
 	state.Sys["user_id"] = "tenant-1"
 

internal/agent/component/stagehand_runtime_test.go

@@ -46,6 +46,7 @@ func cacheSize(r *stagehandRuntime) int {
 func TestStagehandRuntime_ValidatesRequiredFields(t *testing.T) {
 	r := newStagehandRuntime(time.Hour, 0, time.Minute) // TTL large → no sweeper interference
 	t.Cleanup(func() { _ = r.Close() })
+	ctx := t.Context()
 
 	cases := []struct {
 		name string
@@ -58,7 +59,7 @@ func TestStagehandRuntime_ValidatesRequiredFields(t *testing.T) {
 	}
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
-			_, err := r.RunTask(context.Background(), tc.req)
+			_, err := r.RunTask(ctx, tc.req)
 			if err == nil {
 				t.Fatalf("expected error for %s, got nil", tc.name)
 			}
@@ -535,7 +536,8 @@ func TestStagehandRuntime_SetDefaultStagehandInvoker(t *testing.T) {
 	if got == nil {
 		t.Fatal("getDefaultStagehandInvoker returned nil after swap")
 	}
-	out, err := got.RunTask(context.Background(), RunTaskRequest{Instruction: "x", ModelName: "m", APIKey: "k"})
+	ctx := t.Context()
+	out, err := got.RunTask(ctx, RunTaskRequest{Instruction: "x", ModelName: "m", APIKey: "k"})
 	if err != nil {
 		t.Fatalf("RunTask: %v", err)
 	}

internal/agent/component/streaming_test.go

@@ -27,9 +27,10 @@ import (
 func TestLLM_Stream_HappyPath(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "hello", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
+	ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}
@@ -61,9 +62,10 @@ func TestLLM_Stream_HappyPath(t *testing.T) {
 func TestLLM_Stream_Error(t *testing.T) {
 	stub := &stubInvoker{err: context.DeadlineExceeded}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
+	ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}
@@ -86,7 +88,7 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) {
 	withStubInvoker(t, stub)
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(t.Context())
 	cancel() // pre-cancel
 
 	ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
@@ -111,9 +113,10 @@ func TestLLM_Stream_RespectsCancellation(t *testing.T) {
 func TestLLM_Stream_BufferDoesNotBlock(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	ch, err := c.Stream(context.Background(), nil, map[string]any{"user_prompt": "hi"})
+	ch, err := c.Stream(ctx, nil, map[string]any{"user_prompt": "hi"})
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}

internal/agent/component/string_transform_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"reflect"
 	"testing"
 
@@ -34,7 +33,7 @@ func TestStringTransform_SplitBasic(t *testing.T) {
 		t.Fatalf("NewStringTransformComponent: %v", err)
 	}
 	state := canvas.NewCanvasState("run-1", "task-1")
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"line": "a,b;c"})
 	if err != nil {
@@ -54,7 +53,7 @@ func TestStringTransform_SplitNoDelim(t *testing.T) {
 		"delimiters": []string{","},
 	})
 	state := canvas.NewCanvasState("run-2", "task-2")
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"line": "abc"})
 	if err != nil {
@@ -75,7 +74,7 @@ func TestStringTransform_Merge(t *testing.T) {
 		"script":     "{{x}} and {{y}}",
 	})
 	state := canvas.NewCanvasState("run-3", "task-3")
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{"x": "foo", "y": "bar"})
 	if err != nil {
@@ -95,7 +94,7 @@ func TestStringTransform_MergeIterationAliases(t *testing.T) {
 	state := canvas.NewCanvasState("run-iter", "task-iter")
 	state.Globals["__item__"] = "beta"
 	state.Globals["__index__"] = 1
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{})
 	if err != nil {
@@ -116,7 +115,7 @@ func TestStringTransform_SplitFromStateRef(t *testing.T) {
 	})
 	state := canvas.NewCanvasState("run-4", "task-4")
 	state.Outputs["cpn_0"] = map[string]any{"x": "alpha,beta,gamma"}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, nil)
 	if err != nil {
@@ -138,7 +137,7 @@ func TestStringTransform_MergeMissingPlaceholder(t *testing.T) {
 		"script":     "hello {{name}}",
 	})
 	state := canvas.NewCanvasState("run-5", "task-5")
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{})
 	if err != nil {

internal/agent/component/structured_output_test.go

@@ -102,12 +102,13 @@ func TestLLM_Invoke_OutputStructure_ValidFirstTry(t *testing.T) {
 		Model:   "echo",
 	}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{
 		ModelID:         "echo",
 		OutputStructure: map[string]any{"name": "", "age": 0},
 	})
-	out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"})
+	out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "who?"})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -136,12 +137,13 @@ func TestLLM_Invoke_OutputStructure_RetryOnInvalid(t *testing.T) {
 		onCall: func() { calls++ },
 	}
 	withStubInvoker(t, inv)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{
 		ModelID:         "echo",
 		OutputStructure: map[string]any{"name": ""},
 	})
-	out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "who?"})
+	out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "who?"})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -173,12 +175,13 @@ func TestLLM_Invoke_OutputStructure_RetryStillFails(t *testing.T) {
 		onCall: func() { calls++ },
 	}
 	withStubInvoker(t, inv)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{
 		ModelID:         "echo",
 		OutputStructure: map[string]any{"x": 0},
 	})
-	out, err := c.Invoke(context.Background(), nil, map[string]any{"user_prompt": "go"})
+	out, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "go"})
 	if err != nil {
 		t.Fatalf("Invoke should not error on parse failure: %v", err)
 	}

internal/agent/component/switch_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"testing"
 
 	"ragflow/internal/agent/canvas"
@@ -56,7 +55,7 @@ func TestSwitch_AndMatches(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-1", "task-1")
 	state.Sys["x"] = "yes"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -87,7 +86,7 @@ func TestSwitch_OrMatches(t *testing.T) {
 	state := canvas.NewCanvasState("run-2", "task-2")
 	state.Sys["score"] = "85"
 	state.Sys["flag"] = "no"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -129,7 +128,7 @@ func TestSwitch_DefaultFallback(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-3", "task-3")
 	state.Sys["x"] = "no"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -156,7 +155,7 @@ func TestSwitch_LegacyEndCpnIDsFallback(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-end-cpn", "task-end-cpn")
 	state.Sys["x"] = "no"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -186,7 +185,7 @@ func TestSwitch_ContainsAndEmpty(t *testing.T) {
 	state := canvas.NewCanvasState("run-4", "task-4")
 	state.Sys["body"] = "hello world"
 	state.Sys["opt"] = ""
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -229,7 +228,7 @@ func TestSwitch_LegacyConditionsAndArrayTo(t *testing.T) {
 	})
 	state := canvas.NewCanvasState("run-legacy", "task-legacy")
 	state.SetVar("UserFillUp:Menu", "demo", "loop")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := s.Invoke(ctx, nil, nil)
 	if err != nil {
@@ -253,7 +252,7 @@ func TestSwitch_NilUpstreamContainsEmptyNeedleMatches(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-nil-contains", "task-nil-contains")
 	state.Sys["answer"] = nil
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -284,7 +283,7 @@ func TestSwitch_NilUpstreamContainsNonEmptyDoesNotMatch(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-nil-needle", "task-nil-needle")
 	state.Sys["answer"] = nil
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -316,7 +315,7 @@ func TestSwitch_NilValueContainsDoesNotRaise(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-nil-value", "task-nil-value")
 	state.Sys["answer"] = "foobar"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -350,7 +349,7 @@ func TestSwitch_NilUpstreamStartWithEndWithDoNotCrash(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-nil-start-end", "task-nil-start-end")
 	state.Sys["answer"] = nil
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	for _, tc := range []struct {
 		name string
@@ -393,7 +392,7 @@ func TestSwitch_MultiTargetTo(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-multi", "task-multi")
 	state.SetVar("UserFillUp:Menu", "demo", "data_ops")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -430,7 +429,7 @@ func TestSwitch_MultiTargetTo(t *testing.T) {
 func TestSwitch_EmptyAndConditionFallsThrough(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-empty-and", "task-1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	// Empty clauses: must not match. Should fall through to default.
 	inputs := map[string]any{
@@ -462,7 +461,7 @@ func TestSwitch_EmptyAndConditionFallsThrough(t *testing.T) {
 func TestSwitch_LegacyEmptyItemsFallsThrough(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-legacy-empty", "task-1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{
@@ -496,7 +495,7 @@ func TestSwitch_SatisfiedAndConditionStillRoutes(t *testing.T) {
 	s, _ := NewSwitchComponent(nil)
 	state := canvas.NewCanvasState("run-and-ok", "task-1")
 	state.Sys["greeting"] = "hello world"
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	inputs := map[string]any{
 		"conditions": []any{

internal/agent/component/tool_call_memory_test.go

@@ -31,8 +31,9 @@ import (
 func TestAddToolCallMemory_NoToolCalls(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
-	got, err := addToolCallMemory(context.Background(), nil, AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"})
+	got, err := addToolCallMemory(ctx, nil, AgentParam{ModelID: "echo"}, &schema.Message{Content: "no tools"})
 	if err != nil {
 		t.Fatalf("err: %v", err)
 	}
@@ -70,7 +71,7 @@ func TestAddToolCallMemory_SummarizesAndAppendsToState(t *testing.T) {
 
 	state := runtime.NewCanvasState("rid", "tid")
 	c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1})
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 	_, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"})
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
@@ -112,7 +113,7 @@ func TestAddToolCallMemory_LLMFailure(t *testing.T) {
 
 	state := runtime.NewCanvasState("rid", "tid")
 	c := NewAgentComponent(AgentParam{ModelID: "echo", MaxRounds: 1})
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 	_, err := c.Invoke(ctx, nil, map[string]any{"user_prompt": "do it"})
 	if err != nil {
 		t.Fatalf("Invoke should not error when memory summary fails: %v", err)

internal/agent/component/tool_dispatch_test.go

@@ -64,12 +64,13 @@ func TestPhase3_6_ToolDSLLoading(t *testing.T) {
 		return &schema.Message{Role: schema.Assistant, Content: "ok"}, nil
 	})
 
+	ctx := t.Context()
 	c := NewAgentComponent(AgentParam{
 		ModelID:   "stub",
 		MaxRounds: 1,
 		Tools:     []string{"retrieval"}, // known tool
 	})
-	_, err := c.Invoke(context.Background(), nil, map[string]any{
+	_, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt": "test",
 	})
 	if err != nil {

internal/agent/component/userfillup_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"testing"
 
 	"ragflow/internal/agent/canvas"
@@ -33,7 +32,7 @@ func TestUserFillUp_RendersTips(t *testing.T) {
 		"tips":        "Hello {{name}}",
 	})
 	state := canvas.NewCanvasState("run-1", "task-1")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"inputs": map[string]any{
@@ -57,7 +56,7 @@ func TestUserFillUp_DisableTips(t *testing.T) {
 		"tips":        "Should not render",
 	})
 	state := canvas.NewCanvasState("run-2", "task-2")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"inputs": map[string]any{
@@ -83,7 +82,7 @@ func TestUserFillUp_DisableTips(t *testing.T) {
 func TestUserFillUp_PassesThroughInputs(t *testing.T) {
 	c, _ := New(componentNameUserFillUp, map[string]any{"enable_tips": false})
 	state := canvas.NewCanvasState("run-3", "task-3")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"inputs": map[string]any{
@@ -116,7 +115,7 @@ func TestUserFillUp_FileInputStub(t *testing.T) {
 		"tips":        "Upload {{cv}} please",
 	})
 	state := canvas.NewCanvasState("run-4", "task-4")
-	ctx := withStateForTest(context.Background(), state)
+	ctx := withStateForTest(t.Context(), state)
 
 	out, err := c.Invoke(ctx, nil, map[string]any{
 		"inputs": map[string]any{

internal/agent/component/variable_aggregator_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"testing"
 
 	"ragflow/internal/agent/canvas"
@@ -31,7 +30,7 @@ func TestVariableAggregator_FirstNonEmpty(t *testing.T) {
 	state.Outputs["cpn_1"] = map[string]any{"y": "second-a"}
 	state.Outputs["cpn_2"] = map[string]any{"y": "second-b"}
 	state.Outputs["cpn_3"] = map[string]any{"y": "second-c"}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	groups := []map[string]any{
 		{
@@ -84,7 +83,7 @@ func TestVariableAggregator_SkipsEmptyString(t *testing.T) {
 	state := canvas.NewCanvasState("run-2", "task-2")
 	state.Outputs["cpn_0"] = map[string]any{"x": ""}
 	state.Outputs["cpn_1"] = map[string]any{"y": "picked"}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	groups := []map[string]any{
 		{
@@ -115,7 +114,7 @@ func TestVariableAggregator_MultipleGroups(t *testing.T) {
 	state.Sys["a"] = "alpha"
 	state.Sys["b"] = ""
 	state.Env["c"] = "gamma"
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	groups := []map[string]any{
 		{
@@ -163,7 +162,7 @@ func TestVariableAggregator_MultipleGroups(t *testing.T) {
 func TestVariableAggregator_AllEmpty(t *testing.T) {
 	state := canvas.NewCanvasState("run-4", "task-4")
 	state.Outputs["cpn_0"] = map[string]any{}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	groups := []map[string]any{
 		{

internal/agent/component/variable_assigner_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"reflect"
 	"testing"
 
@@ -28,7 +27,7 @@ import (
 func TestVariableAssigner_Append(t *testing.T) {
 	state := canvas.NewCanvasState("run-1", "task-1")
 	state.Outputs["cpn_0"] = map[string]any{"xs": []any{1, 2}}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{
@@ -62,7 +61,7 @@ func TestVariableAssigner_Overwrite(t *testing.T) {
 	state := canvas.NewCanvasState("run-2", "task-2")
 	state.Outputs["cpn_0"] = map[string]any{"x": "old"}
 	state.Outputs["cpn_1"] = map[string]any{"y": "fresh"}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{
@@ -88,7 +87,7 @@ func TestVariableAssigner_Overwrite(t *testing.T) {
 func TestVariableAssigner_DivideByZero(t *testing.T) {
 	state := canvas.NewCanvasState("run-3", "task-3")
 	state.Outputs["cpn_0"] = map[string]any{"n": 6.0}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{
@@ -133,7 +132,7 @@ func TestVariableAssigner_Clear(t *testing.T) {
 		"c": map[string]any{"k": "v"},
 		"d": 42,
 	}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{"variable": "cpn_0@a", "operator": "clear", "parameter": "x"},
@@ -166,7 +165,7 @@ func TestVariableAssigner_Clear(t *testing.T) {
 func TestVariableAssigner_Arithmetic(t *testing.T) {
 	state := canvas.NewCanvasState("run-5", "task-5")
 	state.Outputs["cpn_0"] = map[string]any{"n": 10.0}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{"variable": "cpn_0@n", "operator": "+=", "parameter": 5},
@@ -191,7 +190,7 @@ func TestVariableAssigner_Arithmetic(t *testing.T) {
 func TestVariableAssigner_RemoveFirstLast(t *testing.T) {
 	state := canvas.NewCanvasState("run-6", "task-6")
 	state.Outputs["cpn_0"] = map[string]any{"xs": []any{"a", "b", "c", "d"}}
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{"variable": "cpn_0@xs", "operator": "remove_first", "parameter": "x"},
@@ -219,7 +218,7 @@ func TestVariableAssigner_RemoveFirstLast(t *testing.T) {
 // TestVariableAssigner_SysTarget: variable="sys.x" → state.Sys is written.
 func TestVariableAssigner_SysTarget(t *testing.T) {
 	state := canvas.NewCanvasState("run-7", "task-7")
-	ctx := canvas.WithState(context.Background(), state)
+	ctx := canvas.WithState(t.Context(), state)
 
 	vars := []map[string]any{
 		{"variable": "sys.x", "operator": "set", "parameter": "hello"},

internal/agent/component/vision_test.go

@@ -17,7 +17,6 @@
 package component
 
 import (
-	"context"
 	"reflect"
 	"testing"
 
@@ -285,10 +284,11 @@ func TestBuildMessagesWithImages_WithImages_UsesUserInputMultiContent(t *testing
 func TestLLM_Invoke_ForwardsImagesToInvoker(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	uri := "data:image/png;base64,iVBORw0KGgo="
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	_, err := c.Invoke(context.Background(), nil, map[string]any{
+	_, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt":  "what is this?",
 		"visual_files": []string{uri},
 	})
@@ -318,9 +318,10 @@ func TestLLM_Invoke_ForwardsImagesToInvoker(t *testing.T) {
 func TestLLM_Invoke_NoVisualFiles_BackwardCompat(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	_, err := c.Invoke(context.Background(), nil, map[string]any{
+	_, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt": "hi",
 	})
 	if err != nil {
@@ -347,10 +348,11 @@ func TestLLM_Invoke_NoVisualFiles_BackwardCompat(t *testing.T) {
 func TestLLM_Invoke_VisualFilesAsString(t *testing.T) {
 	stub := &stubInvoker{resp: &ChatInvokeResponse{Content: "ok", Model: "echo"}}
 	withStubInvoker(t, stub)
+	ctx := t.Context()
 
 	uri := "data:image/jpeg;base64,/9j/4AAQ"
 	c := NewLLMComponent(LLMParam{ModelID: "echo"})
-	_, err := c.Invoke(context.Background(), nil, map[string]any{
+	_, err := c.Invoke(ctx, nil, map[string]any{
 		"user_prompt":  "describe",
 		"visual_files": "see " + uri,
 	})