fix(agent): review fixes for tool call prefix stripping (#259)

- 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
This commit is contained in:
viettranx committed 2026-03-20 09:14:28 +07:00
1 parent ffb536861f
commit 6e94c5fdf8
7 files changed
+88 -53

No files matched your search

+28 -28
View File
@@ -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
+18 -15
View File
@@ -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 {
+1 -1
View File
@@ -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
}
+38
View File
@@ -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)
}
})
}
}
+1 -3
View File
@@ -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",
+1 -3
View File
@@ -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",
+1 -3
View File
@@ -147,9 +147,7 @@
"contextWindow": "上下文窗口",
"contextWindowHint": "模型上下文的令牌限制。",
"maxToolIterations": "最大工具迭代次数",
"maxToolIterationsHint": "每次请求最大工具调用次数。",
"stripAssistantPrefill": "移除 Assistant Prefill",
"stripAssistantPrefillHint": "发送到 LLM 前移除末尾的 assistant 消息(适用于代理提供商)"
"maxToolIterationsHint": "每次请求最大工具调用次数。"
},
"workspace": {
"title": "工作区",