Go: fix context (#18118)

Signed-off-by: Jin Hai <haijin.chn@gmail.com>
2026-08-15 21:16:03 +08:00 · 2026-08-11 19:19:29 +08:00
parent 8bd5768ebc
commit c75edbfbe8
36 changed files with 492 additions and 352 deletions
--- a/internal/agent/tool/akshare_test.go
+++ b/internal/agent/tool/akshare_test.go
@@ -67,8 +67,10 @@ func TestAkShare_FetchesStockNews(t *testing.T) {
 	akshareStockNewsEndpoint = srv.URL + "/search/jsonp"
 	defer func() { akshareStockNewsEndpoint = oldEndpoint }()

+	ctx := t.Context()
+
 	tool := NewAkShareToolWithTopN(NewHTTPHelper(), 2)
-	out, err := tool.InvokableRun(context.Background(), `{"query":"600519"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"600519"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v (out=%s)", err, out)
 	}
@@ -105,9 +107,10 @@ func TestAkShare_ParseTruncatesToTopN(t *testing.T) {

 func TestAkShare_RejectsMalformedJSON(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewAkShareTool()
-	_, err := tool.InvokableRun(context.Background(), `{not json`)
+	_, err := tool.InvokableRun(ctx, `{not json`)
 	if err == nil {
 		t.Fatal("expected error for malformed JSON, got nil")
 	}
@@ -118,9 +121,10 @@ func TestAkShare_RejectsMalformedJSON(t *testing.T) {

 func TestAkShare_RejectsMissingQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewAkShareTool()
-	out, err := tool.InvokableRun(context.Background(), `{}`)
+	out, err := tool.InvokableRun(ctx, `{}`)
 	if err == nil {
 		t.Fatalf("expected error for missing query, got nil (out=%s)", out)
 	}
--- a/internal/agent/tool/arxiv_test.go
+++ b/internal/agent/tool/arxiv_test.go
@@ -148,9 +148,10 @@ func TestArxiv_Info(t *testing.T) {

 func TestArxiv_EmptyQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewArxivTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty): %v", err)
 	}
@@ -180,6 +181,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) {
 		_, _ = w.Write([]byte(canned))
 	}))
 	defer srv.Close()
+	ctx := t.Context()

 	// rewriteHostTransport points the hard-coded export.arxiv.org at the
 	// test server.
@@ -187,7 +189,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewArxivToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(), `{"query":"rag"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"rag"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -212,6 +214,7 @@ func TestArxiv_FullRoundtrip(t *testing.T) {

 func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	built, err := BuildByName("arxiv", map[string]any{
 		"top_n":   float64(7),
@@ -232,7 +235,7 @@ func TestArxiv_ComponentReferencesAndDefaults(t *testing.T) {
 	envelope := map[string]any{"results": []any{map[string]any{
 		"title": "Paper", "summary": "Paper summary.", "pdf_url": "https://arxiv.org/pdf/1", "entry_id": "kept",
 	}}}
-	chunks, docAggs := arxiv.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := arxiv.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Paper summary." {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
--- a/internal/agent/tool/code_exec_test.go
+++ b/internal/agent/tool/code_exec_test.go
@@ -26,9 +26,10 @@ import (

 func TestCodeExec_StubsErrorWhenClientMissing(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCodeExecTool()
-	out, err := c.InvokableRun(context.Background(), `{"language":"python","code":"def main(): return {}"}`)
+	out, err := c.InvokableRun(ctx, `{"language":"python","code":"def main(): return {}"}`)
 	if !errors.Is(err, ErrCodeExecSandboxMissing) {
 		t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err)
 	}
@@ -47,9 +48,10 @@ func TestCodeExec_StubsErrorWhenClientMissing(t *testing.T) {

 func TestCodeExec_RejectsEmptyCode(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCodeExecTool()
-	_, err := c.InvokableRun(context.Background(), `{"language":"python","code":""}`)
+	_, err := c.InvokableRun(ctx, `{"language":"python","code":""}`)
 	if err == nil || !strings.Contains(err.Error(), "code") {
 		t.Fatalf("err = %v, want to mention empty code", err)
 	}
@@ -57,9 +59,10 @@ func TestCodeExec_RejectsEmptyCode(t *testing.T) {

 func TestCodeExec_RejectsBadLanguage(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCodeExecTool()
-	_, err := c.InvokableRun(context.Background(), `{"language":"brainfuck","code":"x"}`)
+	_, err := c.InvokableRun(ctx, `{"language":"brainfuck","code":"x"}`)
 	if err == nil || !strings.Contains(err.Error(), "language") {
 		t.Fatalf("err = %v, want to reject unsupported language", err)
 	}
@@ -67,11 +70,12 @@ func TestCodeExec_RejectsBadLanguage(t *testing.T) {

 func TestCodeExec_AcceptsLangAlias(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCodeExecTool()
 	// Python tool also accepts "lang" as the field name; the Go shell
 	// should still reach the stub branch.
-	_, err := c.InvokableRun(context.Background(), `{"lang":"nodejs","script":"async function main() {}"}`)
+	_, err := c.InvokableRun(ctx, `{"lang":"nodejs","script":"async function main() {}"}`)
 	if !errors.Is(err, ErrCodeExecSandboxMissing) {
 		t.Fatalf("err = %v, want ErrCodeExecSandboxMissing", err)
 	}
@@ -79,9 +83,10 @@ func TestCodeExec_AcceptsLangAlias(t *testing.T) {

 func TestCodeExec_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCodeExecTool()
-	info, err := c.Info(context.Background())
+	info, err := c.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -101,7 +106,7 @@ func TestCodeExec_Info(t *testing.T) {
 		t.Fatalf("marshal Info schema: %v", err)
 	}
 	var schema map[string]any
-	if err := json.Unmarshal(encoded, &schema); err != nil {
+	if err = json.Unmarshal(encoded, &schema); err != nil {
 		t.Fatalf("decode Info schema: %v", err)
 	}
 	properties, ok := schema["properties"].(map[string]any)
@@ -109,12 +114,12 @@ func TestCodeExec_Info(t *testing.T) {
 		t.Fatalf("Info schema properties = %#v, want object", schema["properties"])
 	}
 	for _, name := range []string{"lang", "script"} {
-		if _, ok := properties[name]; !ok {
+		if _, ok = properties[name]; !ok {
 			t.Errorf("Info schema missing %q", name)
 		}
 	}
 	for _, name := range []string{"language", "code", "arguments", "outputs"} {
-		if _, ok := properties[name]; ok {
+		if _, ok = properties[name]; ok {
 			t.Errorf("Info schema unexpectedly exposes node field %q", name)
 		}
 	}
@@ -344,6 +349,7 @@ func TestCodeExec_ResultFallsBackToStdoutJSON(t *testing.T) {
 // parallel with the other CodeExec tests that depend on the
 // default (loud-fail) stub.
 func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) {
+	ctx := t.Context()
 	var captured SandboxRequest
 	prev := GetSandboxClient()
 	SetSandboxClient(stubSandbox(func(_ context.Context, req SandboxRequest) (*SandboxResponse, error) {
@@ -353,7 +359,7 @@ func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) {
 	t.Cleanup(func() { SetSandboxClient(prev) })

 	c := NewCodeExecTool()
-	_, err := c.InvokableRun(context.Background(),
+	_, err := c.InvokableRun(ctx,
 		`{"language":"python","code":"def main(): return {}","timeout":42}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -368,6 +374,7 @@ func TestCodeExec_PassesTimeoutToSandbox(t *testing.T) {
 // timeout test, this mutates the global sandbox client and must
 // not run in parallel with sibling CodeExec tests.
 func TestCodeExec_PassesArgumentsToSandbox(t *testing.T) {
+	ctx := t.Context()
 	var captured SandboxRequest
 	prev := GetSandboxClient()
 	SetSandboxClient(stubSandbox(func(_ context.Context, req SandboxRequest) (*SandboxResponse, error) {
@@ -377,7 +384,7 @@ func TestCodeExec_PassesArgumentsToSandbox(t *testing.T) {
 	t.Cleanup(func() { SetSandboxClient(prev) })

 	c := NewCodeExecTool()
-	_, err := c.InvokableRun(context.Background(),
+	_, err := c.InvokableRun(ctx,
 		`{"language":"python","code":"def main(**kw): return kw","arguments":{"x":1,"y":"z"}}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
--- a/internal/agent/tool/crawler_test.go
+++ b/internal/agent/tool/crawler_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"net"
@@ -41,6 +40,7 @@ const sampleHTML = `<!DOCTYPE html>

 func TestCrawler_FetchesAndExtractsText(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/html; charset=utf-8")
@@ -63,7 +63,7 @@ func TestCrawler_FetchesAndExtractsText(t *testing.T) {
 		return host, net.ParseIP(host), nil
 	}
 	c := NewCrawlerTool().WithResolver(loopbackResolver)
-	out, err := c.InvokableRun(context.Background(),
+	out, err := c.InvokableRun(ctx,
 		`{"query":`+jsonString(srv.URL)+`,"max_depth":0}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -105,9 +105,10 @@ func TestCrawler_FetchesAndExtractsText(t *testing.T) {

 func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCrawlerTool()
-	_, err := c.InvokableRun(context.Background(), `{"query":"https://example.com","max_depth":1}`)
+	_, err := c.InvokableRun(ctx, `{"query":"https://example.com","max_depth":1}`)
 	if !errors.Is(err, ErrCrawlerDepthUnsupported) {
 		t.Fatalf("err = %v, want ErrCrawlerDepthUnsupported", err)
 	}
@@ -115,9 +116,10 @@ func TestCrawler_RejectsMaxDepthGreaterThanZero(t *testing.T) {

 func TestCrawler_RejectsMissingQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCrawlerTool()
-	_, err := c.InvokableRun(context.Background(), `{"query":""}`)
+	_, err := c.InvokableRun(ctx, `{"query":""}`)
 	if err == nil {
 		t.Fatal("expected error for empty query")
 	}
@@ -125,9 +127,10 @@ func TestCrawler_RejectsMissingQuery(t *testing.T) {

 func TestCrawler_RejectsNonHTTPScheme(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCrawlerTool()
-	_, err := c.InvokableRun(context.Background(), `{"query":"file:///etc/passwd"}`)
+	_, err := c.InvokableRun(ctx, `{"query":"file:///etc/passwd"}`)
 	if err == nil || !strings.Contains(err.Error(), "scheme") {
 		t.Fatalf("err = %v, want to reject file:// scheme", err)
 	}
@@ -135,6 +138,7 @@ func TestCrawler_RejectsNonHTTPScheme(t *testing.T) {

 func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	sentinel := errors.New("stop after legacy url normalization")
 	c := NewCrawlerTool().WithResolver(func(rawURL string) (string, net.IP, error) {
@@ -144,7 +148,7 @@ func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) {
 		return "example.com", net.ParseIP("93.184.216.34"), sentinel
 	})

-	_, err := c.InvokableRun(context.Background(), `{"url":"https://example.com"}`)
+	_, err := c.InvokableRun(ctx, `{"url":"https://example.com"}`)
 	if !errors.Is(err, sentinel) {
 		t.Fatalf("err = %v, want resolver error after accepting legacy url", err)
 	}
@@ -152,9 +156,10 @@ func TestCrawler_AcceptsLegacyURLArgument(t *testing.T) {

 func TestCrawler_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	c := NewCrawlerTool()
-	info, err := c.Info(context.Background())
+	info, err := c.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
--- a/internal/agent/tool/deepl_test.go
+++ b/internal/agent/tool/deepl_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -28,6 +27,7 @@ import (

 func TestDeepL_BuildRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	var gotMethod, gotAuth, gotCT, gotPath string
 	var gotForm url.Values
@@ -52,7 +52,7 @@ func TestDeepL_BuildRequest(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewDeepLToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"api_key":"key-xyz:fx","text":"Hello world","source_lang":"en","target_lang":"de"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -98,6 +98,7 @@ func TestDeepL_BuildRequest(t *testing.T) {

 func TestDeepL_DefaultLanguages(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	var gotForm url.Values
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -112,7 +113,7 @@ func TestDeepL_DefaultLanguages(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewDeepLToolWith(helper)
-	if _, err := tool.InvokableRun(context.Background(),
+	if _, err := tool.InvokableRun(ctx,
 		`{"api_key":"x:fx","text":"Hello"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -126,13 +127,14 @@ func TestDeepL_DefaultLanguages(t *testing.T) {

 func TestDeepL_RequiresAPIKeyAndText(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewDeepLTool()
-	if _, err := tool.InvokableRun(context.Background(),
+	if _, err := tool.InvokableRun(ctx,
 		`{"api_key":"","text":"Hello"}`); err == nil {
 		t.Error("expected error for missing api_key")
 	}
-	if _, err := tool.InvokableRun(context.Background(),
+	if _, err := tool.InvokableRun(ctx,
 		`{"api_key":"x","text":""}`); err == nil {
 		t.Error("expected error for empty text")
 	}
@@ -140,9 +142,10 @@ func TestDeepL_RequiresAPIKeyAndText(t *testing.T) {

 func TestDeepL_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewDeepLTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -178,6 +181,7 @@ func TestDeepL_Info(t *testing.T) {
 // when both tests run in the same package.
 func TestDeepL_TranslationFailureReturnsError(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	// 500 Internal Server Error from a stub DeepL endpoint.
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -189,7 +193,7 @@ func TestDeepL_TranslationFailureReturnsError(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewDeepLToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"api_key":"key-xyz:fx","text":"hello","source_lang":"EN","target_lang":"ZH"}`)
 	if err == nil {
 		t.Fatalf("expected non-nil error, got nil; out=%s", out)
--- a/internal/agent/tool/duckduckgo_test.go
+++ b/internal/agent/tool/duckduckgo_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -92,6 +91,7 @@ func TestDuckDuckGo_BuildNewsURLWithVQD(t *testing.T) {
 }

 func TestDuckDuckGo_ParseGeneralResults(t *testing.T) {
+	ctx := t.Context()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Content-Type", "text/html; charset=utf-8")
 		_, _ = w.Write([]byte(`<!doctype html><html><body>
@@ -114,7 +114,7 @@ func TestDuckDuckGo_ParseGeneralResults(t *testing.T) {
 	t.Cleanup(func() { duckduckgoSearchEndpoint = prevSearch })

 	tool := NewDuckDuckGoTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow","top_n":5}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"ragflow","top_n":5}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -141,6 +141,7 @@ func TestDuckDuckGo_ParseGeneralResults(t *testing.T) {
 }

 func TestDuckDuckGo_ParseNewsResults(t *testing.T) {
+	ctx := t.Context()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		switch r.URL.Path {
 		case "/bootstrap":
@@ -171,7 +172,7 @@ func TestDuckDuckGo_ParseNewsResults(t *testing.T) {
 	t.Cleanup(func() { duckduckgoNewsBootstrapEndpoint = prevBootstrap })

 	tool := NewDuckDuckGoTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow","channel":"news","top_n":1}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"ragflow","channel":"news","top_n":1}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -195,6 +196,7 @@ func TestDuckDuckGo_ParseNewsResults(t *testing.T) {
 }

 func TestDuckDuckGo_DefaultChannelUsesGeneralSearch(t *testing.T) {
+	ctx := t.Context()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if got := r.URL.Path; got != "/" {
 			// keep old behavior impossible to hit if search endpoint override works incorrectly
@@ -214,15 +216,16 @@ func TestDuckDuckGo_DefaultChannelUsesGeneralSearch(t *testing.T) {
 	t.Cleanup(func() { duckduckgoSearchEndpoint = prevSearch })

 	tool := NewDuckDuckGoTool()
-	_, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`)
+	_, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 }

 func TestDuckDuckGo_Info(t *testing.T) {
+	ctx := t.Context()
 	tool := NewDuckDuckGoTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -253,18 +256,20 @@ func TestDuckDuckGo_Info(t *testing.T) {
 }

 func TestDuckDuckGo_EmptyQuery(t *testing.T) {
+	ctx := t.Context()
 	tool := NewDuckDuckGoTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty): %v", err)
 	}
 	var envelope duckduckgoEnvelope
-	if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 {
+	if err = json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 {
 		t.Fatalf("empty result = %s / %v", out, err)
 	}
 }

 func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) {
+	ctx := t.Context()
 	var hitCount int
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		hitCount++
@@ -290,7 +295,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) {
 		"RAGFlow is an open-source RAG engine.",
 	)

-	agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
+	agent, err := react.NewAgent(ctx, &react.AgentConfig{
 		ToolCallingModel: mdl,
 		ToolsConfig: compose.ToolsNodeConfig{
 			Tools: []einotool.BaseTool{realTool},
@@ -301,7 +306,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) {
 		t.Fatalf("react.NewAgent: %v", err)
 	}

-	out, err := agent.Generate(context.Background(), []*schema.Message{
+	out, err := agent.Generate(ctx, []*schema.Message{
 		schema.UserMessage("What is RAGFlow?"),
 	})
 	if err != nil {
@@ -339,6 +344,7 @@ func TestDuckDuckGo_RealReactAgent_ExecutesTool(t *testing.T) {
 }

 func TestDuckDuckGo_ComponentReferencesAndDefaults(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	built, err := BuildByName("duckduckgo", map[string]any{
@@ -365,7 +371,7 @@ func TestDuckDuckGo_ComponentReferencesAndDefaults(t *testing.T) {
 	envelope := map[string]any{"results": []any{map[string]any{
 		"title": "Story", "url": "https://news.example/story", "body": "Breaking update",
 	}}}
-	chunks, docAggs := duck.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := duck.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["content"] != "Breaking update" {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
--- a/internal/agent/tool/email_test.go
+++ b/internal/agent/tool/email_test.go
@@ -64,6 +64,7 @@ func TestEmail_BuildMessage(t *testing.T) {
 }

 func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) {
+	ctx := t.Context()
 	originalSendEmail := sendEmail
 	t.Cleanup(func() { sendEmail = originalSendEmail })
 	var sentParams emailParams
@@ -92,12 +93,12 @@ func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) {
 	argsJSON, _ := json.Marshal(args)
 	state := runtime.NewCanvasState("run-email", "task-email")
 	state.Sys["date"] = "2026-07-15"
-	out, err := built.(*EmailTool).InvokableRun(runtime.WithState(context.Background(), state), string(argsJSON))
+	out, err := built.(*EmailTool).InvokableRun(runtime.WithState(ctx, state), string(argsJSON))
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var env emailEnvelope
-	if err := json.Unmarshal([]byte(out), &env); err != nil || !env.OK || env.Error != "" {
+	if err = json.Unmarshal([]byte(out), &env); err != nil || !env.OK || env.Error != "" {
 		t.Fatalf("output = %s, decode error = %v", out, err)
 	}

@@ -119,6 +120,7 @@ func TestEmail_SendBuildsDistinctHeadersAndEnvelopeRecipients(t *testing.T) {
 }

 func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -161,7 +163,7 @@ func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) {
 	}
 	var portNumber int
 	_, _ = fmt.Sscanf(port, "%d", &portNumber)
-	err = sendEmailSTARTTLS(context.Background(), emailParams{
+	err = sendEmailSTARTTLS(ctx, emailParams{
 		SMTPServer: host, SMTPPort: portNumber, Email: "alice@example.com",
 		ToEmail: "bob@example.com",
 	}, []byte("message"))
@@ -182,6 +184,7 @@ func TestEmail_STARTTLSRequiredBeforeSubmission(t *testing.T) {
 }

 func TestEmail_RequiresFields(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	cases := []struct {
@@ -212,7 +215,7 @@ func TestEmail_RequiresFields(t *testing.T) {
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			_, err := tc.tool.InvokableRun(context.Background(), tc.args)
+			_, err := tc.tool.InvokableRun(ctx, tc.args)
 			if err == nil {
 				t.Fatalf("expected error for %s", tc.name)
 			}
@@ -224,10 +227,11 @@ func TestEmail_RequiresFields(t *testing.T) {
 }

 func TestEmail_Info(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	tool := NewEmailTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
--- a/internal/agent/tool/exesql_test.go
+++ b/internal/agent/tool/exesql_test.go
@@ -65,6 +65,7 @@ func sqlmockDialer(t *testing.T) (exesqlDialer, sqlmock.Sqlmock, func()) {
 }

 func TestExeSQL_NoCredentials(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(exesqlConnParams{}).
@@ -72,7 +73,7 @@ func TestExeSQL_NoCredentials(t *testing.T) {
 			t.Fatal("dialer should not be called when credentials are missing")
 			return nil, nil
 		})
-	_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	_, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if !errors.Is(err, ErrExeSQLNoCredentials) {
 		t.Fatalf("err = %v, want ErrExeSQLNoCredentials", err)
 	}
@@ -108,6 +109,7 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) {
 		{"merge cte", `WITH changed AS (MERGE INTO users USING incoming ON users.id = incoming.id WHEN MATCHED THEN UPDATE SET name = incoming.name RETURNING *) SELECT * FROM changed`},
 	}

+	ctx := t.Context()
 	for _, c := range cases {
 		t.Run(c.name, func(t *testing.T) {
 			t.Parallel()
@@ -116,7 +118,7 @@ func TestExeSQL_RejectsNonSelect(t *testing.T) {
 					t.Fatal("dialer called for rejected SQL")
 					return nil, nil
 				})
-			_, err := e.InvokableRun(context.Background(),
+			_, err := e.InvokableRun(ctx,
 				`{"sql":`+jsonString(c.sql)+`}`)
 			if !errors.Is(err, ErrExeSQLNotSelect) {
 				t.Fatalf("err = %v, want ErrExeSQLNotSelect", err)
@@ -131,8 +133,9 @@ func TestExeSQL_RejectsMixedBatchBeforeDatabaseAccess(t *testing.T) {
 			t.Fatal("dialer called before every SQL statement was validated")
 			return nil, nil
 		})
+	ctx := t.Context()

-	_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1; DROP TABLE users"}`)
+	_, err := e.InvokableRun(ctx, `{"sql":"SELECT 1; DROP TABLE users"}`)
 	if !errors.Is(err, ErrExeSQLNotSelect) {
 		t.Fatalf("err = %v, want ErrExeSQLNotSelect", err)
 	}
@@ -158,6 +161,7 @@ func TestExeSQL_AllowsSelect(t *testing.T) {
 		// Block comment.
 		`/* DROP TABLE foo */ SELECT 1`,
 	}
+	ctx := t.Context()
 	for _, sql := range cases {
 		t.Run(sql, func(t *testing.T) {
 			t.Parallel()
@@ -172,7 +176,7 @@ func TestExeSQL_AllowsSelect(t *testing.T) {
 				sqlmock.NewRows([]string{"1"}),
 			)
 			e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-			_, err := e.InvokableRun(context.Background(),
+			_, err := e.InvokableRun(ctx,
 				`{"sql":`+jsonString(sql)+`}`)
 			// Two acceptable outcomes:
 			//   1. SQL is the literal `SELECT 1` and matches the
@@ -196,20 +200,22 @@ func TestExeSQL_AllowsSelect(t *testing.T) {
 }

 func TestExeSQL_RejectsEmptySQL(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(testConn())
-	_, err := e.InvokableRun(context.Background(), `{"sql":""}`)
+	_, err := e.InvokableRun(ctx, `{"sql":""}`)
 	if err == nil || !strings.Contains(err.Error(), "sql") {
 		t.Fatalf("err = %v, want to mention empty sql", err)
 	}
 }

 func TestExeSQL_RejectsEmptyArgs(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(testConn())
-	_, err := e.InvokableRun(context.Background(), "")
+	_, err := e.InvokableRun(ctx, "")
 	if err == nil {
 		t.Fatal("expected error for empty args")
 	}
@@ -272,6 +278,7 @@ func TestExeSQL_ReadOnlyValidationIgnoresQuotedAndCommentedKeywords(t *testing.T
 }

 func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -283,7 +290,7 @@ func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) {
 		WillReturnRows(sqlmock.NewRows([]string{"value"}).AddRow(2))

 	e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-	if _, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 'hello; world'; SELECT 2"}`); err != nil {
+	if _, err := e.InvokableRun(ctx, `{"sql":"SELECT 'hello; world'; SELECT 2"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	if err := mock.ExpectationsWereMet(); err != nil {
@@ -292,6 +299,7 @@ func TestExeSQL_ExecutesStatementsWithQuotedSemicolonsIntact(t *testing.T) {
 }

 func TestExeSQL_RejectsMySQLExecutableComment(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(testConn()).
@@ -299,17 +307,18 @@ func TestExeSQL_RejectsMySQLExecutableComment(t *testing.T) {
 			t.Fatal("dialer called for an executable comment")
 			return nil, nil
 		})
-	_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1 /*!; DROP TABLE users */"}`)
+	_, err := e.InvokableRun(ctx, `{"sql":"SELECT 1 /*!; DROP TABLE users */"}`)
 	if !errors.Is(err, ErrExeSQLNotSelect) {
 		t.Fatalf("err = %v, want ErrExeSQLNotSelect", err)
 	}
 }

 func TestExeSQL_Info(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(testConn())
-	info, err := e.Info(context.Background())
+	info, err := e.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -337,6 +346,7 @@ func TestExeSQL_Info(t *testing.T) {
 }

 func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) {
+	ctx := t.Context()
 	dialer, mock, cleanup := sqlmockDialer(t)
 	defer cleanup()
 	mock.ExpectPing()
@@ -347,7 +357,7 @@ func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) {
 	conn.SQL = "SELECT 1"
 	e := NewExeSQLTool(conn).WithExeSQLDialer(dialer)

-	out, err := e.InvokableRun(context.Background(), `{}`)
+	out, err := e.InvokableRun(ctx, `{}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -360,6 +370,7 @@ func TestExeSQL_UsesConfiguredSQLDefault(t *testing.T) {
 }

 func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) {
+	ctx := t.Context()
 	dialer, mock, cleanup := sqlmockDialer(t)
 	defer cleanup()
 	mock.ExpectPing()
@@ -370,12 +381,12 @@ func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) {
 	exesql := NewExeSQLTool(conn).WithExeSQLDialer(dialer)
 	state := runtime.NewCanvasState("run", "task")
 	state.SetVar("Agent:Result", "content", "SELECT id FROM orders WHERE status = 'Completed'")
-	out, err := exesql.InvokableRun(runtime.WithState(context.Background(), state), `{}`)
+	out, err := exesql.InvokableRun(runtime.WithState(ctx, state), `{}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var envelope map[string]any
-	if err := json.Unmarshal([]byte(out), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(out), &envelope); err != nil {
 		t.Fatalf("decode output: %v", err)
 	}
 	outputs := exesql.BuildComponentOutputs(envelope)
@@ -390,7 +401,7 @@ func TestExeSQL_ComponentContractAndTemplateResolution(t *testing.T) {
 	if sqlInput, ok := spec.InputForm["sql"].(map[string]any); !ok || sqlInput["type"] != "line" {
 		t.Fatalf("sql input form = %#v", spec.InputForm["sql"])
 	}
-	if err := mock.ExpectationsWereMet(); err != nil {
+	if err = mock.ExpectationsWereMet(); err != nil {
 		t.Fatalf("sql expectations: %v", err)
 	}
 }
@@ -410,6 +421,7 @@ func TestExeSQL_BuildByNameAcceptsCanvasShape(t *testing.T) {
 }

 func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -423,7 +435,7 @@ func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) {
 			AddRow(8, "bob"))

 	e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-	out, err := e.InvokableRun(context.Background(),
+	out, err := e.InvokableRun(ctx,
 		`{"sql":"SELECT id, name FROM t WHERE id = 7"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -444,6 +456,7 @@ func TestExeSQL_ExecuteSelect_ReturnsRows(t *testing.T) {
 }

 func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -453,12 +466,12 @@ func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) {
 		WillReturnRows(sqlmock.NewRows([]string{"x"}))

 	e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-	out, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	out, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var got exesqlResult
-	if err := json.Unmarshal([]byte(out), &got); err != nil {
+	if err = json.Unmarshal([]byte(out), &got); err != nil {
 		t.Fatalf("unmarshal: %v\nout=%s", err, out)
 	}
 	// The Python tool's "No record in the database!" sentinel must
@@ -470,6 +483,7 @@ func TestExeSQL_ExecuteSelect_NoRowsReturnsSentinel(t *testing.T) {
 }

 func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -484,13 +498,13 @@ func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) {
 		WillReturnError(errors.New("syntax error at or near BOGUS"))

 	e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-	out, err := e.InvokableRun(context.Background(),
+	out, err := e.InvokableRun(ctx,
 		`{"sql":"SELECT 1; SELECT * FROM bogus"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun should not abort on a per-statement error: %v", err)
 	}
 	var got exesqlResult
-	if err := json.Unmarshal([]byte(out), &got); err != nil {
+	if err = json.Unmarshal([]byte(out), &got); err != nil {
 		t.Fatalf("unmarshal: %v\nout=%s", err, out)
 	}
 	if len(got.Rows) != 2 {
@@ -507,6 +521,7 @@ func TestExeSQL_ExecuteSelect_PerStatementErrorIsolated(t *testing.T) {
 }

 func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -517,13 +532,13 @@ func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) {
 			AddRow("2024-06-12T03:04:05Z", []byte("hello")))

 	e := NewExeSQLTool(testConn()).WithExeSQLDialer(dialer)
-	out, err := e.InvokableRun(context.Background(),
+	out, err := e.InvokableRun(ctx,
 		`{"sql":"SELECT ts, blob_col FROM t"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var got exesqlResult
-	if err := json.Unmarshal([]byte(out), &got); err != nil {
+	if err = json.Unmarshal([]byte(out), &got); err != nil {
 		t.Fatalf("unmarshal: %v\nout=%s", err, out)
 	}
 	if len(got.Rows) != 1 {
@@ -540,6 +555,7 @@ func TestExeSQL_ExecuteSelect_NormalizesTimeAndBytes(t *testing.T) {
 }

 func TestExeSQL_UnsupportedDB(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	e := NewExeSQLTool(exesqlConnParams{
@@ -547,7 +563,7 @@ func TestExeSQL_UnsupportedDB(t *testing.T) {
 		Host:   "1.1.1.1", Port: 8080, Database: "catalog",
 		Username: "u", Password: "p",
 	})
-	_, err := e.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	_, err := e.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if err == nil {
 		t.Fatal("expected non-nil error for trino without registered driver")
 	}
@@ -747,6 +763,7 @@ func (m *reactScriptedModel) Stream(_ context.Context, _ []*schema.Message, _ ..
 // and the resulting JSON is passed back as a ToolMessage. Replacing
 // the model with a hand-rolled stub would skip all of that.
 func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -763,7 +780,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
 		"the answer is 42",
 	)

-	agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
+	agent, err := react.NewAgent(ctx, &react.AgentConfig{
 		ToolCallingModel: mdl,
 		ToolsConfig: compose.ToolsNodeConfig{
 			Tools: []einotool.BaseTool{realTool},
@@ -774,7 +791,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
 		t.Fatalf("react.NewAgent: %v", err)
 	}

-	out, err := agent.Generate(context.Background(), []*schema.Message{
+	out, err := agent.Generate(ctx, []*schema.Message{
 		schema.UserMessage("What is 42?"),
 	})
 	if err != nil {
@@ -806,7 +823,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
 	if !sawToolResult {
 		t.Errorf("round 2 input did not contain a ToolMessage carrying the tool result; got %d messages", len(mdl.rounds[1]))
 	}
-	if err := mock.ExpectationsWereMet(); err != nil {
+	if err = mock.ExpectationsWereMet(); err != nil {
 		t.Errorf("sqlmock expectations not met: %v", err)
 	}
 }
@@ -818,6 +835,7 @@ func TestExeSQL_RealReactAgent_ExecutesTool(t *testing.T) {
 // the model on round 2 without crashing the ReAct loop, so the model
 // can ground its final answer in the surfaced error.
 func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()

 	dialer, mock, cleanup := sqlmockDialer(t)
@@ -834,7 +852,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) {
 		"the query had a syntax error",
 	)

-	agent, err := react.NewAgent(context.Background(), &react.AgentConfig{
+	agent, err := react.NewAgent(ctx, &react.AgentConfig{
 		ToolCallingModel: mdl,
 		ToolsConfig: compose.ToolsNodeConfig{
 			Tools: []einotool.BaseTool{realTool},
@@ -845,7 +863,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) {
 		t.Fatalf("react.NewAgent: %v", err)
 	}

-	out, err := agent.Generate(context.Background(), []*schema.Message{
+	out, err := agent.Generate(ctx, []*schema.Message{
 		schema.UserMessage("Find bogus rows"),
 	})
 	if err != nil {
@@ -869,7 +887,7 @@ func TestExeSQL_RealReactAgent_ToolErrorIsolated(t *testing.T) {
 	if !sawErrorResult {
 		t.Errorf("round 2 input did not contain a ToolMessage with the DB error; got %d messages", len(mdl.rounds[1]))
 	}
-	if err := mock.ExpectationsWereMet(); err != nil {
+	if err = mock.ExpectationsWereMet(); err != nil {
 		t.Errorf("sqlmock expectations: %v", err)
 	}
 }
--- a/internal/agent/tool/exesql_trino_test.go
+++ b/internal/agent/tool/exesql_trino_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"strings"
 	"testing"

@@ -208,7 +207,8 @@ func TestExeSQL_Trino_HappyPath(t *testing.T) {
 		MaxRecords: 100,
 	}).WithExeSQLDialer(dialer)

-	out, err := tool.InvokableRun(context.Background(),
+	ctx := t.Context()
+	out, err := tool.InvokableRun(ctx,
 		`{"sql":"SELECT id, name FROM catalog.tiny.users"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
--- a/internal/agent/tool/exesql_unsupported_test.go
+++ b/internal/agent/tool/exesql_unsupported_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"errors"
 	"testing"
 )
@@ -27,9 +26,10 @@ import (
 // database/sql driver, so InvokableRun should fail at sql.Open with an
 // unknown-driver error rather than the old unsupported-db sentinel.
 func TestExeSQL_TrinoDriverMissing(t *testing.T) {
+	ctx := t.Context()
 	conn := exesqlConnParams{DBType: "trino", Host: "1.1.1.1", Port: 8080, Database: "d", Username: "u"}
 	tool := NewExeSQLTool(conn)
-	_, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	_, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if err == nil {
 		t.Fatal("expected driver error for trino")
 	}
@@ -40,9 +40,10 @@ func TestExeSQL_TrinoDriverMissing(t *testing.T) {

 // TestExeSQL_IBMDB2Unsupported: same as above for IBM DB2.
 func TestExeSQL_IBMDB2Unsupported(t *testing.T) {
+	ctx := t.Context()
 	conn := exesqlConnParams{DBType: "ibm db2", Host: "1.1.1.1", Port: 50000, Database: "d", Username: "u"}
 	tool := NewExeSQLTool(conn)
-	_, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	_, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if err == nil {
 		t.Fatal("expected ErrExeSQLUnsupportedDB for ibm db2")
 	}
@@ -57,9 +58,10 @@ func TestExeSQL_IBMDB2Unsupported(t *testing.T) {
 // follow-up should normalize the error. The regression guard here
 // is "doesn't panic, returns a non-nil error".
 func TestExeSQL_UnknownDB(t *testing.T) {
+	ctx := t.Context()
 	conn := exesqlConnParams{DBType: "fake-db", Host: "1.1.1.1", Port: 1234, Database: "d", Username: "u"}
 	tool := NewExeSQLTool(conn)
-	_, err := tool.InvokableRun(context.Background(), `{"sql":"SELECT 1"}`)
+	_, err := tool.InvokableRun(ctx, `{"sql":"SELECT 1"}`)
 	if err == nil {
 		t.Fatal("expected error for unknown db_type")
 	}
--- a/internal/agent/tool/github_test.go
+++ b/internal/agent/tool/github_test.go
@@ -17,7 +17,6 @@
 package tool

 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -91,6 +90,7 @@ func TestGitHub_BuildURL(t *testing.T) {

 func TestGitHub_ParseResponse(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	var gotContentType, gotAPIVersion, gotPerPage string
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -111,7 +111,7 @@ func TestGitHub_ParseResponse(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewGitHubToolWithDefaults(helper, githubParams{TopN: 17})
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"query":"ragflow"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -149,14 +149,15 @@ func TestGitHub_ParseResponse(t *testing.T) {

 func TestGitHub_EmptyQueryReturnsEmptyResults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewGitHubTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty query): %v", err)
 	}
 	var envelope githubEnvelope
-	if err := json.Unmarshal([]byte(out), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(out), &envelope); err != nil {
 		t.Fatalf("decode empty result: %v", err)
 	}
 	if len(envelope.Results) != 0 || envelope.Error != "" {
@@ -181,19 +182,19 @@ func TestGitHub_BuildByNameUsesPythonNodeParams(t *testing.T) {
 	if github.defaults.TopN != 17 {
 		t.Errorf("defaults.TopN = %d, want 17", github.defaults.TopN)
 	}
-	if _, err := BuildByName("github", map[string]any{"top_n": 100}); err != nil {
+	if _, err = BuildByName("github", map[string]any{"top_n": 100}); err != nil {
 		t.Errorf("BuildByName(github) rejected GitHub's maximum top_n: %v", err)
 	}
-	if _, err := BuildByName("github", map[string]any{"top_n": 0}); err == nil {
+	if _, err = BuildByName("github", map[string]any{"top_n": 0}); err == nil {
 		t.Fatal("BuildByName(github) accepted non-positive top_n")
 	}
-	if _, err := BuildByName("github", map[string]any{"top_n": 1.5}); err == nil {
+	if _, err = BuildByName("github", map[string]any{"top_n": 1.5}); err == nil {
 		t.Fatal("BuildByName(github) accepted fractional top_n")
 	}
-	if _, err := BuildByName("github", map[string]any{"top_n": "10"}); err == nil {
+	if _, err = BuildByName("github", map[string]any{"top_n": "10"}); err == nil {
 		t.Fatal("BuildByName(github) accepted string top_n")
 	}
-	if _, err := BuildByName("github", map[string]any{"top_n": 101}); err == nil {
+	if _, err = BuildByName("github", map[string]any{"top_n": 101}); err == nil {
 		t.Fatal("BuildByName(github) accepted top_n above GitHub's per_page limit")
 	}
 	ignored, err := BuildByName("github", map[string]any{"max_results": 5})
@@ -220,6 +221,7 @@ func TestGitHub_ComponentContractMatchesPython(t *testing.T) {
 }

 func TestGitHub_ReferencesAndOutputsPreserveRawResults(t *testing.T) {
+	ctx := t.Context()
 	github := NewGitHubTool()
 	results := []any{map[string]any{
 		"name":        "ragflow",
@@ -234,7 +236,7 @@ func TestGitHub_ReferencesAndOutputsPreserveRawResults(t *testing.T) {
 	}
 	envelope := map[string]any{"results": results}

-	chunks, docAggs := github.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := github.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
@@ -269,9 +271,10 @@ func TestGitHub_LimitReferencesKeepsBoundaryChunk(t *testing.T) {

 func TestGitHub_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()

 	tool := NewGitHubTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}