Add tool binary file detection
This commit is contained in:
@@ -5,6 +5,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"encoding/json"
|
"encoding/json"
|
||||||
"fmt"
|
"fmt"
|
||||||
|
"net/http"
|
||||||
"os"
|
"os"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -175,33 +176,50 @@ func (r *Registry) ToAPI() []APITool {
|
|||||||
return tools
|
return tools
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// binaryOutputPrefix is the message shown when tool output is detected as binary.
|
||||||
|
const binaryOutputPrefix = "error: binary output detected: "
|
||||||
|
|
||||||
// ToolResult holds the structured output from a tool execution.
|
// ToolResult holds the structured output from a tool execution.
|
||||||
type ToolResult struct {
|
type ToolResult struct {
|
||||||
Stdout string
|
Stdout []byte
|
||||||
Stderr string
|
Stderr []byte
|
||||||
ProcessState *os.ProcessState
|
ProcessState *os.ProcessState
|
||||||
}
|
}
|
||||||
|
|
||||||
// FormatResult formats the tool result using the structured label-based
|
// FormatResult formats the tool result using the structured label-based
|
||||||
// format with stdout, stderr, and exit code sections.
|
// format with stdout, stderr, and exit code sections.
|
||||||
// FormatResult formats the tool result using the structured label-based
|
//
|
||||||
// format with stdout, stderr, and exit code sections.
|
// http.DetectContentType implements the WHATWG MIME-sniffing spec and
|
||||||
|
// returns text/plain for any content consisting entirely of printable
|
||||||
|
// characters and whitespace. Text-based formats like JSON, JavaScript,
|
||||||
|
// and XML will be detected as text/plain; charset=utf-8. We assume
|
||||||
|
// checking for the "text/" prefix is sufficient.
|
||||||
func (r *ToolResult) FormatResult() string {
|
func (r *ToolResult) FormatResult() string {
|
||||||
var buf strings.Builder
|
var buf strings.Builder
|
||||||
|
|
||||||
fmt.Fprintln(&buf, "▸ stdout:")
|
fmt.Fprintln(&buf, "▸ stdout:")
|
||||||
if r.Stdout != "" {
|
if len(r.Stdout) > 0 {
|
||||||
fmt.Fprint(&buf, r.Stdout)
|
ctype := http.DetectContentType(r.Stdout)
|
||||||
if !strings.HasSuffix(r.Stdout, "\n") {
|
if !strings.HasPrefix(ctype, "text/") {
|
||||||
fmt.Fprint(&buf, "\n")
|
fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
|
||||||
|
} else {
|
||||||
|
buf.Write(r.Stdout)
|
||||||
|
if r.Stdout[len(r.Stdout)-1] != '\n' {
|
||||||
|
fmt.Fprint(&buf, "\n")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fmt.Fprintln(&buf, "▸ stderr:")
|
fmt.Fprintln(&buf, "▸ stderr:")
|
||||||
if r.Stderr != "" {
|
if len(r.Stderr) > 0 {
|
||||||
fmt.Fprint(&buf, r.Stderr)
|
ctype := http.DetectContentType(r.Stderr)
|
||||||
if !strings.HasSuffix(r.Stderr, "\n") {
|
if !strings.HasPrefix(ctype, "text/") {
|
||||||
fmt.Fprint(&buf, "\n")
|
fmt.Fprintf(&buf, binaryOutputPrefix+"%s\n", ctype)
|
||||||
|
} else {
|
||||||
|
buf.Write(r.Stderr)
|
||||||
|
if r.Stderr[len(r.Stderr)-1] != '\n' {
|
||||||
|
fmt.Fprint(&buf, "\n")
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -301,8 +319,8 @@ func (r *Registry) Execute(
|
|||||||
// ProcessState is set; return it along with captured output.
|
// ProcessState is set; return it along with captured output.
|
||||||
default:
|
default:
|
||||||
return &ToolResult{
|
return &ToolResult{
|
||||||
Stdout: stdout.String(),
|
Stdout: stdout.Bytes(),
|
||||||
Stderr: stderr.String(),
|
Stderr: stderr.Bytes(),
|
||||||
ProcessState: cmd.ProcessState,
|
ProcessState: cmd.ProcessState,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
@@ -310,8 +328,8 @@ func (r *Registry) Execute(
|
|||||||
|
|
||||||
// No error. Process exited with status 0.
|
// No error. Process exited with status 0.
|
||||||
return &ToolResult{
|
return &ToolResult{
|
||||||
Stdout: stdout.String(),
|
Stdout: stdout.Bytes(),
|
||||||
Stderr: stderr.String(),
|
Stderr: stderr.Bytes(),
|
||||||
ProcessState: cmd.ProcessState,
|
ProcessState: cmd.ProcessState,
|
||||||
}, nil
|
}, nil
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
package llm
|
package llm
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"bytes"
|
||||||
"context"
|
"context"
|
||||||
"os/exec"
|
"os/exec"
|
||||||
"strings"
|
"strings"
|
||||||
@@ -221,24 +222,54 @@ func TestToolResult_FormatResult(t *testing.T) {
|
|||||||
}{
|
}{
|
||||||
{
|
{
|
||||||
name: "success with stdout",
|
name: "success with stdout",
|
||||||
result: ToolResult{Stdout: "hello\n", Stderr: "", ProcessState: cmdOK.ProcessState},
|
result: ToolResult{Stdout: []byte("hello\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
wantContains: []string{"▸ stdout:", "hello", "▸ stderr:", "┗ exit: 0"},
|
wantContains: []string{"▸ stdout:", "hello", "▸ stderr:", "┗ exit: 0"},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "failure with stderr",
|
name: "failure with stderr",
|
||||||
result: ToolResult{Stdout: "", Stderr: "error occurred\n", ProcessState: cmdFail.ProcessState},
|
result: ToolResult{Stdout: nil, Stderr: []byte("error occurred\n"), ProcessState: cmdFail.ProcessState},
|
||||||
wantContains: []string{"▸ stdout:", "▸ stderr:", "error occurred", "┗ exit: 3"},
|
wantContains: []string{"▸ stdout:", "▸ stderr:", "error occurred", "┗ exit: 3"},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "both stdout and stderr",
|
name: "both stdout and stderr",
|
||||||
result: ToolResult{Stdout: "output", Stderr: "warning", ProcessState: cmdOK.ProcessState},
|
result: ToolResult{Stdout: []byte("output"), Stderr: []byte("warning"), ProcessState: cmdOK.ProcessState},
|
||||||
wantContains: []string{"▸ stdout:", "output", "▸ stderr:", "warning", "┗ exit: 0"},
|
wantContains: []string{"▸ stdout:", "output", "▸ stderr:", "warning", "┗ exit: 0"},
|
||||||
},
|
},
|
||||||
{
|
{
|
||||||
name: "empty output",
|
name: "empty output",
|
||||||
result: ToolResult{Stdout: "", Stderr: "", ProcessState: cmdOK.ProcessState},
|
result: ToolResult{Stdout: nil, Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
wantContains: []string{"▸ stdout:", "▸ stderr:", "┗ exit: 0"},
|
wantContains: []string{"▸ stdout:", "▸ stderr:", "┗ exit: 0"},
|
||||||
},
|
},
|
||||||
|
{
|
||||||
|
name: "binary stdout (PNG)",
|
||||||
|
result: ToolResult{Stdout: []byte("\x89PNG\r\n\x1a\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", binaryOutputPrefix, "image/png", "▸ stderr:", "┗ exit: 0"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "binary stderr",
|
||||||
|
result: ToolResult{Stdout: nil, Stderr: []byte("\x89PNG\r\n\x1a\n"), ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", "▸ stderr:", binaryOutputPrefix, "image/png", "┗ exit: 0"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "HTML",
|
||||||
|
result: ToolResult{Stdout: []byte("<HTML><BODY>Hello</BODY></HTML>\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", "<HTML>", "▸ stderr:", "┗ exit: 0"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "XML",
|
||||||
|
result: ToolResult{Stdout: []byte("<?xml version=\"1.0\"?><root/>\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", "<?xml", "▸ stderr:", "┗ exit: 0"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "Go source",
|
||||||
|
result: ToolResult{Stdout: []byte("package main\n\nfunc main() {}\n"), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", "package main", "▸ stderr:", "┗ exit: 0"},
|
||||||
|
},
|
||||||
|
{
|
||||||
|
name: "JSON",
|
||||||
|
result: ToolResult{Stdout: []byte(`{"key":"value"}`), Stderr: nil, ProcessState: cmdOK.ProcessState},
|
||||||
|
wantContains: []string{"▸ stdout:", `{"key":"value"}`, "▸ stderr:", "┗ exit: 0"},
|
||||||
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
for _, tt := range tests {
|
for _, tt := range tests {
|
||||||
@@ -268,7 +299,7 @@ func TestRegistry_Execute_Success(t *testing.T) {
|
|||||||
if result.ProcessState.ExitCode() != 0 {
|
if result.ProcessState.ExitCode() != 0 {
|
||||||
t.Errorf("ExitCode = %d, want 0", result.ProcessState.ExitCode())
|
t.Errorf("ExitCode = %d, want 0", result.ProcessState.ExitCode())
|
||||||
}
|
}
|
||||||
if !strings.Contains(result.Stdout, "hello") {
|
if !bytes.Contains(result.Stdout, []byte("hello")) {
|
||||||
t.Errorf("Stdout = %q, want to contain hello", result.Stdout)
|
t.Errorf("Stdout = %q, want to contain hello", result.Stdout)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user