Skip to content

feat: MCP rate-limit retries, optional servers, task-scoped resource lifecycle#23

Merged
demianbrecht merged 5 commits into
mainfrom
feat/mcp-resilience-resource-lifecycle
Jun 16, 2026
Merged

feat: MCP rate-limit retries, optional servers, task-scoped resource lifecycle#23
demianbrecht merged 5 commits into
mainfrom
feat/mcp-resilience-resource-lifecycle

Conversation

@demianbrecht

Copy link
Copy Markdown
Contributor

Summary

  • Task-scoped resource lifecycle. MCP/checkpointer scopes are now bound to the task (_agent_resources), so a masked transport fault surfaces with its real cause at agent_main's reporting boundary instead of an opaque CancelledError. The boundary classifies on durable cancel signals and unwraps BaseExceptionGroup to the most specific leaf.
  • 429 retries for the OAuth MCP client. RetryTransport transparently retries HTTP 429s (honoring Retry-After, draining the response before resend) for both the initialize handshake and tool calls. Wired via httpx's public transport= constructor param so TLS verify and the retry layer don't depend on httpx private internals.
  • Optional MCP servers. Servers marked optional degrade gracefully on startup failure rather than aborting the task. Startup errors are now name-keyed (server, message) tuples instead of fragile substring matching of error text against config names.
  • No terminal-event overwrite. A teardown fault raised after a task already emitted its terminal event is logged (resource_teardown_failed_after_terminal), not re-reported — a successful task can no longer be recorded as FAILED.

Testing

make test → 718 passed, 9 skipped. Adds end-to-end agent_main boundary tests (clean completion, wrapped-cause leaf reporting, operator cancel vs. teardown fault, success-path teardown fault), RetryTransport/oauth coverage (drain-before-resend, Retry-After, max-retries), and optional-server filtering tests.

Notes

The two most recent additions (H1 terminal-event guard, M1 public-transport= wiring) came out of an independent quality + security review of the diff.

…lifecycle

Bind MCP/checkpointer resource scopes to the task (`_agent_resources`) so a
masked transport fault surfaces with its real cause at the reporting boundary
instead of an opaque CancelledError. The boundary classifies on durable cancel
signals and unwraps BaseExceptionGroup to the most specific leaf.

- oauth: RetryTransport transparently retries HTTP 429s (with Retry-After
  support) for the OAuth-built MCP client, draining before resend. Wired via
  the public `transport=` constructor param so TLS verify and the retry layer
  don't depend on httpx private internals.
- app/mcp: `optional` MCP servers degrade gracefully on startup failure;
  errors are name-keyed (server, message) tuples instead of fragile substring
  matching against config names.
- agent_runtime: a teardown fault raised after the task already emitted a
  terminal event is logged, not re-reported, so a successful task can't be
  recorded as FAILED.

Adds end-to-end agent_main boundary tests, RetryTransport/oauth coverage, and
optional-server filtering tests.

@demianbrecht demianbrecht left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 [Security] No findings. Reviewed the OAuth HTTP retry transport, MCP lifecycle/optional-server handling, and the agent_runtime exception-classification boundary.

Critical: 0
High: 0
Medium: 0
Low: 0

Notes (no action needed):

  • TLS verification is preserved through the new retry path: _build_transport applies verify=ssl_verify on httpx.AsyncHTTPTransport and passes it via the public transport= param. Since httpx ignores the client-level verify= once transport= is set, putting verify on the transport (as done here) is the correct wiring — this is a posture improvement over the prior private-internals approach.
  • RetryTransport logs only server name / attempt / delay — no URL, headers, body, or token content. Backoff is bounded (max_retries × 60s cap), and Retry-After is parsed as a guarded, clamped float — no DoS or injection vector. Retries re-send the already-signed request below the auth layer, so no credential re-mint concern.
  • Token storage perms (0o600 file / 0o700 dir) confirmed unchanged.
  • Consistent with the project's single-user, local-only trust model; MCP servers are operator-configured, not attacker-controlled.

