1package modes
2
3import (
4 "context"
5 "fmt"
6 "io"
7 "os"
8 "strings"
9
10 "github.com/bnema/zut/packages/core"
11 "github.com/bnema/zut/packages/provider"
12)
13
14// RunStream runs the agent and writes assistant text to out as it
15// arrives (EvTextDelta). Tool activity is written to stderr so stdout
16// stays the live assistant transcript. When a provider does not emit
17// deltas, the full assistant text is written when the message completes.
18func RunStream(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out io.Writer) error {
19 return RunStreamWithDiag(ctx, ag, prompt, images, out, os.Stderr)
20}
21
22// RunStreamWithDiag is RunStream with an explicit diagnostics writer
23// for tool activity (tests pass a buffer).
24func RunStreamWithDiag(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out, diag io.Writer) error {
25 _, err := RunStreamWithContextRecovery(ctx, ag, prompt, images, out, diag, nil)
26 return err
27}
28
29// RunStreamWithContextRecovery is RunStreamWithDiag with an optional
30// checkpoint writer for a one-shot context-overflow compaction. OutputStart
31// identifies the transcript suffix that callers must persist after a
32// successful checkpoint.
33func RunStreamWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, out, diag io.Writer, persistCompaction func([]provider.Message) error) (ContextRecoveryResult, error) {
34 if diag == nil {
35 diag = io.Discard
36 }
37 var (
38 runErr error
39 streamed bool
40 wroteText bool
41 lastWasNL bool
42 )
43 flush := func() {
44 if f, ok := out.(interface{ Sync() error }); ok {
45 _ = f.Sync()
46 }
47 }
48 writeText := func(s string) {
49 if s == "" {
50 return
51 }
52 _, _ = fmt.Fprint(out, s)
53 flush()
54 wroteText = true
55 lastWasNL = strings.HasSuffix(s, "\n")
56 }
57 assistantText := func(m provider.Message) string {
58 var sb strings.Builder
59 for _, c := range m.Content {
60 if tb, ok := c.(provider.TextBlock); ok {
61 if sb.Len() > 0 {
62 sb.WriteString("\n")
63 }
64 sb.WriteString(tb.Text)
65 }
66 }
67 return sb.String()
68 }
69
70 sink := func(ev core.AgentEvent) {
71 switch e := ev.(type) {
72 case core.EvAssistantStart:
73 streamed = false
74 case core.EvTextDelta:
75 streamed = true
76 writeText(e.Delta)
77 case core.EvAssistantMessage:
78 if !streamed {
79 writeText(assistantText(e.Message))
80 }
81 case core.EvToolCall:
82 fmt.Fprintf(diag, "\n[tool] %s\n", e.Name)
83 if len(e.Args) > 0 {
84 fmt.Fprintf(diag, "%s\n", string(e.Args))
85 }
86 case core.EvToolProgress:
87 if e.Text != "" {
88 fmt.Fprint(diag, e.Text)
89 }
90 case core.EvToolResult:
91 for _, c := range e.Result.Content {
92 if tb, ok := c.(provider.TextBlock); ok && tb.Text != "" {
93 fmt.Fprint(diag, tb.Text)
94 if !strings.HasSuffix(tb.Text, "\n") {
95 fmt.Fprint(diag, "\n")
96 }
97 }
98 }
99 if e.Result.IsError {
100 fmt.Fprintln(diag, "[tool error]")
101 }
102 case core.EvTurnEnd:
103 if e.Err != nil {
104 runErr = e.Err
105 }
106 case core.EvError:
107 if e.Err != nil {
108 runErr = e.Err
109 }
110 }
111 }
112
113 recovery, err := PromptWithContextRecovery(ctx, ag, prompt, images, sink, ContextRecoveryOptions{
114 PersistCompaction: persistCompaction,
115 SuppressInitialOverflowEvent: true,
116 })
117 if err != nil {
118 return recovery, err
119 }
120 if runErr != nil {
121 return recovery, runErr
122 }
123 if wroteText && !lastWasNL {
124 _, _ = fmt.Fprint(out, "\n")
125 flush()
126 }
127 return recovery, nil
128}