loveuer/agentu
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
loveuer
2026-08-13 13:56:24 +08:00
parent 28195519d1
commit 47671f3098
8 changed files with 302 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/agent/agent.go
+63 -19
View File
@@ -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()
internal/agent/agent_test.go
+80
View File
@@ -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