From ffb536861ff88f9743ab85c66288a26f0d09917d Mon Sep 17 00:00:00 2001 From: "HXD.VN" Date: Fri, 20 Mar 2026 09:10:06 +0700 Subject: [PATCH] feat(agent): tool call prefix stripping for proxy providers (#259) * feat(agent): support tool call prefix stripping for proxy providers Proxy providers like LiteLLM and OpenRouter may prepend a prefix to tool call names returned by the model (e.g. "proxy_exec" instead of "exec"). This broke tool policy validation, registry lookup, and hardcoded name checks for "team_tasks" and "spawn" in both serial and parallel execution paths. Add per-agent toolCallPrefix configuration that strips the configured prefix from incoming tool call names before registry resolution. The stripping is applied at resolveToolCallName() which is called before permission checks, registry execution, and spawn/team_tasks detection. - Add StripToolPrefix() supporting literal ("proxy_") and template ("{tool_name}") patterns - Add toolCallPrefix to ToolPolicySpec with backward compat from old "toolPrefix" JSON key - Fix config save using spread operator to prevent dropping new fields - Add UI input in Tool Policy section with i18n (en/vi/zh) * chore: ignore AI tool config directories Add .gemini/, .claude/, .opencode/ to .gitignore to prevent committing user-specific AI tool configurations. --- .gitignore | 3 ++ internal/agent/loop.go | 52 ++++++++++++------- internal/config/config_channels.go | 1 + internal/store/agent_store.go | 27 ++++++++++ internal/tools/policy.go | 33 ++++++++++++ ui/web/src/i18n/locales/en/agents.json | 8 ++- ui/web/src/i18n/locales/vi/agents.json | 8 ++- ui/web/src/i18n/locales/zh/agents.json | 8 ++- .../config-sections/tool-policy-section.tsx | 11 ++++ ui/web/src/types/agent.ts | 1 + 10 files changed, 126 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index 3317a008..09090946 100644 --- a/.gitignore +++ b/.gitignore @@ -46,3 +46,6 @@ k8s-*/* tests/**/creds.json ui/simple-saas *.tar +# AI tool config directories (user-specific) +.gemini/ +.opencode/ diff --git a/internal/agent/loop.go b/internal/agent/loop.go index 86744794..9e39bd1d 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -646,13 +646,14 @@ 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: messages, + Messages: chatMessages, Tools: toolDefs, Model: model, Options: map[string]any{ @@ -845,7 +846,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) // Track team_tasks create for orphan detection (argument-based, pre-execution). // Spawn counting is done post-execution so failed spawns don't get counted. for _, tc := range resp.ToolCalls { - if tc.Name == "team_tasks" { + if l.resolveToolCallName(tc.Name) == "team_tasks" { if action, _ := tc.Arguments["action"].(string); action == "create" { teamTaskCreates++ } @@ -901,29 +902,30 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) argsHash := loopDetector.record(tc.Name, 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) var result *tools.Result - if allowedTools != nil && !allowedTools[tc.Name] { + if allowedTools != nil && !allowedTools[registryName] { // Attempt lazy activation: deferred MCP tools can be activated on first call // so the LLM can call them by name directly without mcp_tool_search. - if l.tools.TryActivateDeferred(tc.Name) { + if l.tools.TryActivateDeferred(registryName) { // Verify tool isn't explicitly denied by policy before allowing. - if l.toolPolicy != nil && l.toolPolicy.IsDenied(tc.Name, l.agentToolPolicy) { - slog.Warn("security.tool_policy_denied_lazy", "agent", l.id, "tool", tc.Name) + if l.toolPolicy != nil && l.toolPolicy.IsDenied(registryName, l.agentToolPolicy) { + slog.Warn("security.tool_policy_denied_lazy", "agent", l.id, "tool", tc.Name, "resolved", registryName) result = tools.ErrorResult("tool not allowed by policy: " + tc.Name) } else { - allowedTools[tc.Name] = true - slog.Info("mcp.tool.lazy_activated", "agent", l.id, "tool", tc.Name) + allowedTools[registryName] = true + slog.Info("mcp.tool.lazy_activated", "agent", l.id, "tool", tc.Name, "resolved", registryName) } } else { - slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name) + slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name, "resolved", registryName) result = tools.ErrorResult("tool not allowed by policy: " + tc.Name) } } if result == nil { - result = l.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) + result = l.tools.ExecuteWithContext(ctx, registryName, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) } stopSlowTimer() @@ -949,7 +951,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } // Count successful spawn calls for orphan detection (post-execution). - if tc.Name == "spawn" && !result.IsError { + if registryName == "spawn" && !result.IsError { if tid, _ := tc.Arguments["team_task_id"].(string); tid != "" { teamTaskSpawns++ } @@ -1073,31 +1075,32 @@ 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), "parallel", true) spanStart := time.Now().UTC() + registryName := l.resolveToolCallName(tc.Name) // Emit running span inside goroutine — goroutine-safe (channel send only). // End is also emitted here to prevent orphans on ctx cancellation. spanID := l.emitToolSpanStart(ctx, spanStart, tc.Name, tc.ID, string(argsJSON)) stopSlowTimer := toolTiming.StartSlowTimer(tc.Name, l.id, req.RunID, slowToolEnabled, emitRun) var result *tools.Result - if allowedTools != nil && !allowedTools[tc.Name] { + if allowedTools != nil && !allowedTools[registryName] { // Attempt lazy activation for deferred MCP tools. // Note: don't write back to allowedTools — concurrent goroutines share // the map and writes would race. TryActivateDeferred is idempotent. - if l.tools.TryActivateDeferred(tc.Name) { + if l.tools.TryActivateDeferred(registryName) { // Verify tool isn't explicitly denied by policy before allowing. - if l.toolPolicy != nil && l.toolPolicy.IsDenied(tc.Name, l.agentToolPolicy) { - slog.Warn("security.tool_policy_denied_lazy", "agent", l.id, "tool", tc.Name) + if l.toolPolicy != nil && l.toolPolicy.IsDenied(registryName, l.agentToolPolicy) { + slog.Warn("security.tool_policy_denied_lazy", "agent", l.id, "tool", tc.Name, "resolved", registryName) result = tools.ErrorResult("tool not allowed by policy: " + tc.Name) } else { - slog.Info("mcp.tool.lazy_activated", "agent", l.id, "tool", tc.Name) + slog.Info("mcp.tool.lazy_activated", "agent", l.id, "tool", tc.Name, "resolved", registryName) } } else { - slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name) + slog.Warn("security.tool_policy_blocked", "agent", l.id, "tool", tc.Name, "resolved", registryName) result = tools.ErrorResult("tool not allowed by policy: " + tc.Name) } } if result == nil { - result = l.tools.ExecuteWithContext(ctx, tc.Name, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) + result = l.tools.ExecuteWithContext(ctx, registryName, tc.Arguments, req.Channel, req.ChatID, req.PeerKind, req.SessionKey, nil) } stopSlowTimer() l.emitToolSpanEnd(ctx, spanID, spanStart, result) @@ -1144,7 +1147,7 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) } // Count successful spawn calls for orphan detection (post-execution). - if r.tc.Name == "spawn" && !r.result.IsError { + if l.resolveToolCallName(r.tc.Name) == "spawn" && !r.result.IsError { if tid, _ := r.tc.Arguments["team_task_id"].(string); tid != "" { teamTaskSpawns++ } @@ -1401,7 +1404,16 @@ func (l *Loop) runLoop(ctx context.Context, req RunRequest) (*RunResult, error) }, nil } -// truncateToolArgs returns a copy of arguments with string values truncated to maxLen. +// resolveToolCallName strips the configured tool call prefix from a name +// returned by the model, returning the original registry name. +// Example: prefix "proxy_" + model calls "proxy_exec" → returns "exec". +func (l *Loop) resolveToolCallName(name string) string { + if l.agentToolPolicy != nil && l.agentToolPolicy.ToolCallPrefix != "" { + return tools.StripToolPrefix(l.agentToolPolicy.ToolCallPrefix, name) + } + return name +} + func truncateToolArgs(args map[string]any, maxLen int) map[string]any { out := make(map[string]any, len(args)) for k, v := range args { diff --git a/internal/config/config_channels.go b/internal/config/config_channels.go index fc3f1b5f..206f4812 100644 --- a/internal/config/config_channels.go +++ b/internal/config/config_channels.go @@ -406,6 +406,7 @@ type ToolPolicySpec struct { Deny []string `json:"deny,omitempty"` AlsoAllow []string `json:"alsoAllow,omitempty"` ByProvider map[string]*ToolPolicySpec `json:"byProvider,omitempty"` + ToolCallPrefix string `json:"toolCallPrefix,omitempty"` // prefix to strip from model's tool call names before registry lookup } type WebToolsConfig struct { diff --git a/internal/store/agent_store.go b/internal/store/agent_store.go index da1576a9..c695534e 100644 --- a/internal/store/agent_store.go +++ b/internal/store/agent_store.go @@ -61,6 +61,18 @@ func (a *AgentData) ParseToolsConfig() *config.ToolPolicySpec { if json.Unmarshal(a.ToolsConfig, &c) != nil { return nil } + // Backward compat: migrate old "toolPrefix" key to "toolCallPrefix" + if c.ToolCallPrefix == "" { + var raw map[string]json.RawMessage + if json.Unmarshal(a.ToolsConfig, &raw) == nil { + if v, ok := raw["toolPrefix"]; ok { + var s string + if json.Unmarshal(v, &s) == nil && s != "" { + c.ToolCallPrefix = s + } + } + } + } return &c } @@ -154,6 +166,21 @@ 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 6b276754..1c72b25e 100644 --- a/internal/tools/policy.go +++ b/internal/tools/policy.go @@ -398,6 +398,39 @@ func matchDenySpec(name string, spec []string) bool { return false } +// StripToolPrefix removes a prefix pattern from a tool name returned by the LLM. +// The template uses {tool_name} as placeholder. Example: template "proxy_{tool_name}" +// strips "proxy_" from "proxy_exec" → "exec". +// If template has no {tool_name}, it's treated as a literal prefix to strip. +func StripToolPrefix(tmpl, name string) string { + const placeholder = "{tool_name}" + if strings.Contains(tmpl, placeholder) { + parts := strings.SplitN(tmpl, placeholder, 2) + prefix, suffix := parts[0], parts[1] + if strings.HasPrefix(name, prefix) && strings.HasSuffix(name, suffix) { + result := name[len(prefix):] + if suffix != "" { + result = result[:len(result)-len(suffix)] + } + if result != "" { + return result + } + } + return name + } + // Plain prefix: strip literal prefix and any leading underscore separator + stripped := strings.TrimPrefix(name, tmpl) + if stripped == name { + return name // prefix didn't match + } + stripped = strings.TrimLeft(stripped, "_") + if stripped == "" { + return name // nothing left after stripping + } + slog.Debug("tool_prefix.stripped", "from", name, "to", stripped, "template", tmpl) + return stripped +} + func resolveAlias(name string) string { if canonical, ok := legacyToolAliases[name]; ok { return canonical diff --git a/ui/web/src/i18n/locales/en/agents.json b/ui/web/src/i18n/locales/en/agents.json index 7fa0fe5e..1be37366 100644 --- a/ui/web/src/i18n/locales/en/agents.json +++ b/ui/web/src/i18n/locales/en/agents.json @@ -147,7 +147,9 @@ "contextWindow": "Context Window", "contextWindowHint": "Token limit for the model context.", "maxToolIterations": "Max Tool Iterations", - "maxToolIterationsHint": "Max tool calls per request." + "maxToolIterationsHint": "Max tool calls per request.", + "stripAssistantPrefill": "Strip Assistant Prefill", + "stripAssistantPrefillHint": "Remove trailing assistant messages before LLM call (for proxy providers)" }, "workspace": { "title": "Workspace", @@ -319,7 +321,9 @@ "alsoAllow": "Also Allow", "selectToolsAllow": "Select tools to allow...", "selectToolsDeny": "Select tools to deny...", - "selectToolsAlsoAllow": "Select additional tools..." + "selectToolsAlsoAllow": "Select additional tools...", + "toolCallPrefix": "Tool Call Prefix", + "toolCallPrefixHint": "Strips this prefix from model's tool call names before registry lookup." }, "workspaceSharing": { "title": "Workspace Sharing", diff --git a/ui/web/src/i18n/locales/vi/agents.json b/ui/web/src/i18n/locales/vi/agents.json index 0e076195..cb6613b8 100644 --- a/ui/web/src/i18n/locales/vi/agents.json +++ b/ui/web/src/i18n/locales/vi/agents.json @@ -147,7 +147,9 @@ "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." + "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)" }, "workspace": { "title": "Workspace", @@ -319,7 +321,9 @@ "alsoAllow": "Cũng cho phép", "selectToolsAllow": "Chọn công cụ để cho phép...", "selectToolsDeny": "Chọn công cụ để từ chối...", - "selectToolsAlsoAllow": "Chọn công cụ bổ sung..." + "selectToolsAlsoAllow": "Chọn công cụ bổ sung...", + "toolCallPrefix": "Tiền tố Tool Call", + "toolCallPrefixHint": "Loại bỏ tiền tố này từ tên tool call của model trước khi tra cứu registry." }, "workspaceSharing": { "title": "Chia sẻ Workspace", diff --git a/ui/web/src/i18n/locales/zh/agents.json b/ui/web/src/i18n/locales/zh/agents.json index 315c4fdf..4d4ed0dc 100644 --- a/ui/web/src/i18n/locales/zh/agents.json +++ b/ui/web/src/i18n/locales/zh/agents.json @@ -147,7 +147,9 @@ "contextWindow": "上下文窗口", "contextWindowHint": "模型上下文的令牌限制。", "maxToolIterations": "最大工具迭代次数", - "maxToolIterationsHint": "每次请求最大工具调用次数。" + "maxToolIterationsHint": "每次请求最大工具调用次数。", + "stripAssistantPrefill": "移除 Assistant Prefill", + "stripAssistantPrefillHint": "发送到 LLM 前移除末尾的 assistant 消息（适用于代理提供商）" }, "workspace": { "title": "工作区", @@ -319,7 +321,9 @@ "alsoAllow": "同时允许", "selectToolsAllow": "选择要允许的工具...", "selectToolsDeny": "选择要拒绝的工具...", - "selectToolsAlsoAllow": "选择附加工具..." + "selectToolsAlsoAllow": "选择附加工具...", + "toolCallPrefix": "工具调用前缀", + "toolCallPrefixHint": "从模型的工具调用名称中去除此前缀后再查找注册表。" }, "workspaceSharing": { "title": "工作区共享", diff --git a/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx b/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx index 89845999..4c825cff 100644 --- a/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx +++ b/ui/web/src/pages/agents/agent-detail/config-sections/tool-policy-section.tsx @@ -6,6 +6,7 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { Input } from "@/components/ui/input"; import type { ToolPolicyConfig } from "@/types/agent"; import { ConfigSection, InfoLabel } from "./config-section"; import { ToolNameSelect } from "@/components/shared/tool-name-select"; @@ -42,6 +43,16 @@ export function ToolPolicySection({ enabled, value, onToggle, onChange }: ToolPo +
+ {t(`${s}.toolCallPrefix`)} + onChange({ ...value, toolCallPrefix: e.target.value.replace(/[^a-z0-9_{}/]/g, "") || undefined })} + placeholder="e.g. proxy_" + className="font-mono text-sm" + /> +

{t(`${s}.toolCallPrefixHint`)}

+
{t(`${s}.allow`)} ; + toolCallPrefix?: string; // prefix to strip from model's tool call names } export interface SubagentsConfig {