fix: eliminate data races on shared tool instances (#1080)

* fix: eliminate data races on shared tool instances Signed-off-by: Boris Bliznioukov <blib@mail.com> * fix: remove unused indirect dependency on github.com/gdamore/tcell/v2 Signed-off-by: Boris Bliznioukov <blib@mail.com> * fix: reviewer comments improve context handling for tool execution and ensure defaults for non-conversation callers Signed-off-by: Boris Bliznioukov <blib@mail.com> --------- Signed-off-by: Boris Bliznioukov <blib@mail.com>
2026-06-12 18:08:54 +00:00 · 2026-03-05 02:57:33 +01:00
parent 204038ec60
commit aef1e8e8c4
12 changed files with 181 additions and 252 deletions
@@ -10,11 +10,38 @@ type Tool interface {
 	Execute(ctx context.Context, args map[string]any) *ToolResult
 }

-// ContextualTool is an optional interface that tools can implement
-// to receive the current message context (channel, chatID)
-type ContextualTool interface {
-	Tool
-	SetContext(channel, chatID string)
+// --- Request-scoped tool context (channel / chatID) ---
+//
+// Carried via context.Value so that concurrent tool calls each receive
+// their own immutable copy — no mutable state on singleton tool instances.
+//
+// Keys are unexported pointer-typed vars — guaranteed collision-free,
+// and only accessible through the helper functions below.
+
+type toolCtxKey struct{ name string }
+
+var (
+	ctxKeyChannel = &toolCtxKey{"channel"}
+	ctxKeyChatID  = &toolCtxKey{"chatID"}
+)
+
+// WithToolContext returns a child context carrying channel and chatID.
+func WithToolContext(ctx context.Context, channel, chatID string) context.Context {
+	ctx = context.WithValue(ctx, ctxKeyChannel, channel)
+	ctx = context.WithValue(ctx, ctxKeyChatID, chatID)
+	return ctx
+}
+
+// ToolChannel extracts the channel from ctx, or "" if unset.
+func ToolChannel(ctx context.Context) string {
+	v, _ := ctx.Value(ctxKeyChannel).(string)
+	return v
+}
+
+// ToolChatID extracts the chatID from ctx, or "" if unset.
+func ToolChatID(ctx context.Context) string {
+	v, _ := ctx.Value(ctxKeyChatID).(string)
+	return v
 }

 // AsyncCallback is a function type that async tools use to notify completion.
@@ -22,51 +49,36 @@ type ContextualTool interface {
 //
 // The ctx parameter allows the callback to be canceled if the agent is shutting down.
 // The result parameter contains the tool's execution result.
-//
-// Example usage in an async tool:
-//
-//	func (t *MyAsyncTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
-//	    // Start async work in background
-//	    go func() {
-//	        result := doAsyncWork()
-//	        if t.callback != nil {
-//	            t.callback(ctx, result)
-//	        }
-//	    }()
-//	    return AsyncResult("Async task started")
-//	}
 type AsyncCallback func(ctx context.Context, result *ToolResult)

-// AsyncTool is an optional interface that tools can implement to support
+// AsyncExecutor is an optional interface that tools can implement to support
 // asynchronous execution with completion callbacks.
 //
-// Async tools return immediately with an AsyncResult, then notify completion
-// via the callback set by SetCallback.
+// Unlike the old AsyncTool pattern (SetCallback + Execute), AsyncExecutor
+// receives the callback as a parameter of ExecuteAsync. This eliminates the
+// data race where concurrent calls could overwrite each other's callbacks
+// on a shared tool instance.
 //
 // This is useful for:
-// - Long-running operations that shouldn't block the agent loop
-// - Subagent spawns that complete independently
-// - Background tasks that need to report results later
+//   - Long-running operations that shouldn't block the agent loop
+//   - Subagent spawns that complete independently
+//   - Background tasks that need to report results later
 //
 // Example:
 //
-//	type SpawnTool struct {
-//	    callback AsyncCallback
-//	}
-//
-//	func (t *SpawnTool) SetCallback(cb AsyncCallback) {
-//	    t.callback = cb
-//	}
-//
-//	func (t *SpawnTool) Execute(ctx context.Context, args map[string]interface{}) *ToolResult {
-//	    go t.runSubagent(ctx, args)
+//	func (t *SpawnTool) ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
+//	    go func() {
+//	        result := t.runSubagent(ctx, args)
+//	        if cb != nil { cb(ctx, result) }
+//	    }()
 //	    return AsyncResult("Subagent spawned, will report back")
 //	}
-type AsyncTool interface {
+type AsyncExecutor interface {
 	Tool
-	// SetCallback registers a callback function to be invoked when the async operation completes.
-	// The callback will be called from a goroutine and should handle thread-safety if needed.
-	SetCallback(cb AsyncCallback)
+	// ExecuteAsync runs the tool asynchronously. The callback cb will be
+	// invoked (possibly from another goroutine) when the async operation
+	// completes. cb is guaranteed to be non-nil by the caller (registry).
+	ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult
 }

 func ToolToSchema(tool Tool) map[string]any {
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"strings"
-	"sync"
 	"time"

 	"github.com/sipeed/picoclaw/pkg/bus"
@@ -24,9 +23,6 @@ type CronTool struct {
 	executor    JobExecutor
 	msgBus      *bus.MessageBus
 	execTool    *ExecTool
-	channel     string
-	chatID      string
-	mu          sync.RWMutex
 }

 // NewCronTool creates a new CronTool
@@ -102,14 +98,6 @@ func (t *CronTool) Parameters() map[string]any {
 	}
 }

-// SetContext sets the current session context for job creation
-func (t *CronTool) SetContext(channel, chatID string) {
-	t.mu.Lock()
-	defer t.mu.Unlock()
-	t.channel = channel
-	t.chatID = chatID
-}
-
 // Execute runs the tool with the given arguments
 func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
 	action, ok := args["action"].(string)
@@ -119,7 +107,7 @@ func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult

 	switch action {
 	case "add":
-		return t.addJob(args)
+		return t.addJob(ctx, args)
 	case "list":
 		return t.listJobs()
 	case "remove":
@@ -133,11 +121,9 @@ func (t *CronTool) Execute(ctx context.Context, args map[string]any) *ToolResult
 	}
 }

-func (t *CronTool) addJob(args map[string]any) *ToolResult {
-	t.mu.RLock()
-	channel := t.channel
-	chatID := t.chatID
-	t.mu.RUnlock()
+func (t *CronTool) addJob(ctx context.Context, args map[string]any) *ToolResult {
+	channel := ToolChannel(ctx)
+	chatID := ToolChatID(ctx)

 	if channel == "" || chatID == "" {
 		return ErrorResult("no session context (channel/chat_id not set). Use this tool in an active conversation.")
@@ -9,10 +9,8 @@ import (
 type SendCallback func(channel, chatID, content string) error

 type MessageTool struct {
-	sendCallback   SendCallback
-	defaultChannel string
-	defaultChatID  string
-	sentInRound    atomic.Bool // Tracks whether a message was sent in the current processing round
+	sendCallback SendCallback
+	sentInRound  atomic.Bool // Tracks whether a message was sent in the current processing round
 }

 func NewMessageTool() *MessageTool {
@@ -48,10 +46,10 @@ func (t *MessageTool) Parameters() map[string]any {
 	}
 }

-func (t *MessageTool) SetContext(channel, chatID string) {
-	t.defaultChannel = channel
-	t.defaultChatID = chatID
-	t.sentInRound.Store(false) // Reset send tracking for new processing round
+// ResetSentInRound resets the per-round send tracker.
+// Called by the agent loop at the start of each inbound message processing round.
+func (t *MessageTool) ResetSentInRound() {
+	t.sentInRound.Store(false)
 }

 // HasSentInRound returns true if the message tool sent a message during the current round.
@@ -73,10 +71,10 @@ func (t *MessageTool) Execute(ctx context.Context, args map[string]any) *ToolRes
 	chatID, _ := args["chat_id"].(string)

 	if channel == "" {
-		channel = t.defaultChannel
+		channel = ToolChannel(ctx)
 	}
 	if chatID == "" {
-		chatID = t.defaultChatID
+		chatID = ToolChatID(ctx)
 	}

 	if channel == "" || chatID == "" {
@@ -8,7 +8,6 @@ import (

 func TestMessageTool_Execute_Success(t *testing.T) {
 	tool := NewMessageTool()
-	tool.SetContext("test-channel", "test-chat-id")

 	var sentChannel, sentChatID, sentContent string
 	tool.SetSendCallback(func(channel, chatID, content string) error {
@@ -18,7 +17,7 @@ func TestMessageTool_Execute_Success(t *testing.T) {
 		return nil
 	})

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id")
 	args := map[string]any{
 		"content": "Hello, world!",
 	}
@@ -60,7 +59,6 @@ func TestMessageTool_Execute_Success(t *testing.T) {

 func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
 	tool := NewMessageTool()
-	tool.SetContext("default-channel", "default-chat-id")

 	var sentChannel, sentChatID string
 	tool.SetSendCallback(func(channel, chatID, content string) error {
@@ -69,7 +67,7 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {
 		return nil
 	})

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "default-channel", "default-chat-id")
 	args := map[string]any{
 		"content": "Test message",
 		"channel": "custom-channel",
@@ -96,14 +94,13 @@ func TestMessageTool_Execute_WithCustomChannel(t *testing.T) {

 func TestMessageTool_Execute_SendFailure(t *testing.T) {
 	tool := NewMessageTool()
-	tool.SetContext("test-channel", "test-chat-id")

 	sendErr := errors.New("network error")
 	tool.SetSendCallback(func(channel, chatID, content string) error {
 		return sendErr
 	})

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id")
 	args := map[string]any{
 		"content": "Test message",
 	}
@@ -133,9 +130,8 @@ func TestMessageTool_Execute_SendFailure(t *testing.T) {

 func TestMessageTool_Execute_MissingContent(t *testing.T) {
 	tool := NewMessageTool()
-	tool.SetContext("test-channel", "test-chat-id")

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id")
 	args := map[string]any{} // content missing

 	result := tool.Execute(ctx, args)
@@ -151,7 +147,7 @@ func TestMessageTool_Execute_MissingContent(t *testing.T) {

 func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {
 	tool := NewMessageTool()
-	// No SetContext called, so defaultChannel and defaultChatID are empty
+	// No WithToolContext — channel/chatID are empty

 	tool.SetSendCallback(func(channel, chatID, content string) error {
 		return nil
@@ -175,10 +171,9 @@ func TestMessageTool_Execute_NoTargetChannel(t *testing.T) {

 func TestMessageTool_Execute_NotConfigured(t *testing.T) {
 	tool := NewMessageTool()
-	tool.SetContext("test-channel", "test-chat-id")
 	// No SetSendCallback called

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "test-channel", "test-chat-id")
 	args := map[string]any{
 		"content": "Test message",
 	}
@@ -45,8 +45,9 @@ func (r *ToolRegistry) Execute(ctx context.Context, name string, args map[string
 }

 // ExecuteWithContext executes a tool with channel/chatID context and optional async callback.
-// If the tool implements AsyncTool and a non-nil callback is provided,
-// the callback will be set on the tool before execution.
+// If the tool implements AsyncExecutor and a non-nil callback is provided,
+// ExecuteAsync is called instead of Execute — the callback is a parameter,
+// never stored as mutable state on the tool.
 func (r *ToolRegistry) ExecuteWithContext(
 	ctx context.Context,
 	name string,
@@ -69,22 +70,23 @@ func (r *ToolRegistry) ExecuteWithContext(
 		return ErrorResult(fmt.Sprintf("tool %q not found", name)).WithError(fmt.Errorf("tool not found"))
 	}

-	// If tool implements ContextualTool, set context
-	if contextualTool, ok := tool.(ContextualTool); ok && channel != "" && chatID != "" {
-		contextualTool.SetContext(channel, chatID)
-	}
+	// Inject channel/chatID into ctx so tools read them via ToolChannel(ctx)/ToolChatID(ctx).
+	// Always inject — tools validate what they require.
+	ctx = WithToolContext(ctx, channel, chatID)

-	// If tool implements AsyncTool and callback is provided, set callback
-	if asyncTool, ok := tool.(AsyncTool); ok && asyncCallback != nil {
-		asyncTool.SetCallback(asyncCallback)
-		logger.DebugCF("tool", "Async callback injected",
+	// If tool implements AsyncExecutor and callback is provided, use ExecuteAsync.
+	// The callback is a call parameter, not mutable state on the tool instance.
+	var result *ToolResult
+	start := time.Now()
+	if asyncExec, ok := tool.(AsyncExecutor); ok && asyncCallback != nil {
+		logger.DebugCF("tool", "Executing async tool via ExecuteAsync",
 			map[string]any{
 				"tool": name,
 			})
+		result = asyncExec.ExecuteAsync(ctx, args, asyncCallback)
+	} else {
+		result = tool.Execute(ctx, args)
 	}
-
-	start := time.Now()
-	result := tool.Execute(ctx, args)
 	duration := time.Since(start)

 	// Log based on result type
@@ -25,24 +25,24 @@ func (m *mockRegistryTool) Execute(_ context.Context, _ map[string]any) *ToolRes
 	return m.result
 }

-type mockCtxTool struct {
+type mockContextAwareTool struct {
 	mockRegistryTool
-	channel string
-	chatID  string
+	lastCtx context.Context
 }

-func (m *mockCtxTool) SetContext(channel, chatID string) {
-	m.channel = channel
-	m.chatID = chatID
+func (m *mockContextAwareTool) Execute(ctx context.Context, _ map[string]any) *ToolResult {
+	m.lastCtx = ctx
+	return m.result
 }

 type mockAsyncRegistryTool struct {
 	mockRegistryTool
-	cb AsyncCallback
+	lastCB AsyncCallback
 }

-func (m *mockAsyncRegistryTool) SetCallback(cb AsyncCallback) {
-	m.cb = cb
+func (m *mockAsyncRegistryTool) ExecuteAsync(_ context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
+	m.lastCB = cb
+	return m.result
 }

 // --- helpers ---
@@ -136,34 +136,44 @@ func TestToolRegistry_Execute_NotFound(t *testing.T) {
 	}
 }

-func TestToolRegistry_ExecuteWithContext_ContextualTool(t *testing.T) {
+func TestToolRegistry_ExecuteWithContext_InjectsToolContext(t *testing.T) {
 	r := NewToolRegistry()
-	ct := &mockCtxTool{
+	ct := &mockContextAwareTool{
 		mockRegistryTool: *newMockTool("ctx_tool", "needs context"),
 	}
 	r.Register(ct)

 	r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "telegram", "chat-42", nil)

-	if ct.channel != "telegram" {
-		t.Errorf("expected channel 'telegram', got %q", ct.channel)
+	if ct.lastCtx == nil {
+		t.Fatal("expected Execute to be called")
 	}
-	if ct.chatID != "chat-42" {
-		t.Errorf("expected chatID 'chat-42', got %q", ct.chatID)
+	if got := ToolChannel(ct.lastCtx); got != "telegram" {
+		t.Errorf("expected channel 'telegram', got %q", got)
+	}
+	if got := ToolChatID(ct.lastCtx); got != "chat-42" {
+		t.Errorf("expected chatID 'chat-42', got %q", got)
 	}
 }

-func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) {
+func TestToolRegistry_ExecuteWithContext_EmptyContext(t *testing.T) {
 	r := NewToolRegistry()
-	ct := &mockCtxTool{
+	ct := &mockContextAwareTool{
 		mockRegistryTool: *newMockTool("ctx_tool", "needs context"),
 	}
 	r.Register(ct)

 	r.ExecuteWithContext(context.Background(), "ctx_tool", nil, "", "", nil)

-	if ct.channel != "" || ct.chatID != "" {
-		t.Error("SetContext should not be called with empty channel/chatID")
+	if ct.lastCtx == nil {
+		t.Fatal("expected Execute to be called")
+	}
+	// Empty values are still injected; tools decide what to do with them.
+	if got := ToolChannel(ct.lastCtx); got != "" {
+		t.Errorf("expected empty channel, got %q", got)
+	}
+	if got := ToolChatID(ct.lastCtx); got != "" {
+		t.Errorf("expected empty chatID, got %q", got)
 	}
 }

@@ -179,14 +189,14 @@ func TestToolRegistry_ExecuteWithContext_AsyncCallback(t *testing.T) {
 	cb := func(_ context.Context, _ *ToolResult) { called = true }

 	result := r.ExecuteWithContext(context.Background(), "async_tool", nil, "", "", cb)
-	if at.cb == nil {
-		t.Error("expected SetCallback to have been called")
+	if at.lastCB == nil {
+		t.Error("expected ExecuteAsync to have received a callback")
 	}
 	if !result.Async {
 		t.Error("expected async result")
 	}

-	at.cb(context.Background(), SilentResult("done"))
+	at.lastCB(context.Background(), SilentResult("done"))
 	if !called {
 		t.Error("expected callback to be invoked")
 	}
@@ -8,25 +8,18 @@ import (

 type SpawnTool struct {
 	manager        *SubagentManager
-	originChannel  string
-	originChatID   string
 	allowlistCheck func(targetAgentID string) bool
-	callback       AsyncCallback // For async completion notification
 }

+// Compile-time check: SpawnTool implements AsyncExecutor.
+var _ AsyncExecutor = (*SpawnTool)(nil)
+
 func NewSpawnTool(manager *SubagentManager) *SpawnTool {
 	return &SpawnTool{
-		manager:       manager,
-		originChannel: "cli",
-		originChatID:  "direct",
+		manager: manager,
 	}
 }

-// SetCallback implements AsyncTool interface for async completion notification
-func (t *SpawnTool) SetCallback(cb AsyncCallback) {
-	t.callback = cb
-}
-
 func (t *SpawnTool) Name() string {
 	return "spawn"
 }
@@ -56,16 +49,21 @@ func (t *SpawnTool) Parameters() map[string]any {
 	}
 }

-func (t *SpawnTool) SetContext(channel, chatID string) {
-	t.originChannel = channel
-	t.originChatID = chatID
-}
-
 func (t *SpawnTool) SetAllowlistChecker(check func(targetAgentID string) bool) {
 	t.allowlistCheck = check
 }

 func (t *SpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
+	return t.execute(ctx, args, nil)
+}
+
+// ExecuteAsync implements AsyncExecutor. The callback is passed through to the
+// subagent manager as a call parameter — never stored on the SpawnTool instance.
+func (t *SpawnTool) ExecuteAsync(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
+	return t.execute(ctx, args, cb)
+}
+
+func (t *SpawnTool) execute(ctx context.Context, args map[string]any, cb AsyncCallback) *ToolResult {
 	task, ok := args["task"].(string)
 	if !ok || strings.TrimSpace(task) == "" {
 		return ErrorResult("task is required and must be a non-empty string")
@@ -85,8 +83,20 @@ func (t *SpawnTool) Execute(ctx context.Context, args map[string]any) *ToolResul
 		return ErrorResult("Subagent manager not configured")
 	}

+	// Read channel/chatID from context (injected by registry).
+	// Fall back to "cli"/"direct" for non-conversation callers (e.g., CLI, tests)
+	// to preserve the same defaults as the original NewSpawnTool constructor.
+	channel := ToolChannel(ctx)
+	if channel == "" {
+		channel = "cli"
+	}
+	chatID := ToolChatID(ctx)
+	if chatID == "" {
+		chatID = "direct"
+	}
+
 	// Pass callback to manager for async completion notification
-	result, err := t.manager.Spawn(ctx, task, label, agentID, t.originChannel, t.originChatID, t.callback)
+	result, err := t.manager.Spawn(ctx, task, label, agentID, channel, chatID, cb)
 	if err != nil {
 		return ErrorResult(fmt.Sprintf("failed to spawn subagent: %v", err))
 	}
@@ -252,16 +252,12 @@ func (sm *SubagentManager) ListTasks() []*SubagentTask {
 // Unlike SpawnTool which runs tasks asynchronously, SubagentTool waits for completion
 // and returns the result directly in the ToolResult.
 type SubagentTool struct {
-	manager       *SubagentManager
-	originChannel string
-	originChatID  string
+	manager *SubagentManager
 }

 func NewSubagentTool(manager *SubagentManager) *SubagentTool {
 	return &SubagentTool{
-		manager:       manager,
-		originChannel: "cli",
-		originChatID:  "direct",
+		manager: manager,
 	}
 }

@@ -290,11 +286,6 @@ func (t *SubagentTool) Parameters() map[string]any {
 	}
 }

-func (t *SubagentTool) SetContext(channel, chatID string) {
-	t.originChannel = channel
-	t.originChatID = chatID
-}
-
 func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolResult {
 	task, ok := args["task"].(string)
 	if !ok {
@@ -341,13 +332,24 @@ func (t *SubagentTool) Execute(ctx context.Context, args map[string]any) *ToolRe
 		}
 	}

+	// Fall back to "cli"/"direct" for non-conversation callers (e.g., CLI, tests)
+	// to preserve the same defaults as the original NewSubagentTool constructor.
+	channel := ToolChannel(ctx)
+	if channel == "" {
+		channel = "cli"
+	}
+	chatID := ToolChatID(ctx)
+	if chatID == "" {
+		chatID = "direct"
+	}
+
 	loopResult, err := RunToolLoop(ctx, ToolLoopConfig{
 		Provider:      sm.provider,
 		Model:         sm.defaultModel,
 		Tools:         tools,
 		MaxIterations: maxIter,
 		LLMOptions:    llmOptions,
-	}, messages, t.originChannel, t.originChatID)
+	}, messages, channel, chatID)
 	if err != nil {
 		return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err)
 	}
@@ -50,9 +50,8 @@ func TestSubagentManager_SetLLMOptions_AppliesToRunToolLoop(t *testing.T) {
 	manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil)
 	manager.SetLLMOptions(2048, 0.6)
 	tool := NewSubagentTool(manager)
-	tool.SetContext("cli", "direct")

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "cli", "direct")
 	args := map[string]any{"task": "Do something"}
 	result := tool.Execute(ctx, args)

@@ -147,28 +146,14 @@ func TestSubagentTool_Parameters(t *testing.T) {
 	}
 }

-// TestSubagentTool_SetContext verifies context setting
-func TestSubagentTool_SetContext(t *testing.T) {
-	provider := &MockLLMProvider{}
-	manager := NewSubagentManager(provider, "test-model", "/tmp/test", nil)
-	tool := NewSubagentTool(manager)
-
-	tool.SetContext("test-channel", "test-chat")
-
-	// Verify context is set (we can't directly access private fields,
-	// but we can verify it doesn't crash)
-	// The actual context usage is tested in Execute tests
-}
-
 // TestSubagentTool_Execute_Success tests successful execution
 func TestSubagentTool_Execute_Success(t *testing.T) {
 	provider := &MockLLMProvider{}
 	msgBus := bus.NewMessageBus()
 	manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus)
 	tool := NewSubagentTool(manager)
-	tool.SetContext("telegram", "chat-123")

-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), "telegram", "chat-123")
 	args := map[string]any{
 		"task":  "Write a haiku about coding",
 		"label": "haiku-task",
@@ -297,12 +282,9 @@ func TestSubagentTool_Execute_ContextPassing(t *testing.T) {
 	manager := NewSubagentManager(provider, "test-model", "/tmp/test", msgBus)
 	tool := NewSubagentTool(manager)

-	// Set context
 	channel := "test-channel"
 	chatID := "test-chat"
-	tool.SetContext(channel, chatID)
-
-	ctx := context.Background()
+	ctx := WithToolContext(context.Background(), channel, chatID)
 	args := map[string]any{
 		"task": "Test context passing",
 	}