Go: add context to test (#18070)

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

internal/agent/audio/tts_dispatch_test.go

@@ -70,6 +70,7 @@ func (f *fakeTTSDispatcher) AudioSpeech(
 }
 
 func TestNewTTSDispatchFunc_HappyPath(t *testing.T) {
+	ctx := t.Context()
 	fake := &fakeTTSDispatcher{
 		resp: &modelModule.TTSResponse{Audio: []byte("mp3bytes")},
 		code: common.CodeSuccess,
@@ -78,7 +79,7 @@ func TestNewTTSDispatchFunc_HappyPath(t *testing.T) {
 	if fn == nil {
 		t.Fatal("NewTTSDispatchFunc returned nil for non-nil dispatcher")
 	}
-	resp, err := fn(context.Background(), ModelProviderRequest{
+	resp, err := fn(ctx, ModelProviderRequest{
 		TenantID:  "tenant-1",
 		ModelName: "tts-fish",
 		Text:      "hello world",
@@ -126,12 +127,13 @@ func TestNewTTSDispatchFunc_HappyPath(t *testing.T) {
 }
 
 func TestNewTTSDispatchFunc_EmptyModelName(t *testing.T) {
+	ctx := t.Context()
 	fake := &fakeTTSDispatcher{
 		resp: &modelModule.TTSResponse{Audio: []byte("x")},
 		code: common.CodeSuccess,
 	}
 	fn := NewTTSDispatchFunc(fake)
-	_, err := fn(context.Background(), ModelProviderRequest{
+	_, err := fn(ctx, ModelProviderRequest{
 		TenantID: "t1",
 		Text:     "no model hint",
 	})
@@ -144,6 +146,7 @@ func TestNewTTSDispatchFunc_EmptyModelName(t *testing.T) {
 }
 
 func TestNewTTSDispatchFunc_EmptyVoiceAndLang(t *testing.T) {
+	ctx := t.Context()
 	// Voice and Lang empty → TTSConfig.Params should be nil (not
 	// an empty map) so the model's default voice/lang take effect.
 	fake := &fakeTTSDispatcher{
@@ -151,7 +154,7 @@ func TestNewTTSDispatchFunc_EmptyVoiceAndLang(t *testing.T) {
 		code: common.CodeSuccess,
 	}
 	fn := NewTTSDispatchFunc(fake)
-	_, err := fn(context.Background(), ModelProviderRequest{TenantID: "t1", Text: "hi"})
+	_, err := fn(ctx, ModelProviderRequest{TenantID: "t1", Text: "hi"})
 	if err != nil {
 		t.Fatalf("dispatch: %v", err)
 	}
@@ -164,10 +167,11 @@ func TestNewTTSDispatchFunc_EmptyVoiceAndLang(t *testing.T) {
 }
 
 func TestNewTTSDispatchFunc_DispatcherError(t *testing.T) {
+	ctx := t.Context()
 	sentinel := errors.New("dispatch boom")
 	fake := &fakeTTSDispatcher{err: sentinel}
 	fn := NewTTSDispatchFunc(fake)
-	_, err := fn(context.Background(), ModelProviderRequest{TenantID: "t1", Text: "hi"})
+	_, err := fn(ctx, ModelProviderRequest{TenantID: "t1", Text: "hi"})
 	if err == nil {
 		t.Fatal("expected error from dispatcher, got nil")
 	}
@@ -177,18 +181,20 @@ func TestNewTTSDispatchFunc_DispatcherError(t *testing.T) {
 }
 
 func TestNewTTSDispatchFunc_NonSuccessCode(t *testing.T) {
+	ctx := t.Context()
 	fake := &fakeTTSDispatcher{
 		resp: &modelModule.TTSResponse{Audio: []byte("ignored")},
 		code: common.CodeNotFound,
 	}
 	fn := NewTTSDispatchFunc(fake)
-	_, err := fn(context.Background(), ModelProviderRequest{TenantID: "t1", Text: "hi"})
+	_, err := fn(ctx, ModelProviderRequest{TenantID: "t1", Text: "hi"})
 	if err == nil {
 		t.Fatal("expected error for non-CodeSuccess, got nil")
 	}
 }
 
 func TestNewTTSDispatchFunc_EmptyAudioFromModel(t *testing.T) {
+	ctx := t.Context()
 	// Some buggy model drivers return nil error + nil TTSResponse
 	// (or empty audio). The dispatch must surface that as the
 	// ErrSynthesizeEmpty sentinel so the audio package's caller
@@ -209,7 +215,7 @@ func TestNewTTSDispatchFunc_EmptyAudioFromModel(t *testing.T) {
 				code: common.CodeSuccess,
 			}
 			fn := NewTTSDispatchFunc(fake)
-			_, err := fn(context.Background(), ModelProviderRequest{TenantID: "t1", Text: "hi"})
+			_, err := fn(ctx, ModelProviderRequest{TenantID: "t1", Text: "hi"})
 			if !errors.Is(err, ErrSynthesizeEmpty) {
 				t.Errorf("err = %v, want ErrSynthesizeEmpty", err)
 			}

internal/agent/audio/tts_test.go

@@ -25,11 +25,12 @@ import (
 // TestStubSynth_EmptyEngine: an empty engine returns
 // ErrTTSEngineNotConfigured (the deferred-state sentinel).
 func TestStubSynth_EmptyEngine(t *testing.T) {
+	ctx := t.Context()
 	// Ensure the stub is installed (in case a previous test
 	// registered a different one).
 	SetSynthesizer(nil)
 	synth := GetSynthesizer()
-	_, err := synth.Synthesize(context.Background(), SynthesizeRequest{
+	_, err := synth.Synthesize(ctx, SynthesizeRequest{
 		Engine: EngineEmpty,
 		Text:   "hi",
 	})
@@ -41,9 +42,10 @@ func TestStubSynth_EmptyEngine(t *testing.T) {
 // TestStubSynth_UnknownEngine: a non-empty unknown engine
 // returns ErrTTSUnsupportedEngine.
 func TestStubSynth_UnknownEngine(t *testing.T) {
+	ctx := t.Context()
 	SetSynthesizer(nil)
 	synth := GetSynthesizer()
-	_, err := synth.Synthesize(context.Background(), SynthesizeRequest{
+	_, err := synth.Synthesize(ctx, SynthesizeRequest{
 		Engine: Engine("unknown-engine"),
 		Text:   "hi",
 	})
@@ -55,6 +57,7 @@ func TestStubSynth_UnknownEngine(t *testing.T) {
 // TestSetSynthesizer_Roundtrip: a custom synthesizer set via
 // SetSynthesizer is returned by GetSynthesizer.
 func TestSetSynthesizer_Roundtrip(t *testing.T) {
+	ctx := t.Context()
 	var called bool
 	custom := &fakeSynth{called: &called}
 	SetSynthesizer(custom)
@@ -63,7 +66,7 @@ func TestSetSynthesizer_Roundtrip(t *testing.T) {
 	if got != custom {
 		t.Fatalf("synthesizer not registered")
 	}
-	resp, err := got.Synthesize(context.Background(), SynthesizeRequest{
+	resp, err := got.Synthesize(ctx, SynthesizeRequest{
 		Engine: EngineGTTS,
 		Text:   "hi",
 	})

internal/agent/canvas/cancel_test.go

@@ -77,8 +77,9 @@ func TestWatchCancel_FiresAfterRequest(t *testing.T) {
 }
 
 func TestWatchCancel_StopsOnContextCancel(t *testing.T) {
+	ctx := t.Context()
 	withCancelClient(t)
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 
 	sessionID := "session_test_ctx"
 	done := make(chan struct{})
@@ -102,8 +103,9 @@ func TestWatchCancel_StopsOnContextCancel(t *testing.T) {
 }
 
 func TestWatchCancel_OnCancelNotInvokedForEmptyKey(t *testing.T) {
+	ctx := t.Context()
 	withCancelClient(t)
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
 	invoked := atomic.Int32{}
@@ -132,8 +134,7 @@ func TestRequestCancel_EmptyValueStillFires(t *testing.T) {
 	// we rely on RequestCancel to always set "x" so this test is just
 	// a sanity check that the value round-trips.
 	mr := withCancelClient(t)
-	ctx := context.Background()
-
+	ctx := t.Context()
 	if err := RequestCancel(ctx, "session_value"); err != nil {
 		t.Fatalf("RequestCancel: %v", err)
 	}
@@ -151,7 +152,7 @@ func TestRequestCancel_EmptyValueStillFires(t *testing.T) {
 
 func TestCancelRequested(t *testing.T) {
 	withCancelClient(t)
-	ctx := context.Background()
+	ctx := t.Context()
 
 	requested, err := CancelRequested(ctx, "session-check")
 	if err != nil || requested {

internal/agent/canvas/canvas_test.go

@@ -16,7 +16,6 @@
 package canvas
 
 import (
-	"context"
 	"strings"
 	"testing"
 )
@@ -44,7 +43,8 @@ func TestBeginToMessage_Smoke(t *testing.T) {
 		Path: []string{"begin_0", "message_0"},
 	}
 
-	cc, err := Compile(context.Background(), dsl)
+	ctx := t.Context()
+	cc, err := Compile(ctx, dsl)
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}
@@ -63,13 +63,13 @@ func TestBeginToMessage_Smoke(t *testing.T) {
 
 	// Stash runState on the context so the canvas runner can extract
 	// it via GetStateFromContext.
-	ctx := withState(context.Background(), runState)
+	newCtx := withState(ctx, runState)
 
 	// Invoke with the seed input. The "query" key flows into Begin's
 	// Invoke and is written to state.Sys["query"], where Message's
 	// ResolveTemplate of "{{sys.query}}" will read it.
 	in := map[string]any{"query": "world"}
-	out, err := cc.Workflow.Invoke(ctx, in)
+	out, err := cc.Workflow.Invoke(newCtx, in)
 	if err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
@@ -113,7 +113,8 @@ func TestBuildWorkflow_PreservesRequestSysValues(t *testing.T) {
 		Path: []string{"begin_0", "message_0"},
 	}
 
-	cc, err := Compile(context.Background(), dsl)
+	ctx := t.Context()
+	cc, err := Compile(ctx, dsl)
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}
@@ -121,8 +122,8 @@ func TestBuildWorkflow_PreservesRequestSysValues(t *testing.T) {
 	state := NewCanvasState("run-request-sys", "task-request-sys")
 	state.Sys["files"] = []string{"File: notes.txt\nContent as following:\nhello"}
 	state.Sys["user_id"] = "user-1"
-	ctx := withState(context.Background(), state)
-	if _, err := cc.Workflow.Invoke(ctx, map[string]any{"query": "hello"}); err != nil {
+	newCtx := withState(ctx, state)
+	if _, err = cc.Workflow.Invoke(newCtx, map[string]any{"query": "hello"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 

internal/agent/canvas/checkpoint_store_test.go

@@ -17,7 +17,6 @@
 package canvas
 
 import (
-	"context"
 	"testing"
 	"time"
 
@@ -46,7 +45,7 @@ func newTestStore(t *testing.T, ttl time.Duration) (*RedisCheckPointStore, *mini
 
 func TestRedisCheckPointStore_RoundTrip(t *testing.T) {
 	store, _ := newTestStore(t, 30*24*time.Hour)
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// missing key → (nil, false, nil)
 	got, ok, err := store.Get(ctx, "absent")
@@ -83,7 +82,7 @@ func TestRedisCheckPointStore_RoundTrip(t *testing.T) {
 
 func TestRedisCheckPointStore_TTL(t *testing.T) {
 	store, mr := newTestStore(t, 2*time.Second)
-	ctx := context.Background()
+	ctx := t.Context()
 
 	if err := store.Set(ctx, "cpn_ttl", []byte("x")); err != nil {
 		t.Fatalf("Set: %v", err)
@@ -105,13 +104,13 @@ func TestRedisCheckPointStore_TTL(t *testing.T) {
 
 func TestRedisCheckPointStore_Delete(t *testing.T) {
 	store, _ := newTestStore(t, time.Minute)
-	ctx := context.Background()
+	ctx := t.Context()
 
 	// Delete on missing key is a no-op (no error).
 	if err := store.Delete(ctx, "absent"); err != nil {
 		t.Fatalf("Delete absent: %v", err)
 	}
-	// Set then Delete then Get → missing.
+	// Set then Delete Get → missing.
 	if err := store.Set(ctx, "cpn_del", []byte("payload")); err != nil {
 		t.Fatalf("Set: %v", err)
 	}
@@ -127,7 +126,7 @@ func TestRedisCheckPointStore_NilClient(t *testing.T) {
 	// Cache uninitialized → NewRedisCheckPointStore returns a store with
 	// nil client. Operations must error rather than panic.
 	store := &RedisCheckPointStore{client: nil, ttl: time.Minute}
-	ctx := context.Background()
+	ctx := t.Context()
 
 	if _, _, err := store.Get(ctx, "x"); err == nil {
 		t.Fatal("Get with nil client: err = nil, want error")

internal/agent/canvas/compile_test.go

@@ -17,7 +17,6 @@ package canvas
 
 import (
 	"bytes"
-	"context"
 	"sort"
 	"strings"
 	"testing"
@@ -69,9 +68,11 @@ func TestCompile_LogsWhenLegacyNodesPresent(t *testing.T) {
 		},
 	}
 
+	ctx := t.Context()
+
 	// Compile may return an error from downstream BuildWorkflow —
 	// we ignore it; the assertion is on the log line.
-	_, _ = Compile(context.Background(), c)
+	_, _ = Compile(ctx, c)
 
 	got := buf.String()
 	if !strings.Contains(got, "LoopItem/IterationItem") {
@@ -110,10 +111,11 @@ func TestCompile_NoLogOnCleanCanvas(t *testing.T) {
 		},
 	}
 
+	ctx := t.Context()
 	// We don't fail on Compile's own error (it may fail for many
 	// reasons unrelated to legacy names); the assertion is on the
 	// absence of the legacy log line.
-	_, _ = Compile(context.Background(), c)
+	_, _ = Compile(ctx, c)
 
 	got := buf.String()
 	if strings.Contains(got, "LoopItem/IterationItem") {
@@ -181,7 +183,8 @@ func TestCompile_RejectsUserFillUpInResumeMode(t *testing.T) {
 			"end":   {Obj: CanvasComponentObj{ComponentName: "Answer", Params: map[string]any{}}},
 		},
 	}
-	_, err := Compile(context.Background(), c, WithInterruptAfterNonTerminalCpn())
+	ctx := t.Context()
+	_, err := Compile(ctx, c, WithInterruptAfterNonTerminalCpn())
 	if err == nil {
 		t.Fatal("expected Compile to reject UserFillUp in resume mode, got nil error")
 	}
@@ -204,7 +207,8 @@ func TestCompile_PropagatesCheckPointID(t *testing.T) {
 			"answer:0": {Obj: CanvasComponentObj{ComponentName: "Answer", Params: map[string]any{}}},
 		},
 	}
-	compiled, err := Compile(context.Background(), c, WithCheckPointID("task-9"))
+	ctx := t.Context()
+	compiled, err := Compile(ctx, c, WithCheckPointID("task-9"))
 	if err != nil {
 		t.Skipf("skipping propagation assertion: canvas did not compile in unit scope: %v", err)
 	}

internal/agent/canvas/fixture_compile_test.go

@@ -1,7 +1,6 @@
 package canvas
 
 import (
-	"context"
 	"encoding/json"
 	"os"
 	"path/filepath"
@@ -85,7 +84,8 @@ func TestAllFixture_NormalizeAndCompile(t *testing.T) {
 			Upstream:   stringSliceFromAny(comp["upstream"]),
 		}
 	}
-	cc, err := Compile(context.Background(), c)
+	ctx := t.Context()
+	cc, err := Compile(ctx, c)
 	if err != nil {
 		t.Fatalf("Compile(all.json): %v", err)
 	}

internal/agent/canvas/interrupt_resume_test.go

@@ -26,8 +26,9 @@ import (
 	"strings"
 	"testing"
 
-	"github.com/cloudwego/eino/compose"
 	"ragflow/internal/agent/workflowx"
+
+	"github.com/cloudwego/eino/compose"
 )
 
 // TestBuildInputSpec_BasicFields passes enable_tips/tips/inputs and
@@ -56,7 +57,7 @@ func TestBuildInputSpec_BasicFields(t *testing.T) {
 func TestUserFillUpNodeBody_ResolvesTipsFromCanvasState(t *testing.T) {
 	state := NewCanvasState("run-1", "task-1")
 	state.SetVar("Agent:MoodyIdeasMarry", "content", "How old are you?")
-	ctx := WithState(context.Background(), state)
+	ctx := WithState(t.Context(), state)
 
 	body := UserFillUpNodeBody("UserFillUp:TwelveBadgersRescue", map[string]any{
 		"enable_tips": true,
@@ -85,7 +86,7 @@ func TestUserFillUpNodeBody_OmitsDisabledTips(t *testing.T) {
 		"enable_tips": false,
 		"tips":        "should not be shown",
 	})
-	_, err := body(context.Background(), nil)
+	_, err := body(t.Context(), nil)
 	if err == nil {
 		t.Fatal("UserFillUp should interrupt while waiting for input")
 	}
@@ -165,10 +166,10 @@ func TestExtractInterruptContexts_FlattensSubGraphs(t *testing.T) {
 	}))
 	subNode.AddInput(compose.START)
 	sub.End().AddInput("waiter")
-
+	ctx := t.Context()
 	outer := compose.NewWorkflow[int, int]()
 	loopNode, err := workflowx.AddLoopNode(
-		context.Background(),
+		ctx,
 		outer,
 		"loop",
 		sub,
@@ -180,11 +181,11 @@ func TestExtractInterruptContexts_FlattensSubGraphs(t *testing.T) {
 	loopNode.AddInput(compose.START)
 	outer.End().AddInput("loop")
 
-	compiled, err := outer.Compile(context.Background())
+	compiled, err := outer.Compile(ctx)
 	if err != nil {
 		t.Fatalf("compile: %v", err)
 	}
-	_, err = compiled.Invoke(context.Background(), 0)
+	_, err = compiled.Invoke(ctx, 0)
 	if err == nil {
 		t.Fatal("expected interrupt error, got nil")
 	}
@@ -194,8 +195,8 @@ func TestExtractInterruptContexts_FlattensSubGraphs(t *testing.T) {
 		t.Fatalf("len(ExtractInterruptContexts) = %d; want >= 1", len(got))
 	}
 	foundUserFillUp := false
-	for _, ctx := range got {
-		if info, ok := ctx.Info.(map[string]any); ok {
+	for _, interruptCtx := range got {
+		if info, ok := interruptCtx.Info.(map[string]any); ok {
 			if kind, _ := info["kind"].(string); kind == "user_fill_up" {
 				foundUserFillUp = true
 				break
@@ -279,7 +280,7 @@ func TestUserFillUpNodeBody_FirstCallInterrupts(t *testing.T) {
 		"enable_tips": true,
 		"tips":        "hello",
 	})
-	_, err := body(context.Background(), map[string]any{"x": 1})
+	_, err := body(t.Context(), map[string]any{"x": 1})
 	if err == nil {
 		t.Fatalf("UserFillUpNodeBody first call returned nil err; want interrupt signal")
 	}
@@ -312,7 +313,7 @@ func TestUserFillUpNodeBody_ResumeReturnsInput(t *testing.T) {
 	// string form of the node's address. We pass the cpnID as the
 	// address — that's what UserFillUpNodeBody advertises when it
 	// composes its output.
-	ctx := compose.ResumeWithData(context.Background(), "ufu_1", "user typed this")
+	ctx := compose.ResumeWithData(t.Context(), "ufu_1", "user typed this")
 
 	_, err := body(ctx, map[string]any{"x": 1})
 	// Outside an engine runner, GetResumeContext cannot match the
@@ -387,7 +388,7 @@ func TestBuildUserFillUpResumeOutput_ValueFieldMirrorsResumeData(t *testing.T) {
 func TestUserFillUpNodeBody_DoesNotConsumeSysQuery(t *testing.T) {
 	state := NewCanvasState("run-1", "task-1")
 	state.Sys["query"] = "loop"
-	ctx := WithState(context.Background(), state)
+	ctx := WithState(t.Context(), state)
 
 	body := UserFillUpNodeBody("ufu_1", map[string]any{
 		"inputs": map[string]any{

internal/agent/canvas/loop_semantics_test.go

@@ -52,13 +52,14 @@ import (
 // assert per-iteration writes landed.
 func runLoopCanvas(t *testing.T, dsl *Canvas) (*CanvasState, error) {
 	t.Helper()
-	cc, err := Compile(context.Background(), dsl)
+	ctx := t.Context()
+	cc, err := Compile(ctx, dsl)
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}
 	state := NewCanvasState("run-loop", "task-loop")
-	ctx := withState(context.Background(), state)
-	_, runErr := cc.Workflow.Invoke(ctx, map[string]any{"query": "go"})
+	newCtx := withState(ctx, state)
+	_, runErr := cc.Workflow.Invoke(newCtx, map[string]any{"query": "go"})
 	return state, runErr
 }
 
@@ -174,13 +175,14 @@ func TestLoop_MaxCount(t *testing.T) {
 // canvas must therefore install workflowx loops in every_iteration mode rather
 // than buffering and exposing only the final body result.
 func TestLoop_StreamEmitsEveryIteration(t *testing.T) {
-	cc, err := Compile(context.Background(), counterLoopDSL(1, 3, 50))
+	ctx := t.Context()
+	cc, err := Compile(ctx, counterLoopDSL(1, 3, 50))
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}
 	state := NewCanvasState("run-loop-stream", "task-loop-stream")
-	ctx := withState(context.Background(), state)
-	sr, err := cc.Workflow.Stream(ctx, map[string]any{"query": "go"})
+	newCtx := withState(ctx, state)
+	sr, err := cc.Workflow.Stream(newCtx, map[string]any{"query": "go"})
 	if err != nil {
 		t.Fatalf("Stream: %v", err)
 	}
@@ -210,16 +212,17 @@ func TestLoop_StreamEmitsEveryIteration(t *testing.T) {
 }
 
 func TestLoop_EmitsLifecycleEventsForMacroAndBody(t *testing.T) {
-	cc, err := Compile(context.Background(), counterLoopDSL(1, 3, 50))
+	ctx := t.Context()
+	cc, err := Compile(ctx, counterLoopDSL(1, 3, 50))
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}
 	state := NewCanvasState("run-loop-events", "task-loop-events")
 	events := make(chan RunEvent, 32)
-	ctx := withState(context.Background(), state)
-	ctx = WithRunMeta(ctx, &RunMeta{Events: events})
+	newCtx := withState(ctx, state)
+	newCtx = WithRunMeta(newCtx, &RunMeta{Events: events})
 
-	if _, err := cc.Workflow.Invoke(ctx, map[string]any{"query": "go"}); err != nil {
+	if _, err = cc.Workflow.Invoke(newCtx, map[string]any{"query": "go"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 
@@ -239,6 +242,7 @@ func TestLoop_EmitsLifecycleEventsForMacroAndBody(t *testing.T) {
 }
 
 func TestLoop_MessageEmitsEveryIteration(t *testing.T) {
+	ctx := t.Context()
 	dsl := counterLoopDSL(1, 3, 50)
 	bump := dsl.Components["bump"]
 	bump.Downstream = []string{"msg"}
@@ -258,13 +262,13 @@ func TestLoop_MessageEmitsEveryIteration(t *testing.T) {
 		t.Fatalf("Compile: %v", err)
 	}
 	state := NewCanvasState("run-loop-message", "task-loop-message")
-	ctx := withState(context.Background(), state)
+	newCtx := withState(ctx, state)
 	var emitted []string
-	ctx = runtime.WithCanvasMessageEmitter(ctx, func(content string) {
+	newCtx = runtime.WithCanvasMessageEmitter(newCtx, func(content string) {
 		emitted = append(emitted, content)
 	})
 
-	if _, err := cc.Workflow.Invoke(ctx, map[string]any{"query": "go"}); err != nil {
+	if _, err = cc.Workflow.Invoke(newCtx, map[string]any{"query": "go"}); err != nil {
 		t.Fatalf("Invoke: %v", err)
 	}
 	if len(emitted) != 3 {
@@ -341,7 +345,8 @@ func TestLoop_FactoryErrorSurfaces(t *testing.T) {
 			},
 		},
 	}
-	_, err := Compile(context.Background(), dsl)
+	ctx := t.Context()
+	_, err := Compile(ctx, dsl)
 	if err == nil {
 		t.Fatal("expected factory error, got nil")
 	}
@@ -373,7 +378,8 @@ func TestLoop_LegacyExitLoopStaysNoOp(t *testing.T) {
 			},
 		},
 	}
-	if _, err := Compile(context.Background(), dsl); err != nil {
+	ctx := t.Context()
+	if _, err := Compile(ctx, dsl); err != nil {
 		t.Fatalf("Compile with legacy ExitLoop (factory registered): %v", err)
 	}
 	// Also verify the factory IS registered — otherwise this test

internal/agent/canvas/loop_subgraph_test.go

@@ -33,7 +33,6 @@
 package canvas
 
 import (
-	"context"
 	"strings"
 	"testing"
 )
@@ -341,7 +340,7 @@ func TestTranslateLoopCondition_SingleOp(t *testing.T) {
 	}
 	state := NewCanvasState("", "")
 	state.SetVar("loop_0", "counter", 3)
-	ctx := WithState(context.Background(), state)
+	ctx := WithState(t.Context(), state)
 	quit, err := cond(ctx, 3, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
@@ -352,7 +351,7 @@ func TestTranslateLoopCondition_SingleOp(t *testing.T) {
 	// counter=2 should NOT quit.
 	state2 := NewCanvasState("", "")
 	state2.SetVar("loop_0", "counter", 2)
-	ctx2 := WithState(context.Background(), state2)
+	ctx2 := WithState(t.Context(), state2)
 	quit, err = cond(ctx2, 2, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
@@ -379,7 +378,7 @@ func TestTranslateLoopCondition_OrQuitsEarly(t *testing.T) {
 	state := NewCanvasState("", "")
 	state.SetVar("L", "a", 1)
 	state.SetVar("L", "b", 0)
-	quit, err := cond(WithState(context.Background(), state), 1, nil, nil)
+	quit, err := cond(WithState(t.Context(), state), 1, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -390,7 +389,7 @@ func TestTranslateLoopCondition_OrQuitsEarly(t *testing.T) {
 	state2 := NewCanvasState("", "")
 	state2.SetVar("L", "a", 0)
 	state2.SetVar("L", "b", 2)
-	quit, err = cond(WithState(context.Background(), state2), 1, nil, nil)
+	quit, err = cond(WithState(t.Context(), state2), 1, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -401,7 +400,7 @@ func TestTranslateLoopCondition_OrQuitsEarly(t *testing.T) {
 	state3 := NewCanvasState("", "")
 	state3.SetVar("L", "a", 0)
 	state3.SetVar("L", "b", 0)
-	quit, err = cond(WithState(context.Background(), state3), 1, nil, nil)
+	quit, err = cond(WithState(t.Context(), state3), 1, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -425,7 +424,7 @@ func TestTranslateLoopCondition_AndRequiresAll(t *testing.T) {
 	state := NewCanvasState("", "")
 	state.SetVar("L", "a", 1)
 	state.SetVar("L", "b", 2)
-	quit, _ := cond(WithState(context.Background(), state), 1, nil, nil)
+	quit, _ := cond(WithState(t.Context(), state), 1, nil, nil)
 	if !quit {
 		t.Errorf("AND with both true should quit")
 	}
@@ -433,7 +432,7 @@ func TestTranslateLoopCondition_AndRequiresAll(t *testing.T) {
 	state2 := NewCanvasState("", "")
 	state2.SetVar("L", "a", 1)
 	state2.SetVar("L", "b", 0)
-	quit, _ = cond(WithState(context.Background(), state2), 1, nil, nil)
+	quit, _ = cond(WithState(t.Context(), state2), 1, nil, nil)
 	if quit {
 		t.Errorf("AND with one false should not quit")
 	}
@@ -448,7 +447,7 @@ func TestTranslateLoopCondition_EmptyConditionsNeverQuit(t *testing.T) {
 		t.Fatalf("translate: %v", err)
 	}
 	state := NewCanvasState("", "")
-	quit, err := cond(WithState(context.Background(), state), 1, nil, nil)
+	quit, err := cond(WithState(t.Context(), state), 1, nil, nil)
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -502,7 +501,7 @@ func TestTranslateLoopCondition_VariableInputMode(t *testing.T) {
 	state := NewCanvasState("", "")
 	state.SetVar("L", "counter", 10)
 	state.SetVar("Begin", "threshold", 5)
-	quit, _ := cond(WithState(context.Background(), state), 1, nil, nil)
+	quit, _ := cond(WithState(t.Context(), state), 1, nil, nil)
 	if !quit {
 		t.Errorf("counter(10) >= threshold(5) should quit")
 	}
@@ -588,7 +587,7 @@ func TestBuildWorkflow_LoopInstallsOneNode(t *testing.T) {
 				Upstream: []string{"loop"}},
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow: %v", err)
 	}
 }
@@ -605,7 +604,7 @@ func TestBuildWorkflow_LegacyExitLoop(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow with ExitLoop: %v", err)
 	}
 }
@@ -685,7 +684,7 @@ func TestBuildWorkflow_LoopExitLoopDoesNotBecomeTerminal(t *testing.T) {
 			"continue":  "loop",
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow with grouped loop ExitLoop: %v", err)
 	}
 }
@@ -703,7 +702,7 @@ func TestBuildWorkflow_UnknownComponentErrors(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	_, err := BuildWorkflow(context.Background(), c)
+	_, err := BuildWorkflow(t.Context(), c)
 	if err == nil {
 		t.Fatal("expected error on unknown component name, got nil")
 	}
@@ -725,7 +724,7 @@ func TestBuildWorkflow_EmptyComponentNameErrors(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	_, err := BuildWorkflow(context.Background(), c)
+	_, err := BuildWorkflow(t.Context(), c)
 	if err == nil {
 		t.Fatal("expected error on empty component_name, got nil")
 	}
@@ -779,7 +778,7 @@ func TestBuildWorkflow_LoopSharesOuterCanvasState(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	exp, err := buildLoopExpansion(context.Background(), c, "loop")
+	exp, err := buildLoopExpansion(t.Context(), c, "loop")
 	if err != nil {
 		t.Fatalf("buildLoopExpansion: %v", err)
 	}
@@ -800,7 +799,7 @@ func TestBuildWorkflow_LoopSharesOuterCanvasState(t *testing.T) {
 	// performs, and confirm the mutation is visible to a
 	// LoopCondition-style reader on the SAME *CanvasState.
 	state := NewCanvasState("run-1", "task-1")
-	ctx := WithState(context.Background(), state)
+	ctx := WithState(t.Context(), state)
 
 	got, _, err := GetStateFromContext[*CanvasState](ctx)
 	if err != nil {
@@ -861,7 +860,7 @@ func TestBuildWorkflow_LoopWithBody(t *testing.T) {
 				Upstream: []string{"a"}},
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow: %v", err)
 	}
 }
@@ -914,7 +913,7 @@ func TestBuildWorkflow_LoopBodyWithMultiTerminalCompiles(t *testing.T) {
 		},
 	}
 
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow with loop multi-terminal body: %v", err)
 	}
 }
@@ -933,7 +932,7 @@ func TestBuildWorkflow_LoopMissingParams(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err != nil {
+	if _, err := BuildWorkflow(t.Context(), c); err != nil {
 		t.Fatalf("BuildWorkflow: %v", err)
 	}
 }
@@ -954,7 +953,7 @@ func TestBuildWorkflow_LoopIncompleteCondition(t *testing.T) {
 				Upstream: []string{"begin"}},
 		},
 	}
-	if _, err := BuildWorkflow(context.Background(), c); err == nil {
+	if _, err := BuildWorkflow(t.Context(), c); err == nil {
 		t.Errorf("expected error on incomplete condition")
 	}
 }

internal/agent/canvas/multibranch_test.go

@@ -39,7 +39,6 @@
 package canvas
 
 import (
-	"context"
 	"testing"
 
 	"github.com/cloudwego/eino/compose"
@@ -51,7 +50,7 @@ import (
 // no chosen end-nodes and skips routing.
 func TestMakeSwitchBranchCondition_MissingField(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true, "b": true})
-	got, err := cond(context.Background(), map[string]any{"other": "x"})
+	got, err := cond(t.Context(), map[string]any{"other": "x"})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -64,7 +63,7 @@ func TestMakeSwitchBranchCondition_MissingField(t *testing.T) {
 // the same as missing.
 func TestMakeSwitchBranchCondition_EmptyString(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true})
-	got, err := cond(context.Background(), map[string]any{"_next": ""})
+	got, err := cond(t.Context(), map[string]any{"_next": ""})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -79,7 +78,7 @@ func TestMakeSwitchBranchCondition_EmptyString(t *testing.T) {
 // a []any (list of strings).
 func TestMakeSwitchBranchCondition_WrongType(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true})
-	got, err := cond(context.Background(), map[string]any{"_next": 42})
+	got, err := cond(t.Context(), map[string]any{"_next": 42})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -95,7 +94,7 @@ func TestMakeSwitchBranchCondition_WrongType(t *testing.T) {
 // end node" at runtime and crash the run.
 func TestMakeSwitchBranchCondition_UnknownKey(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true, "b": true})
-	got, err := cond(context.Background(), map[string]any{"_next": "ghost"})
+	got, err := cond(t.Context(), map[string]any{"_next": "ghost"})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -108,7 +107,7 @@ func TestMakeSwitchBranchCondition_UnknownKey(t *testing.T) {
 // cpn_id is passed through as a single-entry map.
 func TestMakeSwitchBranchCondition_KnownKey(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true, "b": true})
-	got, err := cond(context.Background(), map[string]any{"_next": "b"})
+	got, err := cond(t.Context(), map[string]any{"_next": "b"})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -123,7 +122,7 @@ func TestMakeSwitchBranchCondition_KnownKey(t *testing.T) {
 // dropped.
 func TestMakeSwitchBranchCondition_MultiTargetList(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true, "b": true})
-	got, err := cond(context.Background(), map[string]any{"_next": []any{"a", "b", "ghost"}})
+	got, err := cond(t.Context(), map[string]any{"_next": []any{"a", "b", "ghost"}})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -136,7 +135,7 @@ func TestMakeSwitchBranchCondition_MultiTargetList(t *testing.T) {
 // treated as no branch chosen.
 func TestMakeSwitchBranchCondition_EmptyList(t *testing.T) {
 	cond := makeSwitchBranchCondition(map[string]bool{"a": true})
-	got, err := cond(context.Background(), map[string]any{"_next": []any{}})
+	got, err := cond(t.Context(), map[string]any{"_next": []any{}})
 	if err != nil {
 		t.Fatalf("cond: %v", err)
 	}
@@ -328,7 +327,7 @@ func TestMultiBranch_CompileSucceeds(t *testing.T) {
 			},
 		},
 	}
-	cc, err := Compile(context.Background(), dsl)
+	cc, err := Compile(t.Context(), dsl)
 	if err != nil {
 		t.Fatalf("Compile: %v", err)
 	}