if not ctx._completed:
ctx.fail(f"{type(e).__name__}: {e}", traceback.format_exc())
raise
except BaseException as e:

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 [Quality] Edge case worth noting (non-blocking): _unwrap_cause returns the first non-cancellation leaf, so a BaseExceptionGroup that wraps a KeyboardInterrupt/SystemExit (e.g. a TaskGroup where one child raised KeyboardInterrupt) lands in this except BaseException branch rather than the (KeyboardInterrupt, SystemExit) branch above — the group itself is neither type. It would then be reported via ctx.fail(...) as an ordinary failure leaf and not re-raised, so the interpreter teardown the upper branch is designed to preserve would be skipped. Rare in practice for this agent, but if you want the guarantee to hold you could check isinstance(leaf, (KeyboardInterrupt, SystemExit)) after unwrapping and re-raise. Fine to leave as-is given how unlikely it is.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — addressed in ea3aa0b. I unwrap once at the top of the except BaseException handler and, when the leaf is a KeyboardInterrupt/SystemExit, report it (respecting the _completed guard) and raise leaf from e so interpreter teardown is preserved. Added a regression test (test_grouped_process_signal_reraised) that raises a BaseExceptionGroup([SystemExit(...)]) through the real agent_main boundary and asserts both the re-raise and the reported leaf. Note CPython's asyncio.TaskGroup actually extracts base exceptions and re-raises them bare (so they'd hit the dedicated branch); the wrapped case applies to anyio/nested groups, which the test exercises directly.

@demianbrecht demianbrecht left a comment

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🤖 [Quality] Reviewed all 7 changed files. This is a high-quality, carefully-reasoned PR — the docstrings explain the why behind each design decision and the test coverage is thorough and meaningful (the agent_main boundary tests genuinely exercise the masking-fix integration, not just framework behavior).

Strengths

  • Replacing fragile substring matching of error text with name-keyed (server, message) tuples (mcp.py / agent_runtime.py) is a real robustness improvement, not just a refactor.
  • Wiring the retry transport via httpx’s public transport= param instead of mutating client._transport/_mounts is the right call and well-justified — and AsyncClient.aclose() does delegate to the custom transport, so the RetryTransport.aclose delegation is correct.
  • The cancel-classification logic keying on durable signals (ctx._cancelled.is_set() and task_handle.cancelling() > 0) rather than exception shape is sound, and the operator-cancel-wins-over-teardown-fault case is correctly handled and tested.
  • _retry_after_seconds taking the max of header and exponential backoff (so a constant Retry-After: 1 can’t defeat escalation) is a nice detail with a test that pins it.

One non-blocking edge (inline on agent_runtime.py:948): a BaseExceptionGroup wrapping a KeyboardInterrupt/SystemExit would be reported as an ordinary failure leaf and not re-raised, skipping interpreter teardown. Rare for this agent; documented inline.

CI: test/e2e/SAST all green across 3.12–3.14. The lint check is failing on the Check formatting step (ruff format) — 2 ** attempt2**attempt in oauth.py and a couple of collection-literal wraps in test_agent_runtime.py. Looks like a ruff-version formatting drift; make fmt should clear it. Worth resolving before merge so the branch is green.

Verdict: comment (not blocking). Nice work.

…ent boundary

A BaseExceptionGroup wrapping a KeyboardInterrupt/SystemExit (e.g. an anyio
task group or nested group whose child raised one) lands in the generic
'except BaseException' branch, not the dedicated process-signal branch — the
group itself is neither type. It would then be reported via ctx.fail() and
swallowed, skipping the interpreter teardown the dedicated branch preserves.
Unwrap once at the top of the handler and re-raise when the leaf is a process
signal. Adds a regression test.
@demianbrecht demianbrecht merged commit ad1fce5 into main Jun 16, 2026
9 checks passed
@demianbrecht demianbrecht deleted the feat/mcp-resilience-resource-lifecycle branch June 16, 2026 16:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant