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 // SuppressInitialOverflowEvent hides the recoverable first-attempt
21 // EvTurnEnd overflow from sinks whose protocol cannot retract a terminal
22 // error frame, such as JSONL. A retry overflow is never suppressed.
23 SuppressInitialOverflowEvent bool
24}
25
26// ContextRecoveryResult describes a completed host-level context recovery.
27type ContextRecoveryResult struct {
28 // OutputStart is the first message to append after a persisted compaction.
29 // When Compacted is false, it is the message count before Prompt.
30 OutputStart int
31 Compacted bool
32}
33
34// PromptWithContextRecovery appends prompt and runs it once. If that provider
35// request exceeds its context window, it compacts the existing transcript and
36// continues the already-appended prompt exactly once. Cancellation, a
37// one-message transcript, a compaction failure, and a second overflow remain
38// terminal outcomes.
39func PromptWithContextRecovery(ctx context.Context, ag *core.Agent, prompt string, images []provider.ImageBlock, sink func(core.AgentEvent), options ContextRecoveryOptions) (result ContextRecoveryResult, err error) {
40 result.OutputStart = len(ag.Messages())
41 if sink == nil {
42 sink = func(core.AgentEvent) {}
43 }
44
45 firstSink := sink
46 if options.SuppressInitialOverflowEvent {
47 firstSink = func(event core.AgentEvent) {
48 if turnEnd, ok := event.(core.EvTurnEnd); ok && provider.IsContextOverflowError(turnEnd.Err) {
49 return
50 }
51 sink(event)
52 }
53 }
54
55 if err = ag.Prompt(ctx, prompt, images, firstSink); err == nil || ctx.Err() != nil || !provider.IsContextOverflowError(err) {
56 return result, err
57 }
58
59 messages := ag.Messages()
60 if len(messages) <= 1 {
61 // The failed prompt is the only message. Summarizing it would replace
62 // the user's task with an LLM-generated paraphrase.
63 return result, err
64 }
65 keepTail := min(4, len(messages)-1)
66 if _, compactErr := ag.Compact(ctx, keepTail, nil); compactErr != nil {
67 return result, fmt.Errorf("compact transcript after initial overflow: %w", compactErr)
68 }
69
70 compacted := ag.Messages()
71 if options.PersistCompaction != nil {
72 if persistErr := options.PersistCompaction(compacted); persistErr != nil {
73 // The session did not receive the compaction checkpoint, so restore
74 // the original history plus the already-appended prompt rather than
75 // let a later transcript flush make the session inconsistent.
76 ag.SetMessages(messages)
77 return result, fmt.Errorf("persist compacted transcript: %w", persistErr)
78 }
79 }
80
81 result.OutputStart = len(compacted)
82 result.Compacted = true
83 return result, ag.Continue(ctx, sink)
84}