fix(tools): harden tool_call dispatch, filesystem feedback, and shell deny

pkg/tools/call.go
- normalizeToolName() recovers from malformed tool_name values the LLM
  occasionally emits: boundary-safe regex match, substring fallback, and
  comma/space token split; logs a warning when normalization fires
- Pre-flight required-arg check before dispatch; returns schemaHintError
  on missing required parameters
- schemaHintError() includes the tool's full properties schema and required
  list in the error message so the LLM can self-correct without another
  tool_search round-trip
- tool_call now resolves the Tool object before argument parsing so
  not-found errors are returned early with a tool_search suggestion
- Description updated: tool_call is a fallback; prefer direct calls after
  tool_search promotes tools to native callables
- Removed dead inner if-found guard around ResourceProvider loading

pkg/tools/call_test.go
- TestToolCallTool_ToolNotFoundSuggestsSearch: error message must mention
  tool_search
- TestToolCallTool_MissingRequiredArgs_IncludesSchemaHint: error must
  contain "missing required" and the parameter name
- TestToolCallTool_InvalidJSONArguments: updated to use a real registered
  tool so the schema hint path is exercised
- TestToolCallTool_NormalizesCommaSeparatedToolName: "write_file, read_file"
  normalizes to write_file (first valid token)
- TestToolCallTool_NormalizesEmbeddedToolName: long noisy string containing
  "exec" normalizes to exec

pkg/tools/filesystem.go
- write_file now returns a content preview (first 80 chars) in the silent
  result message so the LLM can confirm the write without a read_file

pkg/tools/shell.go
- Add `hostname` to the exec deny-list to block host identity exfiltration
This commit is contained in:
ZanzyTHEbar 2026-02-21 00:40:44 +00:00
parent 3d940883a6
commit d8c5592388
4 changed files with 233 additions and 27 deletions

View file

