| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "errors" |
| 6 | "math/rand" |
| 7 | "time" |
| 8 | |
| 9 | "reasonix/internal/event" |
| 10 | "reasonix/internal/provider" |
| 11 | ) |
| 12 | |
| 13 | // defaultRecoveryWaitBudget bounds continuous waiting on an unreachable |
| 14 | // provider (#9889): mainstream agents stop after ~10 attempts or ~10 minutes. |
| 15 | const defaultRecoveryWaitBudget = 10 * time.Minute |
| 16 | |
| 17 | var recoveryWaitBudget = defaultRecoveryWaitBudget |
| 18 | |
| 19 | type samplingRecoveryState struct { |
| 20 | frozen samplingRequest |
| 21 | context contextRecoveryBudget |
| 22 | replay reasoningReplayRecoveryBudget |
| 23 | output, protocol, partial, missing bool |
| 24 | billable *provider.Usage |
| 25 | waited time.Duration |
| 26 | } |
| 27 | |
| 28 | func (a *Agent) samplingDeadline(ctx context.Context) (context.Context, context.CancelFunc, TaskBudget) { |
| 29 | limit := a.taskBudgetLimit(ctx) |
| 30 | if a.turn.graceRound { |
| 31 | limit = TaskBudget{} |
| 32 | } |
| 33 | if limit.Wall <= 0 { |
| 34 | return ctx, func() {}, limit |
| 35 | } |
| 36 | started := a.task.budget.started |
| 37 | if started.IsZero() { |
| 38 | started = a.turn.budget.started |
| 39 | } |
| 40 | if started.IsZero() { |
| 41 | started = time.Now() |
| 42 | } |
| 43 | next, cancel := context.WithDeadline(ctx, started.Add(limit.Wall)) |
| 44 | return next, cancel, limit |
| 45 | } |
| 46 | |
| 47 | func (a *Agent) streamWithSamplingRecovery(parent context.Context, turn int) (terminal streamedTurn) { |
| 48 | ctx, cancel, limit := a.samplingDeadline(parent) |
| 49 | defer cancel() |
| 50 | state := samplingRecoveryState{} |
| 51 | defer func() { |
| 52 | if limit.Wall > 0 && errors.Is(terminal.err, context.DeadlineExceeded) && parent.Err() == nil { |
| 53 | terminal.err = &taskBudgetPause{axis: "time", detail: "recovery reached the task deadline"} |
| 54 | } |
| 55 | if terminal.err == nil && state.replay.retries > 0 { |
| 56 | a.activateReasoningReplayStrongProjection(state.replay) |
| 57 | } |
| 58 | }() |
| 59 | var err error |
| 60 | state.frozen, err = a.prepareSamplingRequest(ctx) |
| 61 | if err != nil { |
| 62 | return streamedTurn{err: err} |
| 63 | } |
| 64 | if err := a.consumeManualProtocolRecovery(ctx, &state); err != nil { |
| 65 | return streamedTurn{err: err} |
| 66 | } |
| 67 | ctx = provider.WithManagedRecovery(provider.WithRequestAttemptCounter(ctx)) |
| 68 | for attempt := 1; ; attempt++ { |
| 69 | if err := a.samplingRecoveryStop(ctx, limit, state.billable, attempt); err != nil { |
| 70 | return streamedTurn{err: err, usage: state.billable} |
| 71 | } |
| 72 | if state.protocol && !state.replay.persisted { |
| 73 | record := a.protocolRecord(state.frozen, "consumed") |
| 74 | if state.replay.cutoff > 0 { |
| 75 | record.Projected = true |
| 76 | record.Prefix, record.Anchor = state.replay.cutoff, state.replay.anchor |
| 77 | } |
| 78 | if err := a.saveProtocolRecord(record); err != nil { |
| 79 | return streamedTurn{err: err, usage: state.billable} |
| 80 | } |
| 81 | state.replay.persisted = true |
| 82 | } |
| 83 | id := newStreamAttemptID(attempt) |
| 84 | a.emitStreamAttempt(id, event.StreamAttemptBegin, attempt, "", nil) |
| 85 | sink, attemptSink := a.samplingAttemptSinks() |
| 86 | a.freezeVisibleReads(state.frozen.req.Messages) |
| 87 | result := a.runSamplingAttempt(ctx, turn, attemptSink, &state.frozen, id) |
| 88 | state.billable, _ = a.recordSamplingAttempt(state.billable, result) |
| 89 | if ctx.Err() != nil { |
| 90 | // A user cancellation settles the visible prefix as local display |
| 91 | // history. Dropping the attempt here loses the only complete prefix. |
| 92 | sink.Flush() |
| 93 | result.err, result.interrupted, result.usage = ctx.Err(), true, state.billable |
| 94 | return result |
| 95 | } |
| 96 | if result.err == nil { |
| 97 | retry, done := a.handleSamplingCandidate(&state, result, sink, attempt, id) |
| 98 | if retry { |
| 99 | continue |
| 100 | } |
| 101 | return done |
| 102 | } |
| 103 | state.partial = state.partial || sawSpeculativeSamplingOutput(result) || len(result.responsesItems) > 0 || len(result.serverSearch) > 0 |
| 104 | if attempt < maxSamplingAttempts && a.trySamplingRepair(ctx, &state, result, sink, attempt, id) { |
| 105 | continue |
| 106 | } |
| 107 | if a.waitSamplingRetry(ctx, &state, &result, sink, attempt, id) { |
| 108 | continue |
| 109 | } |
| 110 | sink.Flush() |
| 111 | if !state.protocol { |
| 112 | if err := a.offerProtocolRecovery(state.frozen, result.err); err != nil { |
| 113 | result.err = err |
| 114 | } |
| 115 | } |
| 116 | if provider.AsContextLimitError(result.err) != nil { |
| 117 | a.setLastRecovery(contextRecoveryFailed) |
| 118 | } |
| 119 | result.usage = finalizeSamplingUsage(state.billable, result.usage) |
| 120 | if ctx.Err() != nil { |
| 121 | result.err = ctx.Err() |
| 122 | result.interrupted = true |
| 123 | } |
| 124 | return result |
| 125 | } |
| 126 | } |
| 127 | |
| 128 | func (a *Agent) samplingRecoveryStop(ctx context.Context, limit TaskBudget, usage *provider.Usage, attempt int) error { |
| 129 | if ctx.Err() != nil { |
| 130 | return ctx.Err() |
| 131 | } |
| 132 | if attempt <= 1 { |
| 133 | return nil |
| 134 | } |
| 135 | shadow := a.task.budget |
| 136 | if usage != nil { |
| 137 | shadow.observe(usage, a.svc.pricing) |
| 138 | } |
| 139 | if axis, detail := shadow.exceeded(limit); axis != "" { |
| 140 | return &taskBudgetPause{axis: axis, detail: detail} |
| 141 | } |
| 142 | return nil |
| 143 | } |
| 144 | |
| 145 | func (a *Agent) handleSamplingCandidate(s *samplingRecoveryState, result streamedTurn, sink *deferredStreamSink, attempt int, id string) (bool, streamedTurn) { |
| 146 | issue := a.reasoningReplayIssue(result) |
| 147 | if issue == "" { |
| 148 | a.observeMissingAssistantReasoning(result.assistantMessage(), result.reasoningComplete) |
| 149 | if s.missing { |
| 150 | a.recordRecoveredCandidate(result) |
| 151 | } |
| 152 | sink.Flush() |
| 153 | result.settledAttemptID, result.settledAttempt = id, attempt |
| 154 | result.usage = finalizeSamplingUsage(s.billable, result.usage) |
| 155 | return false, result |
| 156 | } |
| 157 | s.partial = true |
| 158 | _, claimed := a.observeMissingAssistantReasoning(result.assistantMessage(), result.reasoningComplete) |
| 159 | if (issue != ReasoningReplayMissing && issue != ReasoningReplayIncomplete) || s.protocol || a.protocolRecoverySpent() || !claimed || attempt >= maxSamplingAttempts { |
| 160 | return false, a.finishReasoningReplayOverflow(result, sink, issue, s.billable, id, attempt) |
| 161 | } |
| 162 | s.protocol, s.missing = true, true |
| 163 | event.RecordProtocolRecovery(a.svc.sink, event.ProtocolRecoveryAudit{Kind: event.ProtocolRecoveryMissingReasoningRetryAttempted}) |
| 164 | if next, ok := a.recoverReasoningReplayHistory(s.frozen, &s.replay); ok { |
| 165 | s.frozen = next |
| 166 | s.replay.local = true |
| 167 | } |
| 168 | sink.Discard() |
| 169 | a.emitStreamAttempt(id, event.StreamAttemptDiscard, attempt, "reasoning_replay", nil) |
| 170 | a.emitProtocolRetry(attempt, false) |
| 171 | return true, streamedTurn{} |
| 172 | } |
| 173 | |
| 174 | func (a *Agent) recordRecoveredCandidate(result streamedTurn) { |
| 175 | kind := event.ProtocolRecoveryMissingReasoningRetryRecovered |
| 176 | if len(result.calls) == 0 && len(result.serverSearch) == 0 { |
| 177 | kind = event.ProtocolRecoveryMissingReasoningRetryReplaced |
| 178 | } |
| 179 | event.RecordProtocolRecovery(a.svc.sink, event.ProtocolRecoveryAudit{Kind: kind}) |
| 180 | } |
| 181 | |
| 182 | func (a *Agent) trySamplingRepair(ctx context.Context, s *samplingRecoveryState, result streamedTurn, sink *deferredStreamSink, attempt int, id string) bool { |
| 183 | if limit := provider.AsOutputLimitError(result.err); !s.output && limit != nil && s.frozen.req.MaxTokens > limit.MaxOutputTokens { |
| 184 | s.output = true |
| 185 | a.learnOutputBudget(limit.MaxOutputTokens) |
| 186 | s.frozen.req.MaxTokens = limit.MaxOutputTokens |
| 187 | sink.Discard() |
| 188 | a.emitStreamAttempt(id, event.StreamAttemptDiscard, attempt, "output_limit", result.err) |
| 189 | return true |
| 190 | } |
| 191 | if next, ok, _ := a.recoverContextLimit(ctx, s.frozen, result.err, &s.context); ok { |
| 192 | sink.Discard() |
| 193 | a.emitStreamAttempt(id, event.StreamAttemptDiscard, attempt, "context_limit", result.err) |
| 194 | s.frozen = next |
| 195 | return true |
| 196 | } |
| 197 | if s.protocol { |
| 198 | return false |
| 199 | } |
| 200 | next, ok := a.tryRecoverReasoningReplay400(sink, s.frozen, id, attempt, result.err, &s.replay) |
| 201 | if ok { |
| 202 | s.protocol = true |
| 203 | s.frozen = next |
| 204 | } |
| 205 | return ok |
| 206 | } |
| 207 | |
| 208 | func (a *Agent) canWaitSampling(ctx context.Context, s *samplingRecoveryState, f provider.RecoveryFailure) bool { |
| 209 | role, _ := ctx.Value(turnContextRoleKey{}).(turnContextRole) |
| 210 | if role == turnContextPlanner { |
| 211 | return false |
| 212 | } |
| 213 | if SubagentDepth(ctx) != 0 || a.turn.graceRound || s.partial { |
| 214 | return false |
| 215 | } |
| 216 | return f.Retryable && (f.Phase == "connect" || (f.Phase == "headers" && (f.Status == 408 || f.Status == 429 || f.Status >= 500))) |
| 217 | } |
| 218 | |
| 219 | func (a *Agent) waitSamplingRetry(ctx context.Context, s *samplingRecoveryState, result *streamedTurn, sink *deferredStreamSink, attempt int, id string) bool { |
| 220 | failure := provider.ClassifyRecovery(result.err) |
| 221 | waiting := attempt >= maxSamplingAttempts && a.canWaitSampling(ctx, s, failure) |
| 222 | if !failure.Retryable || (attempt >= maxSamplingAttempts && !waiting) { |
| 223 | return false |
| 224 | } |
| 225 | base := time.Duration(1<<min(attempt-1, 2)) * 2 * time.Second |
| 226 | delay := base |
| 227 | if waiting { |
| 228 | delay = time.Minute + time.Duration(rand.Intn(6001))*time.Millisecond |
| 229 | } |
| 230 | delay = max(delay, failure.RetryAfter) |
| 231 | if waiting && s.waited+delay > recoveryWaitBudget { |
| 232 | result.err = &provider.RecoveryWaitExhaustedError{Phase: failure.Phase, Code: failure.Code, Status: failure.Status, Waited: s.waited, Attempts: attempt, Cause: result.err} |
| 233 | return false |
| 234 | } |
| 235 | sink.Discard() |
| 236 | reason := failure.Phase |
| 237 | if provider.IsStreamInterrupted(result.err) { |
| 238 | reason = provider.StreamInterruptReason(result.err) |
| 239 | } |
| 240 | a.emitStreamAttempt(id, event.StreamAttemptDiscard, attempt, reason, result.err) |
| 241 | status := &event.RecoveryStatus{Phase: failure.Phase, Reason: failure.Code, NextAttemptAt: time.Now().Add(delay).UnixMilli(), WaitedMs: s.waited.Milliseconds(), Waiting: waiting} |
| 242 | if waiting { |
| 243 | status.WaitBudgetMs = recoveryWaitBudget.Milliseconds() |
| 244 | } |
| 245 | a.svc.sink.Emit(event.Event{Kind: event.Retrying, RetryAttempt: attempt, RetryMax: maxStreamRecoveries, RetryScope: event.RetryScopeStream, Recovery: status}) |
| 246 | s.waited += delay |
| 247 | if !waiting && failure.RetryAfter <= base { |
| 248 | return streamRetrySleep(ctx, attempt) |
| 249 | } |
| 250 | return recoverySleep(ctx, delay) |
| 251 | } |
| 252 | |
| 253 | func unmeteredHeaderFailure(result streamedTurn, httpRequests int) bool { |
| 254 | if httpRequests <= 0 || sawSpeculativeSamplingOutput(result) { |
| 255 | return false |
| 256 | } |
| 257 | failure := provider.ClassifyRecovery(result.err) |
| 258 | return failure.Phase == "headers" || failure.Phase == "connect" |
| 259 | } |
| 260 | |
| 261 | func unmeteredUsage(usage *provider.Usage, result streamedTurn, httpRequests int) *provider.Usage { |
| 262 | if usage == nil && unmeteredHeaderFailure(result, httpRequests) { |
| 263 | return &provider.Usage{Unknown: true, RequestCount: httpRequests} |
| 264 | } |
| 265 | return usage |
| 266 | } |
| 267 |