From b900af334f47632284864ca8de32c2d1b129dbfa Mon Sep 17 00:00:00 2001 From: Boris Bliznioukov Date: Wed, 4 Mar 2026 13:28:41 +0100 Subject: [PATCH] fix: reviewer comments improve context handling for tool execution and ensure defaults for non-conversation callers Signed-off-by: Boris Bliznioukov --- pkg/agent/loop.go | 4 ++-- pkg/tools/base.go | 19 +++++++++++-------- pkg/tools/registry.go | 6 ++---- pkg/tools/registry_test.go | 3 ++- pkg/tools/spawn.go | 10 +++++++++- pkg/tools/subagent.go | 13 ++++++++++++- 6 files changed, 38 insertions(+), 17 deletions(-) diff --git a/pkg/agent/loop.go b/pkg/agent/loop.go index f952ca082..ff107673b 100644 --- a/pkg/agent/loop.go +++ b/pkg/agent/loop.go @@ -480,8 +480,8 @@ func (al *AgentLoop) processMessage(ctx context.Context, msg bus.InboundMessage) // Reset message-tool state for this round so we don't skip publishing due to a previous round. if tool, ok := agent.Tools.Get("message"); ok { - if mt, ok := tool.(*tools.MessageTool); ok { - mt.ResetSentInRound() + if resetter, ok := tool.(interface{ ResetSentInRound() }); ok { + resetter.ResetSentInRound() } } diff --git a/pkg/tools/base.go b/pkg/tools/base.go index 19afa192e..ec743e164 100644 --- a/pkg/tools/base.go +++ b/pkg/tools/base.go @@ -14,30 +14,33 @@ type Tool interface { // // 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 ToolContextKey string +type toolCtxKey struct{ name string } -const ( - CtxKeyChannel ToolContextKey = "channel" - CtxKeyChatID ToolContextKey = "chatID" +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) + 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) + 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) + v, _ := ctx.Value(ctxKeyChatID).(string) return v } diff --git a/pkg/tools/registry.go b/pkg/tools/registry.go index 25e94da2b..ca8436c67 100644 --- a/pkg/tools/registry.go +++ b/pkg/tools/registry.go @@ -71,10 +71,8 @@ func (r *ToolRegistry) ExecuteWithContext( } // Inject channel/chatID into ctx so tools read them via ToolChannel(ctx)/ToolChatID(ctx). - // Immutable per-call — no shared mutable state on tool instances. - if channel != "" && chatID != "" { - ctx = WithToolContext(ctx, channel, chatID) - } + // Always inject — tools validate what they require. + ctx = WithToolContext(ctx, channel, chatID) // If tool implements AsyncExecutor and callback is provided, use ExecuteAsync. // The callback is a call parameter, not mutable state on the tool instance. diff --git a/pkg/tools/registry_test.go b/pkg/tools/registry_test.go index b1cabc34f..92d7d5abd 100644 --- a/pkg/tools/registry_test.go +++ b/pkg/tools/registry_test.go @@ -156,7 +156,7 @@ func TestToolRegistry_ExecuteWithContext_InjectsToolContext(t *testing.T) { } } -func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { +func TestToolRegistry_ExecuteWithContext_EmptyContext(t *testing.T) { r := NewToolRegistry() ct := &mockContextAwareTool{ mockRegistryTool: *newMockTool("ctx_tool", "needs context"), @@ -168,6 +168,7 @@ func TestToolRegistry_ExecuteWithContext_SkipsEmptyContext(t *testing.T) { 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) } diff --git a/pkg/tools/spawn.go b/pkg/tools/spawn.go index 18ea8d185..be40ffda2 100644 --- a/pkg/tools/spawn.go +++ b/pkg/tools/spawn.go @@ -83,9 +83,17 @@ func (t *SpawnTool) execute(ctx context.Context, args map[string]any, cb AsyncCa return ErrorResult("Subagent manager not configured") } - // Read channel/chatID from context (injected by registry) + // 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, channel, chatID, cb) diff --git a/pkg/tools/subagent.go b/pkg/tools/subagent.go index ea011e7d5..429340047 100644 --- a/pkg/tools/subagent.go +++ b/pkg/tools/subagent.go @@ -332,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, ToolChannel(ctx), ToolChatID(ctx)) + }, messages, channel, chatID) if err != nil { return ErrorResult(fmt.Sprintf("Subagent execution failed: %v", err)).WithError(err) }