From 0b2eeaadd0f2d93af233026e906a0214de63ac48 Mon Sep 17 00:00:00 2001 From: Mateusz Date: Fri, 14 Aug 2026 14:37:31 +0200 Subject: [PATCH] fix: retry mid-turn 429s with backoff and native continue Keep the ACP prompt open across Gemini quota errors: wait using retry-after or exponential backoff, resume the same agy conversation, and only return -32003 after the attempt budget is exhausted. Co-authored-by: Cursor --- .release-files | 2 + README.md | 22 +++++ internal/acp/agent.go | 144 ++++++++++++++++++++++++++++--- internal/acp/agent_test.go | 147 +++++++++++++++++++++++++++++++- internal/agy/retry.go | 149 +++++++++++++++++++++++++++++++++ internal/agy/retry_test.go | 95 +++++++++++++++++++++ internal/agy/runner.go | 3 +- internal/config/config.go | 59 ++++++++----- internal/config/config_test.go | 36 ++++++++ 9 files changed, 621 insertions(+), 36 deletions(-) create mode 100644 internal/agy/retry.go create mode 100644 internal/agy/retry_test.go diff --git a/.release-files b/.release-files index c85e7f0..0960105 100644 --- a/.release-files +++ b/.release-files @@ -32,6 +32,8 @@ internal/agy/process_unix.go internal/agy/process_unix_test.go internal/agy/process_windows.go internal/agy/process_windows_test.go +internal/agy/retry.go +internal/agy/retry_test.go internal/agy/runner.go internal/agy/runner_test.go internal/agy/signal_unix.go diff --git a/README.md b/README.md index 70e5ece..d40110d 100644 --- a/README.md +++ b/README.md @@ -49,6 +49,24 @@ Prompts exceeding the configurable byte threshold (default 8KB) are written unde first session for a workdir starts, when the last session in that workdir closes, and on wrapper shutdown. +### Provider 429 / quota continue + +A mid-turn Gemini `429` / `RESOURCE_EXHAUSTED` used to fail the ACP `session/prompt` +immediately (or, before fail-fast, hang until `--timeout-seconds`). The wrapper now +keeps that prompt open and retries the **native** `agy --conversation` turn: + +1. Parse `retryDelay` / "retry in …" / `Retry-After` when present; otherwise use + exponential backoff (`2s` … `60s`). +2. Wait up to the remaining per-turn timeout (or `--quota-retry-max-wait-seconds`). +3. Stream a short "waiting … before retry" notice so the ACP client is not silent. +4. If the failed attempt already produced a conversation ID, send a continue cue + instead of replaying the user prompt. +5. After the attempt budget is exhausted, still return ACP `-32003` and **do not** + dump fallback context (that would burn more quota). + +This is independent of the virtual-context fallback used for other native-resume +failures. + ### Response Extraction When `agy --print` produces no stdout (a known platform-specific issue), the wrapper extracts the model's response directly from agy's transcript log at @@ -116,6 +134,8 @@ Use `go-agy-acp-wrapper --version` for executable validation without starting AC | `AGY_MODEL` | _(empty = catalog default)_ | Default model for new sessions (e.g. `gemini-3.6-flash-high`) | | `AGY_PROMPT_THRESHOLD` | `8000` | Byte threshold above which prompts are written to temp files | | `AGY_TIMEOUT_SECONDS` | `14400` (4 hours) | Per-turn execution timeout in seconds. Forwarded to agy as `--print-timeout` (agy's own print-mode default is 5 minutes). | +| `AGY_QUOTA_RETRY_ATTEMPTS` | `5` | How many `agy --print` invocations to make for one ACP prompt after a 429 / quota error (1 initial + retries). | +| `AGY_QUOTA_RETRY_MAX_WAIT_SECONDS` | `0` (remaining turn timeout) | Cap for a single backoff wait. `0` waits up to the remaining per-turn timeout (so a 3h quota reset can still resume inside a 4h turn). | | `AGY_SKIP_PERMISSIONS` | `true` | Whether to pass `--dangerously-skip-permissions` to `agy`; set to `false` to opt out | | `AGY_ACP_SKIP_ENV_NOTE` | `false` | Set to `true` to skip the once-per-session execution-environment steering note | @@ -127,6 +147,8 @@ Equivalent CLI flags are available and override environment values: | `--model ` | Default model for new sessions | | `--prompt-threshold ` | Byte threshold above which prompts are written to workdir files | | `--timeout-seconds ` | Per-turn execution timeout and agy `--print-timeout` | +| `--quota-retry-attempts ` | Native continue retries after 429 / quota errors | +| `--quota-retry-max-wait-seconds ` | Cap for a single 429 backoff (`0` = remaining turn timeout) | | `--skip-permissions` | Force-enable `--dangerously-skip-permissions` | | `--no-skip-permissions` | Opt out of `--dangerously-skip-permissions` | | `--execution-env-note` | Force-enable the once-per-session execution-environment note | diff --git a/internal/acp/agent.go b/internal/acp/agent.go index 51586f1..b14d293 100644 --- a/internal/acp/agent.go +++ b/internal/acp/agent.go @@ -28,6 +28,7 @@ type AgyAgent struct { modelCatalog *agy.ModelCatalog promptWriter *agy.PromptFileWriter coordinator *agy.RepoCoordinator + sleep func(context.Context, time.Duration) error mu sync.Mutex workdirs map[string]int cancels map[string]activePrompt @@ -55,6 +56,7 @@ func NewAgyAgent(cfg *config.Config) *AgyAgent { modelCatalog: agy.NewModelCatalog(cfg.AgyBinary), promptWriter: agy.NewPromptFileWriter(cfg.PromptThreshold), coordinator: agy.NewRepoCoordinator(cfg.AgyConfigDir()), + sleep: sleepContext, workdirs: make(map[string]int), cancels: make(map[string]activePrompt), finalizers: make(map[string]*sessionFinalizer), @@ -173,6 +175,9 @@ func (a *AgyAgent) Prompt(ctx context.Context, params acp.PromptRequest) (acp.Pr if !ok { return } + if a.conn == nil { + return + } if err := a.conn.SessionUpdate(promptCtx, acp.SessionNotification{ SessionId: params.SessionId, Update: update, @@ -205,7 +210,7 @@ func (a *AgyAgent) Prompt(ctx context.Context, params acp.PromptRequest) (acp.Pr tail = response[len(streamedText):] } - if tail != "" { + if tail != "" && a.conn != nil { if err := a.conn.SessionUpdate(ctx, acp.SessionNotification{ SessionId: params.SessionId, Update: acp.UpdateAgentMessageText(tail), @@ -218,6 +223,99 @@ func (a *AgyAgent) Prompt(ctx context.Context, params acp.PromptRequest) (acp.Pr } func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promptText string, onEvent func(agy.StreamEvent)) (string, error) { + attempts := a.quotaRetryAttempts() + turnDeadline := time.Now().Add(time.Duration(a.cfg.TimeoutSeconds) * time.Second) + if deadline, ok := ctx.Deadline(); ok && deadline.Before(turnDeadline) { + turnDeadline = deadline + } + + prompt := promptText + var native nativeAttemptState + var lastErr error + for attempt := 0; attempt < attempts; attempt++ { + remaining := time.Until(turnDeadline) + if remaining < 0 { + remaining = 0 + } + timeout := remaining + configured := time.Duration(a.cfg.TimeoutSeconds) * time.Second + if configured > 0 && (timeout == 0 || timeout > configured) { + timeout = configured + } + out, err := a.executeTurnOnce(ctx, sess, prompt, timeout, onEvent, &native) + if err == nil { + return out, nil + } + lastErr = err + if ctx.Err() != nil { + return "", err + } + if !agy.IsProviderLimitError(err) || attempt+1 >= attempts { + return "", err + } + wait, ok := agy.NextProviderLimitWait(attempt, err, remaining, a.quotaRetryMaxWait()) + if !ok { + return "", err + } + notice := fmt.Sprintf("\nProvider rate limit; waiting %s before retry %d/%d.\n", wait.Round(time.Second), attempt+2, attempts) + if onEvent != nil { + onEvent(agy.StreamEvent{Kind: agy.StreamEventText, Text: notice}) + } + slog.Warn("provider rate limit, backing off before native continue", + "sessionId", sess.ID, + "wait", wait, + "attempt", attempt+2, + "attempts", attempts, + "error", err, + ) + if sleepErr := a.sleep(ctx, wait); sleepErr != nil { + return "", sleepErr + } + if sess.GetConversationID() != "" { + prompt = providerLimitContinuePrompt + } + } + return "", lastErr +} + +const providerLimitContinuePrompt = "Continue from where you left off. Do not restart the task." + +type nativeAttemptState struct { + previousID string + startedAt time.Time + snapshotDone bool + snapshotFailed bool +} + +func (a *AgyAgent) quotaRetryAttempts() int { + if a.cfg == nil || a.cfg.QuotaRetryAttempts < 1 { + return config.DefaultQuotaRetryAttempts + } + return a.cfg.QuotaRetryAttempts +} + +func (a *AgyAgent) quotaRetryMaxWait() time.Duration { + if a.cfg == nil || a.cfg.QuotaRetryMaxWaitSeconds <= 0 { + return 0 + } + return time.Duration(a.cfg.QuotaRetryMaxWaitSeconds) * time.Second +} + +func sleepContext(ctx context.Context, d time.Duration) error { + if d <= 0 { + return ctx.Err() + } + timer := time.NewTimer(d) + defer timer.Stop() + select { + case <-ctx.Done(): + return ctx.Err() + case <-timer.C: + return nil + } +} + +func (a *AgyAgent) executeTurnOnce(ctx context.Context, sess *session.Context, promptText string, timeout time.Duration, onEvent func(agy.StreamEvent), native *nativeAttemptState) (string, error) { mode := sess.GetMode() convID := sess.GetConversationID() turnCount := sess.GetTurnCount() @@ -230,7 +328,7 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp opts := agy.ExecuteOpts{ Cwd: sess.Cwd, Model: nativeModel, - Timeout: time.Duration(a.cfg.TimeoutSeconds) * time.Second, + Timeout: timeout, SkipPerms: a.cfg.SkipPerms, } @@ -238,7 +336,7 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp case mode == session.ModeFallbackContext: return a.executeFallbackTurn(ctx, sess, opts, promptText, onEvent) - case convID != "" && turnCount > 1: + case convID != "": opts.ConversationID = convID opts.Prompt = promptText if a.promptWriter.NeedsFile(promptText) { @@ -251,6 +349,7 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp } resp, err := a.runner.ExecuteStream(ctx, opts, onEvent) + captureConversationID(sess, resp) out, err := turnOutput(resp, err, opts) if err != nil { if !agy.ShouldFallback(err) { @@ -263,13 +362,22 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp return out, nil default: - previousID, err := a.discoverer.SnapshotConversationID(ctx, sess.Cwd) - if err != nil { - slog.Warn("conversation cache snapshot failed, using fallback context", "error", err, "sessionId", sess.ID) + if native.snapshotFailed { sess.SwitchToFallback() return a.executeFallbackTurn(ctx, sess, opts, promptText, onEvent) } - startedAt := time.Now() + if !native.snapshotDone { + previousID, snapErr := a.discoverer.SnapshotConversationID(ctx, sess.Cwd) + native.snapshotDone = true + if snapErr != nil { + slog.Warn("conversation cache snapshot failed, using fallback context", "error", snapErr, "sessionId", sess.ID) + native.snapshotFailed = true + sess.SwitchToFallback() + return a.executeFallbackTurn(ctx, sess, opts, promptText, onEvent) + } + native.previousID = previousID + native.startedAt = time.Now() + } opts.Prompt = promptText if a.promptWriter.NeedsFile(promptText) { path, err := a.promptWriter.WritePromptFile(sess.Cwd, sess.ID, turnCount, promptText) @@ -281,16 +389,17 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp } resp, err := a.runner.ExecuteStream(ctx, opts, onEvent) + captureConversationID(sess, resp) out, err := turnOutput(resp, err, opts) if err != nil { + if sess.GetConversationID() == "" && !native.startedAt.IsZero() { + _ = a.discoverAndSetConversationID(ctx, sess, native.previousID, native.startedAt) + } return "", err } - if convID == "" { - if resp.ConversationID != "" { - sess.SetConversationID(resp.ConversationID) - slog.Info("conversation ID received from agy stream", "conversationId", resp.ConversationID, "sessionId", sess.ID) - } else if !a.discoverAndSetConversationID(ctx, sess, previousID, startedAt) { + if sess.GetConversationID() == "" { + if !a.discoverAndSetConversationID(ctx, sess, native.previousID, native.startedAt) { sess.SwitchToFallback() } } @@ -299,6 +408,17 @@ func (a *AgyAgent) executeTurn(ctx context.Context, sess *session.Context, promp } } +func captureConversationID(sess *session.Context, resp *agy.Response) { + if resp == nil || strings.TrimSpace(resp.ConversationID) == "" { + return + } + if sess.GetConversationID() == resp.ConversationID { + return + } + sess.SetConversationID(resp.ConversationID) + slog.Info("conversation ID received from agy stream", "conversationId", resp.ConversationID, "sessionId", sess.ID) +} + func (a *AgyAgent) executeFallbackTurn(ctx context.Context, sess *session.Context, opts agy.ExecuteOpts, promptText string, onEvent func(agy.StreamEvent)) (string, error) { transcript := sess.GetTranscript() turnCount := sess.GetTurnCount() diff --git a/internal/acp/agent_test.go b/internal/acp/agent_test.go index f2cca03..ceb2a40 100644 --- a/internal/acp/agent_test.go +++ b/internal/acp/agent_test.go @@ -42,6 +42,9 @@ func newTestAgent(t *testing.T) *AgyAgent { "claude-sonnet-4-6", "claude-opus-4-6-thinking", ) + agent.sleep = func(ctx context.Context, d time.Duration) error { + return ctx.Err() + } t.Cleanup(agent.Close) return agent } @@ -345,6 +348,32 @@ func (r *errorRunner) ExecuteStream(ctx context.Context, opts agy.ExecuteOpts, o return nil, r.err } +type scriptedCall struct { + resp *agy.Response + err error +} + +type scriptedRunner struct { + mu sync.Mutex + calls []agy.ExecuteOpts + script []scriptedCall +} + +func (r *scriptedRunner) Execute(ctx context.Context, opts agy.ExecuteOpts) (*agy.Response, error) { + return r.ExecuteStream(ctx, opts, nil) +} + +func (r *scriptedRunner) ExecuteStream(ctx context.Context, opts agy.ExecuteOpts, onEvent func(agy.StreamEvent)) (*agy.Response, error) { + r.mu.Lock() + defer r.mu.Unlock() + idx := len(r.calls) + r.calls = append(r.calls, opts) + if idx >= len(r.script) { + return nil, fmt.Errorf("unexpected extra ExecuteStream call %d", idx) + } + return r.script[idx].resp, r.script[idx].err +} + type recordingRunner struct { mu sync.Mutex calls []agy.ExecuteOpts @@ -512,14 +541,128 @@ func TestAgyAgent_Prompt_ReturnsDescriptiveQuotaRequestError(t *testing.T) { if !strings.Contains(fmt.Sprint(requestErr.Data), "Gemini quota exceeded") { t.Fatalf("expected provider detail in request error data: %#v", requestErr.Data) } - if runner.calls != 1 { - t.Fatalf("quota failure should not be retried in fallback mode; got %d calls", runner.calls) + if runner.calls != config.DefaultQuotaRetryAttempts { + t.Fatalf("quota failure should retry native conversation, not fallback; got %d calls", runner.calls) } if sess.GetMode() != session.ModeNativeConversation { t.Fatal("quota failure unexpectedly switched session to fallback mode") } } +func TestAgyAgent_Prompt_RetriesQuotaThenContinuesNativeConversation(t *testing.T) { + agent := newTestAgent(t) + var waits []time.Duration + agent.sleep = func(ctx context.Context, d time.Duration) error { + waits = append(waits, d) + return ctx.Err() + } + runner := &scriptedRunner{script: []scriptedCall{ + {resp: &agy.Response{ConversationID: "conv-123"}, err: &agy.ProcessError{ExitCode: 1, Detail: "status 429: overloaded"}}, + {resp: &agy.Response{Output: "resumed", ConversationID: "conv-123"}}, + }} + agent.runner = runner + + sessResp, err := agent.NewSession(context.Background(), acp.NewSessionRequest{Cwd: t.TempDir()}) + if err != nil { + t.Fatalf("NewSession failed: %v", err) + } + sess, ok := agent.store.Get(string(sessResp.SessionId)) + if !ok { + t.Fatal("session not found") + } + sess.SetConversationID("conv-123") + sess.AddUserMessage("previous") + sess.AddAssistantMessage("answer") + + resp, err := agent.Prompt(context.Background(), acp.PromptRequest{ + SessionId: sessResp.SessionId, + Prompt: []acp.ContentBlock{acp.TextBlock("next")}, + }) + if err != nil { + t.Fatalf("Prompt failed after retry: %v", err) + } + if resp.StopReason != acp.StopReasonEndTurn { + t.Fatalf("stop reason = %s", resp.StopReason) + } + if len(runner.calls) != 2 { + t.Fatalf("expected 2 native calls, got %d", len(runner.calls)) + } + if runner.calls[0].Prompt == "" || !strings.Contains(runner.calls[0].Prompt, "next") { + t.Fatalf("first prompt should keep the user text, got %q", runner.calls[0].Prompt) + } + if runner.calls[1].Prompt != providerLimitContinuePrompt { + t.Fatalf("retry prompt = %q, want continue cue", runner.calls[1].Prompt) + } + if runner.calls[1].ConversationID != "conv-123" { + t.Fatalf("retry conversation = %q", runner.calls[1].ConversationID) + } + if len(waits) != 1 || waits[0] != 2*time.Second { + t.Fatalf("backoff waits = %#v", waits) + } +} + +func TestAgyAgent_Prompt_FirstTurnQuotaRetriesOriginalPromptWithoutConversation(t *testing.T) { + agent := newTestAgent(t) + agent.cfg.InjectExecutionEnvNote = false + runner := &scriptedRunner{script: []scriptedCall{ + {err: &agy.ProcessError{ExitCode: 1, Detail: "code 429"}}, + {resp: &agy.Response{Output: "ok", ConversationID: "conv-new"}}, + }} + agent.runner = runner + + sessResp, err := agent.NewSession(context.Background(), acp.NewSessionRequest{Cwd: t.TempDir()}) + if err != nil { + t.Fatalf("NewSession failed: %v", err) + } + _, err = agent.Prompt(context.Background(), acp.PromptRequest{ + SessionId: sessResp.SessionId, + Prompt: []acp.ContentBlock{acp.TextBlock("start work")}, + }) + if err != nil { + t.Fatalf("Prompt failed after retry: %v", err) + } + if len(runner.calls) != 2 { + t.Fatalf("expected 2 calls, got %d", len(runner.calls)) + } + if runner.calls[0].ConversationID != "" || runner.calls[1].ConversationID != "" { + t.Fatalf("first-turn retries without a conversation id must not pass --conversation: %#v", runner.calls) + } + if runner.calls[0].Prompt != "start work" || runner.calls[1].Prompt != "start work" { + t.Fatalf("expected original prompt on both attempts, got %#v %#v", runner.calls[0].Prompt, runner.calls[1].Prompt) + } +} + +func TestAgyAgent_Prompt_FirstTurnQuotaContinuesAfterStreamConversationID(t *testing.T) { + agent := newTestAgent(t) + agent.cfg.InjectExecutionEnvNote = false + runner := &scriptedRunner{script: []scriptedCall{ + {resp: &agy.Response{ConversationID: "conv-from-stream"}, err: &agy.ProcessError{ExitCode: 1, Detail: "RESOURCE_EXHAUSTED"}}, + {resp: &agy.Response{Output: "ok", ConversationID: "conv-from-stream"}}, + }} + agent.runner = runner + + sessResp, err := agent.NewSession(context.Background(), acp.NewSessionRequest{Cwd: t.TempDir()}) + if err != nil { + t.Fatalf("NewSession failed: %v", err) + } + _, err = agent.Prompt(context.Background(), acp.PromptRequest{ + SessionId: sessResp.SessionId, + Prompt: []acp.ContentBlock{acp.TextBlock("start work")}, + }) + if err != nil { + t.Fatalf("Prompt failed after retry: %v", err) + } + if len(runner.calls) != 2 { + t.Fatalf("expected 2 calls, got %d", len(runner.calls)) + } + if runner.calls[0].ConversationID != "" { + t.Fatalf("first attempt should create a conversation, got %q", runner.calls[0].ConversationID) + } + if runner.calls[1].ConversationID != "conv-from-stream" || runner.calls[1].Prompt != providerLimitContinuePrompt { + t.Fatalf("second attempt = %#v", runner.calls[1]) + } +} + func TestAgyAgent_Prompt_TimeoutDoesNotFallback(t *testing.T) { agent := newTestAgent(t) runner := &errorRunner{err: &agy.TimeoutError{Timeout: 4 * time.Hour, Detail: "print timeout"}} diff --git a/internal/agy/retry.go b/internal/agy/retry.go new file mode 100644 index 0000000..75ab940 --- /dev/null +++ b/internal/agy/retry.go @@ -0,0 +1,149 @@ +package agy + +import ( + "errors" + "regexp" + "strconv" + "strings" + "time" +) + +const ( + DefaultQuotaRetryAttempts = 5 + minProviderLimitBackoff = 2 * time.Second + maxProviderLimitBackoff = 60 * time.Second + maxParsedRetryAfter = 24 * time.Hour + retryRunReserve = 2 * time.Second +) + +var ( + retryDelayJSON = regexp.MustCompile(`(?i)retrydelay"\s*:\s*"?([0-9]+(?:\.[0-9]+)?)(ms|s|m|h)?"?`) + retryInGoDur = regexp.MustCompile(`(?i)(?:please\s+)?retry\s+in\s+((?:[0-9]+(?:\.[0-9]+)?(?:ms|h|m|s))+)\b`) + retryInWords = regexp.MustCompile(`(?i)(?:please\s+)?retry\s+in\s+([0-9]+(?:\.[0-9]+)?)\s*(milliseconds?|seconds?|minutes?|hours?|ms|h|m|s)\b`) + retryAfterHdr = regexp.MustCompile(`(?i)retry-after[:\s]+([0-9]+)`) + resetInGoDur = regexp.MustCompile(`(?i)reset(?:s)?\s+in\s+((?:[0-9]+(?:\.[0-9]+)?(?:ms|h|m|s))+)\b`) +) + +// ParseRetryAfter extracts a provider-suggested wait from quota or 429 text. +func ParseRetryAfter(text string) (time.Duration, bool) { + text = strings.TrimSpace(text) + if text == "" { + return 0, false + } + if wait, ok := parseRetryAfterMatch(retryDelayJSON.FindStringSubmatch(text), true); ok { + return wait, true + } + if wait, ok := parseGoDurationCapture(retryInGoDur.FindStringSubmatch(text)); ok { + return wait, true + } + if wait, ok := parseRetryAfterMatch(retryInWords.FindStringSubmatch(text), false); ok { + return wait, true + } + if wait, ok := parseGoDurationCapture(resetInGoDur.FindStringSubmatch(text)); ok { + return wait, true + } + if matches := retryAfterHdr.FindStringSubmatch(text); len(matches) == 2 { + seconds, err := strconv.Atoi(matches[1]) + if err == nil && seconds > 0 { + return clampRetryAfter(time.Duration(seconds) * time.Second), true + } + } + return 0, false +} + +func parseGoDurationCapture(matches []string) (time.Duration, bool) { + if len(matches) < 2 { + return 0, false + } + wait, err := time.ParseDuration(strings.ToLower(matches[1])) + if err != nil || wait <= 0 { + return 0, false + } + return clampRetryAfter(wait), true +} + +func parseRetryAfterMatch(matches []string, jsonDelay bool) (time.Duration, bool) { + if len(matches) < 2 { + return 0, false + } + amount, err := strconv.ParseFloat(matches[1], 64) + if err != nil || amount <= 0 { + return 0, false + } + unit := "s" + if len(matches) >= 3 && matches[2] != "" { + unit = strings.ToLower(matches[2]) + } else if !jsonDelay { + return 0, false + } + wait := durationFromUnit(amount, unit) + if wait <= 0 { + return 0, false + } + return clampRetryAfter(wait), true +} + +func durationFromUnit(amount float64, unit string) time.Duration { + switch unit { + case "ms", "millisecond", "milliseconds": + return time.Duration(amount * float64(time.Millisecond)) + case "s", "second", "seconds": + return time.Duration(amount * float64(time.Second)) + case "m", "minute", "minutes": + return time.Duration(amount * float64(time.Minute)) + case "h", "hour", "hours": + return time.Duration(amount * float64(time.Hour)) + default: + return 0 + } +} + +func clampRetryAfter(wait time.Duration) time.Duration { + if wait > maxParsedRetryAfter { + return maxParsedRetryAfter + } + return wait +} + +// NextProviderLimitWait chooses the backoff after a provider-limit failure. +// remaining is time left on the ACP turn. maxWait 0 means "no extra cap". +func NextProviderLimitWait(failedAttempt int, err error, remaining, maxWait time.Duration) (time.Duration, bool) { + if !IsProviderLimitError(err) { + return 0, false + } + if remaining < minProviderLimitBackoff+retryRunReserve { + return 0, false + } + wait := minProviderLimitBackoff << failedAttempt + if wait <= 0 || wait > maxProviderLimitBackoff { + wait = maxProviderLimitBackoff + } + if parsed, ok := ParseRetryAfter(errorDetail(err)); ok { + wait = parsed + } + if maxWait > 0 && wait > maxWait { + wait = maxWait + } + budget := remaining - retryRunReserve + if budget < minProviderLimitBackoff { + return 0, false + } + if wait > budget { + wait = budget + } + if wait < minProviderLimitBackoff { + wait = minProviderLimitBackoff + if wait > budget { + return 0, false + } + } + return wait, true +} + +func errorDetail(err error) string { + var processErr *ProcessError + if errors.As(err, &processErr) { + return processErr.Detail + "\n" + processErr.Error() + } + return err.Error() +} diff --git a/internal/agy/retry_test.go b/internal/agy/retry_test.go new file mode 100644 index 0000000..9c7a867 --- /dev/null +++ b/internal/agy/retry_test.go @@ -0,0 +1,95 @@ +package agy + +import ( + "testing" + "time" +) + +func TestParseRetryAfter(t *testing.T) { + tests := []struct { + name string + in string + want time.Duration + ok bool + }{ + {name: "empty", in: "", ok: false}, + {name: "json retryDelay seconds", in: `"retryDelay":"58s"`, want: 58 * time.Second, ok: true}, + {name: "json retryDelay number as seconds", in: `retryDelay": 12`, want: 12 * time.Second, ok: true}, + {name: "please retry in go duration", in: "RESOURCE_EXHAUSTED Please retry in 3h25m", want: 3*time.Hour + 25*time.Minute, ok: true}, + {name: "retry in words", in: "please retry in 2 hours", want: 2 * time.Hour, ok: true}, + {name: "retry-after header seconds", in: "Retry-After: 120", want: 120 * time.Second, ok: true}, + {name: "reset in duration", in: "individual quota resets in 2h59m1s", want: 2*time.Hour + 59*time.Minute + time.Second, ok: true}, + {name: "unrelated", in: "authentication failed", ok: false}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, ok := ParseRetryAfter(tt.in) + if ok != tt.ok { + t.Fatalf("ok=%v want %v (got %s)", ok, tt.ok, got) + } + if got != tt.want { + t.Fatalf("wait=%s want %s", got, tt.want) + } + }) + } +} + +func TestNextProviderLimitWait_ExponentialWhenNoRetryAfter(t *testing.T) { + err := &ProcessError{Detail: "status 429: model API overloaded"} + remaining := 10 * time.Minute + wait, ok := NextProviderLimitWait(0, err, remaining, 0) + if !ok || wait != 2*time.Second { + t.Fatalf("attempt 0: wait=%s ok=%v", wait, ok) + } + wait, ok = NextProviderLimitWait(1, err, remaining, 0) + if !ok || wait != 4*time.Second { + t.Fatalf("attempt 1: wait=%s ok=%v", wait, ok) + } + wait, ok = NextProviderLimitWait(8, err, remaining, 0) + if !ok || wait != maxProviderLimitBackoff { + t.Fatalf("capped wait=%s ok=%v", wait, ok) + } +} + +func TestNextProviderLimitWait_UsesRetryAfterWithinBudget(t *testing.T) { + err := &ProcessError{Detail: `RESOURCE_EXHAUSTED Please retry in 3h0m0s`} + remaining := 4 * time.Hour + wait, ok := NextProviderLimitWait(0, err, remaining, 0) + if !ok || wait != 3*time.Hour { + t.Fatalf("wait=%s ok=%v", wait, ok) + } +} + +func TestNextProviderLimitWait_ClampsToRemainingBudget(t *testing.T) { + err := &ProcessError{Detail: `RESOURCE_EXHAUSTED Please retry in 3h0m0s`} + remaining := 10 * time.Minute + wait, ok := NextProviderLimitWait(0, err, remaining, 0) + if !ok { + t.Fatal("expected a clamped wait") + } + if wait != remaining-retryRunReserve { + t.Fatalf("wait=%s want %s", wait, remaining-retryRunReserve) + } +} + +func TestNextProviderLimitWait_HonorsMaxWaitCap(t *testing.T) { + err := &ProcessError{Detail: `RESOURCE_EXHAUSTED Please retry in 3h0m0s`} + wait, ok := NextProviderLimitWait(0, err, 4*time.Hour, 15*time.Minute) + if !ok || wait != 15*time.Minute { + t.Fatalf("wait=%s ok=%v", wait, ok) + } +} + +func TestNextProviderLimitWait_GivesUpWhenTurnAlmostOver(t *testing.T) { + err := &ProcessError{Detail: "status 429"} + if _, ok := NextProviderLimitWait(0, err, 3*time.Second, 0); ok { + t.Fatal("expected no wait when remaining time cannot cover backoff plus retry") + } +} + +func TestNextProviderLimitWait_IgnoresNonQuotaErrors(t *testing.T) { + err := &ProcessError{Detail: "authentication failed"} + if _, ok := NextProviderLimitWait(0, err, time.Hour, 0); ok { + t.Fatal("unexpected wait for non-quota error") + } +} diff --git a/internal/agy/runner.go b/internal/agy/runner.go index f44bf76..a4bc83e 100644 --- a/internal/agy/runner.go +++ b/internal/agy/runner.go @@ -158,7 +158,8 @@ func IsTimeoutError(err error) bool { // ShouldFallback reports whether a failed native --conversation turn should be // retried with a dumped virtual context. Timeouts and provider limits must not -// be retried: a retry would start another long wait or burn quota. +// dump context: timeouts would start another long wait, and quota retries are +// handled as wait-and-continue on the native conversation instead. func ShouldFallback(err error) bool { if err == nil { return false diff --git a/internal/config/config.go b/internal/config/config.go index 25c5d8b..1d10651 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -17,26 +17,33 @@ const ( // is 5m; keep this far above long tool waits such as test suites. DefaultTimeoutSeconds = 4 * 60 * 60 DefaultSessionIDLen = 8 + // DefaultQuotaRetryAttempts is the number of agy invocations per ACP prompt + // when the model provider returns a 429 / quota error (1 initial + retries). + DefaultQuotaRetryAttempts = 5 ) type Config struct { - AgyBinary string - HomeDir string - PromptThreshold int - TimeoutSeconds int - DefaultModel string - SkipPerms bool - InjectExecutionEnvNote bool + AgyBinary string + HomeDir string + PromptThreshold int + TimeoutSeconds int + QuotaRetryAttempts int + QuotaRetryMaxWaitSeconds int + DefaultModel string + SkipPerms bool + InjectExecutionEnvNote bool } type CLIOptions struct { - AgyBinary string - Model string - PromptThreshold int - TimeoutSeconds int - SkipPerms *bool - InjectExecutionEnvNote *bool - ListModels bool + AgyBinary string + Model string + PromptThreshold int + TimeoutSeconds int + QuotaRetryAttempts int + QuotaRetryMaxWaitSeconds int + SkipPerms *bool + InjectExecutionEnvNote *bool + ListModels bool } func Load() (*Config, error) { @@ -50,13 +57,15 @@ func LoadWithOptions(opts CLIOptions) (*Config, error) { } cfg := &Config{ - AgyBinary: getEnvOrDefault("AGY_BINARY", detectAgyBinary()), - HomeDir: home, - PromptThreshold: getEnvIntOrDefault("AGY_PROMPT_THRESHOLD", DefaultPromptThreshold), - TimeoutSeconds: getEnvIntOrDefault("AGY_TIMEOUT_SECONDS", DefaultTimeoutSeconds), - DefaultModel: os.Getenv("AGY_MODEL"), - SkipPerms: getEnvBoolOrDefault("AGY_SKIP_PERMISSIONS", true), - InjectExecutionEnvNote: !getEnvBoolOrDefault("AGY_ACP_SKIP_ENV_NOTE", false), + AgyBinary: getEnvOrDefault("AGY_BINARY", detectAgyBinary()), + HomeDir: home, + PromptThreshold: getEnvIntOrDefault("AGY_PROMPT_THRESHOLD", DefaultPromptThreshold), + TimeoutSeconds: getEnvIntOrDefault("AGY_TIMEOUT_SECONDS", DefaultTimeoutSeconds), + QuotaRetryAttempts: getEnvIntOrDefault("AGY_QUOTA_RETRY_ATTEMPTS", DefaultQuotaRetryAttempts), + QuotaRetryMaxWaitSeconds: getEnvIntOrDefault("AGY_QUOTA_RETRY_MAX_WAIT_SECONDS", 0), + DefaultModel: os.Getenv("AGY_MODEL"), + SkipPerms: getEnvBoolOrDefault("AGY_SKIP_PERMISSIONS", true), + InjectExecutionEnvNote: !getEnvBoolOrDefault("AGY_ACP_SKIP_ENV_NOTE", false), } if opts.AgyBinary != "" { cfg.AgyBinary = opts.AgyBinary @@ -70,6 +79,12 @@ func LoadWithOptions(opts CLIOptions) (*Config, error) { if opts.TimeoutSeconds > 0 { cfg.TimeoutSeconds = opts.TimeoutSeconds } + if opts.QuotaRetryAttempts > 0 { + cfg.QuotaRetryAttempts = opts.QuotaRetryAttempts + } + if opts.QuotaRetryMaxWaitSeconds > 0 { + cfg.QuotaRetryMaxWaitSeconds = opts.QuotaRetryMaxWaitSeconds + } if opts.SkipPerms != nil { cfg.SkipPerms = *opts.SkipPerms } @@ -94,6 +109,8 @@ func ParseCLIOptions(args []string) (CLIOptions, bool, error) { fs.StringVar(&opts.Model, "model", "", "default agy model") fs.IntVar(&opts.PromptThreshold, "prompt-threshold", 0, "prompt file threshold") fs.IntVar(&opts.TimeoutSeconds, "timeout-seconds", 0, "agy execution and --print-timeout (seconds)") + fs.IntVar(&opts.QuotaRetryAttempts, "quota-retry-attempts", 0, "agy invocations per prompt after 429/quota errors") + fs.IntVar(&opts.QuotaRetryMaxWaitSeconds, "quota-retry-max-wait-seconds", 0, "cap for a single 429 backoff (0 = remaining turn timeout)") fs.BoolVar(&skipPerms, "skip-permissions", false, "pass --dangerously-skip-permissions to agy") fs.BoolVar(&noSkipPerms, "no-skip-permissions", false, "do not pass --dangerously-skip-permissions to agy") fs.BoolVar(&envNote, "execution-env-note", false, "prepend the execution-environment steering note (default)") diff --git a/internal/config/config_test.go b/internal/config/config_test.go index b3b7693..da1c553 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -58,6 +58,9 @@ func TestParseCLIOptions(t *testing.T) { if opts.AgyBinary != "custom-agy" || opts.Model != "gemini-test" || opts.PromptThreshold != 123 || opts.TimeoutSeconds != 45 { t.Fatalf("unexpected parsed options: %+v", opts) } + if opts.QuotaRetryAttempts != 0 || opts.QuotaRetryMaxWaitSeconds != 0 { + t.Fatalf("expected quota retry flags to stay unset, got %+v", opts) + } if opts.SkipPerms == nil || *opts.SkipPerms { t.Fatalf("expected skip perms opt-out, got %+v", opts.SkipPerms) } @@ -169,3 +172,36 @@ func TestLoadWithOptions_ExecutionEnvNoteOverridesEnv(t *testing.T) { t.Fatal("expected CLI override to re-enable env-note injection") } } + +func TestLoad_QuotaRetryDefaults(t *testing.T) { + t.Setenv("AGY_QUOTA_RETRY_ATTEMPTS", "") + t.Setenv("AGY_QUOTA_RETRY_MAX_WAIT_SECONDS", "") + + cfg, err := Load() + if err != nil { + t.Fatalf("Load failed: %v", err) + } + if cfg.QuotaRetryAttempts != DefaultQuotaRetryAttempts { + t.Fatalf("attempts = %d, want %d", cfg.QuotaRetryAttempts, DefaultQuotaRetryAttempts) + } + if cfg.QuotaRetryMaxWaitSeconds != 0 { + t.Fatalf("max wait = %d, want 0 (remaining timeout)", cfg.QuotaRetryMaxWaitSeconds) + } +} + +func TestParseCLIOptions_QuotaRetry(t *testing.T) { + opts, showVersion, err := ParseCLIOptions([]string{ + "--quota-retry-attempts", "3", + "--quota-retry-max-wait-seconds", "900", + }) + if err != nil { + t.Fatalf("ParseCLIOptions failed: %v", err) + } + if showVersion { + t.Fatal("did not expect version mode") + } + if opts.QuotaRetryAttempts != 3 || opts.QuotaRetryMaxWaitSeconds != 900 { + t.Fatalf("unexpected quota retry options: %+v", opts) + } +} +