diff --git a/internal/agent/loop.go b/internal/agent/loop.go index fe691ac4c..85e84ed16 100644 --- a/internal/agent/loop.go +++ b/internal/agent/loop.go @@ -25,14 +25,18 @@ const maxTurnsAnswer = "Agent reached maximum number of turns without a final an const maxTurnsFinalAnswerPrompt = "You have reached the tool-turn limit. Do not call tools. Give a concise final answer now: summarize what you completed, what you found, and any remaining blockers." // maxStreamStallRetries bounds how many times a turn that timed out (idle/stall) -// WITH NO OUTPUT yet is re-issued on a fresh connection before giving up. Only -// the no-output case is retried (a partial turn would duplicate), so this is a -// safe recovery for a stalled/dead pooled connection. +// OR hit a mid-stream transport abort (#973) with no answer text yet is re-issued +// on a fresh connection before giving up. Incomplete tool-call previews are +// allowed through the gate (they are never dispatched before the error return). +// The safety claim is only that visible answer prose is not duplicated on retry — +// not that every partial stream shape is excluded. Covers a stalled/dead pooled +// connection or a socket abort before tool dispatch. // // Set to 1 (not 2): each attempt can itself idle for the full stream timeout // (~5min) before the stall is even detected, so 2 retries left an interactive // session frozen for ~15min. One retry keeps the common single-hiccup recovery -// while bounding the worst case to ~2× the idle timeout. +// while bounding the worst case to ~2× the idle timeout. Mid-stream transport +// aborts share this bound (do not raise it for #973). const maxStreamStallRetries = 1 const ( @@ -374,7 +378,11 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) rec.StampFirstToken() } if onText != nil { - forwardedVisibleText = true + // Zero-length chunks are not committed answer prose; counting + // them would block an eligible mid-stream abort retry. + if s != "" { + forwardedVisibleText = true + } onText(s) } } @@ -450,9 +458,14 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, ctx.Err() } - // A stream idle/stall timeout is safely re-issued when the turn committed NO - // answer text — no forwarded visible prose (forwardedVisibleText) and no - // collected final text (collected.Text). This covers two cases: + // A stream idle/stall timeout OR mid-stream transport abort (#973 — wsarecv / + // connection reset / unexpected EOF / connection closed / broken pipe) is + // safely re-issued when the turn committed NO answer text — no forwarded + // visible prose (forwardedVisibleText) and no collected final text + // (collected.Text). The abort gate matches midStreamAbortNeedles only, not + // the full shouldReconnect list: connect-phase timeouts and connection + // refused are NOT retried here (a slow healthy server must not cost a + // second prefill). This covers three cases: // 1. Nothing streamed at all before the connection died (the original // macOS stale-pooled-connection hang past the response-header timeout). // 2. The model streamed transient reasoning and began a tool call (e.g. a @@ -463,13 +476,27 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) // re-render is transient previews, not duplicated answer text. This is // why the gate no longer excludes collected.ToolCalls: an incomplete // tool call from a timed-out stream is discard-and-retry, not output. + // 3. Mid-stream socket abort after connect succeeded (Windows WSAECONNABORTED / + // wsarecv, connection reset by peer). Same no-tool-executed safety: the + // error returns before tool dispatch, so a bounded retry is safe. // A turn that forwarded real prose is NOT retried (it would duplicate visible // answer text) and falls through to the error return below. Capped + // exponential backoff, with a user-visible notice per attempt. for attempt := 1; attempt <= maxStreamStallRetries && - isStreamTimeoutError(collected.Error) && !forwardedVisibleText && + (isStreamTimeoutError(collected.Error) || isMidStreamTransportAbort(collected.Error)) && + !forwardedVisibleText && collected.Text == ""; attempt++ { - if notify := stallRetryNoticeFor(options); notify != nil { + var notify reconnectNotifier + if isStreamTimeoutError(collected.Error) { + notify = stallRetryNoticeFor(options) + } else { + notify = reconnectNoticeFor(options) + // Count this post-connect reissue even when the replacement + // connect succeeds immediately. streamWithReconnect only + // increments reconnect_count after a failed connect attempt. + options.Trace.Counter(trace.CounterReconnectCount, 1) + } + if notify != nil { notify(attempt, maxStreamStallRetries) } if err := sleepWithContext(ctx, backoffFor(attempt)); err != nil { @@ -482,12 +509,22 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) retryStream, retryErr := streamWithReconnect(ctx, provider, retryRequest, reconnectNoticeFor(options)) if retryErr != nil { result.Messages = copyMessages(messages) + if ctx.Err() != nil { + return result, ctx.Err() + } return result, retryErr } stallGenSpan := options.Trace.Span(trace.SpanGeneration) collected = zeroruntime.CollectStreamWithOptions(ctx, retryStream, forwardingOpts) stallGenSpan.End() } + // Recheck ctx after a retried CollectStream: helpers.go stores + // ctx.Err().Error() in collected.Error, and errors.New of that string + // drops the context.Canceled sentinel ACP branches on. + if ctx.Err() != nil { + result.Messages = copyMessages(messages) + return result, ctx.Err() + } if collected.Error != "" { // Route a reissued stream's non-stall error through the SAME recovery as // the initial stream (image-rejection wrapping / context-limit compaction) @@ -498,6 +535,10 @@ func Run(ctx context.Context, prompt string, provider Provider, options Options) result.Messages = copyMessages(messages) return result, stop } + if ctx.Err() != nil { + result.Messages = copyMessages(messages) + return result, ctx.Err() + } if collected.Error != "" { result.Messages = copyMessages(messages) return result, errors.New(collected.Error) diff --git a/internal/agent/midstream_retry_test.go b/internal/agent/midstream_retry_test.go new file mode 100644 index 000000000..13f0839f6 --- /dev/null +++ b/internal/agent/midstream_retry_test.go @@ -0,0 +1,454 @@ +package agent + +import ( + "context" + "errors" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/Gitlawb/zero/internal/tools" + "github.com/Gitlawb/zero/internal/trace" + "github.com/Gitlawb/zero/internal/zeroruntime" +) + +// midStreamAbortProvider connects successfully but emits a transport-abort +// StreamEventError on the first abortBefore calls, then succeeds with "done". +// hangOnCall, if > 0, makes that 1-based call block until ctx is done so a +// cancel during the retried CollectStream can be reproduced. +// connectFailFrom, if > 0, makes that 1-based call and later return a +// reconnectable connect error so streamWithReconnect enters its backoff. +type midStreamAbortProvider struct { + calls int32 + abortBefore int32 + abortError string + partialText string + emptyTextEvent bool + partialToolCall string + hangOnCall int32 + connectFailFrom int32 + connectErr string + started chan struct{} +} + +func (p *midStreamAbortProvider) StreamCompletion(ctx context.Context, _ zeroruntime.CompletionRequest) (<-chan zeroruntime.StreamEvent, error) { + n := atomic.AddInt32(&p.calls, 1) + if p.connectFailFrom > 0 && n >= p.connectFailFrom { + if p.started != nil { + select { + case <-p.started: + default: + close(p.started) + } + } + errMsg := p.connectErr + if errMsg == "" { + errMsg = "connection reset by peer" + } + return nil, errors.New(errMsg) + } + if p.hangOnCall > 0 && n == p.hangOnCall { + if p.started != nil { + close(p.started) + } + ch := make(chan zeroruntime.StreamEvent) + go func() { + <-ctx.Done() + close(ch) + }() + return ch, nil + } + ch := make(chan zeroruntime.StreamEvent, 5) + if n <= p.abortBefore { + if p.emptyTextEvent { + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: ""} + } + if p.partialText != "" { + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventText, Content: p.partialText} + } + if p.partialToolCall != "" { + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallStart, ToolCallID: "tc_1", ToolName: p.partialToolCall} + ch <- zeroruntime.StreamEvent{Type: zeroruntime.StreamEventToolCallDelta, ToolCallID: "tc_1", ArgumentsFragment: `{"path":"x.html","content":"= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' +} + // backoffFor is the deterministic exponential base delay for a 1-based attempt, // capped at streamReconnectMax. Jitter is layered on separately (jitteredBackoff). func backoffFor(attempt int) time.Duration { diff --git a/internal/agent/reconnect_test.go b/internal/agent/reconnect_test.go index 84bd3a3ff..cc4d762bf 100644 --- a/internal/agent/reconnect_test.go +++ b/internal/agent/reconnect_test.go @@ -92,6 +92,10 @@ func TestShouldReconnectClassification(t *testing.T) { disconnects := []string{ "unexpected EOF", "connection reset by peer", "broken pipe", "i/o timeout", "server closed the connection", "connection refused", + // Windows socket aborts also remain reconnectable at connect time. + "wsarecv: An established connection was aborted by the software in your host machine", + "read: connection was aborted", + "An existing connection was forcibly closed by the remote host", } for _, m := range disconnects { if !shouldReconnect(ctx, errors.New(m)) {