1package modes
2
3import (
4 "context"
5 "fmt"
6
7 "github.com/bnema/zut/packages/core"
8 "github.com/bnema/zut/packages/provider"
9)
10
11// ContextRecoveryOptions configures the host-owned, one-shot recovery for a
12// provider context-overflow error. It deliberately does not add a terminal
13// continuation policy to core.Agent.
14type ContextRecoveryOptions struct {
15 // PersistCompaction records the compacted transcript before the existing
16 // prompt is continued. A persistence failure restores the pre-compaction
17 // transcript and prevents the continuation request.
18 PersistCompaction func([]provider.Message) error
19
20 // ContextWindow and AutoCompactThreshold enable proactive compaction before
21 // Prompt. The threshold uses the same supported percentage values and 85%
22 // default as interactive mode; zero disables it. Context usage includes
23 // uncached input, cache reads, and cache writes from the latest request.
24 ContextWindow int
25 AutoCompactThreshold *int
26
27 // SuppressInitialOverflowEvent hides the recoverable first-attempt
28 // EvTurnEnd overflow from sinks whose protocol cannot retract a terminal
29 // error frame, such as JSONL. A retry overflow is never suppressed.
30 SuppressInitialOverflowEvent bool
31}
32
33// ContextRecoveryResult describes a completed host-level context recovery.
34type ContextRecoveryResult struct {
35 // OutputStart is the first message to append after a persisted compaction.
36 // When Compacted is false, it is the message count before Prompt.
37 OutputStart int
38 Compacted bool
39}
40
41// PromptWithContextRecovery proactively compacts a transcript whose latest
42// request reached the configured context percentage, then appends prompt and
43// runs it once. If that provider request still exceeds its context window, it
44// compacts the existing transcript and continues the already-appended prompt
45// exactly once. Cancellation, a compaction failure, and a second overflow
46// remain terminal outcomes.
47func PromptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, sink func(core.AgentEvent), options ContextRecoveryOptions) (result ContextRecoveryResult, err error) {
48 result.OutputStart = len(ag.Messages())
49 if sink == nil {
50 sink = func(core.AgentEvent) {}
51 }
52
53 usage := ag.LastTurnUsage()
54 contextUsed := usage.InputTokens + usage.CacheReadTokens + usage.CacheWriteTokens
55 threshold := NormalizeAutoCompactThreshold(options.AutoCompactThreshold)
56 proactiveSummaryOverflow := false
57 if ShouldAutoCompact(contextUsed, options.ContextWindow, threshold) {
58 messages := ag.Messages()
59 // A sole transcript message is commonly the worker's only task.
60 // Replacing it before it has run loses the exact instruction.
61 if len(messages) > 1 {
62 if err := ctx.Err(); err != nil {
63 return result, err
64 }
65 compacted, compactErr := compactContext(ctx, ag, messages, options.PersistCompaction)
66 if compactErr != nil {
67 if ctxErr := ctx.Err(); ctxErr != nil {
68 return result, ctxErr
69 }
70 if !provider.IsContextOverflowError(compactErr) {
71 return result, fmt.Errorf("proactively compact transcript: %w", compactErr)
72 }
73 // The summary request can itself exceed a provider limit. The
74 // transcript is unchanged, so make one normal prompt attempt. If
75 // that also overflows, do not repeat the doomed summary request.
76 proactiveSummaryOverflow = true
77 } else {
78 result.OutputStart = len(compacted)
79 result.Compacted = true
80 if err := ctx.Err(); err != nil {
81 return result, err
82 }
83 }
84 }
85 }
86
87 firstSink := sink
88 if options.SuppressInitialOverflowEvent && !proactiveSummaryOverflow {
89 firstSink = func(event core.AgentEvent) {
90 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
91 return
92 }
93 sink(event)
94 }
95 }
96
97 if err = ag.Prompt(ctx, prompt, images, firstSink); err == nil || ctx.Err() != nil || !provider.IsContextOverflowError(err) || proactiveSummaryOverflow {
98 return result, err
99 }
100
101 if ctxErr := ctx.Err(); ctxErr != nil {
102 return result, ctxErr
103 }
104 messages := ag.Messages()
105 if len(messages) <= 1 {
106 // The failed prompt is the only message. Summarizing it would replace
107 // the user's task with an LLM-generated paraphrase.
108 return result, err
109 }
110 compacted, compactErr := compactContext(ctx, ag, messages, options.PersistCompaction)
111 if compactErr != nil {
112 return result, fmt.Errorf("compact transcript after initial overflow: %w", compactErr)
113 }
114
115 result.OutputStart = len(compacted)
116 result.Compacted = true
117 if ctxErr := ctx.Err(); ctxErr != nil {
118 return result, ctxErr
119 }
120 return result, ag.Continue(ctx, sink)
121}
122
123func compactContext(ctx context.Context, ag *core.Agent, messages []provider.Message, persist func([]provider.Message) error) ([]provider.Message, error) {
124 if err := ctx.Err(); err != nil {
125 return nil, err
126 }
127 keepTail := min(4, len(messages)-1)
128 if _, err := ag.Compact(ctx, keepTail, nil); err != nil {
129 return nil, err
130 }
131
132 compacted := ag.Messages()
133 if err := ctx.Err(); err != nil {
134 ag.SetMessages(messages)
135 return nil, err
136 }
137 if persist != nil {
138 if err := ctx.Err(); err != nil {
139 ag.SetMessages(messages)
140 return nil, err
141 }
142 if err := persist(compacted); err != nil {
143 // The session did not receive the compaction checkpoint. Restore the
144 // exact pre-compaction transcript so a later suffix flush cannot mix
145 // the old durable generation with the new in-memory generation.
146 ag.SetMessages(messages)
147 return nil, fmt.Errorf("persist compacted transcript: %w", err)
148 }
149 }
150 return compacted, nil
151}