feat: built-in large tool output handling #2728
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,101 @@ | ||
| package builtins | ||
|
|
||
| import ( | ||
| "context" | ||
| "encoding/json" | ||
| "fmt" | ||
| "os" | ||
| "path/filepath" | ||
| "strings" | ||
|
|
||
| "github.com/docker/docker-agent/pkg/hooks" | ||
| ) | ||
|
|
||
| const HandleLargeToolOutput = "handle_large_tool_output" | ||
|
|
||
| func handleLargeToolOutput(ctx context.Context, in *hooks.Input, args []string) (*hooks.Output, error) { | ||
| if in == nil { | ||
| return nil, nil | ||
| } | ||
|
|
||
| if in.HookEventName != hooks.EventToolResponseTransform { | ||
| return nil, nil | ||
| } | ||
|
|
||
| response, ok := in.ToolResponse.(string) | ||
| if !ok || response == "" { | ||
| return nil, nil | ||
| } | ||
|
|
||
| cfg := parseArgs(args) | ||
| threshold := cfg.Threshold | ||
| if threshold == 0 { | ||
| threshold = 30000 | ||
| } | ||
|
|
||
| if len(response) <= threshold { | ||
| return nil, nil | ||
| } | ||
|
|
||
| outputDir := cfg.OutputDir | ||
| if outputDir == "" { | ||
| outputDir = os.TempDir() | ||
| } | ||
|
|
||
| if err := os.MkdirAll(outputDir, 0o750); err != nil { | ||
| return nil, fmt.Errorf("create output directory: %w", err) | ||
| } | ||
|
|
||
| filename := fmt.Sprintf("%s_%s.txt", | ||
| sanitizeFilename(in.SessionID), | ||
| sanitizeFilename(in.ToolUseID)) | ||
| path := filepath.Join(outputDir, filename) | ||
|
|
||
| if err := os.WriteFile(path, []byte(response), 0o600); err != nil { | ||
| return nil, fmt.Errorf("write output file: %w", err) | ||
| } | ||
|
|
||
| previewSize := cfg.PreviewSize | ||
| if previewSize == 0 { | ||
| previewSize = 3000 | ||
| } | ||
| preview := response | ||
| if len(preview) > previewSize { | ||
| preview = response[:previewSize] | ||
| } | ||
|
|
||
| pointer := fmt.Sprintf("[%s response: %d chars, full output saved to %s]\n\nFirst %d chars:\n%s\n\n[Use shell tool to read: cat %s]", | ||
|
There was a problem hiding this comment. [LOW] Shell metacharacters in The pointer message shown to the agent includes: pointer := fmt.Sprintf("...[Use shell tool to read: cat %s]", path)
Fix: Either quote the path in the suggested command, or expand // Option A: shell-quote the path in the message
pointer := fmt.Sprintf("...[Use shell tool to read: cat '%s']", path)
// Option B: allowlist in sanitizeFilename
func sanitizeFilename(name string) string {
var b strings.Builder
for _, r := range name {
if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') ||
(r >= '0' && r <= '9') || r == '-' || r == '.' {
b.WriteRune(r)
} else {
b.WriteRune('_')
}
}
return b.String()
} |
||
| in.ToolName, len(response), path, previewSize, preview, path) | ||
|
|
||
| return &hooks.Output{ | ||
| HookSpecificOutput: &hooks.HookSpecificOutput{ | ||
| HookEventName: hooks.EventToolResponseTransform, | ||
| UpdatedToolResponse: &pointer, | ||
| }, | ||
| }, nil | ||
| } | ||
|
|
||
| type toolOutputConfig struct { | ||
|
There was a problem hiding this comment. [HIGH]
Threshold int `json:"threshold,omitempty"`
OutputDir string `json:"output_dir,omitempty"`
PreviewSize int `json:"preview_size,omitempty"`But type toolOutputConfig struct {
Threshold int
OutputDir string
PreviewSize int
}In
Result: Any Fix: Add json struct tags to type toolOutputConfig struct {
Threshold int `json:"threshold"`
OutputDir string `json:"output_dir"`
PreviewSize int `json:"preview_size"`
} |
||
| Threshold int | ||
| OutputDir string | ||
| PreviewSize int | ||
| } | ||
|
|
||
| func parseArgs(args []string) toolOutputConfig { | ||
| if len(args) == 0 { | ||
| return toolOutputConfig{} | ||
| } | ||
|
|
||
| var cfg toolOutputConfig | ||
| if err := json.Unmarshal([]byte(args[0]), &cfg); err != nil { | ||
| return toolOutputConfig{} | ||
| } | ||
| return cfg | ||
| } | ||
|
|
||
| func sanitizeFilename(name string) string { | ||
| name = strings.ReplaceAll(name, "..", "__") | ||
| name = strings.ReplaceAll(name, "/", "_") | ||
| name = strings.ReplaceAll(name, "\\", "_") | ||
| return name | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[MEDIUM] Filename collision risk when
SessionIDorToolUseIDis emptyThe output filename is built as:
If either
SessionIDorToolUseIDis empty (e.g.sanitizeFilename("") == ""), multiple tool calls can produce the same filename (e.g.,_.txtor_<toolUseID>.txt). Each subsequentos.WriteFilecall silently overwrites the previous content. When the agent later reads the file, it gets the last response only — the earlier ones are lost without any error.Fix: Validate that both IDs are non-empty before constructing the filename, or incorporate a timestamp/random component to guarantee uniqueness: