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:
parent
3d940883a6
commit
d8c5592388
4 changed files with 233 additions and 27 deletions
|
|
@ -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
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -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"})
|
||||||
|
|
|
||||||
|
|
@ -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 {
|
||||||
|
|
|
||||||
|
|
@ -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))
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue