[andrewpark-scaleai]

Stacked on #695 (base is that branch, not main). #695 adds the model_key_budget and model_key_auth error types this PR reads.

Why

Goal: give agentenv enough signal to debug every failure and heal itself where code can. This PR stops runs from spending tokens on a model key that cannot work.

Today nothing knows a budget or auth refusal is permanent. The rubrics judge retries a refused LiteLLM call DEFAULT_MAX_RETRIES (3) times, so each grader step spends about 2.7 judge calls on a dead key. A step with retry_config rolls back and runs the same agent turn again with the same key. Each retry also redeploys sandboxes.

What

• ModelKeyUnusable(cause, where) in task_step/context.py, next to VerifierNotPassingError. cause is "budget" or "auth". The message is fixed text with no key material, and it drops the original exception chain, because a 401 body can echo the key.
• _prompt_llm_judge classifies a LiteLLM error once in _model_key_refusal: a budget-exceeded error or message marker gives budget, a 401 gives auth. It raises ModelKeyUnusable on the first such error. Every other error, including 429 rate limits and 5xx, retries as before.
• _invoke_judge_a2a raises ModelKeyUnusable when the judge task failed with error_type model_key_budget or model_key_auth.
• prompt_agent records the failed turn, then raises ModelKeyUnusable for those two error_type values instead of the plain RuntimeError.
• The task scheduler does not retry a step that raised ModelKeyUnusable, the same way it skips VerifierNotPassingError.

ModelKeyUnusable subclasses RuntimeError, so callers that catch RuntimeError still see it. Behavior changes only for a refused key.

Design decisions and tradeoffs

One attempt, even if the key gets fixed mid-run. A key can be topped up or rotated while a run is in flight, and one attempt gives that up. I chose it because a budget or auth refusal does not clear within a retry window of seconds, and the cost of being wrong is one failed step that a person can rerun. The cost of the old behavior is every step in every run burning its full retry budget. Waiting out a key reset belongs to a later breaker, not to blind retries.

Parse at the boundary. The A2A paths branch on the typed error_type from #695 and never read message text. A direct LiteLLM call has no typed field, so _model_key_refusal parses the exception once at the call site, and the retry loop sees only a cause. A message that merely mentions "budget" in another context can match the marker, so the markers are LiteLLM's own budget phrases. The auth case uses the 401 status or AuthenticationError, and a 403 or 429 does not count.

Not done here. The hub breaker, preflight checks and the pass@k pause are separate work. This PR only stops the retries.

How to verify

PYTHONPATH=src:packages/agentenv-protocol/src \
  python -m pytest tst/unit/task_step/verifiers/test_model_key_unusable.py \
  tst/unit/task_step/test_prompt_agent_model_key.py tst/unit/task_step/retry_test.py

Before the change (with only the exception class present), 6 tests fail: a refused judge key is called 3 times and raises RuntimeError, the A2A judge and prompt_agent raise RuntimeError, and the step with retry_config runs 4 times. After, all 52 pass. The transient-error tests (503 then 429 then success, and a plain agent_error) pass both before and after, which pins that nothing else changed.

Follow-up for downstream workers

A worker that retries activities in place needs to treat ModelKeyUnusable as non-retryable once this ships in a release. In the worker that consumes this package, add it to the retry_cannot_help exception set, next to VerifierNotPassingError.