Go: fix context, part2 (#18133)

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

internal/agent/tool/google_scholar_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -180,6 +179,7 @@ func TestGoogleScholar_BuildURL(t *testing.T) {
 
 func TestGoogleScholar_ParseResults(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")
@@ -191,7 +191,7 @@ func TestGoogleScholar_ParseResults(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewGoogleScholarToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"query":"transformer","top_n":5,"sort_by":"relevance","year_low":2020,"year_high":2024,"patents":true}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -229,23 +229,25 @@ func TestGoogleScholar_ParseResults(t *testing.T) {
 
 func TestGoogleScholar_EmptyQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewGoogleScholarTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty): %v", err)
 	}
 	var envelope googleScholarEnvelope
-	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 TestGoogleScholar_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewGoogleScholarTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -259,6 +261,7 @@ func TestGoogleScholar_Info(t *testing.T) {
 
 func TestGoogleScholar_ComponentReferencesAndValidation(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	built, err := BuildByName("google_scholar", map[string]any{
 		"top_n":    float64(7),
@@ -286,7 +289,7 @@ func TestGoogleScholar_ComponentReferencesAndValidation(t *testing.T) {
 	envelope := map[string]any{"results": []any{map[string]any{
 		"title": "Paper", "link": "https://paper.example", "authors": "A Author", "year": "2024", "snippet": "Abstract",
 	}}}
-	chunks, docAggs := scholar.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := scholar.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 || !strings.Contains(chunks[0]["content"].(string), "Authors: A Author") {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
@@ -325,6 +328,7 @@ func TestRenderGoogleScholarReferencesStopsBeforeOverBudgetBlock(t *testing.T) {
 
 func TestGoogleScholar_MergesNodeLevelDefaults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		q := r.URL.Query()
@@ -362,7 +366,7 @@ func TestGoogleScholar_MergesNodeLevelDefaults(t *testing.T) {
 		YearHigh: 2024,
 		Patents:  boolPtr(false),
 	})
-	_, err := tool.InvokableRun(context.Background(), `{}`)
+	_, err := tool.InvokableRun(ctx, `{}`)
 	if err != nil {
 		t.Fatalf("InvokableRun with defaults: %v", err)
 	}

internal/agent/tool/google_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -86,6 +85,7 @@ func TestGoogle_BuildURL(t *testing.T) {
 
 func TestGoogle_ParseResults(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", "application/json")
@@ -102,7 +102,7 @@ func TestGoogle_ParseResults(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewGoogleToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"q":"ragflow","api_key":"K","num":5,"country":"us","language":"en"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -128,9 +128,10 @@ func TestGoogle_ParseResults(t *testing.T) {
 
 func TestGoogle_RequiresAPIKey(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewGoogleTool()
-	_, err := tool.InvokableRun(context.Background(), `{"q":"x","api_key":""}`)
+	_, err := tool.InvokableRun(ctx, `{"q":"x","api_key":""}`)
 	if err == nil {
 		t.Fatal("expected error for missing api_key")
 	}
@@ -141,9 +142,10 @@ func TestGoogle_RequiresAPIKey(t *testing.T) {
 
 func TestGoogle_InfoAndInputForm(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewGoogleTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -246,6 +248,7 @@ func TestGoogle_BuildByNameIgnoresUnrelatedCanvasParams(t *testing.T) {
 
 func TestGoogle_ComponentReferencesAndOutputs(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	google := NewGoogleTool()
 	spec := google.ComponentSpec()
@@ -270,7 +273,7 @@ func TestGoogle_ComponentReferencesAndOutputs(t *testing.T) {
 			}},
 		},
 	}}
-	chunks, docAggs := google.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := google.BuildReferences(ctx, envelope)
 	if len(chunks) != 2 || len(docAggs) != 2 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}

internal/agent/tool/http_helper_test.go

@@ -49,6 +49,7 @@ func newTestHelper(maxAttempts int, base, max time.Duration) *HTTPHelper {
 // attempt with no retry, and the body / content-type round-trip cleanly.
 func TestHTTPHelper_HappyPath(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", "application/json")
@@ -58,7 +59,7 @@ func TestHTTPHelper_HappyPath(t *testing.T) {
 	defer srv.Close()
 
 	h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
-	resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
+	resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
 	if err != nil {
 		t.Fatalf("Do returned error: %v", err)
 	}
@@ -80,6 +81,7 @@ func TestHTTPHelper_HappyPath(t *testing.T) {
 // returns the first 2xx response. Server returns 503 twice, then 200.
 func TestHTTPHelper_RetriesOn5xx(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var hits int32
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -94,7 +96,7 @@ func TestHTTPHelper_RetriesOn5xx(t *testing.T) {
 	defer srv.Close()
 
 	h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
-	resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
+	resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
 	if err != nil {
 		t.Fatalf("Do returned error: %v", err)
 	}
@@ -116,6 +118,7 @@ func TestHTTPHelper_RetriesOn5xx(t *testing.T) {
 // no retry — the caller is responsible for fixing 4xx, retrying won't help.
 func TestHTTPHelper_NoRetryOn4xx(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var hits int32
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -126,7 +129,7 @@ func TestHTTPHelper_NoRetryOn4xx(t *testing.T) {
 	defer srv.Close()
 
 	h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
-	resp, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
+	resp, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
 	if err != nil {
 		t.Fatalf("Do returned error: %v", err)
 	}
@@ -158,6 +161,7 @@ func TestHTTPHelper_NoRetryOn4xx(t *testing.T) {
 // and flaked under load.
 func TestHTTPHelper_Timeout(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var hits int32
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -176,7 +180,7 @@ func TestHTTPHelper_Timeout(t *testing.T) {
 	})
 	// Tight 50ms deadline. The server takes 500ms, so this call must
 	// abort due to the context, not the server finishing.
-	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
+	ctx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
 	defer cancel()
 
 	_, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
@@ -196,6 +200,7 @@ func TestHTTPHelper_Timeout(t *testing.T) {
 // 5xx error is returned.
 func TestHTTPHelper_5xxExhaustion(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var hits int32
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -205,7 +210,7 @@ func TestHTTPHelper_5xxExhaustion(t *testing.T) {
 	defer srv.Close()
 
 	h := newTestHelper(3, 1*time.Millisecond, 5*time.Millisecond)
-	_, err := h.Do(context.Background(), http.MethodGet, srv.URL, "", "", nil)
+	_, err := h.Do(ctx, http.MethodGet, srv.URL, "", "", nil)
 	if err == nil {
 		t.Fatal("expected error after 5xx exhaustion, got nil")
 	}
@@ -218,6 +223,7 @@ func TestHTTPHelper_5xxExhaustion(t *testing.T) {
 // custom headers and a non-empty content-type on POST bodies.
 func TestHTTPHelper_HeadersAndContentType(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if got := r.Header.Get("X-Token"); got != "abc" {
@@ -234,7 +240,7 @@ func TestHTTPHelper_HeadersAndContentType(t *testing.T) {
 	defer srv.Close()
 
 	h := newTestHelper(1, 1*time.Millisecond, 5*time.Millisecond)
-	resp, err := h.Do(context.Background(), http.MethodPost, srv.URL, `{"k":1}`, "application/json", map[string]string{"X-Token": "abc"})
+	resp, err := h.Do(ctx, http.MethodPost, srv.URL, `{"k":1}`, "application/json", map[string]string{"X-Token": "abc"})
 	if err != nil {
 		t.Fatalf("Do: %v", err)
 	}
@@ -251,20 +257,20 @@ func TestBackoffExponential(t *testing.T) {
 	t.Parallel()
 
 	base := 50 * time.Millisecond
-	max := 300 * time.Millisecond
+	maxDuration := 300 * time.Millisecond
 
-	got1 := backoff(base, max, 1)
+	got1 := backoff(base, maxDuration, 1)
 	if got1 < 0 || got1 > base {
 		t.Fatalf("backoff(attempt=1) = %s, want [0, %s]", got1, base)
 	}
-	got3 := backoff(base, max, 3)
-	if got3 < 0 || got3 > max {
-		t.Fatalf("backoff(attempt=3) = %s, want [0, %s] (capped)", got3, max)
+	got3 := backoff(base, maxDuration, 3)
+	if got3 < 0 || got3 > maxDuration {
+		t.Fatalf("backoff(attempt=3) = %s, want [0, %s] (capped)", got3, maxDuration)
 	}
 	// With base=50ms, attempt=10 should be capped at 300ms.
-	got10 := backoff(base, max, 10)
-	if got10 > max {
-		t.Fatalf("backoff(attempt=10) = %s, want <= %s (cap)", got10, max)
+	got10 := backoff(base, maxDuration, 10)
+	if got10 > maxDuration {
+		t.Fatalf("backoff(attempt=10) = %s, want <= %s (cap)", got10, maxDuration)
 	}
 }
 
@@ -289,7 +295,7 @@ func TestRetryConfigDefaults(t *testing.T) {
 // test for the M1-rebinding fix as hardened by the post-Phase-7
 // review: DNS pinning MUST happen at the transport layer, not by
 // rewriting the request URL. If the URL host were rewritten to the IP,
-// the TLS ServerName (auto-populated by Go from req.URL.Host) would
+// the TLS ServerName (autopopulated by Go from req.URL.Host) would
 // become the IP, the SNI would send the IP, and cert verification
 // would target the IP — which is not what real HTTPS sites have, and
 // would manifest as x509 errors against any host-cert-only target.
@@ -297,12 +303,13 @@ func TestRetryConfigDefaults(t *testing.T) {
 // This test stands up a real TLS server with a cert whose DNS SAN is
 // "example.test" and whose IP SAN covers the loopback address. The
 // pinned dialer connects to 127.0.0.1, but the request URL host stays
-// as "example.test". The server observes the SNI the client sent and
+// as "example.test". The server observes the SNI the client sent, and
 // we assert it equals "example.test" (not the IP), and the request
 // completes successfully (cert verification passes because the URL
 // host matches the SAN).
 func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	// Cert valid for "example.test" (DNS SAN) and 127.0.0.1, ::1
 	// (IP SANs, just so the test environment itself can resolve).
@@ -353,7 +360,7 @@ func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) {
 	h := NewHTTPHelper()
 	h.baseTransport.TLSClientConfig = &tls.Config{RootCAs: pool}
 
-	resp, err := h.DoPinned(context.Background(),
+	resp, err := h.DoPinned(ctx,
 		http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP)
 	if err != nil {
 		t.Fatalf("DoPinned: %v", err)
@@ -378,9 +385,10 @@ func TestHTTPHelper_DoPinnedHTTPS_PreservesSNIAndCert(t *testing.T) {
 // refuse rather than silently deliver a broken connection.
 func TestHTTPHelper_DoPinnedRefusesMismatchedURLHost(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	h := NewHTTPHelper()
-	_, err := h.DoPinned(context.Background(),
+	_, err := h.DoPinned(ctx,
 		http.MethodGet,
 		"https://attacker.example/foo",
 		"", "", nil,
@@ -414,6 +422,7 @@ func TestHTTPHelper_DoPinnedRefusesMismatchedURLHost(t *testing.T) {
 // bypasses the proxy, the direct dial to 127.0.0.1 succeeds.
 func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	cert := generateTestCert(t, "example.test")
 	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -467,7 +476,7 @@ func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) {
 	}
 	h.baseTransport.Proxy = http.ProxyURL(closedProxy)
 
-	resp, err := h.DoPinned(context.Background(),
+	resp, err := h.DoPinned(ctx,
 		http.MethodGet, targetURL, "", "", nil, "example.test", pinnedIP)
 	if err != nil {
 		t.Fatalf("DoPinned: %v (proxy may not have been bypassed)", err)
@@ -486,6 +495,7 @@ func TestHTTPHelper_DoPinnedBypassesProxy(t *testing.T) {
 // transport-layer primitive that DoPinned uses.
 func TestPinnedDialer_RewritesAddress(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	// Stand up a TCP server on 127.0.0.1:<random> and capture the
 	// accepted conn to confirm the pinned dialer actually dialed it.
@@ -511,7 +521,7 @@ func TestPinnedDialer_RewritesAddress(t *testing.T) {
 	// Pass a deliberately misleading host in the addr — the dialer
 	// must ignore it and dial 127.0.0.1:<ln port> instead.
 	misleading := net.JoinHostPort("203.0.113.99", fmt.Sprint(ln.Addr().(*net.TCPAddr).Port))
-	conn, derr := d.DialContext(context.Background(), "tcp", misleading)
+	conn, derr := d.DialContext(ctx, "tcp", misleading)
 	if derr != nil {
 		t.Fatalf("DialContext: %v", derr)
 	}

internal/agent/tool/jin10_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"strings"
 	"testing"
@@ -25,6 +24,7 @@ import (
 
 func TestJin10_StubsUnsupported(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	cases := []struct {
 		name string
@@ -47,7 +47,7 @@ func TestJin10_StubsUnsupported(t *testing.T) {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
 			tool := NewJin10Tool()
-			out, err := tool.InvokableRun(context.Background(), tc.args)
+			out, err := tool.InvokableRun(ctx, tc.args)
 			if err == nil {
 				t.Fatalf("expected error, got nil (out=%s)", out)
 			}
@@ -70,9 +70,10 @@ func TestJin10_StubsUnsupported(t *testing.T) {
 
 func TestJin10_RejectsMalformedJSON(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewJin10Tool()
-	_, 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")
 	}
@@ -83,9 +84,10 @@ func TestJin10_RejectsMalformedJSON(t *testing.T) {
 
 func TestJin10_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewJin10Tool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}

internal/agent/tool/keenable_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -30,6 +29,7 @@ import (
 // without an X-API-Key header.
 func TestKeenable_KeylessPath(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotMethod, gotPath, gotUA, gotTitle, gotAPIKey, gotCT string
 	var gotBody map[string]any
@@ -52,7 +52,7 @@ func TestKeenable_KeylessPath(t *testing.T) {
 	})
 	tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] })
 
-	if _, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil {
+	if _, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 
@@ -87,6 +87,7 @@ func TestKeenable_KeylessPath(t *testing.T) {
 // the request.
 func TestKeenable_KeyedPath(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotPath, gotAPIKey string
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -103,7 +104,7 @@ func TestKeenable_KeyedPath(t *testing.T) {
 	tool := NewKeenableToolWithAPIKey(helper, "key-xyz")
 	tool.envBaseURL = func() string { return "https://" + srv.URL[len("http://"):] }
 
-	if _, err := tool.InvokableRun(context.Background(),
+	if _, err := tool.InvokableRun(ctx,
 		`{"query":"ragflow","mode":"realtime"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -120,6 +121,7 @@ func TestKeenable_KeyedPath(t *testing.T) {
 // that the result list is truncated to top_n.
 func TestKeenable_SiteAndTopN(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotBody map[string]any
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -139,7 +141,7 @@ func TestKeenable_SiteAndTopN(t *testing.T) {
 	})
 	tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] })
 
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"query":"x","site":"example.com","top_n":2}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -171,6 +173,7 @@ func TestKeenable_SiteAndTopN(t *testing.T) {
 // results from the upstream response (the default in the Python tool).
 func TestKeenable_DefaultTopN(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		// 12 results; default top_n is 10, so we expect 10 in the envelope.
@@ -193,7 +196,7 @@ func TestKeenable_DefaultTopN(t *testing.T) {
 	})
 	tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] })
 
-	out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"x"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -209,9 +212,10 @@ func TestKeenable_DefaultTopN(t *testing.T) {
 
 func TestKeenable_MissingQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableTool()
-	out, err := tool.InvokableRun(context.Background(), `{}`)
+	out, err := tool.InvokableRun(ctx, `{}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -228,9 +232,10 @@ func TestKeenable_MissingQuery(t *testing.T) {
 // of realtime mode without a configured api_key.
 func TestKeenable_RealtimeRequiresAPIKey(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableTool()
-	_, err := tool.InvokableRun(context.Background(), `{"query":"x","mode":"realtime"}`)
+	_, err := tool.InvokableRun(ctx, `{"query":"x","mode":"realtime"}`)
 	if err == nil {
 		t.Fatal("expected error for realtime mode without api_key")
 	}
@@ -243,9 +248,10 @@ func TestKeenable_RealtimeRequiresAPIKey(t *testing.T) {
 // up front instead of being forwarded to the upstream.
 func TestKeenable_InvalidMode(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableTool()
-	_, err := tool.InvokableRun(context.Background(), `{"query":"x","mode":"bogus"}`)
+	_, err := tool.InvokableRun(ctx, `{"query":"x","mode":"bogus"}`)
 	if err == nil {
 		t.Fatal("expected error for invalid mode")
 	}
@@ -302,6 +308,7 @@ func TestKeenable_ResolveBaseURL(t *testing.T) {
 // the test does not depend on the host environment.
 func TestKeenable_BaseURLFromEnv(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotPath string
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -318,7 +325,7 @@ func TestKeenable_BaseURLFromEnv(t *testing.T) {
 		return "https://" + srv.URL[len("http://"):]
 	})
 
-	if _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`); err != nil {
+	if _, err := tool.InvokableRun(ctx, `{"query":"x"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	if gotPath != "/v1/search/public" {
@@ -330,9 +337,10 @@ func TestKeenable_BaseURLFromEnv(t *testing.T) {
 // reported back to the caller instead of being silently sent.
 func TestKeenable_BadBaseURL(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableToolWithEnvBaseURL(NewHTTPHelper(), func() string { return "http://example.com" })
-	_, err := tool.InvokableRun(context.Background(), `{"query":"x"}`)
+	_, err := tool.InvokableRun(ctx, `{"query":"x"}`)
 	if err == nil {
 		t.Fatal("expected error for non-https non-loopback base URL")
 	}
@@ -345,6 +353,7 @@ func TestKeenable_BadBaseURL(t *testing.T) {
 // is surfaced as an error and an _ERROR envelope.
 func TestKeenable_UpstreamError(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, "boom", http.StatusInternalServerError)
@@ -356,7 +365,7 @@ func TestKeenable_UpstreamError(t *testing.T) {
 	})
 	tool := NewKeenableToolWithEnvBaseURL(helper, func() string { return "https://" + srv.URL[len("http://"):] })
 
-	out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"x"}`)
 	if err == nil {
 		t.Fatal("expected error for 5xx response")
 	}
@@ -372,9 +381,10 @@ func TestKeenable_UpstreamError(t *testing.T) {
 // TestKeenable_Info verifies the model-facing metadata.
 func TestKeenable_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -411,6 +421,7 @@ func TestKeenable_Info(t *testing.T) {
 
 func TestKeenable_ComponentContractReferencesAndOutputs(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewKeenableTool()
 	spec := tool.ComponentSpec()
@@ -430,7 +441,7 @@ func TestKeenable_ComponentContractReferencesAndOutputs(t *testing.T) {
 		"url":         "https://example.com/item",
 		"description": "Fresh search result",
 	}}}
-	chunks, docAggs := tool.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := tool.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}

internal/agent/tool/mcp_test.go

@@ -63,8 +63,9 @@ func TestMCPToolAdapter_InfoReturnsMCPDescriptor(t *testing.T) {
 // mcpclient is discovery-only; InvokableRun must return a clear error
 // until tools/call lands.
 func TestMCPToolAdapter_InvokableRunNotYetImplemented(t *testing.T) {
+	ctx := t.Context()
 	a := NewMCPToolAdapter(mcpclient.Tool{Name: "x"})
-	out, err := a.InvokableRun(context.Background(), `{"q":"hi"}`)
+	out, err := a.InvokableRun(ctx, `{"q":"hi"}`)
 	if err == nil {
 		t.Fatal("expected error from unimplemented tools/call")
 	}
@@ -162,6 +163,7 @@ func TestMarshalArguments_ValidJSON(t *testing.T) {
 // envelope (string result) and the session lifecycle
 // (initialize → tools/call).
 func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) {
+	ctx := t.Context()
 	defer mcpLoopbackOverride(t)()
 	var sawCall bool
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -185,7 +187,7 @@ func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) {
 	defer srv.Close()
 
 	a := NewMCPToolAdapterFull(mcpclient.Tool{Name: "echo"}, srv.URL, nil, 2*time.Second, srv.Client())
-	out, err := a.InvokableRun(context.Background(), `{"msg":"hi"}`)
+	out, err := a.InvokableRun(ctx, `{"msg":"hi"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -200,6 +202,7 @@ func TestMCPToolAdapter_InvokableRunDispatchesCallTool(t *testing.T) {
 // TestMCPToolAdapter_InvokableRunIsError: a tools/call response
 // with isError=true surfaces as a Go error.
 func TestMCPToolAdapter_InvokableRunIsError(t *testing.T) {
+	ctx := t.Context()
 	defer mcpLoopbackOverride(t)()
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		body, _ := io.ReadAll(r.Body)
@@ -220,7 +223,7 @@ func TestMCPToolAdapter_InvokableRunIsError(t *testing.T) {
 	defer srv.Close()
 
 	a := NewMCPToolAdapterFull(mcpclient.Tool{Name: "echo"}, srv.URL, nil, 2*time.Second, srv.Client())
-	_, err := a.InvokableRun(context.Background(), `{}`)
+	_, err := a.InvokableRun(ctx, `{}`)
 	if err == nil {
 		t.Fatalf("expected error for isError response")
 	}

internal/agent/tool/pubmed_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -76,6 +75,7 @@ func TestPubMed_BuildURL(t *testing.T) {
 }
 
 func TestPubMed_InvokableRunParsesXML(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	var esearchHits, efetchHits int32
@@ -123,7 +123,7 @@ func TestPubMed_InvokableRunParsesXML(t *testing.T) {
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
 	tool := NewPubMedToolWithDefaults(helper, pubmedParams{TopN: 3, Email: "tester@example.com"})
-	out, err := tool.InvokableRun(context.Background(), `{"query":"ragflow"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"ragflow"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -168,6 +168,7 @@ func TestPubMed_InvokableRunParsesXML(t *testing.T) {
 }
 
 func TestPubMed_InvokableRunEmptyResults(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -178,7 +179,7 @@ func TestPubMed_InvokableRunEmptyResults(t *testing.T) {
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
 	tool := NewPubMedToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(), `{"query":"missing"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"missing"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -192,10 +193,11 @@ func TestPubMed_InvokableRunEmptyResults(t *testing.T) {
 }
 
 func TestPubMed_InvokableRunEmptyQuery(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	tool := NewPubMedTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty): %v", err)
 	}
@@ -206,10 +208,11 @@ func TestPubMed_InvokableRunEmptyQuery(t *testing.T) {
 }
 
 func TestPubMed_InfoOnlyExposesQuery(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	tool := NewPubMedTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -258,6 +261,7 @@ func TestPubMed_BuildByNameAcceptsNodeParams(t *testing.T) {
 }
 
 func TestPubMed_ComponentReferencesAndOutputs(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	pubmed := NewPubMedTool()
@@ -268,7 +272,7 @@ func TestPubMed_ComponentReferencesAndOutputs(t *testing.T) {
 	envelope := map[string]any{"results": []any{map[string]any{
 		"title": "Paper", "url": "https://pubmed.ncbi.nlm.nih.gov/1", "content": "Title: Paper\nAbstract: Evidence.",
 	}}}
-	chunks, docAggs := pubmed.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := pubmed.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "Paper" {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}

internal/agent/tool/qweather_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -88,6 +87,7 @@ func TestQWeather_BuildURL(t *testing.T) {
 }
 
 func TestQWeather_ParseResponse(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	var gotQuery url.Values
@@ -117,7 +117,7 @@ func TestQWeather_ParseResponse(t *testing.T) {
 	})
 	tool := NewQWeatherToolWith(helper)
 
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"api_key":"K-abc","location":"101010100"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -157,6 +157,7 @@ func TestQWeather_ParseResponse(t *testing.T) {
 }
 
 func TestQWeather_UpstreamBusinessError(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -174,7 +175,7 @@ func TestQWeather_UpstreamBusinessError(t *testing.T) {
 	})
 	tool := NewQWeatherToolWith(helper)
 
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"api_key":"K-abc","location":"000000000"}`)
 	if err == nil {
 		t.Fatal("expected error for non-200 upstream code, got nil")
@@ -185,10 +186,11 @@ func TestQWeather_UpstreamBusinessError(t *testing.T) {
 }
 
 func TestQWeather_RejectsMissingAPIKey(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	tool := NewQWeatherTool()
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"location":"101010100"}`)
 	if err == nil {
 		t.Fatal("expected error for missing api_key")
@@ -199,10 +201,11 @@ func TestQWeather_RejectsMissingAPIKey(t *testing.T) {
 }
 
 func TestQWeather_RejectsMissingLocation(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	tool := NewQWeatherTool()
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"api_key":"K-abc"}`)
 	if err == nil {
 		t.Fatal("expected error for missing location")
@@ -213,10 +216,11 @@ func TestQWeather_RejectsMissingLocation(t *testing.T) {
 }
 
 func TestQWeather_Info(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	tool := NewQWeatherTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}

internal/agent/tool/retrieval_test.go

@@ -29,10 +29,11 @@ import (
 )
 
 func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	rt := NewRetrievalTool()
-	out, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-1"]}`)
+	out, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
 	if err == nil {
 		t.Fatal("expected stub error, got nil")
 	}
@@ -54,9 +55,10 @@ func TestRetrieval_StubsErrorWhenServiceMissing(t *testing.T) {
 }
 
 func TestRetrieval_RejectsUseKG(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 	rt := NewRetrievalTool()
-	out, err := rt.InvokableRun(context.Background(), `{"query":"x","use_kg":true}`)
+	out, err := rt.InvokableRun(ctx, `{"query":"x","use_kg":true}`)
 	if !errors.Is(err, ErrGraphRAGNotSupported) {
 		t.Fatalf("err = %v, want ErrGraphRAGNotSupported", err)
 	}
@@ -66,10 +68,11 @@ func TestRetrieval_RejectsUseKG(t *testing.T) {
 }
 
 func TestRetrieval_InfoMatchesPythonMeta(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	rt := NewRetrievalTool()
-	info, err := rt.Info(context.Background())
+	info, err := rt.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -102,10 +105,11 @@ func TestRetrieval_InfoMatchesPythonMeta(t *testing.T) {
 }
 
 func TestRetrieval_EmptyArgsIsHandled(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	rt := NewRetrievalTool()
-	out, err := rt.InvokableRun(context.Background(), "")
+	out, err := rt.InvokableRun(ctx, "")
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -126,7 +130,7 @@ func TestRetrieval_PassesTenantIDFromCanvasState(t *testing.T) {
 
 	state := runtime.NewCanvasState("run-1", "task-1")
 	state.Sys["tenant_id"] = "tenant-1"
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 
 	rt := NewRetrievalTool()
 	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
@@ -146,7 +150,7 @@ func TestRetrieval_PassesUserIDWhenTenantIDMissing(t *testing.T) {
 
 	state := runtime.NewCanvasState("run-1", "task-1")
 	state.Sys["user_id"] = "user-1"
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 
 	rt := NewRetrievalTool()
 	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)
@@ -183,7 +187,9 @@ func TestRetrieval_UsesNodeParamsAsDefaults(t *testing.T) {
 		t.Fatalf("BuildByName(retrieval) returned %T, want *RetrievalTool", built)
 	}
 
-	_, err = rt.InvokableRun(context.Background(), `{"query":"hello"}`)
+	ctx := t.Context()
+
+	_, err = rt.InvokableRun(ctx, `{"query":"hello"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -226,7 +232,8 @@ func TestRetrieval_ExplicitZeroSimilarityArgsOverrideDefaults(t *testing.T) {
 		SimilarityThreshold:      &similarityThreshold,
 		KeywordsSimilarityWeight: &keywordsSimilarityWeight,
 	})
-	_, err := retrievalTool.InvokableRun(context.Background(), `{"query":"hello","similarity_threshold":0,"keywords_similarity_weight":0}`)
+	ctx := t.Context()
+	_, err := retrievalTool.InvokableRun(ctx, `{"query":"hello","similarity_threshold":0,"keywords_similarity_weight":0}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -274,9 +281,10 @@ func TestRetrieval_ParsesEnhancementNodeParams(t *testing.T) {
 
 func TestRetrieval_UsesEmptyResponseForEmptyQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	rt := NewRetrievalToolWithDefaults(retrievalArgs{EmptyResponse: "No query or result."})
-	out, err := rt.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := rt.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -314,7 +322,7 @@ func TestRetrieval_ResolvesCanvasVariables(t *testing.T) {
 	state.SetVar("source", "ids", []any{"kb-1", "kb-2"})
 	state.SetVar("source", "query", "semantic question")
 	state.SetVar("source", "value", "2026")
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 
 	ids, err := resolveRetrievalDatasetIDs(ctx, []string{"source@ids"})
 	if err != nil {
@@ -347,17 +355,18 @@ func TestRetrieval_OmitsUnsetEmptyResponseFromArguments(t *testing.T) {
 }
 
 func TestRetrieval_UsesEmptyResponseWhenSearchHasNoChunks(t *testing.T) {
+	ctx := t.Context()
 	prev := GetRetrievalService()
 	SetRetrievalService(staticRetrievalService{})
 	t.Cleanup(func() { SetRetrievalService(prev) })
 
 	rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-1"}, EmptyResponse: "No matching chunk."})
-	out, err := rt.InvokableRun(context.Background(), `{"query":"love"}`)
+	out, err := rt.InvokableRun(ctx, `{"query":"love"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var result retrievalResult
-	if err := json.Unmarshal([]byte(out), &result); err != nil {
+	if err = json.Unmarshal([]byte(out), &result); err != nil {
 		t.Fatalf("unmarshal result: %v", err)
 	}
 	if result.FormalizedContent != "No matching chunk." {
@@ -384,13 +393,14 @@ func TestRetrieval_IgnoresCanvasMetadataNodeParams(t *testing.T) {
 }
 
 func TestRetrieval_ModelArgsOverrideNodeDatasetIDs(t *testing.T) {
+	ctx := t.Context()
 	prev := GetRetrievalService()
 	svc := &capturingRetrievalService{}
 	SetRetrievalService(svc)
 	t.Cleanup(func() { SetRetrievalService(prev) })
 
 	rt := NewRetrievalToolWithDefaults(retrievalArgs{DatasetIDs: []string{"kb-default"}, TopN: 3})
-	_, err := rt.InvokableRun(context.Background(), `{"query":"hello","dataset_ids":["kb-call"],"top_n":5}`)
+	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-call"],"top_n":5}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -421,7 +431,7 @@ func TestRetrieval_RecordsFrontendReferencePayload(t *testing.T) {
 	t.Cleanup(func() { SetRetrievalService(prev) })
 
 	state := runtime.NewCanvasState("run-1", "task-1")
-	ctx := runtime.WithState(context.Background(), state)
+	ctx := runtime.WithState(t.Context(), state)
 
 	rt := NewRetrievalTool()
 	_, err := rt.InvokableRun(ctx, `{"query":"hello","dataset_ids":["kb-1"]}`)

internal/agent/tool/searxng_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"errors"
 	"net"
@@ -59,6 +58,7 @@ func TestSearXNGBuildURLMatchesPythonQuery(t *testing.T) {
 }
 
 func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, request *http.Request) {
@@ -83,7 +83,7 @@ func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) {
 	defaults.SearXNGURL = server.URL
 	defaults.TopN = 2
 	searchTool := newLocalSearXNGTool(t, defaults)
-	out, err := searchTool.InvokableRun(context.Background(), `{"query":"  ragflow search  ","searxng_url":"http://127.0.0.1:1"}`)
+	out, err := searchTool.InvokableRun(ctx, `{"query":"  ragflow search  ","searxng_url":"http://127.0.0.1:1"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -106,9 +106,10 @@ func TestSearXNGInvokableRunPreservesRawResultsAndTopN(t *testing.T) {
 }
 
 func TestSearXNGInfoMatchesPythonModelSchema(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
-	info, err := NewSearXNGTool().Info(context.Background())
+	info, err := NewSearXNGTool().Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -138,6 +139,7 @@ func TestSearXNGInfoMatchesPythonModelSchema(t *testing.T) {
 }
 
 func TestSearXNGEmptyTryRunInputsSkipRequest(t *testing.T) {
+	ctx := t.Context()
 	t.Parallel()
 
 	searchTool := NewSearXNGTool()
@@ -151,12 +153,12 @@ func TestSearXNGEmptyTryRunInputsSkipRequest(t *testing.T) {
 		`{"query":"   ","searxng_url":"https://example.com"}`,
 		`{"query":"ragflow"}`,
 	} {
-		out, err := searchTool.InvokableRun(context.Background(), args)
+		out, err := searchTool.InvokableRun(ctx, args)
 		if err != nil {
 			t.Fatalf("InvokableRun(%s): %v", args, err)
 		}
 		var envelope searxngEnvelope
-		if err := json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 || envelope.Error != "" {
+		if err = json.Unmarshal([]byte(out), &envelope); err != nil || len(envelope.Results) != 0 || envelope.Error != "" {
 			t.Fatalf("InvokableRun(%s) = %s, want empty results", args, out)
 		}
 	}
@@ -202,6 +204,7 @@ func TestSearXNGBuildByNameRejectsInvalidNodeParams(t *testing.T) {
 
 func TestSearXNGComponentContractReferencesAndOutputs(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewSearXNGTool()
 	spec := tool.ComponentSpec()
@@ -225,7 +228,7 @@ func TestSearXNGComponentContractReferencesAndOutputs(t *testing.T) {
 		map[string]any{"title": "RAGFlow\nDocs", "url": "https://ragflow.io", "content": content, "engine": "bing", "score": 0.9},
 		map[string]any{"title": "Empty", "url": "https://example.com", "content": ""},
 	}}
-	chunks, docAggs := tool.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := tool.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
@@ -283,6 +286,7 @@ func TestRenderSearXNGReferencesStopsBeforeOverBudgetBlock(t *testing.T) {
 
 func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var calls atomic.Int32
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
@@ -294,7 +298,7 @@ func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) {
 	defaults := defaultSearXNGParams()
 	defaults.SearXNGURL = server.URL
 	searchTool := newLocalSearXNGTool(t, defaults)
-	if _, err := searchTool.InvokableRun(context.Background(), `{"query":"single attempt"}`); err == nil {
+	if _, err := searchTool.InvokableRun(ctx, `{"query":"single attempt"}`); err == nil {
 		t.Fatal("InvokableRun succeeded, want upstream error")
 	}
 	if calls.Load() != 1 {
@@ -304,11 +308,12 @@ func TestSearXNGDoesNotRetryFailedRequest(t *testing.T) {
 
 func TestSearXNGSSRFGuardRejectsLoopback(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	defaults := defaultSearXNGParams()
 	defaults.SearXNGURL = "http://127.0.0.1:4000"
 	searchTool := newSearXNGToolWithDefaults(nil, defaults)
-	out, err := searchTool.InvokableRun(context.Background(), `{"query":"metadata"}`)
+	out, err := searchTool.InvokableRun(ctx, `{"query":"metadata"}`)
 	if err == nil || !errors.Is(err, ErrSSRFBlocked) {
 		t.Fatalf("err = %v, want ErrSSRFBlocked", err)
 	}

internal/agent/tool/tavily_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -29,6 +28,7 @@ import (
 
 func TestTavily_BuildRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotPath, gotAuth, gotCT, gotMethod string
 	var gotBody map[string]any
@@ -44,14 +44,14 @@ func TestTavily_BuildRequest(t *testing.T) {
 	}))
 	defer srv.Close()
 
-	// Point the hard-coded tavily endpoint at the test server via a
+	// Point the hard-coded tavily endpoint at the test server via
 	// transport that rewrites the host. Avoids the global-package-var
 	// race that breaks parallel tests.
 	helper := NewHTTPHelper().WithClient(&http.Client{
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewTavilyToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"query":"ragflow","api_key":"key-xyz","max_results":3,"search_depth":"advanced","include_raw_content":true,"include_images":true}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -91,6 +91,7 @@ func TestTavily_BuildRequest(t *testing.T) {
 
 func TestTavily_ParseResponse(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", "application/json")
@@ -108,7 +109,7 @@ func TestTavily_ParseResponse(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewTavilyToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"query":"x","api_key":"k"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -134,11 +135,12 @@ func TestTavily_ParseResponse(t *testing.T) {
 
 func TestTavily_RequiresAPIKey(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	// envKey always returns "" so we know the failure is from the
 	// missing api_key, not from a stray process env var.
 	tool := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" })
-	out, err := tool.InvokableRun(context.Background(), `{"query":"x"}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"x"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun should not return a Go error for missing api_key: %v", err)
 	}
@@ -149,6 +151,7 @@ func TestTavily_RequiresAPIKey(t *testing.T) {
 
 func TestTavily_APIKeyFromEnv(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotAuth string
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -161,7 +164,7 @@ func TestTavily_APIKeyFromEnv(t *testing.T) {
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewTavilyToolWithEnvKey(helper, func() string { return "from-env" })
-	if _, err := tool.InvokableRun(context.Background(), `{"query":"x"}`); err != nil {
+	if _, err := tool.InvokableRun(ctx, `{"query":"x"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	if gotAuth != "Bearer from-env" {
@@ -171,9 +174,10 @@ func TestTavily_APIKeyFromEnv(t *testing.T) {
 
 func TestTavily_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTavilyTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -187,15 +191,16 @@ func TestTavily_Info(t *testing.T) {
 
 func TestTavily_InfoArraySchemasIncludeStringItems(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	info, err := NewTavilyTool().Info(context.Background())
+	info, err := NewTavilyTool().Info(ctx)
 	if err != nil {
 		t.Fatalf("Tavily Info: %v", err)
 	}
 	assertToolArrayItemType(t, info, "include_domains")
 	assertToolArrayItemType(t, info, "exclude_domains")
 
-	extractInfo, err := NewTavilyExtractTool().Info(context.Background())
+	extractInfo, err := NewTavilyExtractTool().Info(ctx)
 	if err != nil {
 		t.Fatalf("Tavily Extract Info: %v", err)
 	}
@@ -233,14 +238,15 @@ func assertToolArrayItemType(t *testing.T, info *schema.ToolInfo, fieldName stri
 
 func TestTavily_EmptyQueryReturnsEmptyResults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tavily := NewTavilyToolWithEnvKey(NewHTTPHelper(), func() string { return "" })
-	out, err := tavily.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tavily.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty query): %v", err)
 	}
 	var envelope tavilyEnvelope
-	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 != "" {
@@ -250,6 +256,7 @@ func TestTavily_EmptyQueryReturnsEmptyResults(t *testing.T) {
 
 func TestTavily_PreservesRawResults(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", "application/json")
@@ -257,12 +264,12 @@ func TestTavily_PreservesRawResults(t *testing.T) {
 	}))
 	defer srv.Close()
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
-	out, err := NewTavilyToolWith(helper).InvokableRun(context.Background(), `{"query":"x","api_key":"k"}`)
+	out, err := NewTavilyToolWith(helper).InvokableRun(ctx, `{"query":"x","api_key":"k"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var envelope tavilyEnvelope
-	if err := json.Unmarshal([]byte(out), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(out), &envelope); err != nil {
 		t.Fatalf("decode response: %v", err)
 	}
 	custom, ok := envelope.Results[0]["custom"].(map[string]any)
@@ -303,6 +310,7 @@ func TestTavily_BuildByNameUsesNodeDefaults(t *testing.T) {
 
 func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotBody map[string]any
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -316,7 +324,7 @@ func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) {
 	tavily := newTavilyTool(helper, func() string { return "" }, tavilyParams{
 		APIKey: "stored-key", IncludeAnswer: true,
 	})
-	if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow"}`); err != nil {
+	if _, err := tavily.InvokableRun(ctx, `{"query":"ragflow"}`); err != nil {
 		t.Fatalf("InvokableRun(node defaults): %v", err)
 	}
 	if gotBody["include_answer"] != true {
@@ -325,7 +333,7 @@ func TestTavily_ExplicitFlagsOverrideNodeDefaults(t *testing.T) {
 	if gotBody["include_raw_content"] != false || gotBody["include_images"] != false {
 		t.Fatalf("include_raw_content/images should be forced false: %#v", gotBody)
 	}
-	if _, err := tavily.InvokableRun(context.Background(), `{"query":"ragflow","include_answer":false}`); err != nil {
+	if _, err := tavily.InvokableRun(ctx, `{"query":"ragflow","include_answer":false}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	if gotBody["include_answer"] != false {
@@ -354,6 +362,7 @@ func TestTavily_BuildByNameRejectsInvalidNodeDefaults(t *testing.T) {
 
 func TestTavily_ComponentReferencesAndOutputs(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tavily := NewTavilyTool()
 	spec := tavily.ComponentSpec()
@@ -373,7 +382,7 @@ func TestTavily_ComponentReferencesAndOutputs(t *testing.T) {
 		"score":       float64(0.75),
 		"custom":      "preserved",
 	}}}
-	chunks, docAggs := tavily.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := tavily.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}
@@ -398,6 +407,7 @@ func TestTavily_ComponentReferencesAndOutputs(t *testing.T) {
 
 func TestTavilyExtract_BuildRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotPath, gotAuth, gotMethod string
 	var gotBody map[string]any
@@ -413,7 +423,7 @@ func TestTavilyExtract_BuildRequest(t *testing.T) {
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
 	tool := NewTavilyExtractToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"urls":"https://a.example, https://b.example","api_key":"key-xyz","extract_depth":"advanced","format":"text"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -444,6 +454,7 @@ func TestTavilyExtract_BuildRequest(t *testing.T) {
 
 func TestTavilyExtract_ParseResponse(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", "application/json")
@@ -453,7 +464,7 @@ func TestTavilyExtract_ParseResponse(t *testing.T) {
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(srv.URL)})
 	tool := NewTavilyExtractToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"],"api_key":"k"}`)
+	out, err := tool.InvokableRun(ctx, `{"urls":["https://a.example/"],"api_key":"k"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -478,9 +489,10 @@ func TestTavilyExtract_ParseResponse(t *testing.T) {
 
 func TestTavilyExtract_RequiresAPIKey(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTavilyExtractToolWithEnvKey(NewHTTPHelper(), func() string { return "" })
-	out, err := tool.InvokableRun(context.Background(), `{"urls":["https://a.example/"]}`)
+	out, err := tool.InvokableRun(ctx, `{"urls":["https://a.example/"]}`)
 	if err != nil {
 		t.Fatalf("InvokableRun should not return a Go error for missing api_key: %v", err)
 	}
@@ -491,9 +503,10 @@ func TestTavilyExtract_RequiresAPIKey(t *testing.T) {
 
 func TestTavilyExtract_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTavilyExtractTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -571,6 +584,7 @@ func TestTavilyExtract_BuildByNameAcceptsNodeDefaults(t *testing.T) {
 
 func TestTavilyExtract_UsesNodeDefaults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotAuth string
 	var gotBody map[string]any
@@ -589,7 +603,7 @@ func TestTavilyExtract_UsesNodeDefaults(t *testing.T) {
 		ExtractDepth: "advanced",
 		Format:       "text",
 	})
-	if _, err := tavily.InvokableRun(context.Background(), `{"unrelated":"ignored"}`); err != nil {
+	if _, err := tavily.InvokableRun(ctx, `{"unrelated":"ignored"}`); err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	if gotAuth != "Bearer stored-key" {

internal/agent/tool/tushare_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"io"
 	"net/http"
@@ -88,6 +87,7 @@ func TestTushare_BuildRequest(t *testing.T) {
 
 func TestTushare_ParseResponse(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var (
 		gotMethod  string
@@ -122,7 +122,7 @@ func TestTushare_ParseResponse(t *testing.T) {
 	})
 	tool := NewTushareToolWith(helper)
 
-	out, err := tool.InvokableRun(context.Background(),
+	out, err := tool.InvokableRun(ctx,
 		`{"token":"T-test","api_name":"stock_basic","params":{"list_status":"L"}}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
@@ -169,9 +169,10 @@ func TestTushare_ParseResponse(t *testing.T) {
 
 func TestTushare_RejectsMissingToken(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTushareTool()
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"api_name":"stock_basic"}`)
 	if err == nil {
 		t.Fatal("expected error for missing token")
@@ -183,9 +184,10 @@ func TestTushare_RejectsMissingToken(t *testing.T) {
 
 func TestTushare_RejectsMissingAPIName(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTushareTool()
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"token":"T-abc"}`)
 	if err == nil {
 		t.Fatal("expected error for missing api_name")
@@ -197,6 +199,7 @@ func TestTushare_RejectsMissingAPIName(t *testing.T) {
 
 func TestTushare_UpstreamErrorCode(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", "application/json")
@@ -213,7 +216,7 @@ func TestTushare_UpstreamErrorCode(t *testing.T) {
 	})
 	tool := NewTushareToolWith(helper)
 
-	_, err := tool.InvokableRun(context.Background(),
+	_, err := tool.InvokableRun(ctx,
 		`{"token":"T-abc","api_name":"premium_only"}`)
 	if err == nil {
 		t.Fatal("expected error for non-zero code, got nil")
@@ -225,9 +228,10 @@ func TestTushare_UpstreamErrorCode(t *testing.T) {
 
 func TestTushare_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewTushareTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}

internal/agent/tool/wencai_test.go

@@ -26,6 +26,7 @@ import (
 
 func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	cases := []struct {
 		name string
@@ -39,7 +40,7 @@ func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) {
 	for _, tc := range cases {
 		t.Run(tc.name, func(t *testing.T) {
 			t.Parallel()
-			out, err := NewWencaiTool().InvokableRun(context.Background(), tc.args)
+			out, err := NewWencaiTool().InvokableRun(ctx, tc.args)
 			if err != nil {
 				t.Fatalf("InvokableRun errored: %v (out=%s)", err, out)
 			}
@@ -59,8 +60,9 @@ func TestWencai_InvokeMatchesCurrentPythonResult(t *testing.T) {
 
 func TestWencai_RejectsMalformedJSON(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	out, err := NewWencaiTool().InvokableRun(context.Background(), `{not json`)
+	out, err := NewWencaiTool().InvokableRun(ctx, `{not json`)
 	if err == nil {
 		t.Fatal("expected malformed JSON error")
 	}
@@ -78,8 +80,9 @@ func TestWencai_RejectsMalformedJSON(t *testing.T) {
 
 func TestWencai_RespectsCanceledContext(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(ctx)
 	cancel()
 	out, err := NewWencaiTool().InvokableRun(ctx, `{"query":"商业航天"}`)
 	if !errors.Is(err, context.Canceled) {
@@ -96,8 +99,9 @@ func TestWencai_RespectsCanceledContext(t *testing.T) {
 
 func TestWencai_InfoMatchesPythonMeta(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	info, err := NewWencaiTool().Info(context.Background())
+	info, err := NewWencaiTool().Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}

internal/agent/tool/wikipedia_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -97,6 +96,7 @@ func TestWikipedia_BuildURL(t *testing.T) {
 
 func TestWikipedia_ParseResults(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotUA string
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -114,12 +114,12 @@ func TestWikipedia_ParseResults(t *testing.T) {
 	defer srv.Close()
 
 	// Point the hard-coded en.wikipedia.org endpoint at the test server
-	// by injecting a transport that rewrites the request host.
+	// by injecting transport that rewrites the request host.
 	helper := NewHTTPHelper().WithClient(&http.Client{
 		Transport: rewriteHostTransport(srv.URL),
 	})
 	tool := NewWikipediaToolWith(helper)
-	out, err := tool.InvokableRun(context.Background(), `{"query":"RAG","lang":"en","max_results":5}`)
+	out, err := tool.InvokableRun(ctx, `{"query":"RAG","lang":"en","max_results":5}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -179,9 +179,10 @@ func (t *hostSwapRT) RoundTrip(req *http.Request) (*http.Response, error) {
 
 func TestWikipedia_Info(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewWikipediaTool()
-	info, err := tool.Info(context.Background())
+	info, err := tool.Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}
@@ -195,9 +196,10 @@ func TestWikipedia_Info(t *testing.T) {
 
 func TestWikipedia_EmptyQuery(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tool := NewWikipediaTool()
-	out, err := tool.InvokableRun(context.Background(), `{"query":""}`)
+	out, err := tool.InvokableRun(ctx, `{"query":""}`)
 	if err != nil {
 		t.Fatalf("InvokableRun(empty): %v", err)
 	}
@@ -212,6 +214,7 @@ func TestWikipedia_EmptyQuery(t *testing.T) {
 
 func TestWikipedia_ComponentReferencesAndOutputs(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	wikipedia := NewWikipediaTool()
 	spec := wikipedia.ComponentSpec()
@@ -223,7 +226,7 @@ func TestWikipedia_ComponentReferencesAndOutputs(t *testing.T) {
 		"url":     "https://en.wikipedia.org/wiki/RAG",
 		"content": "RAG is an acronym.",
 	}}}
-	chunks, docAggs := wikipedia.BuildReferences(context.Background(), envelope)
+	chunks, docAggs := wikipedia.BuildReferences(ctx, envelope)
 	if len(chunks) != 1 || len(docAggs) != 1 || chunks[0]["document_name"] != "RAG" || chunks[0]["similarity"] != 1 {
 		t.Fatalf("references = %#v / %#v", chunks, docAggs)
 	}

internal/agent/tool/yahoo_finance_test.go

@@ -17,7 +17,6 @@
 package tool
 
 import (
-	"context"
 	"encoding/json"
 	"net/http"
 	"net/http/httptest"
@@ -50,6 +49,7 @@ func TestYahooFinanceBuildURL(t *testing.T) {
 
 func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotQuery, gotUserAgent string
 	var gotPaths []string
@@ -92,7 +92,7 @@ func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) {
 	defer server.Close()
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
-	raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":" AAPL "}`)
+	raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"stock_code":" AAPL "}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -132,6 +132,7 @@ func TestYahooFinanceInvokableRunBuildsMarkdownReport(t *testing.T) {
 
 func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var gotQuery string
 	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
@@ -158,7 +159,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) {
 	defer server.Close()
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
-	raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"query":"3800.HK"}`)
+	raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"query":"3800.HK"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -166,7 +167,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) {
 		t.Fatalf("q = %q", gotQuery)
 	}
 	var envelope yahooFinanceEnvelope
-	if err := json.Unmarshal([]byte(raw), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(raw), &envelope); err != nil {
 		t.Fatalf("unmarshal output %q: %v", raw, err)
 	}
 	if !strings.Contains(envelope.Report, "| symbol | 3800.HK |") ||
@@ -177,6 +178,7 @@ func TestYahooFinanceInvokableRunAcceptsQueryAlias(t *testing.T) {
 
 func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var paths []string
 	var gotSearchQuery string
@@ -232,7 +234,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T)
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
 	yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{Info: true, Count: true})
-	raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"Apple"}`)
+	raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"Apple"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -249,7 +251,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T)
 		}
 	}
 	var envelope yahooFinanceEnvelope
-	if err := json.Unmarshal([]byte(raw), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(raw), &envelope); err != nil {
 		t.Fatalf("unmarshal output %q: %v", raw, err)
 	}
 	if !strings.Contains(envelope.Report, "| longname | Apple Inc. |") ||
@@ -260,6 +262,7 @@ func TestYahooFinanceUsesSearchResolvedSymbolForDownstreamRequests(t *testing.T)
 
 func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	calls := 0
 	server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
@@ -269,12 +272,12 @@ func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) {
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
 
 	for _, args := range []string{`{"stock_code":""}`, `{"stock_code":"   "}`} {
-		raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), args)
+		raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, args)
 		if err != nil {
 			t.Fatalf("InvokableRun(%s): %v", args, err)
 		}
 		var envelope yahooFinanceEnvelope
-		if err := json.Unmarshal([]byte(raw), &envelope); err != nil {
+		if err = json.Unmarshal([]byte(raw), &envelope); err != nil {
 			t.Fatalf("unmarshal output: %v", err)
 		}
 		if envelope.Report != "" || envelope.Error != "" {
@@ -288,6 +291,7 @@ func TestYahooFinanceEmptyStockCodeSkipsRequest(t *testing.T) {
 
 func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	calls := 0
 	server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {
@@ -297,12 +301,12 @@ func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) {
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
 	yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{})
 
-	raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"AAPL"}`)
+	raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"AAPL"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
 	var envelope yahooFinanceEnvelope
-	if err := json.Unmarshal([]byte(raw), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(raw), &envelope); err != nil {
 		t.Fatalf("unmarshal output: %v", err)
 	}
 	if envelope.Report != "" || calls != 0 {
@@ -312,6 +316,7 @@ func TestYahooFinanceAllSectionsDisabledSkipsRequest(t *testing.T) {
 
 func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	var sawCookieOnCrumb, sawCookieOnSummary bool
 	server := httptest.NewServer(http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
@@ -352,7 +357,7 @@ func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) {
 
 	helper := NewHTTPHelper().WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
 	yahoo := NewYahooFinanceToolWithDefaults(helper, yahooFinanceParams{Count: true})
-	raw, err := yahoo.InvokableRun(context.Background(), `{"stock_code":"AAPL"}`)
+	raw, err := yahoo.InvokableRun(ctx, `{"stock_code":"AAPL"}`)
 	if err != nil {
 		t.Fatalf("InvokableRun: %v", err)
 	}
@@ -360,7 +365,7 @@ func TestYahooFinanceSummaryUsesCrumbAndCookie(t *testing.T) {
 		t.Fatalf("cookie propagation crumb=%v summary=%v", sawCookieOnCrumb, sawCookieOnSummary)
 	}
 	var envelope yahooFinanceEnvelope
-	if err := json.Unmarshal([]byte(raw), &envelope); err != nil {
+	if err = json.Unmarshal([]byte(raw), &envelope); err != nil {
 		t.Fatalf("unmarshal output: %v", err)
 	}
 	if !strings.Contains(envelope.Report, "# Count:") ||
@@ -459,6 +464,7 @@ func TestMergeYahooFinanceParamsKeepsStockCodeAndUsesNodeConfig(t *testing.T) {
 
 func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
 	tests := []struct {
 		name       string
@@ -484,7 +490,7 @@ func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) {
 			defer server.Close()
 			helper := NewHTTPHelperWithRetry(RetryConfig{MaxAttempts: 1}).WithClient(&http.Client{Transport: rewriteHostTransport(server.URL)})
 
-			raw, err := NewYahooFinanceToolWith(helper).InvokableRun(context.Background(), `{"stock_code":"AAPL"}`)
+			raw, err := NewYahooFinanceToolWith(helper).InvokableRun(ctx, `{"stock_code":"AAPL"}`)
 			if err == nil || !strings.Contains(err.Error(), test.wantError) {
 				t.Fatalf("err = %v, want %q", err, test.wantError)
 			}
@@ -501,8 +507,9 @@ func TestYahooFinanceErrorsReturnEnvelope(t *testing.T) {
 
 func TestYahooFinanceMalformedArguments(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	raw, err := NewYahooFinanceTool().InvokableRun(context.Background(), `{`)
+	raw, err := NewYahooFinanceTool().InvokableRun(ctx, `{`)
 	if err == nil || !strings.Contains(err.Error(), "parse arguments") {
 		t.Fatalf("err = %v", err)
 	}
@@ -535,8 +542,9 @@ func TestYahooFinanceComponentContract(t *testing.T) {
 
 func TestYahooFinanceInfoExposesPythonCompatibleParams(t *testing.T) {
 	t.Parallel()
+	ctx := t.Context()
 
-	info, err := NewYahooFinanceTool().Info(context.Background())
+	info, err := NewYahooFinanceTool().Info(ctx)
 	if err != nil {
 		t.Fatalf("Info: %v", err)
 	}