@ -3,6 +3,8 @@ package tools
import ( import (
"context" "context"
"fmt" "fmt"
"regexp"
"strings"
jsonv2 "github.com/go-json-experiment/json" jsonv2 "github.com/go-json-experiment/json"
@ -39,7 +41,7 @@ func NewToolCallTool(registry *ToolRegistry) *ToolCallTool {
func (t *ToolCallTool) Name() string { return "tool_call" } func (t *ToolCallTool) Name() string { return "tool_call" }
func (t *ToolCallTool) Description() string { func (t *ToolCallTool) Description() string {
return "Execute any registered tool by name. Use tool_search first to discover available tools and their parameters, then call them here." return "Fallback tool executor: dispatch any registered tool by name. Prefer calling discovered tools directly — after tool_search, tools become native callables. Use tool_call only if a tool is not yet directly available."
} }
func (t *ToolCallTool) Parameters() map[string]interface{} { func (t *ToolCallTool) Parameters() map[string]interface{} {
@ -69,12 +71,18 @@ func (t *ToolCallTool) Execute(ctx context.Context, args map[string]interface{})
if toolName == "" { if toolName == "" {
return ErrorResult("tool_name is required") return ErrorResult("tool_name is required")
} }
toolName = t.normalizeToolName(toolName)
// Prevent recursive calls to meta-tools // Prevent recursive calls to meta-tools
if toolName == "tool_call" || toolName == "tool_search" { if toolName == "tool_call" || toolName == "tool_search" {
return ErrorResult(fmt.Sprintf("cannot recursively call meta-tool %q", toolName)) return ErrorResult(fmt.Sprintf("cannot recursively call meta-tool %q", toolName))
} }
tool, found := t.registry.Get(toolName)
if !found {
return ErrorResult(fmt.Sprintf("tool %q not found — use tool_search to discover available tools", toolName))
}
// Extract arguments — handle both direct object and JSON string // Extract arguments — handle both direct object and JSON string
var toolArgs map[string]interface{} var toolArgs map[string]interface{}
switch v := args["arguments"].(type) { switch v := args["arguments"].(type) {
@ -86,20 +94,26 @@ func (t *ToolCallTool) Execute(ctx context.Context, args map[string]interface{})
return ErrorResult(fmt.Sprintf("arguments JSON too large: %d bytes (max %d)", len(v), maxArgsJSON)) return ErrorResult(fmt.Sprintf("arguments JSON too large: %d bytes (max %d)", len(v), maxArgsJSON))
} }
if err := jsonv2.Unmarshal([]byte(v), &toolArgs); err != nil { if err := jsonv2.Unmarshal([]byte(v), &toolArgs); err != nil {
return ErrorResult(fmt.Sprintf("invalid arguments JSON: %v", err)) return t.schemaHintError(tool, fmt.Sprintf("invalid arguments JSON: %v", err))
} }
case nil: case nil:
toolArgs = map[string]interface{}{} toolArgs = map[string]interface{}{}
default: default:
return ErrorResult(fmt.Sprintf("arguments must be a JSON object, got %T", v)) return t.schemaHintError(tool, fmt.Sprintf("arguments must be a JSON object, got %T", v))
} }
if len(toolArgs) > 50 { if len(toolArgs) > 50 {
return ErrorResult(fmt.Sprintf("too many arguments: %d (max 50)", len(toolArgs))) return ErrorResult(fmt.Sprintf("too many arguments: %d (max 50)", len(toolArgs)))
} }
// Pre-flight: check required parameters before dispatch
if schema := tool.Parameters(); schema != nil {
if missing := t.checkRequired(schema, toolArgs); len(missing) > 0 {
return t.schemaHintError(tool, fmt.Sprintf("missing required arguments: %s", strings.Join(missing, ", ")))
}
}
// If the target tool declares resources, load them before dispatch. // If the target tool declares resources, load them before dispatch.
if tool, found := t.registry.Get(toolName); found {
if rp, ok := tool.(ResourceProvider); ok { if rp, ok := tool.(ResourceProvider); ok {
resources, err := rp.LoadResources(ctx) resources, err := rp.LoadResources(ctx)
if err != nil { if err != nil {
@ -109,7 +123,122 @@ func (t *ToolCallTool) Execute(ctx context.Context, args map[string]interface{})
ctx = context.WithValue(ctx, ctxKeyResources{}, resources) ctx = context.WithValue(ctx, ctxKeyResources{}, resources)
} }
} }
}
return t.registry.ExecuteWithContext(ctx, toolName, toolArgs, t.channel, t.chatID, nil) return t.registry.ExecuteWithContext(ctx, toolName, toolArgs, t.channel, t.chatID, nil)
} }
// schemaHintError returns an error result that includes the tool's expected
// parameter schema, giving the LLM a clear correction path.
func (t *ToolCallTool) schemaHintError(tool Tool, msg string) *ToolResult {
schema := tool.Parameters()
hint := fmt.Sprintf("%s\n\nExpected schema for %q:\n", msg, tool.Name())
if props, ok := schema["properties"].(map[string]interface{}); ok {
schemaJSON, err := jsonv2.Marshal(props)
if err == nil {
hint += string(schemaJSON)
}
}
if req := t.extractRequired(schema); len(req) > 0 {
hint += fmt.Sprintf("\nRequired: %s", strings.Join(req, ", "))
}
hint += "\n\nNote: discovered tools are directly callable — you do not need tool_call for tools returned by tool_search."
return ErrorResult(hint)
}
// checkRequired returns any required parameters that are missing from the args.
func (t *ToolCallTool) checkRequired(schema, args map[string]interface{}) []string {
required := t.extractRequired(schema)
var missing []string
for _, key := range required {
if _, ok := args[key]; !ok {
missing = append(missing, key)
}
}
return missing
}
func (t *ToolCallTool) extractRequired(schema map[string]interface{}) []string {
switch r := schema["required"].(type) {
case []string:
return r
case []interface{}:
out := make([]string, 0, len(r))
for _, v := range r {
if s, ok := v.(string); ok {
out = append(out, s)
}
}
return out
}
return nil
}
func (t *ToolCallTool) normalizeToolName(raw string) string {
name := strings.TrimSpace(raw)
if name == "" {
return name
}
if _, found := t.registry.Get(name); found {
return name
}
known := t.registry.List()
normalized := strings.ToLower(name)
// First pass: boundary-safe contains (prefers full tool names embedded in noisy strings).
best := ""
bestPos := -1
for _, candidate := range known {
pat := fmt.Sprintf(`(^|[^a-z0-9_])%s([^a-z0-9_]|$)`, regexp.QuoteMeta(strings.ToLower(candidate)))
re, err := regexp.Compile(pat)
if err != nil {
continue
}
loc := re.FindStringIndex(normalized)
if loc == nil {
continue
}
if bestPos == -1 || loc[0] < bestPos {
best = candidate
bestPos = loc[0]
}
}
if best != "" {
logger.WarnCF("tool_call", "Normalized malformed tool_name",
map[string]interface{}{"raw": raw, "normalized": best})
return best
}
// Third pass: raw substring fallback for heavily malformed names
// like "exec_tool_search_query_exec_run_shell_command...".
best = ""
bestPos = -1
for _, candidate := range known {
pos := strings.Index(normalized, strings.ToLower(candidate))
if pos == -1 {
continue
}
if bestPos == -1 || pos < bestPos {
best = candidate
bestPos = pos
}
}
if best != "" {
logger.WarnCF("tool_call", "Normalized tool_name by substring fallback",
map[string]interface{}{"raw": raw, "normalized": best})
return best
}
// Second pass: comma/space separated fragments, choose first valid tool token.
replacer := strings.NewReplacer(",", " ", ";", " ", "|", " ")
for _, token := range strings.Fields(replacer.Replace(name)) {
if _, found := t.registry.Get(token); found {
logger.WarnCF("tool_call", "Normalized split tool_name token",
map[string]interface{}{"raw": raw, "normalized": token})
return token
}
}
return name
}

View file

@ -2,6 +2,7 @@ package tools
import ( import (
"context" "context"
"strings"
"testing" "testing"
) )
@ -66,19 +67,6 @@ func TestToolCallTool_DispatchesToTool(t *testing.T) {
} }
} }
func TestToolCallTool_ToolNotFound(t *testing.T) {
r := NewToolRegistry()
tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "nonexistent",
})
if !result.IsError {
t.Error("expected error for nonexistent tool")
}
}
func TestToolCallTool_PreventRecursion_ToolCall(t *testing.T) { func TestToolCallTool_PreventRecursion_ToolCall(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.RegisterMetaTools() r.RegisterMetaTools()
@ -141,16 +129,58 @@ func TestToolCallTool_NilArguments(t *testing.T) {
func TestToolCallTool_InvalidJSONArguments(t *testing.T) { func TestToolCallTool_InvalidJSONArguments(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(&stubToolWithSchema{name: "read_file", desc: "Read a file"})
tc := NewToolCallTool(r) tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{ result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "anything", "tool_name": "read_file",
"arguments": "not-json", "arguments": "not-json",
}) })
if !result.IsError { if !result.IsError {
t.Error("expected error for invalid JSON arguments") t.Error("expected error for invalid JSON arguments")
} }
// Should include schema hint
if !strings.Contains(result.ForLLM, "path") {
t.Errorf("expected schema hint with 'path' parameter, got: %s", result.ForLLM)
}
}
func TestToolCallTool_MissingRequiredArgs_IncludesSchemaHint(t *testing.T) {
r := NewToolRegistry()
r.Register(&stubToolWithSchema{name: "read_file", desc: "Read a file"})
tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "read_file",
"arguments": map[string]interface{}{},
})
if !result.IsError {
t.Error("expected error for missing required args")
}
if !strings.Contains(result.ForLLM, "missing required") {
t.Errorf("expected 'missing required' in error, got: %s", result.ForLLM)
}
if !strings.Contains(result.ForLLM, "path") {
t.Errorf("expected 'path' in schema hint, got: %s", result.ForLLM)
}
}
func TestToolCallTool_ToolNotFoundSuggestsSearch(t *testing.T) {
r := NewToolRegistry()
tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "nonexistent",
})
if !result.IsError {
t.Error("expected error for nonexistent tool")
}
if !strings.Contains(result.ForLLM, "tool_search") {
t.Errorf("expected suggestion to use tool_search, got: %s", result.ForLLM)
}
} }
func TestToolCallTool_ContextPropagation(t *testing.T) { func TestToolCallTool_ContextPropagation(t *testing.T) {
@ -200,6 +230,43 @@ func TestToolCallTool_ResourceProvider_LoadsResources(t *testing.T) {
} }
} }
func TestToolCallTool_NormalizesCommaSeparatedToolName(t *testing.T) {
r := NewToolRegistry()
r.Register(&stubTool{name: "write_file", desc: "write"})
r.Register(&stubTool{name: "read_file", desc: "read"})
tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "write_file, read_file",
"arguments": map[string]interface{}{"path": "x.txt"},
})
if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM)
}
if result.ForLLM != "executed write_file" {
t.Fatalf("expected normalized dispatch to write_file, got %q", result.ForLLM)
}
}
func TestToolCallTool_NormalizesEmbeddedToolName(t *testing.T) {
r := NewToolRegistry()
r.Register(&stubTool{name: "exec", desc: "exec"})
tc := NewToolCallTool(r)
result := tc.Execute(context.Background(), map[string]interface{}{
"tool_name": "exec_tool_search_query_exec_run_shell_command_return_output_caution",
"arguments": map[string]interface{}{"command": "echo hi"},
})
if result.IsError {
t.Fatalf("unexpected error: %s", result.ForLLM)
}
if result.ForLLM != "executed exec" {
t.Fatalf("expected normalized dispatch to exec, got %q", result.ForLLM)
}
}
func TestToolCallTool_NoResourceProvider_StillWorks(t *testing.T) { func TestToolCallTool_NoResourceProvider_StillWorks(t *testing.T) {
r := NewToolRegistry() r := NewToolRegistry()
r.Register(&stubTool{name: "plain", desc: "No resources"}) r.Register(&stubTool{name: "plain", desc: "No resources"})

View file

@ -244,7 +244,14 @@ func (t *WriteFileTool) Execute(ctx context.Context, args map[string]interface{}
return ErrorResult(fmt.Sprintf("failed to write file: %v", err)) return ErrorResult(fmt.Sprintf("failed to write file: %v", err))
} }
return SilentResult(fmt.Sprintf("File written: %s", path)) preview := strings.TrimSpace(content)
if len(preview) > 80 {
preview = preview[:80] + "..."
}
if preview == "" {
return SilentResult(fmt.Sprintf("File written: %s (empty content)", path))
}
return SilentResult(fmt.Sprintf("File written: %s (content preview: %q)", path, preview))
} }
type ListDirTool struct { type ListDirTool struct {

View file

@ -131,6 +131,9 @@ func buildDenyPatterns() []*regexp.Regexp {
// Command obfuscation via eval // Command obfuscation via eval
`\beval\s+.*\$\(`, `\beval\s+.*\$\(`,
// Host identity exfiltration shortcut
`\bhostname\b`,
} }
compiled := make([]*regexp.Regexp, 0, len(patterns)) compiled := make([]*regexp.Regexp, 0, len(patterns))