From 6e94c5fdf82aadc9ece89b0b90c253f0b4e1273a Mon Sep 17 00:00:00 2001 From: viettranx Date: Fri, 20 Mar 2026 09:14:28 +0700 Subject: [PATCH] fix(agent): review fixes for tool call prefix stripping (#259) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix TrimLeft → TrimPrefix to strip only one underscore separator - Add registryName to indexedResult for parallel path cache - Use registryName for bootstrapToolAllowlist and loopDetector checks - Add server-side sanitization for ToolCallPrefix input - Remove dead chatMessages alias and unrelated ParseStripAssistantPrefill - Remove orphaned stripAssistantPrefill i18n keys - Add 13 unit tests for StripToolPrefix --- internal/agent/loop.go | 56 +++++++++++++------------- internal/store/agent_store.go | 33 ++++++++------- internal/tools/policy.go | 2 +- internal/tools/policy_prefix_test.go | 38 +++++++++++++++++ ui/web/src/i18n/locales/en/agents.json | 4 +- ui/web/src/i18n/locales/vi/agents.json | 4 +- ui/web/src/i18n/locales/zh/agents.json | 4 +- 7 files changed, 88 insertions(+), 53 deletions(-) create mode 100644 internal/tools/policy_prefix_test.go diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 9e39bd1d..0445b5c4 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -646,14 +646,13 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } // Use per-request model override if set (e.g. heartbeat uses cheaper model). - chatMessages := messages model := l.model if req.ModelOverride != "" { model = req.ModelOverride } chatReq := providers.ChatRequest{ - Messages: chatMessages, + Messages: messages, Tools: toolDefs, Model: model, Options: map[string]any{ @@ -899,10 +898,10 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) argsJSON, _ := json.Marshal(tc.Arguments) slog.Info("tool call", "agent", l.id, "tool", tc.Name, "args_len", len(argsJSON)) - argsHash := loopDetector.record(tc.Name, tc.Arguments) + registryName := l.resolveToolCallName(tc.Name) + argsHash := loopDetector.record(registryName, tc.Arguments) toolSpanStart := time.Now().UTC() - registryName := l.resolveToolCallName(tc.Name) toolSpanID := l.emitToolSpanStart(ctx, toolSpanStart, tc.Name, tc.ID, string(argsJSON)) stopSlowTimer := toolTiming.StartSlowTimer(tc.Name, l.id, req.RunID, slowToolEnabled, emitRun) @@ -936,7 +935,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Record result for loop detection. loopDetector.recordResult(argsHash, result.ForLLM) - loopDetector.recordMutation(tc.Name) + loopDetector.recordMutation(registryName) if result.Async { asyncToolCalls = append(asyncToolCalls, tc.Name) @@ -956,7 +955,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) teamTaskSpawns++ } } - if hadBootstrap && bootstrapToolAllowlist[tc.Name] { + if hadBootstrap && bootstrapToolAllowlist[registryName] { bootstrapWriteDetected = true } @@ -1006,14 +1005,14 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) pendingMsgs = append(pendingMsgs, toolMsg) // Check for tool call loop after recording result. - if level, msg := loopDetector.detect(tc.Name, argsHash); level != "" { + if level, msg := loopDetector.detect(registryName, argsHash); level != "" { if level == "critical" { - slog.Warn("tool loop critical", "agent", l.id, "tool", tc.Name, "message", msg) - finalContent = "I was unable to complete this task — I got stuck repeatedly calling " + tc.Name + " without making progress. Please try rephrasing your request." + slog.Warn("tool loop critical", "agent", l.id, "tool", registryName, "message", msg) + finalContent = "I was unable to complete this task — I got stuck repeatedly calling " + registryName + " without making progress. Please try rephrasing your request." break } // Warning: inject message so model knows to change strategy. - slog.Warn("tool loop warning", "agent", l.id, "tool", tc.Name, "message", msg) + slog.Warn("tool loop warning", "agent", l.id, "tool", registryName, "message", msg) messages = append(messages, providers.Message{Role: "user", Content: msg}) } @@ -1032,10 +1031,10 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Check for same tool returning identical results with different args. if rh := hashResult(result.ForLLM); rh != "" { - if level, msg := loopDetector.detectSameResult(tc.Name, rh); level != "" { + if level, msg := loopDetector.detectSameResult(registryName, rh); level != "" { if level == "critical" { slog.Warn("tool loop critical: same result", - "tool", tc.Name, "agent", l.id, "run", req.RunID) + "tool", registryName, "agent", l.id, "run", req.RunID) finalContent = msg break } @@ -1047,11 +1046,12 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Tool instances are immutable (context-based) so concurrent access is safe. // Results are collected then processed sequentially for deterministic ordering. type indexedResult struct { - idx int - tc providers.ToolCall - result *tools.Result - argsJSON string - spanStart time.Time + idx int + tc providers.ToolCall + registryName string + result *tools.Result + argsJSON string + spanStart time.Time } // 1. Emit all tool.call events upfront (client sees all calls starting) @@ -1104,7 +1104,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } stopSlowTimer() l.emitToolSpanEnd(ctx, spanID, spanStart, result) - resultCh <- indexedResult{idx: idx, tc: tc, result: result, argsJSON: string(argsJSON), spanStart: spanStart} + resultCh <- indexedResult{idx: idx, tc: tc, registryName: registryName, result: result, argsJSON: string(argsJSON), spanStart: spanStart} }(i, tc) } @@ -1130,9 +1130,9 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) toolTiming.Record(r.tc.Name, time.Since(r.spanStart).Milliseconds()) // Record for loop detection. - argsHash := loopDetector.record(r.tc.Name, r.tc.Arguments) + argsHash := loopDetector.record(r.registryName, r.tc.Arguments) loopDetector.recordResult(argsHash, r.result.ForLLM) - loopDetector.recordMutation(r.tc.Name) + loopDetector.recordMutation(r.registryName) if r.result.Async { asyncToolCalls = append(asyncToolCalls, r.tc.Name) @@ -1147,12 +1147,12 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } // Count successful spawn calls for orphan detection (post-execution). - if l.resolveToolCallName(r.tc.Name) == "spawn" && !r.result.IsError { + if r.registryName == "spawn" && !r.result.IsError { if tid, _ := r.tc.Arguments["team_task_id"].(string); tid != "" { teamTaskSpawns++ } } - if hadBootstrap && bootstrapToolAllowlist[r.tc.Name] { + if hadBootstrap && bootstrapToolAllowlist[r.registryName] { bootstrapWriteDetected = true } @@ -1202,23 +1202,23 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) pendingMsgs = append(pendingMsgs, toolMsg) // Check for tool call loop. - if level, msg := loopDetector.detect(r.tc.Name, argsHash); level != "" { + if level, msg := loopDetector.detect(r.registryName, argsHash); level != "" { if level == "critical" { - slog.Warn("tool loop critical", "agent", l.id, "tool", r.tc.Name, "message", msg) - finalContent = "I was unable to complete this task — I got stuck repeatedly calling " + r.tc.Name + " without making progress. Please try rephrasing your request." + slog.Warn("tool loop critical", "agent", l.id, "tool", r.registryName, "message", msg) + finalContent = "I was unable to complete this task — I got stuck repeatedly calling " + r.registryName + " without making progress. Please try rephrasing your request." loopStuck = true break } - slog.Warn("tool loop warning", "agent", l.id, "tool", r.tc.Name, "message", msg) + slog.Warn("tool loop warning", "agent", l.id, "tool", r.registryName, "message", msg) messages = append(messages, providers.Message{Role: "user", Content: msg}) } // Check for same tool returning identical results with different args. if rh := hashResult(r.result.ForLLM); rh != "" { - if level, msg := loopDetector.detectSameResult(r.tc.Name, rh); level != "" { + if level, msg := loopDetector.detectSameResult(r.registryName, rh); level != "" { if level == "critical" { slog.Warn("tool loop critical: same result", - "tool", r.tc.Name, "agent", l.id, "run", req.RunID) + "tool", r.registryName, "agent", l.id, "run", req.RunID) finalContent = msg loopStuck = true break diff --git a/internal/store/agent_store.go b/internal/store/agent_store.go index c695534e..735d73a2 100644 --- a/internal/store/agent_store.go +++ b/internal/store/agent_store.go @@ -3,11 +3,27 @@ package store import ( "context" "encoding/json" + "strings" "github.com/google/uuid" "github.com/nextlevelbuilder/goclaw/internal/config" ) +// sanitizeToolCallPrefix strips characters not in [a-z0-9_{}] from the prefix. +// This matches the UI-side regex and prevents injection via direct API calls. +func sanitizeToolCallPrefix(s string) string { + if s == "" { + return "" + } + var b strings.Builder + for _, r := range s { + if (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9') || r == '_' || r == '{' || r == '}' { + b.WriteRune(r) + } + } + return b.String() +} + // Agent type constants. const ( AgentTypeOpen = "open" // per-user context files, seeded on first chat @@ -73,6 +89,8 @@ func (a *AgentData) ParseToolsConfig() *config.ToolPolicySpec { } } } + // Sanitize: only allow [a-z0-9_{}] to prevent injection via API bypass. + c.ToolCallPrefix = sanitizeToolCallPrefix(c.ToolCallPrefix) return &c } @@ -166,21 +184,6 @@ func (a *AgentData) ParseMaxTokens() int { return cfg.MaxTokens } -// ParseStripAssistantPrefill extracts strip_assistant_prefill from other_config JSONB. -// When true, trailing assistant messages are removed before sending to LLM (for proxy providers). -func (a *AgentData) ParseStripAssistantPrefill() bool { - if len(a.OtherConfig) == 0 { - return false - } - var cfg struct { - StripAssistantPrefill bool `json:"strip_assistant_prefill"` - } - if json.Unmarshal(a.OtherConfig, &cfg) != nil { - return false - } - return cfg.StripAssistantPrefill -} - // ParseSelfEvolve extracts self_evolve from other_config JSONB. // When true, predefined agents can update their SOUL.md (style/tone) through chat. func (a *AgentData) ParseSelfEvolve() bool { diff --git a/internal/tools/policy.go b/internal/tools/policy.go index 1c72b25e..31241898 100644 --- a/internal/tools/policy.go +++ b/internal/tools/policy.go @@ -423,7 +423,7 @@ func StripToolPrefix(tmpl, name string) string { if stripped == name { return name // prefix didn't match } - stripped = strings.TrimLeft(stripped, "_") + stripped = strings.TrimPrefix(stripped, "_") if stripped == "" { return name // nothing left after stripping } diff --git a/internal/tools/policy_prefix_test.go b/internal/tools/policy_prefix_test.go new file mode 100644 index 00000000..a07f7972 --- /dev/null +++ b/internal/tools/policy_prefix_test.go @@ -0,0 +1,38 @@ +package tools + +import "testing" + +func TestStripToolPrefix(t *testing.T) { + tests := []struct { + name string + tmpl string + input string + expect string + }{ + // Template-based ({tool_name} placeholder) + {"template match", "proxy_{tool_name}", "proxy_exec", "exec"}, + {"template with suffix", "{tool_name}_v2", "exec_v2", "exec"}, + {"template prefix+suffix", "pre_{tool_name}_suf", "pre_exec_suf", "exec"}, + {"template no match", "proxy_{tool_name}", "other_exec", "other_exec"}, + {"template empty result", "proxy_{tool_name}", "proxy_", "proxy_"}, + {"template exact placeholder", "{tool_name}", "exec", "exec"}, + + // Literal prefix + {"literal prefix with separator", "proxy_", "proxy_exec", "exec"}, + {"literal prefix no match", "proxy_", "other_exec", "other_exec"}, + {"literal prefix underscore join", "proxy", "proxy_exec", "exec"}, + {"literal prefix single underscore strip", "proxy", "proxy__exec", "_exec"}, + {"literal prefix empty after strip", "proxy_", "proxy_", "proxy_"}, + {"literal name shorter than prefix", "proxy_", "pro", "pro"}, + {"literal name equals prefix", "proxy", "proxy", "proxy"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got := StripToolPrefix(tt.tmpl, tt.input) + if got != tt.expect { + t.Errorf("StripToolPrefix(%q, %q) = %q, want %q", tt.tmpl, tt.input, got, tt.expect) + } + }) + } +} diff --git a/ui/web/src/i18n/locales/en/agents.json b/ui/web/src/i18n/locales/en/agents.json index 1be37366..a3773c33 100644 --- a/ui/web/src/i18n/locales/en/agents.json +++ b/ui/web/src/i18n/locales/en/agents.json @@ -147,9 +147,7 @@ "contextWindow": "Context Window", "contextWindowHint": "Token limit for the model context.", "maxToolIterations": "Max Tool Iterations", - "maxToolIterationsHint": "Max tool calls per request.", - "stripAssistantPrefill": "Strip Assistant Prefill", - "stripAssistantPrefillHint": "Remove trailing assistant messages before LLM call (for proxy providers)" + "maxToolIterationsHint": "Max tool calls per request." }, "workspace": { "title": "Workspace", diff --git a/ui/web/src/i18n/locales/vi/agents.json b/ui/web/src/i18n/locales/vi/agents.json index cb6613b8..30328069 100644 --- a/ui/web/src/i18n/locales/vi/agents.json +++ b/ui/web/src/i18n/locales/vi/agents.json @@ -147,9 +147,7 @@ "contextWindow": "Cửa sổ ngữ cảnh", "contextWindowHint": "Giới hạn token cho ngữ cảnh model.", "maxToolIterations": "Số lần lặp công cụ tối đa", - "maxToolIterationsHint": "Số lần gọi công cụ tối đa mỗi yêu cầu.", - "stripAssistantPrefill": "Loại bỏ Assistant Prefill", - "stripAssistantPrefillHint": "Loại bỏ tin nhắn assistant ở cuối trước khi gửi đến LLM (dùng cho proxy provider)" + "maxToolIterationsHint": "Số lần gọi công cụ tối đa mỗi yêu cầu." }, "workspace": { "title": "Workspace", diff --git a/ui/web/src/i18n/locales/zh/agents.json b/ui/web/src/i18n/locales/zh/agents.json index 4d4ed0dc..75900fa5 100644 --- a/ui/web/src/i18n/locales/zh/agents.json +++ b/ui/web/src/i18n/locales/zh/agents.json @@ -147,9 +147,7 @@ "contextWindow": "上下文窗口", "contextWindowHint": "模型上下文的令牌限制。", "maxToolIterations": "最大工具迭代次数", - "maxToolIterationsHint": "每次请求最大工具调用次数。", - "stripAssistantPrefill": "移除 Assistant Prefill", - "stripAssistantPrefillHint": "发送到 LLM 前移除末尾的 assistant 消息（适用于代理提供商）" + "maxToolIterationsHint": "每次请求最大工具调用次数。" }, "workspace": { "title": "工作区",