feat(proxy): lazy input_schema discovery for sandbox-supplied-URL MCP providers (closes #135)

[claudio-michel]

Summary

Closes #135. For handler = "mcp" providers whose upstream URL is sandbox-supplied (the mcp_url_env / X-Ati-Upstream-Url model from #124), the proxy cannot discover tools at boot. So operators declare tool names statically; the gap was the schemas, which surfaced as input_schema: null until they were hand-mirrored.

This PR has the proxy lazily fetch missing schemas on demand and cache them per (provider, upstream_url). The change runs on three surfaces — GET /tools, GET /tools/:name, and MCP JSON-RPC tools/list — so both ati tool info and live agents benefit.

Acceptance criteria from #135 — all green

• ati tool info <provider>:<tool> returns the live inputSchema for a sandbox-URL MCP provider whose manifest declares the tool name but no input_schema. → tool_info_merges_upstream_schema_when_static_is_none
• No regression for OpenAPI-handler tools (schema from spec) or local-mode MCP (already works). → tool_info_without_header_falls_back_to_static_null (header absent = today's behavior) + the existing test suite (1211/1211 green)
• Boot still succeeds when the upstream URL is unresolvable (falls back to the static entry / null, as today). → tool_info_silently_falls_back_when_url_outside_allowlist + negative-result caching avoids retrying a known-broken upstream

What I shipped

New code (all in src/proxy/server.rs):

• ProxyState::lazy_schema_cache: LazySchemaCache — keyed (provider, url), value Option<HashMap<tool_name, Value>>. Negative results cached too.
• lazy_fetch_schemas(state, provider, url) — runs one tools/list via mcp_client::McpClient::connect_with_gen (same path the /call handler uses, so allowlist semantics are unified). Caches positive + negative.
• extract_upstream_url(state, provider, headers) — REST-handler counterpart to the JSON-RPC resolve_upstream_override. An allowlist miss silently becomes None (no 4xx on tool-info).
• resolve_lazy_input_schema(state, provider, tool, headers) — single helper used by handle_tool_info; fast-paths the common case (static schema present).
• The REST GET /tools list handler and the MCP tools/list JSON-RPC arm batch by provider so a multi-tool provider only triggers one upstream tools/list per request.

Test coverage in tests/lazy_schema_discovery_test.rs (six tests):

1. Merge upstream schema when static is none — the acceptance criterion
2. No header → fall back to static null (today's behavior)
3. Tool absent from upstream tools/list → keep static null
4. URL outside allowlist → 200 + null (REST handlers must not 4xx)
5. Cache: second request served without re-dialing upstream
6. MCP JSON-RPC tools/list also serves the merged schema

The tests stand up a wiremock server that impersonates a Streamable HTTP MCP server (initialize / notifications/initialized / tools/list).

Test plan

cargo test               # 1211 / 1211 green
cargo clippy --all-features --lib --tests -- -D warnings  # clean
cargo fmt --all --check  # clean

• CI green
• Confirm ati tool info parcha_tools:firecrawl_parse_document against a Parcha sandbox returns the full schema (issue Proxy: lazy per-call input_schema discovery for sandbox-supplied-URL MCP providers #135 concrete repro)
• Verify parcha PR #6424's hand-maintained [tools.input_schema] mirrors can be removed once this lands

What this does NOT do (deliberately)

• Manifest schema change. None needed — existing mcp_url_env + per-provider allowlist do all the wiring.
• Touch local-mode MCP or stdio transport. Boot-time discovery still handles those; lazy path only fires for handler = "mcp" + mcp_url_env.
• Touch the /call path. Already does lazy discovery via call.rs::execute_local's "tool not in static index" code; this PR is for the info/list surface only.
• Open up the REST handlers to 4xx behavior. An allowlist miss on /call is 403; on /tools/:name it stays 200 with the static fallback. The two paths have different threat models (call dispatches credentials; info just describes).

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR adds lazy inputSchema discovery for handler = "mcp" + mcp_url_env providers, where the proxy can't know tool schemas at boot because the upstream URL is sandbox-supplied. On a cache miss, the three info surfaces (GET /tools, GET /tools/:name, and MCP JSON-RPC tools/list) now dial the upstream via the existing connect_with_gen path, cache the result per (provider, upstream_url) for the process lifetime, and fall back silently to null on any failure.

• src/proxy/server.rs: adds LazySchemaCache to ProxyState; implements lazy_fetch_schemas, extract_upstream_url, and resolve_lazy_input_schema; applies per-provider batching in the two list handlers and the single-tool info handler.
• tests/lazy_schema_discovery_test.rs: six new wiremock-backed integration tests covering all acceptance criteria from Proxy: lazy per-call input_schema discovery for sandbox-supplied-URL MCP providers #135.
• Remaining test files: mechanical addition of lazy_schema_cache field initialization to ProxyState struct literals.

Confidence Score: 4/5

Safe to merge with the minor batching fix applied; the fallback-to-null contract guarantees no new 5xx surface on the info endpoints.

The core logic is sound: allowlist validation happens before discovery, the mutex is never held across an await, and the negative-result cache prevents hammering a broken upstream. When lazy_fetch_schemas returns None (upstream failure), the provider is not recorded in schemas_by_provider, so N-1 remaining tools from that provider each acquire the mutex for a redundant cache-read rather than being short-circuited — a concrete but non-correctness-affecting defect replicated in both list handlers.

src/proxy/server.rs — the two batching loops (handle_tools_list and the tools/list MCP arm) should record a sentinel for failed providers.

Important Files Changed

Filename	Overview
src/proxy/server.rs	Core change: adds LazySchemaCache to ProxyState, implements lazy_fetch_schemas/extract_upstream_url/resolve_lazy_input_schema, applies lazy discovery to all three info surfaces. Logic correct; mutex never held across awaits; allowlist validated before discovery. Minor: failed providers not recorded in schemas_by_provider, causing redundant per-tool cache lookups.
tests/lazy_schema_discovery_test.rs	Six new wiremock-backed integration tests covering all #135 acceptance criteria. build_app docstring has a typo describing both tools as 'without schema'.
tests/common/mod.rs	Mechanical addition of lazy_schema_cache field to shared test-app builders.
tests/proxy_server_test.rs	Mechanical addition of lazy_schema_cache field to inline ProxyState constructions.

Sequence Diagram

sequenceDiagram
    participant Client
    participant Proxy as ATI Proxy
    participant Cache as lazy_schema_cache
    participant Allowlist as extract_upstream_url
    participant MCP as Sandbox MCP Upstream

    Client->>Proxy: GET /tools/:name (X-Ati-Upstream-Url: url)
    Proxy->>Allowlist: resolve_upstream_override(provider, url)
    alt URL rejected or absent
        Allowlist-->>Proxy: None
        Proxy-->>Client: 200 input_schema null
    else URL allowed
        Allowlist-->>Proxy: Allow(url)
        Proxy->>Cache: lock get((provider, url))
        alt Cache hit
            Cache-->>Proxy: Some(map) or None
            Proxy-->>Client: 200 input_schema from map or null
        else Cache miss
            Cache-->>Proxy: absent
            Proxy->>MCP: "connect_with_gen -> initialize -> tools/list"
            MCP-->>Proxy: Vec of Tool with inputSchema
            Proxy->>Cache: lock insert((provider, url), map or None)
            Proxy-->>Client: 200 input_schema from map or null
        end
    end

Loading

_{Reviews (1): Last reviewed commit: "feat(proxy): lazy input_schema discovery..." | Re-trigger Greptile}

[greptile-apps]

Batching short-circuit doesn't fire on upstream failure. When lazy_fetch_schemas returns None (connect or list failed), the provider is not recorded in schemas_by_provider. For a provider with N tools that all lack a static schema, the first tool caches the negative result, but the remaining N-1 tools still re-enter the loop and individually call lazy_fetch_schemas — each grabbing the mutex for a cache read. Replacing the if let Some with unwrap_or_default() inserts an empty map for the failed provider, so the contains_key guard on subsequent tools fires correctly and the N-1 extra mutex acquisitions are avoided.

Suggested change

	if let Some(map) = lazy_fetch_schemas(&state, provider, &upstream_url).await {
	schemas_by_provider.insert(provider.name.clone(), map);
	}
	}
	let tools: Vec<Value> = filtered
	let map = lazy_fetch_schemas(&state, provider, &upstream_url)
	.await
	.unwrap_or_default();
	schemas_by_provider.insert(provider.name.clone(), map);
	}
	let tools: Vec<Value> = filtered

[greptile-apps]

Same batching short-circuit gap in the MCP JSON-RPC tools/list arm. If lazy_fetch_schemas returns None, subsequent tools from the same provider bypass the contains_key guard and each acquire the mutex for an individual negative-cache hit.

Suggested change

	if let Some(map) = lazy_fetch_schemas(&state, provider, &url).await {
	schemas_by_provider.insert(provider.name.clone(), map);
	}
	}
	let mcp_tools: Vec<Value> = visible_tools
	let map = lazy_fetch_schemas(&state, provider, &url)
	.await
	.unwrap_or_default();
	schemas_by_provider.insert(provider.name.clone(), map);
	}
	let mcp_tools: Vec<Value> = visible_tools

[greptile-apps]

Docstring typo: "one without schema, one without" appears to be a copy-paste error. The test body and tool_info_keeps_null_when_upstream_omits_the_tool make clear the intent is to verify both the "present upstream schema" and "absent upstream schema" paths.

Suggested change

	/// Build an axum Router wired with the issue #135 manifest pattern:
	/// MCP provider with `mcp_url_env`, two statically named tools (one
	/// without schema, one without — both meant to be filled by lazy
	/// discovery), and a hostname-glob allowlist that admits the wiremock
	/// server's host.
	/// Build an axum Router wired with the issue #135 manifest pattern:
	/// MCP provider with `mcp_url_env`, two statically named tools (one
	/// with an upstream inputSchema, one without — no static schemas on
	/// either, so the issue #135 lazy-discovery path fires for both),
	/// and a hostname-glob allowlist that admits the wiremock server's host.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!