Compare commits
1 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 47671f3098 |
@@ -167,7 +167,8 @@ Web tools are also available in read-only mode:
|
||||
|
||||
- `file_write` for creating or replacing whole files
|
||||
- `file_edit` for surgical line-range or pattern edits to existing files
|
||||
- `shell_run` for non-interactive shell commands
|
||||
- `shell_run` for non-interactive shell commands; output streams live into the
|
||||
TUI while the command runs instead of appearing only after completion
|
||||
|
||||
Use `--yolo` for prompts that ask agentu to run or inspect local commands, for
|
||||
example:
|
||||
|
||||
+56
-12
@@ -24,6 +24,7 @@ const (
|
||||
toolLogArgumentLimit = 400
|
||||
toolLogOutputPreviewBytes = 4096
|
||||
toolLogOutputMarker = "[output]"
|
||||
toolStreamLogLimit = tools.MaxToolOutputBytes
|
||||
|
||||
summarizerSystemPrompt = "You are a conversation summarizer. Summarize the following conversation history into a concise but detailed summary. Preserve:\n" +
|
||||
"- Key facts, decisions, and conclusions\n" +
|
||||
@@ -182,6 +183,7 @@ func (a *Agent) SetLastUsage(usage *llm.Usage) {
|
||||
copy := *usage
|
||||
a.lastUsage = ©
|
||||
}
|
||||
|
||||
// compactRetainedTokenBudget returns the number of estimated tokens to retain
|
||||
// based on compactRecentRetentionPercent, using integer ceiling division.
|
||||
func compactRetainedTokenBudget(totalTokens int) int {
|
||||
@@ -237,7 +239,6 @@ func compactRecentStart(messages []llm.Message, prefixEnd int) int {
|
||||
return compactTurnBoundary(messages, prefixEnd, start)
|
||||
}
|
||||
|
||||
|
||||
// Compact summarizes older conversation messages while preserving the system
|
||||
// prompt and recent messages intact.
|
||||
func (a *Agent) Compact(ctx context.Context) error {
|
||||
@@ -495,32 +496,75 @@ func (a *Agent) toolDefinitions() []llm.Tool {
|
||||
return a.toolRegistry.Definitions()
|
||||
}
|
||||
|
||||
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall) string {
|
||||
func (a *Agent) executeTool(ctx context.Context, call llm.ToolCall, logs io.Writer) (string, bool) {
|
||||
name := call.Function.Name
|
||||
|
||||
if a.toolRegistry == nil {
|
||||
return tools.FormatError(fmt.Errorf("tool registry is disabled"))
|
||||
return tools.FormatError(fmt.Errorf("tool registry is disabled")), false
|
||||
}
|
||||
tool, ok := a.toolRegistry.Get(name)
|
||||
if !ok {
|
||||
return tools.FormatError(fmt.Errorf("unknown tool: %s", name))
|
||||
return tools.FormatError(fmt.Errorf("unknown tool: %s", name)), false
|
||||
}
|
||||
|
||||
args := strings.TrimSpace(call.Function.Arguments)
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
raw := json.RawMessage(args)
|
||||
|
||||
toolCtx, cancel := context.WithTimeout(ctx, a.toolTimeout)
|
||||
defer cancel()
|
||||
output, err := tool.Execute(toolCtx, json.RawMessage(args))
|
||||
|
||||
if streaming, ok := tool.(tools.StreamingTool); ok {
|
||||
return a.executeStreamingTool(streaming, toolCtx, raw, name, args, logs)
|
||||
}
|
||||
|
||||
output, err := tool.Execute(toolCtx, raw)
|
||||
if err != nil {
|
||||
if output != "" {
|
||||
return output + "\n" + tools.FormatError(err)
|
||||
return output + "\n" + tools.FormatError(err), false
|
||||
}
|
||||
return tools.FormatError(err)
|
||||
return tools.FormatError(err), false
|
||||
}
|
||||
return output
|
||||
return output, false
|
||||
}
|
||||
|
||||
func (a *Agent) executeStreamingTool(streaming tools.StreamingTool, ctx context.Context, raw json.RawMessage, name, args string, logs io.Writer) (string, bool) {
|
||||
var streamed int
|
||||
emit := func(chunk string) error {
|
||||
if chunk == "" || logs == nil {
|
||||
return nil
|
||||
}
|
||||
if streamed >= toolStreamLogLimit {
|
||||
return nil
|
||||
}
|
||||
if len(chunk) > toolStreamLogLimit-streamed {
|
||||
chunk = chunk[:toolStreamLogLimit-streamed]
|
||||
streamed = toolStreamLogLimit
|
||||
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s\n[streaming output truncated after %d bytes]", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk, toolStreamLogLimit)
|
||||
return err
|
||||
}
|
||||
streamed += len(chunk)
|
||||
_, err := fmt.Fprintf(logs, "\n[tool] %s %s\n%s\n%s", name, compact(args, toolLogArgumentLimit), toolLogOutputMarker, chunk)
|
||||
return err
|
||||
}
|
||||
|
||||
output, err := streaming.ExecuteStream(ctx, raw, emit)
|
||||
if logs != nil {
|
||||
if err != nil {
|
||||
_ = emit("\n" + tools.FormatError(err))
|
||||
} else if streamed == 0 {
|
||||
_ = emit("(no output)")
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
if output != "" {
|
||||
return output + "\n" + tools.FormatError(err), true
|
||||
}
|
||||
return tools.FormatError(err), true
|
||||
}
|
||||
return output, true
|
||||
}
|
||||
|
||||
func (a *Agent) hasMutatingCall(calls []llm.ToolCall) bool {
|
||||
@@ -550,9 +594,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
|
||||
if a.hasMutatingCall(calls) {
|
||||
// Sequential: avoid racing mutating tools.
|
||||
for i, call := range calls {
|
||||
output := a.executeTool(ctx, call)
|
||||
output, streamed := a.executeTool(ctx, call, logs)
|
||||
results[i] = toolExecutionResult{call: call, output: output}
|
||||
if logs != nil {
|
||||
if logs != nil && !streamed {
|
||||
logToolOutput(logs, call, output)
|
||||
}
|
||||
}
|
||||
@@ -566,9 +610,9 @@ func (a *Agent) executeTools(ctx context.Context, calls []llm.ToolCall, logs io.
|
||||
wg.Add(1)
|
||||
go func(idx int, c llm.ToolCall) {
|
||||
defer wg.Done()
|
||||
output := a.executeTool(ctx, c)
|
||||
output, streamed := a.executeTool(ctx, c, logs)
|
||||
results[idx] = toolExecutionResult{call: c, output: output}
|
||||
if logs != nil {
|
||||
if logs != nil && !streamed {
|
||||
mu.Lock()
|
||||
logToolOutput(logs, c, output)
|
||||
mu.Unlock()
|
||||
|
||||
@@ -87,6 +87,86 @@ func TestAgentRunsToolLoop(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
type streamToolProvider struct {
|
||||
calls int
|
||||
t *testing.T
|
||||
}
|
||||
|
||||
func (p *streamToolProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
p.calls++
|
||||
switch p.calls {
|
||||
case 1:
|
||||
return emit(llm.StreamEvent{ToolCalls: []llm.ToolCallDelta{
|
||||
{Index: 0, ID: "call_1", Type: "function", Name: "test_stream", Arguments: `{}`},
|
||||
}})
|
||||
case 2:
|
||||
last := req.Messages[len(req.Messages)-1]
|
||||
if last.Role != llm.RoleTool || last.Content != "full result" {
|
||||
p.t.Fatalf("last message = %#v", last)
|
||||
}
|
||||
return emit(llm.StreamEvent{Content: "done"})
|
||||
default:
|
||||
p.t.Fatalf("unexpected call count %d", p.calls)
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
type streamEchoTool struct{}
|
||||
|
||||
func (streamEchoTool) Definition() llm.Tool {
|
||||
return llm.Tool{
|
||||
Type: "function",
|
||||
Function: llm.ToolFunction{
|
||||
Name: "test_stream",
|
||||
Description: "streaming test",
|
||||
Parameters: json.RawMessage(`{"type":"object"}`),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func (streamEchoTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
return "full result", nil
|
||||
}
|
||||
|
||||
func (streamEchoTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
|
||||
if err := emit("one\n"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := emit("two\n"); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return "full result", nil
|
||||
}
|
||||
|
||||
func TestAgentStreamsToolOutputToLogs(t *testing.T) {
|
||||
provider := &streamToolProvider{t: t}
|
||||
a := New(Options{
|
||||
Provider: provider,
|
||||
Model: "test",
|
||||
SystemPrompt: "system",
|
||||
ToolRegistry: tools.NewRegistry(streamEchoTool{}),
|
||||
})
|
||||
|
||||
var out strings.Builder
|
||||
var logs strings.Builder
|
||||
if err := a.RunTurn(context.Background(), "stream", &out, &logs); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.String() != "done" {
|
||||
t.Fatalf("out = %q", out.String())
|
||||
}
|
||||
logText := logs.String()
|
||||
if strings.Count(logText, "[output]") != 2 {
|
||||
t.Fatalf("logs = %q", logText)
|
||||
}
|
||||
if !strings.Contains(logText, "one\n") || !strings.Contains(logText, "two\n") {
|
||||
t.Fatalf("logs missing streamed chunks: %q", logText)
|
||||
}
|
||||
if strings.Contains(logText, "[output]\nfull result") {
|
||||
t.Fatalf("logs contain duplicate full output: %q", logText)
|
||||
}
|
||||
}
|
||||
|
||||
type contentProvider struct {
|
||||
called bool
|
||||
text string
|
||||
|
||||
+40
-5
@@ -1,6 +1,7 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
@@ -8,6 +9,7 @@ import (
|
||||
"os"
|
||||
"os/exec"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
"agentu/pkg/llm"
|
||||
)
|
||||
@@ -46,6 +48,10 @@ type shellRunArgs struct {
|
||||
}
|
||||
|
||||
func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
|
||||
return t.ExecuteStream(ctx, raw, nil)
|
||||
}
|
||||
|
||||
func (t *ShellRunTool) ExecuteStream(ctx context.Context, raw json.RawMessage, emit func(string) error) (string, error) {
|
||||
var args shellRunArgs
|
||||
if err := json.Unmarshal(raw, &args); err != nil {
|
||||
return "", fmt.Errorf("parse arguments: %w", err)
|
||||
@@ -66,16 +72,45 @@ func (t *ShellRunTool) Execute(ctx context.Context, raw json.RawMessage) (string
|
||||
cmd := exec.CommandContext(ctx, shell, "-lc", args.Command)
|
||||
cmd.Dir = workdir
|
||||
|
||||
output, err := cmd.CombinedOutput()
|
||||
text := FormatResult(string(output))
|
||||
var mu sync.Mutex
|
||||
var combined bytes.Buffer
|
||||
writer := &shellStreamWriter{mu: &mu, buf: &combined, emit: emit}
|
||||
cmd.Stdout = writer
|
||||
cmd.Stderr = writer
|
||||
|
||||
runErr := cmd.Run()
|
||||
text := FormatResult(combined.String())
|
||||
if errors.Is(ctx.Err(), context.DeadlineExceeded) {
|
||||
return text, fmt.Errorf("command timed out: %w", ctx.Err())
|
||||
}
|
||||
if err != nil {
|
||||
if runErr != nil {
|
||||
if text != "" {
|
||||
return text, fmt.Errorf("command failed: %w", err)
|
||||
return text, fmt.Errorf("command failed: %w", runErr)
|
||||
}
|
||||
return "", fmt.Errorf("command failed: %w", err)
|
||||
return "", fmt.Errorf("command failed: %w", runErr)
|
||||
}
|
||||
return text, nil
|
||||
}
|
||||
|
||||
// shellStreamWriter captures command output for the final result while
|
||||
// forwarding every write to the streaming callback as it happens.
|
||||
type shellStreamWriter struct {
|
||||
mu *sync.Mutex
|
||||
buf *bytes.Buffer
|
||||
emit func(string) error
|
||||
}
|
||||
|
||||
func (w *shellStreamWriter) Write(p []byte) (int, error) {
|
||||
w.mu.Lock()
|
||||
n, writeErr := w.buf.Write(p)
|
||||
w.mu.Unlock()
|
||||
if writeErr != nil {
|
||||
return n, writeErr
|
||||
}
|
||||
if w.emit != nil {
|
||||
if err := w.emit(string(p)); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
package tools
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestShellRunStreamsOutputWhileRunning(t *testing.T) {
|
||||
tool := NewShellRunTool(t.TempDir())
|
||||
var mu sync.Mutex
|
||||
var chunks []string
|
||||
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"printf 'a\\n'; sleep 0.05; printf 'b\\n'"}`), func(chunk string) error {
|
||||
mu.Lock()
|
||||
chunks = append(chunks, chunk)
|
||||
mu.Unlock()
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(output, "a") || !strings.Contains(output, "b") {
|
||||
t.Fatalf("output = %q", output)
|
||||
}
|
||||
mu.Lock()
|
||||
joined := strings.Join(chunks, "")
|
||||
mu.Unlock()
|
||||
if len(chunks) == 0 {
|
||||
t.Fatal("expected at least one streamed chunk")
|
||||
}
|
||||
if !strings.Contains(joined, "a") || !strings.Contains(joined, "b") {
|
||||
t.Fatalf("streamed chunks = %q", joined)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellRunExecuteCollectsFullOutput(t *testing.T) {
|
||||
tool := NewShellRunTool(t.TempDir())
|
||||
output, err := tool.Execute(context.Background(), json.RawMessage(`{"command":"printf 'hello world'"}`))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if strings.TrimSpace(output) != "hello world" {
|
||||
t.Fatalf("output = %q", output)
|
||||
}
|
||||
}
|
||||
|
||||
func TestShellRunStreamsStderrAndStdout(t *testing.T) {
|
||||
tool := NewShellRunTool(t.TempDir())
|
||||
var mu sync.Mutex
|
||||
var chunks []string
|
||||
output, err := tool.ExecuteStream(context.Background(), json.RawMessage(`{"command":"echo out; echo err >&2"}`), func(chunk string) error {
|
||||
mu.Lock()
|
||||
chunks = append(chunks, chunk)
|
||||
mu.Unlock()
|
||||
return nil
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(output, "out") || !strings.Contains(output, "err") {
|
||||
t.Fatalf("output = %q", output)
|
||||
}
|
||||
mu.Lock()
|
||||
joined := strings.Join(chunks, "")
|
||||
mu.Unlock()
|
||||
if !strings.Contains(joined, "out") || !strings.Contains(joined, "err") {
|
||||
t.Fatalf("streamed chunks = %q", joined)
|
||||
}
|
||||
}
|
||||
@@ -19,6 +19,15 @@ type Tool interface {
|
||||
Execute(ctx context.Context, args json.RawMessage) (string, error)
|
||||
}
|
||||
|
||||
// StreamingTool is optionally implemented by tools that can report partial
|
||||
// output while they execute. emit receives chunks of the tool's output as
|
||||
// they are produced; the returned string is still the complete result that
|
||||
// gets sent back to the model.
|
||||
type StreamingTool interface {
|
||||
Tool
|
||||
ExecuteStream(ctx context.Context, args json.RawMessage, emit func(string) error) (string, error)
|
||||
}
|
||||
|
||||
// MutatingTool is optionally implemented by tools that modify external state.
|
||||
// It is used by the agent to determine whether tool calls can run in parallel.
|
||||
type MutatingTool interface {
|
||||
|
||||
+15
-3
@@ -683,15 +683,27 @@ func (m *model) mergeToolLog(text string) bool {
|
||||
if incoming.output == "" {
|
||||
return false
|
||||
}
|
||||
incomingHeader, _, ok := strings.Cut(text, "\n[output]\n")
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
for i := len(m.messages) - 1; i >= 0; i-- {
|
||||
if m.messages[i].role != roleTool {
|
||||
continue
|
||||
}
|
||||
existing := parseToolLog(m.messages[i].content)
|
||||
if existing.name == incoming.name && existing.args == incoming.args && existing.output == "" {
|
||||
m.messages[i].content = text
|
||||
return true
|
||||
if existing.name != incoming.name || existing.args != incoming.args {
|
||||
continue
|
||||
}
|
||||
if !strings.Contains(m.messages[i].content, "\n[output]\n") {
|
||||
// Placeholder card from the tool-start log; first chunk fills it.
|
||||
m.messages[i].content = text
|
||||
} else {
|
||||
// Streaming tools send incremental chunks; append the raw delta
|
||||
// so line boundaries are preserved inside the same card.
|
||||
m.messages[i].content += strings.TrimPrefix(text, incomingHeader+"\n[output]\n")
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
@@ -395,6 +395,28 @@ func TestToolLogOutputUpdatesExistingToolCard(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestToolLogStreamingChunksAppendToSameCard(t *testing.T) {
|
||||
m := newModel(context.Background(), nil, Options{ModelName: "test-model"})
|
||||
m.messages = append(m.messages, message{role: roleTool, content: `shell_run echo hi`})
|
||||
|
||||
if !m.mergeToolLog("shell_run echo hi\n[output]\none\n") {
|
||||
t.Fatal("expected first mergeToolLog to return true")
|
||||
}
|
||||
if !m.mergeToolLog("shell_run echo hi\n[output]\ntwo\n") {
|
||||
t.Fatal("expected second mergeToolLog to return true")
|
||||
}
|
||||
if len(m.messages) != 1 {
|
||||
t.Fatalf("messages count = %d, want 1", len(m.messages))
|
||||
}
|
||||
plain := stripANSI(m.renderMessages())
|
||||
if !strings.Contains(plain, "one") || !strings.Contains(plain, "two") {
|
||||
t.Fatalf("rendered output missing streamed lines:\n%s", plain)
|
||||
}
|
||||
if strings.Contains(plain, "onetwo") {
|
||||
t.Fatalf("streamed lines were concatenated:\n%s", plain)
|
||||
}
|
||||
}
|
||||
|
||||
type tuiCompactProvider struct{}
|
||||
|
||||
func (tuiCompactProvider) ChatStream(ctx context.Context, req llm.ChatRequest, emit func(llm.StreamEvent) error) error {
|
||||
|
||||
Reference in New Issue
Block a user