Gateway: HTTP MCP backend, hardening (circuit breaker/audit/redaction), parallel reconnect - #108
Conversation
…ay-registry Ported two patterns from forgemax's more mature gateway architecture: - ensure_fresh() discovered each backend's tools sequentially; with N configured servers a single slow one serialized every other one behind it. Now runs concurrently via futures_util::future::join_all (already a transitive dep at 0.3.32, added explicitly). - McpStdioBackend::call/discover only cleared the cached connection on a timeout. A service-level error (Ok(Err(_)) from rmcp's call_tool/ list_all_tools) can equally mean the transport died underneath us (broken pipe, crashed child) since rmcp doesn't distinguish the two at that call site — now drops the cache there too so the next call respawns instead of repeatedly hitting a dead pipe.
…tool sanitization - open a per-backend circuit after repeated consecutive spawn/connection failures, short-circuiting further calls and allowing one probe attempt after a recovery window - append-only audit log for gateway_execute calls (server, tool, args hash, outcome) written alongside the registry db; never logs raw args - redact connection details, file paths, and credential-shaped text from downstream error messages before they reach the LLM - sanitize discovered tool names/descriptions (charset + length bounds) before they're indexed and surfaced via gateway_search
📝 WalkthroughWalkthroughAdds MCP HTTP backend support, shared circuit breaking, audit logging, tool metadata sanitization, and downstream error redaction across gateway-registry and the MCP server. ChangesGateway registry reliability and safety features
Estimated code review effort: 5 (Critical) | ~90 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant Registry
participant Backend
participant AuditLog
Client->>Registry: execute(server, tool, args)
Registry->>Backend: call(args.clone())
Backend-->>Registry: Ok/Err result
Registry->>AuditLog: record(server, tool, hashed args, outcome)
Registry-->>Client: result
Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
crates/gateway-registry/src/audit.rs (1)
35-42: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winWrite each audit entry with a single append
writeln!(f, "{line}")can split a JSONL record across multiple writes, andRegistry::executeis&self, so concurrent calls can splice entries together. Format the line once, add\n, and callwrite_all.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/gateway-registry/src/audit.rs` around lines 35 - 42, The audit write path in the `audit` module currently uses `writeln!` on the opened file handle, which can issue multiple writes and allow concurrent `Registry::execute` calls to interleave JSONL entries. Update the write logic in the `OpenOptions::open` success branch to format the entry once, append a newline, and use a single `write_all` call on `f`, keeping the existing open/error handling in place.crates/gateway-registry/src/mcp_stdio.rs (1)
159-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDuplicated check-then-record wrapper between
discover()andcall().Both methods repeat the identical
check_circuit()→ inner call →record_success/record_failurepattern. Could be extracted into a small private helper (e.g.with_circuit_breaker<T>(&self, fut: impl Future<Output = Result<T, GatewayError>>) -> Result<T, GatewayError>) to avoid the two copies drifting apart if the breaker logic changes again.Also applies to: 212-220
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/gateway-registry/src/mcp_stdio.rs` around lines 159 - 167, Both `discover()` and `call()` in `McpStdioClient` duplicate the same circuit-breaker wrapper logic, so extract that shared pattern into a small private helper (for example a generic `with_circuit_breaker` method) and have both methods delegate to it. Keep the helper responsible for `check_circuit()`, awaiting the inner operation, and calling `record_success()` or `record_failure()` based on the result so the behavior stays identical. This will keep the `discover_inner()` and `call_inner()` paths aligned and prevent the breaker logic from drifting between the two methods.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/gateway-registry/src/mcp_stdio.rs`:
- Around line 91-120: The circuit breaker in check_circuit(), record_success(),
and record_failure() currently releases the mutex before the actual
discover_inner/call_inner probe, so once opened_until expires multiple
concurrent callers can all pass through at once. Update check_circuit() to
re-arm the circuit while still holding the circuit lock as soon as the recovery
window has elapsed, so only the first caller proceeds and the others still see
the circuit as open. Keep the behavior aligned with the existing circuit state
in record_success() and record_failure(), and use the existing symbols
check_circuit, record_success, record_failure, and opened_until to make the fix
local to the current circuit logic.
- Around line 175-187: Treat the `Ok(Err(e))` branch in `mcp_stdio.rs` as a
protocol error rather than an automatic transport failure when `rmcp` can
distinguish them. Update the `Registry::execute`/stdio call path so only real
transport-level failures clear the cached connection and contribute to
circuit-breaker accounting, while clean RPC errors are returned without dropping
`guard` or forcing a reconnect. Use the existing `Ok(Err(e))` handling and
`GatewayError::Upstream` conversion point to preserve protocol errors while
avoiding breaker budget impact.
---
Nitpick comments:
In `@crates/gateway-registry/src/audit.rs`:
- Around line 35-42: The audit write path in the `audit` module currently uses
`writeln!` on the opened file handle, which can issue multiple writes and allow
concurrent `Registry::execute` calls to interleave JSONL entries. Update the
write logic in the `OpenOptions::open` success branch to format the entry once,
append a newline, and use a single `write_all` call on `f`, keeping the existing
open/error handling in place.
In `@crates/gateway-registry/src/mcp_stdio.rs`:
- Around line 159-167: Both `discover()` and `call()` in `McpStdioClient`
duplicate the same circuit-breaker wrapper logic, so extract that shared pattern
into a small private helper (for example a generic `with_circuit_breaker`
method) and have both methods delegate to it. Keep the helper responsible for
`check_circuit()`, awaiting the inner operation, and calling `record_success()`
or `record_failure()` based on the result so the behavior stays identical. This
will keep the `discover_inner()` and `call_inner()` paths aligned and prevent
the breaker logic from drifting between the two methods.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: 1dcc5e58-ec84-490b-9904-1de60c0633f0
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (13)
crates/gateway-registry/Cargo.tomlcrates/gateway-registry/src/audit.rscrates/gateway-registry/src/error.rscrates/gateway-registry/src/lib.rscrates/gateway-registry/src/mcp_stdio.rscrates/gateway-registry/src/redact.rscrates/gateway-registry/src/registry.rscrates/gateway-registry/src/sanitize.rscrates/gateway-registry/tests/fixtures/fixture_server.rscrates/gateway-registry/tests/gateway_audit_log.rscrates/gateway-registry/tests/mcp_stdio_circuit_breaker.rscrates/gateway-registry/tests/mcp_stdio_reconnect.rssrc/mcp_server.rs
A non-object args call previously reached ensure_connected() first, so a caller mistake against a down/slow backend surfaced as Connection/Timeout instead of InvalidArgument -- and held the process-wide gateway lock for up to 30s on a call that was never going to succeed regardless. Validate locally before any connect attempt in both McpStdioBackend and McpHttpBackend, keeping their behavior in parity.
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
crates/gateway-registry/tests/mcp_stdio_reconnect.rs (1)
55-60: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winStrengthen assertion to verify per-call correctness, not just success.
The test only checks
is_ok(); it doesn't confirm each response's text matches its corresponding input index, so a bug that returns another request's response would not be caught.♻️ Suggested strengthening
let calls = (0..20).map(|i| { let backend = backend.clone(); - async move { backend.call("echo", serde_json::json!({"text": i.to_string()})).await } + async move { (i, backend.call("echo", serde_json::json!({"text": i.to_string()})).await) } }); let results = futures_util::future::join_all(calls).await; - assert!(results.iter().all(|r| r.is_ok()), "expected all concurrent calls to succeed: {results:?}"); + for (i, result) in results { + let result = result.unwrap_or_else(|e| panic!("call {i} failed: {e:?}")); + let text = result.get(0).and_then(|c| c.get("text")).and_then(|t| t.as_str()); + assert_eq!(text, Some(format!("echo: {i}").as_str())); + }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/gateway-registry/tests/mcp_stdio_reconnect.rs` around lines 55 - 60, The test in mcp_stdio_reconnect.rs only verifies that backend.call("echo", ...) succeeds, but it does not confirm each response matches the input for that specific request. Strengthen the join_all assertion by checking each result from the echo calls against its corresponding index string so request/response mixing is caught, using the backend.call("echo", ...) and results values to verify per-call correctness rather than only is_ok().crates/gateway-registry/src/mcp_http.rs (1)
35-183: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚖️ Poor tradeoffSignificant structural duplication with
mcp_stdio.rs.
ensure_connected/discover/discover_inner/call/call_innerhere are nearly line-for-line identical in control flow toMcpStdioBackend(circuit wrap, timeout handling, guard invalidation on error, args validation, result mapping) — only the transport construction and error message differ. Consider extracting the shared circuit-wrap-and-timeout skeleton (e.g. a small generic helper taking an async connect/call closure) to avoid maintaining two copies of this error-handling logic in lockstep.Given this doubles the surface for the two backends to drift (e.g. one gets a bugfix the other doesn't), it's a worthwhile — though not urgent — cleanup.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/gateway-registry/src/mcp_http.rs` around lines 35 - 183, Extract the shared circuit-breaker, timeout, guard-invalidation, and result-mapping flow from McpHttpBackend methods like ensure_connected, discover_inner, and call_inner into a reusable helper, so it mirrors the logic already in McpStdioBackend without duplicating the same control flow. Keep the backend-specific parts limited to transport creation and message formatting, and have discover/call use the shared skeleton with closures or a small generic wrapper to prevent the two backends from drifting.crates/gateway-registry/src/config.rs (1)
25-40: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
auth_headerisn't validated againstauth_ref/auth_env.
parse()'s pairing check only inspectsauth_ref/auth_env(Line 62-64);auth_headeris discarded via... So a config that setsauth_headeralone (noauth_ref/auth_env) parses successfully, but perregistry.rs::resolve_mcp_http_auth_header, the header is silently dropped since resolution only proceeds when bothauth_refandauth_envareSome. A user intending to authenticate would get no error at parse time and no header sent at runtime — a silent, hard-to-debug misconfiguration.Consider rejecting (or at least the doc comment on Lines 31-38)
auth_headerset withoutauth_ref/auth_env, mirroring the existingIncompleteAuthConfigpattern.Proposed validation addition
let (auth_ref, auth_env) = match server { ServerConfig::McpStdio { auth_ref, auth_env, .. } => (auth_ref, auth_env), - ServerConfig::McpHttp { auth_ref, auth_env, .. } => (auth_ref, auth_env), + ServerConfig::McpHttp { auth_ref, auth_env, auth_header, .. } => { + if auth_header.is_some() && auth_ref.is_none() { + return Err(ConfigError::IncompleteAuthConfig { + server: name.clone(), + auth_ref: auth_ref.clone(), + auth_env: auth_env.clone(), + }); + } + (auth_ref, auth_env) + } };Also applies to: 61-71
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/gateway-registry/src/config.rs` around lines 25 - 40, `auth_header` is currently accepted even when `auth_ref`/`auth_env` are missing, so the config can parse but never produce an auth header. Update `parse()` in `config.rs` to validate `auth_header` alongside the existing `IncompleteAuthConfig` pairing check, and reject any `McpHttp` entry that sets `auth_header` without both `auth_ref` and `auth_env`. Use the existing `McpHttp`, `IncompleteAuthConfig`, and `resolve_mcp_http_auth_header` flow as the reference points, and align the doc comment on `auth_header` with the new validation behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/gateway-registry/src/mcp_stdio.rs`:
- Around line 167-175: The circuit-breaker accounting in call currently treats
all Err results as backend failures, but GatewayError::InvalidArgument should be
excluded because it is a local validation error. Update the match on result in
call to skip record_failure() when the error is InvalidArgument, and only count
real backend/service failures; apply the same fix in mcp_http::call so both
transport paths behave consistently.
---
Nitpick comments:
In `@crates/gateway-registry/src/config.rs`:
- Around line 25-40: `auth_header` is currently accepted even when
`auth_ref`/`auth_env` are missing, so the config can parse but never produce an
auth header. Update `parse()` in `config.rs` to validate `auth_header` alongside
the existing `IncompleteAuthConfig` pairing check, and reject any `McpHttp`
entry that sets `auth_header` without both `auth_ref` and `auth_env`. Use the
existing `McpHttp`, `IncompleteAuthConfig`, and `resolve_mcp_http_auth_header`
flow as the reference points, and align the doc comment on `auth_header` with
the new validation behavior.
In `@crates/gateway-registry/src/mcp_http.rs`:
- Around line 35-183: Extract the shared circuit-breaker, timeout,
guard-invalidation, and result-mapping flow from McpHttpBackend methods like
ensure_connected, discover_inner, and call_inner into a reusable helper, so it
mirrors the logic already in McpStdioBackend without duplicating the same
control flow. Keep the backend-specific parts limited to transport creation and
message formatting, and have discover/call use the shared skeleton with closures
or a small generic wrapper to prevent the two backends from drifting.
In `@crates/gateway-registry/tests/mcp_stdio_reconnect.rs`:
- Around line 55-60: The test in mcp_stdio_reconnect.rs only verifies that
backend.call("echo", ...) succeeds, but it does not confirm each response
matches the input for that specific request. Strengthen the join_all assertion
by checking each result from the echo calls against its corresponding index
string so request/response mixing is caught, using the backend.call("echo", ...)
and results values to verify per-call correctness rather than only is_ok().
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: a8cff46c-9de7-4a68-969c-6bc69bc5c042
⛔ Files ignored due to path filters (1)
Cargo.lockis excluded by!**/*.lock
📒 Files selected for processing (14)
crates/gateway-registry/Cargo.tomlcrates/gateway-registry/src/backend.rscrates/gateway-registry/src/circuit.rscrates/gateway-registry/src/config.rscrates/gateway-registry/src/lib.rscrates/gateway-registry/src/mcp_http.rscrates/gateway-registry/src/mcp_stdio.rscrates/gateway-registry/src/registry.rscrates/gateway-registry/tests/mcp_http_call.rscrates/gateway-registry/tests/mcp_http_circuit_breaker.rscrates/gateway-registry/tests/mcp_http_discover.rscrates/gateway-registry/tests/mcp_stdio_reconnect.rscrates/gateway-registry/tests/registry.rscrates/gateway-registry/tests/support/mod.rs
🚧 Files skipped from review as they are similar to previous changes (1)
- crates/gateway-registry/src/registry.rs
| pub async fn call(&self, tool: &str, args: Value) -> Result<Value, GatewayError> { | ||
| self.ensure_connected().await?; | ||
| let mut guard = self.running.lock().await; | ||
| let running = guard.as_ref().expect("connected above"); | ||
| self.circuit.check(&self.command).await?; | ||
| let result = self.call_inner(tool, args).await; | ||
| match &result { | ||
| Ok(_) => self.circuit.record_success().await, | ||
| Err(_) => self.circuit.record_failure().await, | ||
| } | ||
| result | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant file and surrounding symbols first.
ast-grep outline crates/gateway-registry/src/mcp_stdio.rs --view expanded || true
printf '\n---\n'
ast-grep outline crates/gateway-registry/src/mcp_http.rs --view expanded || true
printf '\n---\n'
ast-grep outline crates/gateway-registry/src/error.rs --view expanded || true
printf '\n---\n'
# Show the relevant sections with line numbers.
sed -n '1,260p' crates/gateway-registry/src/mcp_stdio.rs | nl -ba | sed -n '1,260p'
printf '\n=== mcp_http.rs ===\n'
sed -n '1,280p' crates/gateway-registry/src/mcp_http.rs | nl -ba | sed -n '1,280p'
printf '\n=== error.rs ===\n'
sed -n '1,220p' crates/gateway-registry/src/error.rs | nl -ba | sed -n '1,220p'
# Search for InvalidArgument and circuit accounting elsewhere.
printf '\n=== InvalidArgument references ===\n'
rg -n "InvalidArgument|record_failure\(|record_success\(" crates/gateway-registry/srcRepository: getappz/agentflare
Length of output: 2468
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline crates/gateway-registry/src/mcp_stdio.rs --view expanded || true
printf '\n---\n'
ast-grep outline crates/gateway-registry/src/mcp_http.rs --view expanded || true
printf '\n---\n'
ast-grep outline crates/gateway-registry/src/error.rs --view expanded || true
printf '\n---\n'
sed -n '1,260p' crates/gateway-registry/src/mcp_stdio.rs | nl -ba | sed -n '1,260p'
printf '\n=== mcp_http.rs ===\n'
sed -n '1,280p' crates/gateway-registry/src/mcp_http.rs | nl -ba | sed -n '1,280p'
printf '\n=== error.rs ===\n'
sed -n '1,220p' crates/gateway-registry/src/error.rs | nl -ba | sed -n '1,220p'
printf '\n=== InvalidArgument references ===\n'
rg -n "InvalidArgument|record_failure\(|record_success\(" crates/gateway-registry/srcRepository: getappz/agentflare
Length of output: 2468
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the relevant function bodies with line numbers.
awk 'NR>=150 && NR<=220 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/mcp_stdio.rs
printf '\n=== mcp_http.rs ===\n'
awk 'NR>=120 && NR<=190 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/mcp_http.rs
printf '\n=== error.rs ===\n'
awk 'NR>=1 && NR<=80 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/error.rs
printf '\n=== InvalidArgument usage ===\n'
rg -n "InvalidArgument" crates/gateway-registry/src
printf '\n=== CircuitBreaker usage/definition ===\n'
rg -n "struct CircuitBreaker|impl CircuitBreaker|record_failure|record_success|check\(" crates/gateway-registry/src crates -g '!target'Repository: getappz/agentflare
Length of output: 17143
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the relevant function bodies with line numbers.
awk 'NR>=150 && NR<=220 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/mcp_stdio.rs
printf '\n=== mcp_http.rs ===\n'
awk 'NR>=120 && NR<=190 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/mcp_http.rs
printf '\n=== error.rs ===\n'
awk 'NR>=1 && NR<=80 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/error.rs
printf '\n=== InvalidArgument usage ===\n'
rg -n "InvalidArgument" crates/gateway-registry/src
printf '\n=== CircuitBreaker usage/definition ===\n'
rg -n "struct CircuitBreaker|impl CircuitBreaker|record_failure|record_success|check\\(" crates/gateway-registry/src crates -g '!target'Repository: getappz/agentflare
Length of output: 17143
🏁 Script executed:
#!/bin/bash
set -euo pipefail
awk 'NR>=1 && NR<=160 {printf "%4d %s\n", NR, $0}' crates/gateway-registry/src/circuit.rsRepository: getappz/agentflare
Length of output: 4694
Skip circuit-breaker failure accounting for InvalidArgument
GatewayError::InvalidArgument is a local, pre-flight validation error, not a backend failure. Counting it in record_failure() lets a few malformed calls open the breaker against a healthy backend. Apply the same guard in mcp_http.rs::call.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/gateway-registry/src/mcp_stdio.rs` around lines 167 - 175, The
circuit-breaker accounting in call currently treats all Err results as backend
failures, but GatewayError::InvalidArgument should be excluded because it is a
local validation error. Update the match on result in call to skip
record_failure() when the error is InvalidArgument, and only count real
backend/service failures; apply the same fix in mcp_http::call so both transport
paths behave consistently.
…t win (#488) self_repair_or_gate only checked the local job queue for an in-flight dispatch, not the item-level claim ledger. A finished job's claim outliving its (#108-capped) in_review TTL looked like "nothing in flight", so the sweep dispatched anyway -- the job died instantly at its own claim-acquire step, still counting the attempt against the 3-try cap. Now it checks claim liveness before dispatching and defers (retrying next sweep) instead of burning a cap slot on an attempt that never had a chance. Agentflare-Agent: claude-code Agentflare-Branch: task/114-ci-self-repair-burns-its-3-attempt-cap-d Agentflare-Item: 114 Co-authored-by: shiva <shiva@gosysinfo.tech>
Summary
McpHttpBackend,kind = "mcp_http"ingateway.toml) alongside the existing stdio one, usingrmcp's Streamable-HTTP client transport. Delete the old unimplementedhttp_apiREST-bridge stub outright — the gateway only ever speaks real MCP.CircuitBreakerused by both backends (was duplicated-in-waiting between stdio and the new HTTP backend).McpStdioBackend:ensure_connectednow returns the still-locked guard instead of releasing and re-acquiring it, closing a window where a concurrent caller could invalidate the "connected" assumption between the two lock acquisitions. The new HTTP backend was built on this corrected pattern from the start.gateway_executecall args before attempting to connect (in both backends), not after — a malformed call now fails instantly instead of first paying for a connect attempt that was never going to succeed.gateway_executecalls (server, tool, args hash, outcome) written alongside the registry db — raw args are never logged, only a hash.gateway_search, since a downstream (or compromised) server fully controls that data.Test plan
cargo build --workspacecargo test --workspace(all green — 207+ tests across the workspace, including new HTTP backend, circuit-breaker, TOCTOU-fix, audit-log, and reconnect integration tests)Summary by CodeRabbit