feat(tools): stream shell_run output live into the TUI

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 = &copy
 }
+
 // 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()