diff --git a/README.md b/README.md index d015567..07ffd82 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/internal/agent/agent.go b/internal/agent/agent.go index f295cc1..bb5bff6 100644 --- a/internal/agent/agent.go +++ b/internal/agent/agent.go @@ -15,15 +15,16 @@ import ( ) const ( - defaultMaxToolRounds = 50 - doomLoopThreshold = 3 + defaultMaxToolRounds = 50 + doomLoopThreshold = 3 compactMinRecentMessages = 6 compactRecentRetentionPercent = 20 - summaryMessageLimit = 8000 - contextSummaryPrefix = "[context summary]" - toolLogArgumentLimit = 400 - toolLogOutputPreviewBytes = 4096 - toolLogOutputMarker = "[output]" + summaryMessageLimit = 8000 + contextSummaryPrefix = "[context summary]" + 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() diff --git a/internal/agent/agent_test.go b/internal/agent/agent_test.go index 5183e90..70afae9 100644 --- a/internal/agent/agent_test.go +++ b/internal/agent/agent_test.go @@ -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 diff --git a/internal/tools/shell.go b/internal/tools/shell.go index 84fd57a..0f81d22 100644 --- a/internal/tools/shell.go +++ b/internal/tools/shell.go @@ -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 +} diff --git a/internal/tools/shell_test.go b/internal/tools/shell_test.go new file mode 100644 index 0000000..1e63dbb --- /dev/null +++ b/internal/tools/shell_test.go @@ -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) + } +} diff --git a/internal/tools/tools.go b/internal/tools/tools.go index 7209fb6..36f9f88 100644 --- a/internal/tools/tools.go +++ b/internal/tools/tools.go @@ -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 { diff --git a/internal/tui/app.go b/internal/tui/app.go index 7bee8a2..a08d394 100644 --- a/internal/tui/app.go +++ b/internal/tui/app.go @@ -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 } diff --git a/internal/tui/app_test.go b/internal/tui/app_test.go index a570ff4..a264911 100644 --- a/internal/tui/app_test.go +++ b/internal/tui/app_test.go @@ -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 {