fix: surface sandbox exit code via proto3 optional exit_code presence

[joseph-shih-wandb]

TL;DR

sandbox.returncode was always None: the old guard read a returncode attribute that never existed on GetSandboxResponse, and mock-based unit tests carried a truthy returncode attribute that hid the field-name drift. The backend now serves the main-process exit code as optional int32 exit_code = 13 on GetSandboxResponse (coreweave/aviato#1490, merged; BSR commit 033ca7a80802), so this PR wires it through: re-vendored stubs, presence-checked read, and a bounded best-effort grace re-poll for the runner's batched report lag. Fixes SBX-PROD-BUG-009 / e2e SBX-LIFE-003.

Merge gate: ✅ satisfied (verified 2026-07-29). The prod gateway runs gateway:v0.132.0 (all pods in sandbox-system @ us-west-09a-core-services), and the v0.132.0 tag contains the aviato#1490 merge commit — prod serves exit_code now. The gate existed because semantic-release publishes on merge, and an SDK released ahead of the gateway would have paid the full ~4s grace window on every natural completion with no chance of a code. (Prod runners are on v0.130.3, which is fine: #1490 is gateway-side only; runners have always reported the code on the internal wire.)

What changed

• Vendored stubs re-generated from buf.build at the new pins (scripts/update-protos.sh). Purely mechanical; also picks up unrelated additive upstream proto changes (StartSandboxRequest.volumes, streaming stub-generator format churn). The SDK never sets the new fields, so requests to current backends are unchanged in meaning.
• Presence-checked read (_exit_code_from_info): HasField("exit_code") distinguishes "exited 0" from "not reported"; getattr fallback for non-proto stand-ins. Read only for source == "poll" and status == COMPLETED, preserving the existing returncode contract.
• Bounded grace re-poll (_grace_repoll_for_exit_code, EXIT_CODE_GRACE_POLLS=2 × 2s): the runner flushes exit codes on a ~5s batched report, so a Get can observe COMPLETED before the code lands; once the poll loop latches the terminal state, returncode is frozen. The re-poll is strictly best-effort:
  • skipped for client-initiated stops (_stop_owned/_is_stopping) and once a terminal state latches concurrently (_is_done, re-checked after the sleep);
  • each bonus poll is a literal single unretried Get (_get_sandbox_once) on a dedicated short timeout (EXIT_CODE_GRACE_RPC_TIMEOUT_SECONDS=2s) — never the primary poll's 15s/30s retry envelope, never _poll_until_stable's transient-status loop;
  • any CWSandboxError from a bonus poll returns the already-in-hand terminal response instead of failing the wait (debug-logged);
  • the bonus response is adopted only when terminal or code-bearing — a stale non-terminal read cannot un-observe an already-observed completion;
  • the in-hand terminal response is published in _grace_pending_response for the duration of the window, and both waiter TimeoutError handlers latch it and resolve through the normal terminal policy — a waiter deadline expiring mid-grace can never convert an already-observed completion into SandboxTimeoutError. The slot is cleared identity-guarded, so a cancelled sibling grace window cannot wipe the other poll task's published rescue response. The waiter-vs-shared-task double-latch is idempotent via _apply_sandbox_info's terminal guard.
• Tests: real-proto presence coverage (TestApplySandboxInfoRealProto) so a field rename fails loudly; grace re-poll coverage including transient-error no-retry (exactly 2 Gets), dedicated-timeout pin, stale-response non-adoption, mid-grace status flip, both stop-guard halves, and deadline-during-grace on both wait paths; mock helpers primed with per-status presence semantics matching the backend contract.

Settled design decisions (do not re-raise)

1. Presence over sentinel: proto3 optional; 0 is meaningful, absent is a legitimate permanent state (older gateways, gateway-initiated stops, containers that never ran). Clients must check presence — hence HasField with a documented returncode is None contract.
2. COMPLETED-only, poll-only: FAILED exit codes (backend sends nonzero-only) are deliberately NOT surfaced; the SDK's public returncode contract is unchanged. Widening it is a separate discussion.
3. Grace re-poll is best-effort enrichment: it must never fail a wait that already holds a terminal response, never add latency to client-initiated stops issued before the window opens, and never let a waiter deadline expiring mid-window surface as a timeout for an observed completion. Bounded (2 extra polls) because "no code" is also a permanent state.
4. Enrich-before-latch, not latch-then-backfill: the grace window runs before the terminal latch so that after wait() resolves, returncode is settled (the primary consumer does wait() then reads returncode). The alternative — latch immediately, backfill returncode on _Terminal afterwards — would retire the residual get_status()-during-grace race and is the recorded follow-up direction if FAILED codes are ever surfaced, but mutating a latched terminal state after wait() returns changes the "returncode is final after wait" contract and is out of scope here.
5. Grace gate keys on raw proto COMPLETED: poll-source UNSPECIFIED (mapped to COMPLETED by _apply_sandbox_info) is only emitted by backends that predate exit codes, so a grace window for it could never produce one (documented at the gate).
6. Grace constants are module constants, not SandboxDefaults knobs: the merge gate above was the kill switch (now satisfied); a config seam for a transitional window was judged not worth the permanent API surface.
7. getattr/ValueError fallback in _exit_code_from_info is deliberate defensive compat for test stand-ins; the stubs are vendored in-package so prod skew is impossible — cheap insurance, not a deployment scenario.

Known limitations / accepted trade-offs

• Concurrent get_status() during the grace window can latch a query-source terminal state (returncode omitted by contract) and permanently forfeit the exit code. Pre-existing microsecond race widened to the bounded window; retired by the backfill follow-up (decision 4).
• Stop-vs-completion race window widened: during the grace window self._state still reads RUNNING, so a concurrent stop() can send a Stop RPC for an already-completed sandbox and a wait_until_complete(raise_on_termination=True) waiter then sees SandboxTerminatedError for a natural completion. Pre-existing race; the stop RPC is benign server-side. A stop() landing mid-grace also waits out the current ≤2s step (guards run at loop boundaries).
• Waiter-deadline latch vs mid-grace FAILED correction: a waiter timing out mid-grace latches the in-hand COMPLETED; a reconciler correction to FAILED arriving on a later bonus poll is then rejected by the terminal guard. Intra-window divergence, bounded; also retired by the backfill follow-up.
• ~5s flush cadence vs 4s window: the lag tail past the window still latches returncode=None permanently. Constants are deliberately conservative; tune after measuring the real flush-lag distribution.
• Sandboxes stopped by other actors (gateway idle/lifetime, another client) pay the grace window once per sandbox even post-deploy.

Follow-ups (out of scope, recorded)

• One-shot None→value returncode backfill on _Terminal (retires the get_status() forfeit and the waiter-latch/FAILED-correction divergence; prerequisite for surfacing FAILED exit codes).
• Consolidate the four presence-mimicking mock helpers into one conftest factory returning real GetSandboxResponse messages; pin the "returncode frozen after wait()" contract and post-latch RPC suppression.
• Pin FAILED/TERMINATED-with-code-present ⇒ returncode is None with real-proto tests; cover the ValueError fallback branch.
• started_at_time truthiness guard yields epoch-0 instead of None on real protos (pre-existing; fix is HasField("started_at_time") + real-proto assertions).
• Narrow _SandboxInfoLike back to fields all real inputs share; type the exit-code read against the poll-path response.
• Hash-verify the vendored proto wheels in update-protos.sh (--require-hashes, --only-binary=:all:).
• Regenerated streaming_pb2.pyi dropped the mypy-protobuf-documented streaming wire contract still referenced by _sandbox.py and the CHANGELOG — regenerate with mypy-protobuf or move the contract docs to a durable home.
• Integration-test sandbox.returncode assertions once the gateway deploy lands.

Verification

• 1266 unit tests green; mypy clean; ruff clean.
• Vendored stub descriptor runtime-checked: field 13, proto3 presence, HasField distinguishes explicit 0 from unset.
• Three rounds of the 6-reviewer heavy review pipeline (Fable 5 reviewers + Codex gpt-5.5 cross-validation) run locally; all confirmed Highs fixed, final round's behavior Mediums in the new grace code fixed (single-Get isolation, stale-response non-adoption, identity-guarded rescue slot), remainder recorded above.

🤖 Generated with Claude Code