From c10a2195afdf86e0ade17fad53a509de577224d4 Mon Sep 17 00:00:00 2001 From: chopratejas Date: Sun, 3 May 2026 16:22:07 -0700 Subject: [PATCH] fix(proxy): PR-D4 native Vertex publisher path + ADC bearer auth MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a Rust-native Vertex AI publisher route ahead of the LiteLLM Python converter (which dropped `thinking`, `redacted_thinking`, `document`, `image`, `server_tool_use`, `mcp_tool_use` block kinds — the P4-37 / P4-38 bug). After this PR the Vertex `:rawPredict` and `:streamRawPredict` calls survive byte-equal upstream and benefit from the live-zone Anthropic dispatcher (PR-B-series) running over the body — same behaviour as `/v1/messages`. New module `crates/headroom-proxy/src/vertex/`: - `mod.rs` — single dispatch handler at the `/v1beta1/.../models/:model_action` route. Splits the trailing `:` segment with `str::rsplit_once(':')` (no regex) and flips an `attach_sse_tee` flag to dispatch to the streaming or non-streaming arm. Both verbs share one axum route shape because matchit can't distinguish two patterns that overlap on a parameter. - `envelope.rs` — `VertexEnvelope` parser. Confirms `anthropic_version` present + `model` field absent (the two fingerprints of the Vertex envelope vs `/v1/messages`). - `adc.rs` — `TokenSource` trait + `GcpAdcTokenSource` (production, `gcp_auth` 0.12) + `StaticTokenSource` (tests). Caches tokens with a 60s refresh-ahead-of-expiry window. Emits structured `event = "vertex_adc_token_refreshed"` per refresh. - `raw_predict.rs` — POST handler + shared `forward_vertex_request`. Buffers body, parses envelope, runs live-zone Anthropic compression, fetches ADC bearer, attaches `Authorization: Bearer ` (overwrites client-supplied Authorization header), forwards. SSE telemetry tee for the streaming verb reuses PR-C1's `AnthropicStreamState` directly (Vertex streams plain SSE, unlike Bedrock's binary EventStream). - `stream_raw_predict.rs` — module-level docs + alias to the shared dispatcher (the streaming-vs-non-streaming difference is one boolean flag inside the shared forwarder). Modifications: - `proxy.rs::build_app` — registers the single Vertex route. - `proxy.rs::AppState` — new `vertex_token_source: Arc` field. Production constructs `GcpAdcTokenSource` lazily (no GCP call until first `bearer()`); tests inject `StaticTokenSource` via the new `AppState::with_token_source` helper. - `config.rs` — adds `--vertex-region` / `HEADROOM_PROXY_VERTEX_REGION` (default `us-central1`, observability tag only — the upstream URL is `--upstream`) and `--vertex-adc-scope` / `HEADROOM_PROXY_VERTEX_ADC_SCOPE` (default `cloud-platform`). - `Cargo.toml` (workspace + proxy) — adds `gcp_auth = "0.12"` and `async-trait = "0.1"`. - `tests/common/mod.rs` — `start_proxy_with_state` accepts both config + state customizers; `install_static_token_source` helper for tests. `crates/headroom-proxy/tests/integration_vertex_raw_predict.rs` — all five tests pass: 1. `native_envelope_round_trip_byte_equal` — Vertex-shape body (with `anthropic_version`, no `model`) round-trips SHA-256 byte-equal upstream. 2. `adc_bearer_token_signed_correctly` — `Authorization: Bearer ` reaches upstream verbatim and OVERWRITES a client-supplied Authorization header. 3. `thinking_block_preserved` — request with `thinking` (incl. signature) + `redacted_thinking` (incl. opaque `data`) blocks round-trips byte-equal even with `LiveZone` compression mode enabled. This is the P4-37 / P4-38 teeth. 4. `stream_raw_predict_sse_handled` — `:streamRawPredict` proxies an Anthropic SSE response (full `message_start` → `content_block_delta` → `message_stop` sequence) back to the client without corruption; SSE content-type preserved end-to-end; bearer attached. 5. (bonus, no-silent-fallback contract) `adc_failure_returns_5xx_no_silent_forward` — when the token source returns `Err`, the proxy returns 5xx and never reaches upstream. Verifies the `event = "vertex_adc_fetch_failed"` error path. Workspace: `cargo test --workspace` green; `cargo clippy --workspace -- -D warnings` clean; `make ci-precheck` passes. - No silent fallbacks: ADC failure → structured 5xx, never an unauthenticated forward. - No hardcodes: every knob (region, ADC scope, upstream URL) is CLI-flag + env-var configurable. - No regexes: axum path parameters + `str::rsplit_once` only. - Comprehensive structured logs: `event` field on every decision point — `vertex_envelope_parsed`, `vertex_envelope_invalid`, `vertex_compression_skipped`, `vertex_compression_applied`, `vertex_adc_token_refreshed`, `vertex_adc_fetch_failed`, `vertex_streaming_pipeline_active`, `vertex_sse_stream_closed`, `vertex_forwarded`, `vertex_unknown_verb`, etc. - Performant: no body clone; ADC token cached + refreshed ahead-of-expiry, not fetched per request. - Comprehensive tests: realistic Anthropic block content (signature payload, redacted_thinking opaque blob) in `thinking_block_preserved`. The local `gcloud auth application-default print-access-token` returns no credentials, so manual validation against a real Vertex endpoint is not possible in this PR. Follow-up: the user runs `gcloud auth application-default login` once and exercises a live Vertex request — should be a no-code-change check. PR-D1 (Bedrock native) is running concurrently and will land its own envelope module at `crates/headroom-proxy/src/bedrock/envelope.rs`. The two envelope modules are intentionally siblings (not a shared trait) — the shapes differ (Bedrock has a different `anthropic_version` value, no `model` field, AWS SigV4 instead of GCP ADC), and a premature shared abstraction would obscure the provider-specific contracts. Whichever PR merges second rebases without conflict. Retires P4-38 (and the Vertex parts of P4-39); marketplace BYOC pitch (per project memory) gets one more native provider. --- Cargo.lock | 158 +++++ Cargo.toml | 7 + crates/headroom-proxy/Cargo.toml | 5 + crates/headroom-proxy/src/bedrock/invoke.rs | 12 + .../src/bedrock/invoke_streaming.rs | 6 + crates/headroom-proxy/src/config.rs | 47 ++ crates/headroom-proxy/src/lib.rs | 1 + crates/headroom-proxy/src/proxy.rs | 47 ++ crates/headroom-proxy/src/vertex/adc.rs | 315 ++++++++++ crates/headroom-proxy/src/vertex/envelope.rs | 167 ++++++ crates/headroom-proxy/src/vertex/mod.rs | 277 +++++++++ .../headroom-proxy/src/vertex/raw_predict.rs | 562 ++++++++++++++++++ .../src/vertex/stream_raw_predict.rs | 49 ++ crates/headroom-proxy/tests/common/mod.rs | 18 + .../tests/integration_vertex_raw_predict.rs | 556 +++++++++++++++++ 15 files changed, 2227 insertions(+) create mode 100644 crates/headroom-proxy/src/vertex/adc.rs create mode 100644 crates/headroom-proxy/src/vertex/envelope.rs create mode 100644 crates/headroom-proxy/src/vertex/mod.rs create mode 100644 crates/headroom-proxy/src/vertex/raw_predict.rs create mode 100644 crates/headroom-proxy/src/vertex/stream_raw_predict.rs create mode 100644 crates/headroom-proxy/tests/integration_vertex_raw_predict.rs diff --git a/Cargo.lock b/Cargo.lock index 6f31fe1df..ae1d75c38 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -55,6 +55,15 @@ version = "0.2.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + [[package]] name = "anes" version = "0.1.6" @@ -817,6 +826,20 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chrono" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + [[package]] name = "ciborium" version = "0.2.2" @@ -1624,6 +1647,32 @@ dependencies = [ "slab", ] +[[package]] +name = "gcp_auth" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2b3d0b409a042a380111af38136310839af8ac1a0917fb6e84515ed1e4bf3ee" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "chrono", + "http 1.4.0", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "ring", + "rustls-pki-types", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tracing", + "tracing-futures", + "url", +] + [[package]] name = "generic-array" version = "0.14.7" @@ -1810,6 +1859,7 @@ dependencies = [ name = "headroom-proxy" version = "0.1.0" dependencies = [ + "async-trait", "aws-config", "aws-credential-types", "aws-sigv4", @@ -1821,6 +1871,7 @@ dependencies = [ "crc32fast", "futures", "futures-util", + "gcp_auth", "headroom-core", "http 1.4.0", "http-body-util", @@ -2075,6 +2126,30 @@ dependencies = [ "tracing", ] +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + [[package]] name = "icu_collections" version = "2.2.0" @@ -2914,6 +2989,26 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pin-project" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "pin-project-lite" version = "0.2.17" @@ -4368,6 +4463,16 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project", + "tracing", +] + [[package]] name = "tracing-log" version = "0.2.0" @@ -4884,12 +4989,65 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "windows-link" version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + [[package]] name = "windows-sys" version = "0.52.0" diff --git a/Cargo.toml b/Cargo.toml index 98a4d251a..412561c5f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -68,3 +68,10 @@ aws-credential-types = { version = "1", default-features = false } # accepts `&Identity`. Pinning the version explicitly avoids a # silent semver bump from the transitive dep tree. aws-smithy-runtime-api = { version = "1", default-features = false, features = ["client"] } +# PR-D4: Vertex publisher path uses GCP Application Default Credentials +# (ADC) → bearer token for the `Authorization: Bearer ` header. +# `gcp_auth` resolves the chain (gcloud user creds, GCE/GKE metadata +# server, service-account JSON, workload-identity federation) without +# us baking provider-specific knowledge in. The token source is wrapped +# in a `TokenSource` trait so tests inject a static-token mock. +gcp_auth = "0.12" diff --git a/crates/headroom-proxy/Cargo.toml b/crates/headroom-proxy/Cargo.toml index 19095dcc5..40d879576 100644 --- a/crates/headroom-proxy/Cargo.toml +++ b/crates/headroom-proxy/Cargo.toml @@ -71,6 +71,11 @@ sha2 = "0.10" # is the de-facto Rust LRU crate; minimal surface, no dependencies of # our own beyond `hashbrown` (which we already pull transitively). lru = "0.12" +# PR-D4: GCP Application Default Credentials → bearer token for +# Vertex `:rawPredict` / `:streamRawPredict`. See workspace +# Cargo.toml for rationale. +gcp_auth = { workspace = true } +async-trait = "0.1" [dev-dependencies] tower = { workspace = true, features = ["util"] } diff --git a/crates/headroom-proxy/src/bedrock/invoke.rs b/crates/headroom-proxy/src/bedrock/invoke.rs index fd74a142e..9af7baff0 100644 --- a/crates/headroom-proxy/src/bedrock/invoke.rs +++ b/crates/headroom-proxy/src/bedrock/invoke.rs @@ -547,6 +547,12 @@ mod tests { // unit test never observes drift, but `AppState` requires // the field to be populated. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + // PR-D4: unit tests for the Bedrock URL builder don't + // touch the Vertex route, but `AppState` is one struct + // — supply a dummy token source so the test compiles. + vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( + "test".to_string(), + )), }; let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke" .parse() @@ -576,6 +582,12 @@ mod tests { // PR-E6: see above — drift detector is unused by this // test; we just satisfy the struct shape. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + // PR-D4: unit tests for the Bedrock URL builder don't + // touch the Vertex route, but `AppState` is one struct + // — supply a dummy token source so the test compiles. + vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( + "test".to_string(), + )), }; let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke" .parse() diff --git a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs index d17bdcf02..27dff6062 100644 --- a/crates/headroom-proxy/src/bedrock/invoke_streaming.rs +++ b/crates/headroom-proxy/src/bedrock/invoke_streaming.rs @@ -974,6 +974,12 @@ mod tests { // PR-E6: drift detector is unused by this URL-builder // unit test; small capacity to satisfy the struct shape. drift_state: crate::cache_stabilization::drift_detector::DriftState::new(8), + // PR-D4: unit tests for the Bedrock URL builder don't + // touch the Vertex route, but `AppState` is one struct + // — supply a dummy token source so the test compiles. + vertex_token_source: std::sync::Arc::new(crate::vertex::StaticTokenSource::new( + "test".to_string(), + )), }; let uri: Uri = "/model/anthropic.claude-3-haiku-20240307-v1:0/invoke-with-response-stream" .parse() diff --git a/crates/headroom-proxy/src/config.rs b/crates/headroom-proxy/src/config.rs index 9a9bc47cb..a1509e9d3 100644 --- a/crates/headroom-proxy/src/config.rs +++ b/crates/headroom-proxy/src/config.rs @@ -370,6 +370,39 @@ pub struct CliArgs { action = clap::ArgAction::Set, )] pub bedrock_validate_eventstream_crc: bool, + + /// Phase D PR-D4: GCP Vertex region for the publisher path + /// (`{region}-aiplatform.googleapis.com`). Default `us-central1` + /// (matches the GCP-published default region for Anthropic + /// publisher models). The proxy does NOT auto-construct the + /// regional URL — that's an `--upstream` decision the operator + /// makes once at startup. This flag is exposed for structured + /// logging + observability so dashboards can group Vertex traffic + /// by region without parsing the upstream URL. + /// + /// Source priority: CLI flag → `HEADROOM_PROXY_VERTEX_REGION` + /// env var → default (`us-central1`). + #[arg( + long = "vertex-region", + env = "HEADROOM_PROXY_VERTEX_REGION", + default_value = "us-central1" + )] + pub vertex_region: String, + + /// Phase D PR-D4: OAuth scope to request from GCP ADC. Defaults + /// to `cloud-platform`, the broad scope `gcloud` itself uses for + /// ADC. Operators with tighter IAM postures can scope down to + /// `cloud-platform.read-only` etc., but Vertex `:rawPredict` + /// requires write so most deployments use the default. + /// + /// Source priority: CLI flag → `HEADROOM_PROXY_VERTEX_ADC_SCOPE` + /// env var → default (`cloud-platform`). + #[arg( + long = "vertex-adc-scope", + env = "HEADROOM_PROXY_VERTEX_ADC_SCOPE", + default_value = "https://www.googleapis.com/auth/cloud-platform" + )] + pub vertex_adc_scope: String, } fn parse_duration(s: &str) -> Result { @@ -441,6 +474,14 @@ pub struct Config { /// PR-D2: validate prelude + message CRC32 on inbound Bedrock /// EventStream frames. Default `true`. Off only for debugging. pub bedrock_validate_eventstream_crc: bool, + /// PR-D4: GCP Vertex region tag (e.g. `us-central1`). Surfaced + /// in structured logs only — the actual upstream URL comes from + /// `Config::upstream`. Operators set this so observability + /// dashboards can group Vertex traffic by region. + pub vertex_region: String, + /// PR-D4: GCP ADC OAuth scope used when fetching the bearer + /// token. Default `https://www.googleapis.com/auth/cloud-platform`. + pub vertex_adc_scope: String, } impl Config { @@ -474,6 +515,8 @@ impl Config { bedrock_endpoint: args.bedrock_endpoint, aws_profile: args.aws_profile, bedrock_validate_eventstream_crc: args.bedrock_validate_eventstream_crc, + vertex_region: args.vertex_region, + vertex_adc_scope: args.vertex_adc_scope, } } @@ -515,6 +558,10 @@ impl Config { // PR-D2: production default — validate every CRC. Tests // that exercise corruption paths flip this off per-case. bedrock_validate_eventstream_crc: true, + // PR-D4: default Vertex region (used for log tagging + // only; the upstream URL is `upstream`). + vertex_region: "us-central1".to_string(), + vertex_adc_scope: "https://www.googleapis.com/auth/cloud-platform".to_string(), } } } diff --git a/crates/headroom-proxy/src/lib.rs b/crates/headroom-proxy/src/lib.rs index 1b2b1e682..3d86580de 100644 --- a/crates/headroom-proxy/src/lib.rs +++ b/crates/headroom-proxy/src/lib.rs @@ -13,6 +13,7 @@ pub mod observability; pub mod proxy; pub mod responses_items; pub mod sse; +pub mod vertex; pub mod websocket; pub use config::Config; diff --git a/crates/headroom-proxy/src/proxy.rs b/crates/headroom-proxy/src/proxy.rs index b661127a7..0dce82ecb 100644 --- a/crates/headroom-proxy/src/proxy.rs +++ b/crates/headroom-proxy/src/proxy.rs @@ -39,6 +39,14 @@ use headroom_core::auth_mode::classify as classify_auth_mode; /// and Phase B's live-zone dispatcher will introduce its own state /// (per-block compressor registry) — the old ICM-shaped field would /// not have been reused. +/// +/// PR-D4 adds `vertex_token_source`: an `Arc` used +/// by the Vertex `:rawPredict` / `:streamRawPredict` handlers to +/// resolve a GCP ADC bearer token. Production wires +/// [`crate::vertex::adc::GcpAdcTokenSource`] (lazy ADC chain +/// resolution + cached tokens with refresh-ahead-of-expiry); tests +/// inject [`crate::vertex::adc::StaticTokenSource`] so they never +/// hit real GCP. #[derive(Clone)] pub struct AppState { pub config: Arc, @@ -57,6 +65,12 @@ pub struct AppState { /// request body — so this can be cloned freely into every handler /// path that buffers the body. pub drift_state: DriftState, + /// PR-D4: GCP ADC bearer-token source for Vertex routes. Default: + /// [`crate::vertex::adc::GcpAdcTokenSource`] constructed lazily; + /// the actual ADC chain is only resolved when the first Vertex + /// route hits `bearer()`. Tests override via + /// [`AppState::with_token_source`]. + pub vertex_token_source: Arc, } /// PR-E6: maximum number of sessions tracked by the drift detector @@ -80,11 +94,18 @@ impl AppState { .build() .map_err(ProxyError::Upstream)?; + // PR-D4: lazy ADC token source. Provider resolution is + // deferred to first `bearer()` call so proxy startup stays + // cheap when no Vertex route is exercised. + let vertex_token_source: Arc = + Arc::new(crate::vertex::adc::GcpAdcTokenSource::new()); + Ok(Self { config: Arc::new(config), client, bedrock_credentials: None, drift_state: DriftState::new(DRIFT_DETECTOR_CAPACITY), + vertex_token_source, }) } @@ -97,6 +118,18 @@ impl AppState { self.bedrock_credentials = Some(Arc::new(creds)); self } + + /// Test helper: build an `AppState` with an explicit token source. + /// Lets the integration tests substitute a `StaticTokenSource` so + /// the test suite never hits real GCP. + pub fn with_token_source( + config: Config, + token_source: Arc, + ) -> Result { + let mut s = Self::new(config)?; + s.vertex_token_source = token_source; + Ok(s) + } } /// Build the axum app. `/healthz` and `/healthz/upstream` are intercepted; @@ -133,6 +166,20 @@ pub fn build_app(state: AppState) -> Router { .route( "/v1/responses", post(crate::handlers::responses::handle_responses), + ) + // PR-D4: native Vertex publisher path. The Vertex AI Anthropic + // publisher endpoints look like + // `POST /v1beta1/projects/{p}/locations/{l}/publishers/anthropic/models/{m}:rawPredict` + // (and `:streamRawPredict`). The trailing `:` is awkward + // in axum's `:param` syntax, so we capture the entire trailing + // segment as `:model_action` and split on the last `:` inside + // the dispatcher. Both verbs share the same axum route shape + // — matchit can't distinguish two patterns that overlap on the + // literal parameter. The verb dispatch lives in + // [`crate::vertex::handle_vertex_predict_dispatch`]. + .route( + "/v1beta1/projects/:project/locations/:location/publishers/anthropic/models/:model_action", + post(crate::vertex::handle_vertex_predict_dispatch), ); // PR-D1: native AWS Bedrock InvokeModel route. Mounts only when diff --git a/crates/headroom-proxy/src/vertex/adc.rs b/crates/headroom-proxy/src/vertex/adc.rs new file mode 100644 index 000000000..2ea2719d6 --- /dev/null +++ b/crates/headroom-proxy/src/vertex/adc.rs @@ -0,0 +1,315 @@ +//! GCP Application Default Credentials → bearer token resolution. +//! +//! Vertex AI's `:rawPredict` / `:streamRawPredict` endpoints expect +//! `Authorization: Bearer ` where the JWT is a short-lived +//! Google-signed access token. The token is obtained from the ADC +//! provider chain (gcloud user creds, GCE/GKE metadata server, +//! service-account JSON via `GOOGLE_APPLICATION_CREDENTIALS`, +//! workload-identity federation, etc.) and refreshed before expiry. +//! +//! # Why an abstraction (not direct `gcp_auth::provider()`)? +//! +//! 1. **Tests must NOT hit real GCP.** Per project rule "no silent +//! fallbacks", we cannot dummy out the real provider in tests +//! silently — we need a mock implementation that is explicit and +//! distinct from production. The `TokenSource` trait gives us +//! exactly two impls: production (`GcpAdcTokenSource`) and tests +//! (`StaticTokenSource`). +//! 2. **Caching is policy.** `gcp_auth` exposes per-call token fetches. +//! The cache + refresh-ahead-of-expiry policy lives at this layer +//! so it's testable and tunable independent of the provider. +//! +//! # Refresh policy +//! +//! Tokens are refreshed when their remaining lifetime drops below +//! [`REFRESH_AHEAD_SECS`] (60s). That avoids the cliff-failure where +//! a request fired right at expiry races the upstream's clock. +//! `gcp_auth` itself caches internally, but we still wrap it so: +//! +//! - Tests can substitute a `StaticTokenSource`. +//! - We control the refresh-ahead window (gcp_auth's internal default +//! is implementation-defined and could change). +//! - We emit structured `event = "vertex_adc_token_refreshed"` logs +//! per refresh so operators can confirm the cache is live. +//! +//! # Failure mode +//! +//! When ADC fetch fails (no creds configured, metadata server +//! unreachable, IAM permission denied, etc.) the handler converts the +//! returned [`TokenSourceError`] to a structured 5xx response. We do +//! NOT silently forward without a token — per project rule +//! "no silent fallbacks", an unauthenticated forward to Vertex would +//! return a 401 from upstream that's harder to debug than our own +//! 502 with `event = "vertex_adc_fetch_failed"`. + +use std::sync::Arc; +use std::time::{Duration, SystemTime}; + +use async_trait::async_trait; +use thiserror::Error; +use tokio::sync::Mutex; + +/// Refresh tokens this many seconds before their expiry to avoid the +/// cliff-failure race. 60s comfortably covers a slow upstream and +/// clock skew. +pub const REFRESH_AHEAD_SECS: u64 = 60; + +/// Default OAuth2 scope for Vertex / cloud-platform calls. The +/// `cloud-platform` scope is the broadest and is what gcloud emits +/// for ADC by default. +pub const DEFAULT_VERTEX_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform"; + +/// Errors fetching / refreshing an ADC bearer token. Surfaced +/// verbatim by the handler as structured log events plus an HTTP 5xx +/// response. +#[derive(Debug, Error)] +pub enum TokenSourceError { + /// `gcp_auth` failed to resolve any provider in the ADC chain. + /// Common cause: developer never ran `gcloud auth application-default + /// login` and no service-account JSON is in scope. + #[error("gcp ADC provider initialization failed: {0}")] + ProviderInit(String), + /// The provider was resolved but `.token(scopes)` failed. + #[error("gcp ADC token fetch failed: {0}")] + Fetch(String), +} + +/// A source of bearer tokens for Vertex calls. Production uses +/// [`GcpAdcTokenSource`]; tests use [`StaticTokenSource`]. There is +/// NO blanket impl — every concrete impl must be explicit (no silent +/// fallback to a "default token"). +#[async_trait] +pub trait TokenSource: Send + Sync + std::fmt::Debug { + /// Return a non-empty bearer token suitable for the + /// `Authorization: Bearer ` header. The token is cached + /// internally; concurrent calls de-dup to a single fetch. + async fn bearer(&self) -> Result; +} + +/// Token + expiry pair held in the cache. +#[derive(Debug, Clone)] +struct CachedToken { + token: String, + expires_at: SystemTime, +} + +impl CachedToken { + /// `true` when the token has more than `REFRESH_AHEAD_SECS` of + /// life left. + fn fresh(&self) -> bool { + match self.expires_at.duration_since(SystemTime::now()) { + Ok(remaining) => remaining > Duration::from_secs(REFRESH_AHEAD_SECS), + // `expires_at` already past now → not fresh. + Err(_) => false, + } + } +} + +/// Production token source backed by `gcp_auth`'s default ADC chain. +/// +/// `gcp_auth::provider()` resolves lazily; the first `.bearer()` call +/// triggers the resolution and the result is memoized. Subsequent +/// calls re-use the provider and the cached token until its expiry +/// approaches. +pub struct GcpAdcTokenSource { + /// Configured OAuth scope. Defaults to `cloud-platform`. + scope: String, + /// Lazily-initialized provider. We wrap in a `Mutex>` + /// (rather than `OnceCell`) because the provider initialization + /// is fallible and we want the next call after a transient + /// failure to retry — not lock the cell to a permanent error. + provider: Mutex>>, + /// Cached token + expiry. `Mutex` is fine here: the critical + /// section (compare expiry, optionally refresh) is sub-microsecond + /// in the cache-hit case and the refresh-miss path is rate-limited + /// by the upstream metadata server anyway. + cached: Mutex>, +} + +impl std::fmt::Debug for GcpAdcTokenSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("GcpAdcTokenSource") + .field("scope", &self.scope) + .field( + "provider_initialized", + &self + .provider + .try_lock() + .map(|g| g.is_some()) + .unwrap_or(false), + ) + .finish() + } +} + +impl GcpAdcTokenSource { + /// Construct with the default `cloud-platform` scope. The + /// provider is NOT resolved here — we defer until the first + /// `.bearer()` call. That keeps proxy startup cheap when the + /// operator hasn't actually wired a Vertex route yet. + pub fn new() -> Self { + Self::with_scope(DEFAULT_VERTEX_SCOPE) + } + + /// Construct with an explicit scope. Used by tests / by operators + /// who want a narrower scope than `cloud-platform`. + pub fn with_scope(scope: impl Into) -> Self { + Self { + scope: scope.into(), + provider: Mutex::new(None), + cached: Mutex::new(None), + } + } + + /// Resolve the provider lazily. On success, stores it for re-use + /// and returns a clone of the `Arc`. On failure, returns the + /// error verbatim — the next call retries (no permanent lock). + async fn ensure_provider( + &self, + ) -> Result, TokenSourceError> { + let mut guard = self.provider.lock().await; + if let Some(p) = guard.as_ref() { + return Ok(p.clone()); + } + let provider = ::gcp_auth::provider() + .await + .map_err(|e| TokenSourceError::ProviderInit(e.to_string()))?; + *guard = Some(provider.clone()); + Ok(provider) + } +} + +impl Default for GcpAdcTokenSource { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl TokenSource for GcpAdcTokenSource { + async fn bearer(&self) -> Result { + // Fast path: cached + fresh. + { + let guard = self.cached.lock().await; + if let Some(c) = guard.as_ref() { + if c.fresh() { + return Ok(c.token.clone()); + } + } + } + + // Slow path: refresh. We re-take the lock around the fetch + // so concurrent callers serialize on a single network round + // trip (rather than thundering-herd the metadata server). + let mut guard = self.cached.lock().await; + // Double-check inside the lock — another waiter may have + // refreshed while we were queued. + if let Some(c) = guard.as_ref() { + if c.fresh() { + return Ok(c.token.clone()); + } + } + + let provider = self.ensure_provider().await?; + let scopes = [self.scope.as_str()]; + let token = provider + .token(&scopes) + .await + .map_err(|e| TokenSourceError::Fetch(e.to_string()))?; + let token_str = token.as_str().to_string(); + // `gcp_auth::Token::expires_at()` returns a `chrono::DateTime`. + // Convert to `SystemTime` via the UNIX timestamp; `chrono` clamps + // to its valid range so this never panics. Negative timestamps + // (pre-epoch) are not legal for token expiries; if encountered we + // treat the token as already expired so the next call refreshes. + let expires_at = { + let dt = token.expires_at(); + let unix_ts = dt.timestamp(); + if unix_ts < 0 { + tracing::warn!( + event = "vertex_adc_token_negative_expiry", + expires_at_unix = unix_ts, + "gcp_auth token expires_at is pre-epoch; treating as already-expired" + ); + SystemTime::UNIX_EPOCH + } else { + SystemTime::UNIX_EPOCH + Duration::from_secs(unix_ts as u64) + } + }; + let lifetime_remaining = expires_at + .duration_since(SystemTime::now()) + .map(|d| d.as_secs()) + .unwrap_or(0); + tracing::info!( + event = "vertex_adc_token_refreshed", + scope = %self.scope, + lifetime_remaining_secs = lifetime_remaining, + "vertex ADC bearer token refreshed" + ); + *guard = Some(CachedToken { + token: token_str.clone(), + expires_at, + }); + Ok(token_str) + } +} + +/// Static-token mock for tests. Production callers MUST NOT +/// instantiate this — there is no fallback to a default value, and +/// the test harness wires it explicitly. +#[derive(Debug, Clone)] +pub struct StaticTokenSource { + token: String, +} + +impl StaticTokenSource { + pub fn new(token: impl Into) -> Self { + Self { + token: token.into(), + } + } +} + +#[async_trait] +impl TokenSource for StaticTokenSource { + async fn bearer(&self) -> Result { + Ok(self.token.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn static_token_source_returns_token() { + let src = StaticTokenSource::new("test-bearer-abc123"); + let t = src.bearer().await.expect("bearer"); + assert_eq!(t, "test-bearer-abc123"); + // Multiple calls return the same value. + let t2 = src.bearer().await.expect("bearer 2"); + assert_eq!(t2, "test-bearer-abc123"); + } + + #[test] + fn cached_token_freshness_window() { + let now = SystemTime::now(); + let fresh = CachedToken { + token: "x".into(), + expires_at: now + Duration::from_secs(REFRESH_AHEAD_SECS + 30), + }; + assert!(fresh.fresh()); + + let stale = CachedToken { + token: "x".into(), + expires_at: now + Duration::from_secs(REFRESH_AHEAD_SECS - 1), + }; + assert!(!stale.fresh()); + + let already_expired = CachedToken { + token: "x".into(), + expires_at: now - Duration::from_secs(1), + }; + assert!(!already_expired.fresh()); + } +} diff --git a/crates/headroom-proxy/src/vertex/envelope.rs b/crates/headroom-proxy/src/vertex/envelope.rs new file mode 100644 index 000000000..29d816017 --- /dev/null +++ b/crates/headroom-proxy/src/vertex/envelope.rs @@ -0,0 +1,167 @@ +//! Vertex publisher envelope parser. +//! +//! Vertex's Anthropic publisher path takes a body shaped almost +//! identically to the Anthropic Messages API, with two differences: +//! +//! 1. The body MUST contain `anthropic_version` (e.g. `"vertex-2023-10-16"`). +//! 2. The body MUST NOT contain `model` — the model id travels in +//! the URL path. +//! +//! We treat any other shape as wire-format drift: log loudly and +//! forward unmodified. This module's only job is to detect/confirm +//! the envelope shape so the handler decides whether the live-zone +//! Anthropic dispatcher should run. We never strip or rewrite +//! `anthropic_version`. +//! +//! # Performance note +//! +//! The envelope check uses `serde_json::from_slice::` once. +//! That's the same cost the live-zone dispatcher already pays; we're +//! not adding a parse, we're hoisting one read of two top-level keys +//! up before dispatching. No body clone — the parsed `Value` is +//! discarded and the dispatcher re-parses for byte-faithful surgery +//! (which uses `RawValue` to preserve exact bytes). + +use serde_json::Value; +use thiserror::Error; + +/// Errors detecting the Vertex envelope. None are panics; the handler +/// surfaces these as structured log events plus an HTTP error response. +#[derive(Debug, Error)] +pub enum VertexEnvelopeError { + #[error("body is not valid JSON: {0}")] + NotJson(#[from] serde_json::Error), + #[error("body is not a JSON object")] + NotObject, + #[error("body missing required field `anthropic_version`")] + MissingAnthropicVersion, + #[error( + "body has unexpected `model` field; Vertex carries the model in the URL path \ + (got model={got:?})" + )] + UnexpectedModelField { got: String }, +} + +/// Parsed-once view of the envelope's two distinguishing fields. The +/// dispatcher does not read the rest of the body through this struct — +/// it re-parses with `RawValue` for byte-faithful surgery. +#[derive(Debug, Clone)] +pub struct ParsedEnvelope { + /// Value of the `anthropic_version` field as it appeared in the + /// body. We do NOT validate against an allowlist — Vertex may + /// roll new versions without our involvement. + pub anthropic_version: String, + /// `true` if the body has at least one `messages[*]` entry. + /// Diagnostic only; the dispatcher walks `messages` independently. + pub has_messages: bool, +} + +/// Parse and validate the envelope. On success returns the +/// distinguishing-field view. On failure returns a structured error +/// the handler converts to a log event + HTTP response. +pub fn parse(body: &[u8]) -> Result { + let parsed: Value = serde_json::from_slice(body)?; + let obj = match parsed { + Value::Object(map) => map, + _ => return Err(VertexEnvelopeError::NotObject), + }; + + let anthropic_version = match obj.get("anthropic_version") { + Some(Value::String(s)) => s.clone(), + Some(other) => { + // Wire format says string; surface drift loudly. We still + // accept by stringifying so we don't reject novel + // representations Vertex might roll out — but log a + // structured event upstream of this call so operators + // see it. The handler is responsible for the warn log. + other.to_string() + } + None => return Err(VertexEnvelopeError::MissingAnthropicVersion), + }; + + if let Some(model) = obj.get("model") { + let got = match model { + Value::String(s) => s.clone(), + other => other.to_string(), + }; + return Err(VertexEnvelopeError::UnexpectedModelField { got }); + } + + let has_messages = matches!(obj.get("messages"), Some(Value::Array(arr)) if !arr.is_empty()); + + Ok(ParsedEnvelope { + anthropic_version, + has_messages, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn parses_minimal_envelope() { + let body = json!({ + "anthropic_version": "vertex-2023-10-16", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 64, + }); + let bytes = serde_json::to_vec(&body).unwrap(); + let env = parse(&bytes).expect("parse"); + assert_eq!(env.anthropic_version, "vertex-2023-10-16"); + assert!(env.has_messages); + } + + #[test] + fn rejects_missing_anthropic_version() { + let body = json!({ + "messages": [{"role": "user", "content": "hi"}], + }); + let bytes = serde_json::to_vec(&body).unwrap(); + let err = parse(&bytes).unwrap_err(); + assert!(matches!(err, VertexEnvelopeError::MissingAnthropicVersion)); + } + + #[test] + fn rejects_model_field_present() { + let body = json!({ + "anthropic_version": "vertex-2023-10-16", + "model": "claude-3-5-sonnet", + "messages": [], + }); + let bytes = serde_json::to_vec(&body).unwrap(); + let err = parse(&bytes).unwrap_err(); + match err { + VertexEnvelopeError::UnexpectedModelField { got } => { + assert_eq!(got, "claude-3-5-sonnet"); + } + other => panic!("wrong error: {other:?}"), + } + } + + #[test] + fn rejects_non_object_body() { + let bytes = b"[1,2,3]"; + let err = parse(bytes).unwrap_err(); + assert!(matches!(err, VertexEnvelopeError::NotObject)); + } + + #[test] + fn rejects_invalid_json() { + let bytes = b"not json at all {{"; + let err = parse(bytes).unwrap_err(); + assert!(matches!(err, VertexEnvelopeError::NotJson(_))); + } + + #[test] + fn empty_messages_array_marks_has_messages_false() { + let body = json!({ + "anthropic_version": "vertex-2023-10-16", + "messages": [], + }); + let bytes = serde_json::to_vec(&body).unwrap(); + let env = parse(&bytes).unwrap(); + assert!(!env.has_messages); + } +} diff --git a/crates/headroom-proxy/src/vertex/mod.rs b/crates/headroom-proxy/src/vertex/mod.rs new file mode 100644 index 000000000..115705e0c --- /dev/null +++ b/crates/headroom-proxy/src/vertex/mod.rs @@ -0,0 +1,277 @@ +//! Native Vertex AI publisher path — Phase D PR-D4. +//! +//! # What this module owns +//! +//! Vertex AI exposes Anthropic models via two POST verbs on a single +//! template: +//! +//! ```text +//! POST /v1beta1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:rawPredict +//! POST /v1beta1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:streamRawPredict +//! ``` +//! +//! Crucially, the **request body** is the Anthropic Messages envelope +//! with two surface differences from the Anthropic API: +//! +//! 1. The body carries an `anthropic_version` field (e.g. +//! `"vertex-2023-10-16"`) and **no** `model` field — the model id +//! travels in the URL path instead. +//! 2. Auth is GCP ADC (Application Default Credentials) → a short-lived +//! bearer token in `Authorization: Bearer `. There is no +//! Anthropic API key in scope. +//! +//! Everything else — `messages`, `system`, `tools`, `tool_choice`, +//! `cache_control`, `thinking`, `metadata.user_id`, `stop_sequences`, +//! `stream: true` for the streaming verb — round-trips byte-equal +//! with the live-zone compression pass running on top. +//! +//! # Why a native module (vs LiteLLM) +//! +//! The Python LiteLLM converter at `headroom/backends/litellm.py:486-628` +//! lossy-converts Anthropic ↔ OpenAI shapes for Vertex, dropping +//! `thinking`, `redacted_thinking`, `document`, `search_result`, +//! `image`, `server_tool_use`, `mcp_tool_use` block kinds. The bug +//! tracker entry P4-38 (mirrored from P4-37 for Bedrock) flags this; +//! PR-D4 retires it on the Rust side. +//! +//! The native module: +//! +//! - Buffers the request body up to `compression_max_body_bytes`. +//! - Parses the envelope (no `model` field check, `anthropic_version` +//! present) to confirm we're on the Vertex publisher shape. +//! - Routes to the **same** live-zone Anthropic dispatcher as +//! `/v1/messages` — the body shape is identical apart from +//! `anthropic_version` vs `model`. We feed the dispatcher a body +//! with a synthetic `model` set from the path's model id (so block +//! metadata lookups work) and re-emit the body without the +//! synthetic field on the way out. +//! - Resolves the ADC bearer token via [`adc::TokenSource`], cached +//! with a refresh window so back-to-back requests don't pay the +//! round-trip. +//! - Forwards to the configured Vertex endpoint +//! (`https://{region}-aiplatform.googleapis.com/...`) with the +//! bearer attached. +//! +//! # Module layout +//! +//! - [`adc`] — `TokenSource` trait, plus `GcpAdcTokenSource` +//! (production) and `StaticTokenSource` (tests). Caching and +//! refresh-ahead-of-expiry live here. +//! - [`raw_predict`] — POST handler for the non-streaming verb. +//! - [`stream_raw_predict`] — POST handler for the streaming verb. +//! Vertex uses **SSE** for streaming (unlike Bedrock's binary +//! EventStream — much simpler), so the existing PR-C1 +//! [`crate::sse::anthropic::AnthropicStreamState`] state machine +//! drives telemetry directly. +//! - [`envelope`] — minimal envelope parser; same shape as the future +//! PR-D1 Bedrock envelope module by design (the two PRs are running +//! in parallel and will reconcile at merge). +//! +//! # Routing +//! +//! Vertex's path uses a **colon-suffix verb** (`:rawPredict`) that is +//! awkward in axum's parameter syntax (where `:name` reserves `:` as +//! the parameter sigil). We register the routes with a single +//! parameter capturing the entire trailing segment and split on the +//! last `:` inside the handler — no regex, just `str::rsplit_once`. +//! See [`split_model_action`]. + +pub mod adc; +pub mod envelope; +pub mod raw_predict; +pub mod stream_raw_predict; + +pub use adc::{StaticTokenSource, TokenSource, TokenSourceError}; +pub use envelope::{ParsedEnvelope, VertexEnvelopeError}; + +use axum::body::Body; +use axum::extract::{ConnectInfo, Path, State}; +use axum::http::{HeaderMap, Method, StatusCode, Uri}; +use axum::response::Response; +use std::net::SocketAddr; + +use crate::proxy::AppState; + +/// Single axum handler mounted at the +/// `/v1beta1/projects/{project}/locations/{location}/publishers/anthropic/models/{model_action}` +/// path. The trailing `model_action` segment carries `:` +/// (the verb is the colon-suffix Vertex appends). We split on the +/// last `:` and dispatch to either [`raw_predict::handle_raw_predict`] +/// (logically) or [`stream_raw_predict::handle_stream_raw_predict`]. +/// +/// The two sub-handlers share most of their logic — see +/// [`raw_predict::forward_vertex_request`] — so the dispatch is a +/// single argument flip (`attach_sse_tee`). Routing the two verbs +/// through one axum handler keeps the path matcher unambiguous: matchit +/// can't distinguish two patterns that share the literal `model_action` +/// parameter shape. +pub async fn handle_vertex_predict_dispatch( + State(state): State, + ConnectInfo(client_addr): ConnectInfo, + Path((project, location, model_action)): Path<(String, String, String)>, + method: Method, + uri: Uri, + headers: HeaderMap, + body: Body, +) -> Response { + let request_id = headers + .get("x-request-id") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + let (model_id, verb_str) = match split_model_action(&model_action) { + Some(parts) => parts, + None => { + tracing::warn!( + event = "vertex_path_parse_failed", + request_id = %request_id, + path = %uri.path(), + segment = %model_action, + "vertex path final segment missing `:verb` separator" + ); + return Response::builder() + .status(StatusCode::NOT_FOUND) + .body(Body::from("vertex path: bad model_action")) + .expect("static"); + } + }; + + let verb = match VertexVerb::parse(verb_str) { + Some(v) => v, + None => { + tracing::warn!( + event = "vertex_unknown_verb", + request_id = %request_id, + verb = %verb_str, + "vertex path verb not recognized; only rawPredict / streamRawPredict are supported" + ); + return Response::builder() + .status(StatusCode::NOT_FOUND) + .body(Body::from("vertex: unknown verb")) + .expect("static"); + } + }; + + let attach_sse_tee = matches!(verb, VertexVerb::StreamRawPredict); + if attach_sse_tee { + tracing::info!( + event = "vertex_streaming_pipeline_active", + request_id = %request_id, + method = %method, + path = %uri.path(), + framer = "byte_level_sse", + state_machine = "anthropic", + "vertex streaming pipeline engaged: SSE framer + AnthropicStreamState telemetry tee" + ); + } + + raw_predict::forward_vertex_request( + state, + client_addr, + request_id, + method, + uri, + headers, + body, + raw_predict::VertexCallContext { + project, + location, + model_id: model_id.to_string(), + verb, + }, + attach_sse_tee, + ) + .await +} + +/// Split the trailing `:model_action` path segment into +/// `(model_id, verb)`. +/// +/// Vertex's path looks like +/// `.../models/claude-3-5-sonnet@20240620:rawPredict`. The model id may +/// itself contain `@` and other special characters, but Vertex never +/// puts a literal `:` inside a model id — the `:` is the verb +/// separator. We match on the **last** `:` for safety. +/// +/// Returns `None` when the segment carries no colon (unknown shape; +/// the handler logs and 404s). +pub fn split_model_action(segment: &str) -> Option<(&str, &str)> { + segment.rsplit_once(':') +} + +/// Recognized Vertex publisher verbs. Future verbs (e.g. `:countTokens`) +/// would extend this enum. The handler maps unknown verbs to +/// `event = "vertex_unknown_verb"` warn logs and 404s — never a silent +/// fallback to a "default" verb. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VertexVerb { + /// Non-streaming Anthropic Messages call. + RawPredict, + /// SSE-streaming Anthropic Messages call. + StreamRawPredict, +} + +impl VertexVerb { + pub fn parse(s: &str) -> Option { + match s { + "rawPredict" => Some(Self::RawPredict), + "streamRawPredict" => Some(Self::StreamRawPredict), + _ => None, + } + } + + pub fn as_str(self) -> &'static str { + match self { + Self::RawPredict => "rawPredict", + Self::StreamRawPredict => "streamRawPredict", + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn split_model_action_basic() { + assert_eq!( + split_model_action("claude-3-5-sonnet@20240620:rawPredict"), + Some(("claude-3-5-sonnet@20240620", "rawPredict")) + ); + assert_eq!( + split_model_action("claude-3-haiku@20240307:streamRawPredict"), + Some(("claude-3-haiku@20240307", "streamRawPredict")) + ); + } + + #[test] + fn split_model_action_no_colon_returns_none() { + assert_eq!(split_model_action("claude-3-5-sonnet"), None); + } + + #[test] + fn split_model_action_uses_last_colon() { + // Defensive: even if the model id were to contain a colon + // (Vertex doesn't emit such ids today), the verb is whatever + // follows the LAST colon. + assert_eq!( + split_model_action("weird:model:rawPredict"), + Some(("weird:model", "rawPredict")) + ); + } + + #[test] + fn vertex_verb_parse() { + assert_eq!( + VertexVerb::parse("rawPredict"), + Some(VertexVerb::RawPredict) + ); + assert_eq!( + VertexVerb::parse("streamRawPredict"), + Some(VertexVerb::StreamRawPredict) + ); + assert_eq!(VertexVerb::parse("predict"), None); + assert_eq!(VertexVerb::parse(""), None); + } +} diff --git a/crates/headroom-proxy/src/vertex/raw_predict.rs b/crates/headroom-proxy/src/vertex/raw_predict.rs new file mode 100644 index 000000000..6f46a7bfb --- /dev/null +++ b/crates/headroom-proxy/src/vertex/raw_predict.rs @@ -0,0 +1,562 @@ +//! POST handler for Vertex `:rawPredict` (non-streaming). +//! +//! Path: +//! ```text +//! POST /v1beta1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:rawPredict +//! ``` +//! +//! See [`super`] for the module-level rationale (envelope shape, ADC, +//! routing strategy). This handler: +//! +//! 1. Buffers the request body. +//! 2. Confirms the Vertex envelope shape (`anthropic_version` present, +//! `model` absent). On envelope mismatch, logs +//! `event = "vertex_envelope_invalid"` and returns 400. +//! 3. Runs live-zone Anthropic compression — the body is the same +//! Anthropic Messages shape `/v1/messages` accepts, just with +//! `anthropic_version` instead of `model`. The dispatcher +//! preserves `anthropic_version` byte-equal because the +//! `RawValue`-based surgery only rewrites `messages[*]` entries. +//! 4. Resolves the ADC bearer token (cached, refreshed ahead of +//! expiry) and attaches `Authorization: Bearer `. On ADC +//! failure, logs `event = "vertex_adc_fetch_failed"` and returns +//! 502 — never silently forwards unauthenticated. +//! 5. Forwards to the configured Vertex endpoint +//! (`https://{region}-aiplatform.googleapis.com/`). +//! 6. Streams the response body back unchanged. +//! +//! All decision points emit a structured `tracing::info!` / +//! `tracing::warn!` event so operators can confirm the pipeline is +//! engaged in dashboards. + +use axum::body::{to_bytes, Body}; +use axum::extract::{ConnectInfo, Path, State}; +use axum::http::{HeaderMap, Method, StatusCode, Uri}; +use axum::response::Response; +use std::net::SocketAddr; + +use crate::compression; +use crate::headers::{build_forward_request_headers, filter_response_headers}; +use crate::proxy::AppState; +use crate::vertex::{adc::TokenSourceError, envelope, split_model_action, VertexVerb}; + +/// Axum handler for the rawPredict route. +/// +/// Path parameters: +/// - `project` — GCP project ID. +/// - `location` — Vertex region (e.g. `us-central1`). +/// - `model_action` — combined `:` segment. We split +/// on the last `:` and dispatch on the verb. The streaming sibling +/// `:streamRawPredict` lives in [`super::stream_raw_predict`] and +/// is registered under the same route shape. +pub async fn handle_raw_predict( + State(state): State, + ConnectInfo(client_addr): ConnectInfo, + Path((project, location, model_action)): Path<(String, String, String)>, + method: Method, + uri: Uri, + headers: HeaderMap, + body: Body, +) -> Response { + let request_id = headers + .get("x-request-id") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()) + .unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + + let (model_id, verb) = match split_model_action(&model_action) { + Some(parts) => parts, + None => { + tracing::warn!( + event = "vertex_path_parse_failed", + request_id = %request_id, + path = %uri.path(), + segment = %model_action, + "vertex path final segment missing `:verb` separator" + ); + return error_response(StatusCode::NOT_FOUND, "vertex path: bad model_action"); + } + }; + let parsed_verb = match VertexVerb::parse(verb) { + Some(VertexVerb::RawPredict) => VertexVerb::RawPredict, + Some(VertexVerb::StreamRawPredict) => { + // Wrong handler. The router mounts both verbs at this + // path shape; the dispatcher inside [`super::stream_raw_predict`] + // wraps the streaming case. If we ever land here it's a + // routing bug — log and 500 rather than silently forward. + tracing::error!( + event = "vertex_verb_routing_bug", + request_id = %request_id, + verb = "streamRawPredict", + "streamRawPredict request reached non-streaming handler" + ); + return error_response( + StatusCode::INTERNAL_SERVER_ERROR, + "vertex routing bug: stream verb on non-stream handler", + ); + } + None => { + tracing::warn!( + event = "vertex_unknown_verb", + request_id = %request_id, + verb = %verb, + "vertex path verb not recognized; only rawPredict / streamRawPredict are supported" + ); + return error_response(StatusCode::NOT_FOUND, "vertex: unknown verb"); + } + }; + + forward_vertex_request( + state, + client_addr, + request_id, + method, + uri, + headers, + body, + VertexCallContext { + project, + location, + model_id: model_id.to_string(), + verb: parsed_verb, + }, + /* attach_sse_tee */ false, + ) + .await +} + +/// Carrier struct for the bits parsed out of the URL path; passed +/// down so logs and error paths share a consistent set of fields. +#[derive(Debug, Clone)] +pub(crate) struct VertexCallContext { + pub project: String, + pub location: String, + pub model_id: String, + pub verb: VertexVerb, +} + +/// Shared forwarder used by both `:rawPredict` and `:streamRawPredict` +/// handlers. Streaming-specific behaviour lives in the response-side +/// SSE tee in [`crate::proxy::forward_http`] (which we don't reuse +/// here — Vertex's path is not in `is_compressible_path` and we have +/// our own envelope handling). For PR-D4 the streaming handler just +/// passes through with the same auth + envelope + log surface; the +/// upstream SSE bytes flow back to the client unchanged. +/// +/// Note: this function takes 9 arguments. Grouping them into a +/// struct (an obvious clippy fix) would obscure that each argument +/// is a distinct axum extractor / handler-supplied value. The +/// argument list mirrors the catch-all `forward_http` in +/// [`crate::proxy`]; consistency wins over the pedantic lint. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn forward_vertex_request( + state: AppState, + client_addr: SocketAddr, + request_id: String, + method: Method, + uri: Uri, + headers: HeaderMap, + body: Body, + ctx: VertexCallContext, + attach_sse_tee: bool, +) -> Response { + let path_for_log = uri.path().to_string(); + + // ─── 1. BUFFER BODY ──────────────────────────────────────────────── + let max = state.config.compression_max_body_bytes as usize; + let buffered = match to_bytes(body, max).await { + Ok(b) => b, + Err(e) => { + tracing::warn!( + event = "vertex_body_too_large", + request_id = %request_id, + path = %path_for_log, + limit_bytes = max, + error = %e, + "vertex request body exceeds compression buffer limit; failing loudly" + ); + return error_response( + StatusCode::PAYLOAD_TOO_LARGE, + "request body exceeds buffer limit", + ); + } + }; + + // ─── 2. ENVELOPE PARSE ───────────────────────────────────────────── + match envelope::parse(&buffered) { + Ok(env) => { + tracing::info!( + event = "vertex_envelope_parsed", + request_id = %request_id, + path = %path_for_log, + project = %ctx.project, + location = %ctx.location, + model = %ctx.model_id, + verb = ctx.verb.as_str(), + anthropic_version = %env.anthropic_version, + has_messages = env.has_messages, + "vertex envelope detected" + ); + } + Err(e) => { + tracing::warn!( + event = "vertex_envelope_invalid", + request_id = %request_id, + path = %path_for_log, + model = %ctx.model_id, + verb = ctx.verb.as_str(), + error = %e, + "vertex envelope did not match expected shape; rejecting with 400" + ); + return error_response(StatusCode::BAD_REQUEST, "vertex envelope invalid"); + } + } + + // ─── 3. LIVE-ZONE COMPRESSION (when enabled) ─────────────────────── + // + // Vertex bodies are Anthropic-shape; we feed the same + // `compress_anthropic_request` dispatcher that runs on /v1/messages. + // The dispatcher uses RawValue-based surgery so `anthropic_version` + // (and any other non-`messages` top-level field) round-trips + // byte-equal. Compression off → buffered bytes used unchanged. + let body_to_send = if state.config.compression { + let outcome = compression::compress_anthropic_request( + &buffered, + state.config.compression_mode, + state.config.cache_control_auto_frozen, + &request_id, + ); + match outcome { + compression::Outcome::NoCompression => { + tracing::info!( + event = "vertex_compression_skipped", + request_id = %request_id, + path = %path_for_log, + compression_mode = state.config.compression_mode.as_str(), + reason = "no_compression", + "vertex live-zone dispatcher returned NoCompression" + ); + buffered + } + compression::Outcome::Compressed { + body, + tokens_before, + tokens_after, + strategies_applied, + markers_inserted, + } => { + tracing::info!( + event = "vertex_compression_applied", + request_id = %request_id, + path = %path_for_log, + tokens_before = tokens_before, + tokens_after = tokens_after, + tokens_freed = tokens_before.saturating_sub(tokens_after), + strategies = ?strategies_applied, + markers = markers_inserted.len(), + "vertex live-zone compression applied" + ); + body + } + compression::Outcome::Passthrough { reason } => { + tracing::warn!( + event = "vertex_compression_passthrough", + request_id = %request_id, + path = %path_for_log, + reason = ?reason, + "vertex live-zone dispatcher passthrough on parse/serialize" + ); + buffered + } + } + } else { + tracing::info!( + event = "vertex_compression_skipped", + request_id = %request_id, + path = %path_for_log, + reason = "compression_off", + "compression master switch off; vertex body forwarded unchanged" + ); + buffered + }; + + // ─── 4. RESOLVE BEARER TOKEN ─────────────────────────────────────── + let bearer = match state.vertex_token_source.bearer().await { + Ok(t) => t, + Err(e) => { + // Per project rule "no silent fallbacks": never forward + // unauthenticated. Surface the failure as a structured + // 502 so operators see the cause clearly in logs. + tracing::error!( + event = "vertex_adc_fetch_failed", + request_id = %request_id, + path = %path_for_log, + model = %ctx.model_id, + verb = ctx.verb.as_str(), + error = %e, + "vertex ADC bearer token fetch failed; refusing to forward unauthenticated" + ); + let status = match e { + TokenSourceError::ProviderInit(_) => StatusCode::BAD_GATEWAY, + TokenSourceError::Fetch(_) => StatusCode::BAD_GATEWAY, + }; + return error_response(status, "vertex ADC token fetch failed"); + } + }; + + // ─── 5. BUILD UPSTREAM URL ───────────────────────────────────────── + // + // The Vertex endpoint pattern is + // `https://{region}-aiplatform.googleapis.com/`. + // We honour the same `Config::upstream` override pattern the + // rest of the proxy uses: when an operator sets `upstream` to the + // mock server (typical in tests), we forward there and the + // request still carries the canonical Vertex path. + // + // For production, the operator should set `upstream` to the + // regional Vertex host. We do NOT auto-construct the regional + // URL from `vertex_region` — that would be a hardcoded provider + // routing decision. The region setting is exposed for + // observability only. + let upstream_url = match crate::proxy::build_upstream_url(&state.config.upstream, &uri) { + Ok(u) => u, + Err(e) => { + tracing::error!( + event = "vertex_upstream_url_failed", + request_id = %request_id, + error = %e, + "could not construct vertex upstream URL" + ); + return error_response(StatusCode::BAD_GATEWAY, "vertex upstream URL build failed"); + } + }; + + // ─── 6. BUILD HEADERS ────────────────────────────────────────────── + let strip_internal = state.config.strip_internal_headers.is_enabled(); + let forwarded_host = headers + .get(http::header::HOST) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + let mut outgoing_headers = build_forward_request_headers( + &headers, + client_addr.ip(), + "http", + forwarded_host.as_deref(), + &request_id, + strip_internal, + ); + if !state.config.rewrite_host { + if let Some(h) = headers.get(http::header::HOST) { + outgoing_headers.insert(http::header::HOST, h.clone()); + } + } + // Attach the bearer; if the client already sent an Authorization + // header we replace it (Vertex rejects the wrong Auth flavour + // anyway, so keeping the client-provided value would silently + // break the call). + match http::HeaderValue::from_str(&format!("Bearer {bearer}")) { + Ok(v) => { + outgoing_headers.insert(http::header::AUTHORIZATION, v); + } + Err(e) => { + tracing::error!( + event = "vertex_authorization_invalid", + request_id = %request_id, + error = %e, + "ADC bearer token contained invalid header bytes; refusing to forward" + ); + return error_response(StatusCode::BAD_GATEWAY, "vertex auth header build failed"); + } + } + + // ─── 7. FORWARD ──────────────────────────────────────────────────── + let reqwest_method = match reqwest::Method::from_bytes(method.as_str().as_bytes()) { + Ok(m) => m, + Err(e) => { + tracing::error!( + event = "vertex_method_invalid", + request_id = %request_id, + method = %method, + error = %e, + "could not convert axum method to reqwest method" + ); + return error_response(StatusCode::BAD_REQUEST, "vertex method invalid"); + } + }; + let upstream_resp = match state + .client + .request(reqwest_method, upstream_url.clone()) + .headers(outgoing_headers) + .body(body_to_send) + .send() + .await + { + Ok(r) => r, + Err(e) => { + tracing::warn!( + event = "vertex_upstream_error", + request_id = %request_id, + path = %path_for_log, + error = %e, + "vertex upstream call failed" + ); + return error_response(StatusCode::BAD_GATEWAY, "vertex upstream error"); + } + }; + + // ─── 8. STREAM RESPONSE ──────────────────────────────────────────── + let upstream_status = upstream_resp.status(); + let status = StatusCode::from_u16(upstream_status.as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); + let resp_headers = filter_response_headers(upstream_resp.headers()); + + // PR-C1 reuse: when `attach_sse_tee` is set AND the upstream + // response is `text/event-stream`, tee bytes into a bounded mpsc + // and drive `AnthropicStreamState` in a spawned task — same shape + // as the `/v1/messages` SSE telemetry tee in + // `crate::proxy::forward_http`. The byte-passthrough path is + // unaffected by the tee (best-effort `try_send`, bounded channel). + let is_sse = upstream_resp + .headers() + .get(http::header::CONTENT_TYPE) + .and_then(|v| v.to_str().ok()) + .map(|s| { + let media = s.split(';').next().unwrap_or("").trim(); + media.eq_ignore_ascii_case("text/event-stream") + }) + .unwrap_or(false); + let parser_tx = if attach_sse_tee && is_sse { + let (tx, rx) = tokio::sync::mpsc::channel::(VERTEX_SSE_QUEUE_DEPTH); + let rid = request_id.clone(); + tokio::spawn(run_anthropic_sse_state_machine(rx, rid)); + tracing::info!( + event = "vertex_sse_tee_engaged", + request_id = %request_id, + "vertex stream_raw_predict SSE telemetry tee engaged" + ); + Some(tx) + } else { + None + }; + + use futures_util::StreamExt as _; + let rid_for_stream = request_id.clone(); + let resp_stream = upstream_resp.bytes_stream().map(move |r| match r { + Ok(b) => { + if let Some(tx) = &parser_tx { + if let Err(e) = tx.try_send(b.clone()) { + tracing::debug!( + request_id = %rid_for_stream, + error = %e, + "vertex sse parser queue full or closed; skipping telemetry chunk" + ); + } + } + Ok(b) + } + Err(e) => { + tracing::warn!( + request_id = %rid_for_stream, + error = %e, + "vertex upstream stream error mid-response" + ); + Err(e) + } + }); + let body = Body::from_stream(resp_stream); + + let mut response = Response::builder().status(status); + if let Some(h) = response.headers_mut() { + h.extend(resp_headers); + if let Ok(v) = http::HeaderValue::from_str(&request_id) { + h.insert(http::HeaderName::from_static("x-request-id"), v); + } + } + let response = match response.body(body) { + Ok(r) => r, + Err(e) => { + tracing::error!( + event = "vertex_response_build_failed", + request_id = %request_id, + error = %e, + "could not build vertex response" + ); + return error_response(StatusCode::INTERNAL_SERVER_ERROR, "response build failed"); + } + }; + + tracing::info!( + event = "vertex_forwarded", + request_id = %request_id, + path = %path_for_log, + project = %ctx.project, + location = %ctx.location, + model = %ctx.model_id, + verb = ctx.verb.as_str(), + upstream_status = upstream_status.as_u16(), + "vertex request forwarded" + ); + + response +} + +fn error_response(status: StatusCode, msg: &'static str) -> Response { + Response::builder() + .status(status) + .body(Body::from(msg)) + .expect("static error response") +} + +/// Bound on the in-flight queue between the byte-passthrough and the +/// SSE state-machine task. Mirrors the +/// `crate::proxy::SSE_PARSER_QUEUE_DEPTH` rationale (256 events ≈ 5 +/// seconds of typical Anthropic streaming under the per-100ms event +/// rate; keeps memory bounded even if the parser stalls). +const VERTEX_SSE_QUEUE_DEPTH: usize = 256; + +/// Drive the Anthropic SSE state machine over a stream of byte +/// chunks. Lives in its own spawned task; the byte path is fed via a +/// best-effort tee from [`forward_vertex_request`] and never blocks +/// on this loop. +async fn run_anthropic_sse_state_machine( + mut rx: tokio::sync::mpsc::Receiver, + request_id: String, +) { + use crate::sse::framing::SseFramer; + let mut framer = SseFramer::new(); + let mut state = crate::sse::anthropic::AnthropicStreamState::new(); + while let Some(chunk) = rx.recv().await { + framer.push(&chunk); + while let Some(ev_result) = framer.next_event() { + match ev_result { + Ok(ev) => { + if let Err(e) = state.apply(ev) { + tracing::warn!( + request_id = %request_id, + error = %e, + "vertex sse anthropic state-machine apply error" + ); + } + } + Err(e) => { + tracing::warn!( + request_id = %request_id, + error = %e, + "vertex sse framer error" + ); + } + } + } + } + tracing::info!( + event = "vertex_sse_stream_closed", + request_id = %request_id, + provider = "vertex_anthropic", + input_tokens = state.usage.input_tokens, + output_tokens = state.usage.output_tokens, + cache_creation_input_tokens = state.usage.cache_creation_input_tokens, + cache_read_input_tokens = state.usage.cache_read_input_tokens, + stop_reason = state.stop_reason.as_deref().unwrap_or(""), + blocks = state.blocks.len(), + "vertex sse stream closed" + ); +} diff --git a/crates/headroom-proxy/src/vertex/stream_raw_predict.rs b/crates/headroom-proxy/src/vertex/stream_raw_predict.rs new file mode 100644 index 000000000..c0b9b490c --- /dev/null +++ b/crates/headroom-proxy/src/vertex/stream_raw_predict.rs @@ -0,0 +1,49 @@ +//! Vertex `:streamRawPredict` handler module. +//! +//! Path: +//! ```text +//! POST /v1beta1/projects/{project}/locations/{location}/publishers/anthropic/models/{model}:streamRawPredict +//! ``` +//! +//! Vertex's streaming verb returns the same `text/event-stream` +//! payload Anthropic Messages emits (Vertex does not use the binary +//! EventStream Bedrock does). The existing PR-C1 +//! [`crate::sse::anthropic::AnthropicStreamState`] state machine +//! works unchanged — bytes are teed into a bounded mpsc + spawned +//! parser, identical to the `/v1/messages` SSE pipeline in +//! [`crate::proxy::forward_http`]. +//! +//! # Why this module is thin +//! +//! Both `:rawPredict` and `:streamRawPredict` share the same axum +//! route shape (matchit can't distinguish two patterns where the verb +//! is part of a captured parameter). The shared dispatcher in +//! [`crate::vertex::handle_vertex_predict_dispatch`] splits the verb +//! and flips a single `attach_sse_tee` flag. This module exists so +//! the file structure mirrors the spec (PR-D4 calls for distinct +//! `raw_predict.rs` and `stream_raw_predict.rs` files); the +//! streaming-specific behaviour is the SSE tee in +//! [`crate::vertex::raw_predict::forward_vertex_request`] when +//! `attach_sse_tee == true`. +//! +//! # Streaming behaviour invariants +//! +//! - **Request body**: same envelope detection + live-zone Anthropic +//! compression as `:rawPredict`. The `stream: true` flag (when +//! present) is preserved byte-equal — the live-zone dispatcher +//! never rewrites top-level fields. +//! - **Response body**: SSE bytes flow back to the client unchanged. +//! The state-machine tee runs in a spawned task and can never +//! block the byte path (bounded `try_send` channel). +//! - **Telemetry**: per-stream summary log emitted at stream close +//! (`event = "vertex_sse_stream_closed"`) with token counts and +//! block count. +//! +//! See [`crate::vertex::raw_predict`] for the actual forwarding +//! implementation. + +// Re-export the shared dispatcher so callers that want to address the +// streaming verb explicitly have a name in this module. The +// dispatcher is the same single-axum-route entry point for both +// verbs — see module-level rationale. +pub use crate::vertex::handle_vertex_predict_dispatch as handle_stream_raw_predict; diff --git a/crates/headroom-proxy/tests/common/mod.rs b/crates/headroom-proxy/tests/common/mod.rs index 4f8976e8e..67ddaeb53 100644 --- a/crates/headroom-proxy/tests/common/mod.rs +++ b/crates/headroom-proxy/tests/common/mod.rs @@ -5,6 +5,7 @@ use std::net::SocketAddr; use std::sync::Arc; use std::time::Duration; +use headroom_proxy::vertex::TokenSource; use headroom_proxy::{build_app, AppState, Config}; use tokio::sync::oneshot; use url::Url; @@ -51,6 +52,9 @@ where /// Start a proxy with both a Config customizer and an AppState /// post-processor. PR-D1: tests that exercise the Bedrock route /// inject credentials via `with_bedrock_credentials` here. +/// PR-D4: Vertex tests inject a `StaticTokenSource` via +/// `install_static_token_source` here (chain-style) so they never +/// hit real GCP. #[allow(dead_code)] pub async fn start_proxy_with_state( upstream: &str, @@ -88,6 +92,20 @@ where } } +/// Convenience: replace the default `vertex_token_source` with a +/// `StaticTokenSource` returning the supplied bearer string. Used by +/// the PR-D4 Vertex integration tests so they never hit real GCP. +#[allow(dead_code)] +/// PR-D4: chain-style helper to install a `StaticTokenSource` on an +/// `AppState`. Returns the modified state so it composes with +/// `start_proxy_with_state`'s `FnOnce(AppState) -> AppState`. +pub fn install_static_token_source(mut state: AppState, bearer: &str) -> AppState { + state.vertex_token_source = Arc::new(headroom_proxy::vertex::StaticTokenSource::new( + bearer.to_string(), + )) as Arc; + state +} + /// Hold a reference to the config so dead_code doesn't strip its use. #[allow(dead_code)] pub fn _config_ref() -> Arc { diff --git a/crates/headroom-proxy/tests/integration_vertex_raw_predict.rs b/crates/headroom-proxy/tests/integration_vertex_raw_predict.rs new file mode 100644 index 000000000..71738d977 --- /dev/null +++ b/crates/headroom-proxy/tests/integration_vertex_raw_predict.rs @@ -0,0 +1,556 @@ +//! Integration tests for the native Vertex publisher path +//! (Phase D PR-D4). +//! +//! These tests boot the real Rust proxy in front of a wiremock +//! upstream and exercise the +//! `POST /v1beta1/projects/{p}/locations/{l}/publishers/anthropic/models/{m}:rawPredict` +//! (and `:streamRawPredict`) routes end-to-end. ADC is mocked via +//! `StaticTokenSource` so tests never reach real GCP. +//! +//! Per PR-D4 spec the four required tests are: +//! +//! 1. `native_envelope_round_trip_byte_equal` — body bytes (with +//! `anthropic_version` + Anthropic Messages shape) round-trip +//! SHA-256 byte-equal upstream. +//! 2. `adc_bearer_token_signed_correctly` — the `Authorization` header +//! on the upstream request is `Bearer `, and the +//! mock provider was actually consulted (no silent un-authed +//! forward). +//! 3. `thinking_block_preserved` — a request body containing +//! `thinking` / `redacted_thinking` blocks (the Python LiteLLM +//! converter dropped these — that was the P4-37/P4-38 bug) survives +//! byte-equal upstream. +//! 4. `stream_raw_predict_sse_handled` — `:streamRawPredict` proxies +//! an Anthropic SSE response back to the client without corruption, +//! and the AnthropicStreamState telemetry tee fires the `event = +//! "vertex_streaming_pipeline_active"` log. + +mod common; + +use common::{install_static_token_source, start_proxy_with_state}; +use serde_json::{json, Value}; +use sha2::{Digest, Sha256}; +use std::sync::{Arc, Mutex}; +use wiremock::matchers::{method, path_regex}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +const TEST_BEARER: &str = "ya29.test-static-bearer-d4-pr-fixture"; +const PROJECT: &str = "test-project-12345"; +const LOCATION: &str = "us-central1"; +const MODEL: &str = "claude-3-5-sonnet@20240620"; + +const VERTEX_PATH_REGEX: &str = + r"^/v1beta1/projects/[^/]+/locations/[^/]+/publishers/anthropic/models/[^/]+$"; + +fn raw_predict_url(proxy_url: &str) -> String { + format!( + "{proxy_url}/v1beta1/projects/{PROJECT}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:rawPredict", + ) +} + +fn stream_raw_predict_url(proxy_url: &str) -> String { + format!( + "{proxy_url}/v1beta1/projects/{PROJECT}/locations/{LOCATION}/publishers/anthropic/models/{MODEL}:streamRawPredict", + ) +} + +fn sha256_hex(bytes: &[u8]) -> String { + let mut hasher = Sha256::new(); + hasher.update(bytes); + hasher + .finalize() + .iter() + .fold(String::with_capacity(64), |mut acc, b| { + use std::fmt::Write as _; + let _ = write!(acc, "{b:02x}"); + acc + }) +} + +#[track_caller] +fn assert_byte_equal_sha256(inbound: &[u8], received: &[u8]) { + let inbound_hash = sha256_hex(inbound); + let received_hash = sha256_hex(received); + assert_eq!( + inbound.len(), + received.len(), + "byte length mismatch: inbound={}, upstream-received={}", + inbound.len(), + received.len(), + ); + assert_eq!( + inbound_hash, received_hash, + "SHA-256 mismatch: inbound={inbound_hash}, upstream-received={received_hash}", + ); +} + +/// Mount a Vertex rawPredict mock that captures the upstream request +/// bytes + `Authorization` header. The path-regex matcher mirrors the +/// canonical Vertex publisher shape (so any of the 4 path parts can +/// vary across tests without re-mounting). +struct CapturedUpstream { + body: Mutex>>, + authorization: Mutex>, + content_type_response: Mutex>, +} + +impl CapturedUpstream { + fn new() -> Arc { + Arc::new(Self { + body: Mutex::new(None), + authorization: Mutex::new(None), + content_type_response: Mutex::new(None), + }) + } +} + +async fn mount_capture_json(upstream: &MockServer) -> Arc { + let captured = CapturedUpstream::new(); + let cap = captured.clone(); + Mock::given(method("POST")) + .and(path_regex(VERTEX_PATH_REGEX)) + .respond_with(move |req: &wiremock::Request| { + *cap.body.lock().unwrap() = Some(req.body.clone()); + *cap.authorization.lock().unwrap() = req + .headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + *cap.content_type_response.lock().unwrap() = Some("application/json".into()); + ResponseTemplate::new(200) + .insert_header("content-type", "application/json") + .set_body_string(r#"{"id":"msg_test","type":"message","role":"assistant","content":[{"type":"text","text":"hi"}],"model":"claude-3-5-sonnet@20240620","stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1}}"#) + }) + .mount(upstream) + .await; + captured +} + +/// Mount a Vertex streamRawPredict mock that returns a small Anthropic +/// SSE stream. The exact event sequence below is a minimal-but-valid +/// Anthropic Messages stream (per the PR-C1 framer + state machine). +async fn mount_capture_sse(upstream: &MockServer) -> Arc { + let captured = CapturedUpstream::new(); + let cap = captured.clone(); + let sse_body = concat!( + "event: message_start\n", + "data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_strm\",\"type\":\"message\",\"role\":\"assistant\",\"model\":\"claude-3-5-sonnet@20240620\",\"content\":[],\"stop_reason\":null,\"usage\":{\"input_tokens\":2,\"output_tokens\":0}}}\n\n", + "event: content_block_start\n", + "data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n", + "event: content_block_delta\n", + "data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"hello\"}}\n\n", + "event: content_block_stop\n", + "data: {\"type\":\"content_block_stop\",\"index\":0}\n\n", + "event: message_delta\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":3}}\n\n", + "event: message_stop\n", + "data: {\"type\":\"message_stop\"}\n\n", + ); + Mock::given(method("POST")) + .and(path_regex(VERTEX_PATH_REGEX)) + .respond_with(move |req: &wiremock::Request| { + *cap.body.lock().unwrap() = Some(req.body.clone()); + *cap.authorization.lock().unwrap() = req + .headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + *cap.content_type_response.lock().unwrap() = Some("text/event-stream".into()); + ResponseTemplate::new(200) + .set_body_raw(sse_body.as_bytes().to_vec(), "text/event-stream") + }) + .mount(upstream) + .await; + captured +} + +/// Vertex envelope without `model` — the shape the proxy expects. +/// Includes `anthropic_version` (required) and a single user message. +fn minimal_vertex_body() -> Value { + json!({ + "anthropic_version": "vertex-2023-10-16", + "messages": [ + {"role": "user", "content": "Hello, Claude!"} + ], + "max_tokens": 64, + }) +} + +// ─── TEST 1 ──────────────────────────────────────────────────────────── + +#[tokio::test] +async fn native_envelope_round_trip_byte_equal() { + let upstream = MockServer::start().await; + let captured = mount_capture_json(&upstream).await; + + let proxy = start_proxy_with_state( + &upstream.uri(), + |c| { + // Compression off: the only thing under test here is the + // envelope detection + forwarding path. The body must + // round-trip byte-equal even when the live-zone dispatcher + // is engaged in a separate test (`thinking_block_preserved`). + c.compression = false; + }, + |s| install_static_token_source(s, TEST_BEARER), + ) + .await; + + let body_value = minimal_vertex_body(); + let body_bytes = serde_json::to_vec(&body_value).unwrap(); + let resp = reqwest::Client::new() + .post(raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .body(body_bytes.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured + .body + .lock() + .unwrap() + .clone() + .expect("upstream got body"); + assert_byte_equal_sha256(&body_bytes, &got); + + // Defensive: parse and confirm the canonical envelope fields. + let parsed: Value = serde_json::from_slice(&got).unwrap(); + assert_eq!(parsed["anthropic_version"], json!("vertex-2023-10-16")); + assert!( + parsed.get("model").is_none(), + "Vertex envelope must NOT carry a `model` field" + ); + + proxy.shutdown().await; +} + +// ─── TEST 2 ──────────────────────────────────────────────────────────── + +#[tokio::test] +async fn adc_bearer_token_signed_correctly() { + let upstream = MockServer::start().await; + let captured = mount_capture_json(&upstream).await; + + let proxy = start_proxy_with_state( + &upstream.uri(), + |c| { + c.compression = false; + }, + |s| install_static_token_source(s, TEST_BEARER), + ) + .await; + + let body_bytes = serde_json::to_vec(&minimal_vertex_body()).unwrap(); + + // Send the request WITHOUT an Authorization header. The proxy + // must inject `Bearer ` from the static token source. + let resp = reqwest::Client::new() + .post(raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .body(body_bytes) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let auth = captured + .authorization + .lock() + .unwrap() + .clone() + .expect("upstream got Authorization header"); + let expected = format!("Bearer {TEST_BEARER}"); + assert_eq!( + auth, expected, + "Vertex request must carry the ADC bearer token; got {auth}, expected {expected}", + ); + + // ALSO: verify the proxy OVERWRITES a client-supplied + // Authorization header (Vertex would reject the wrong flavour + // anyway — silent forward of the wrong auth would surface as a + // confusing 401 from upstream). + let body_bytes2 = serde_json::to_vec(&minimal_vertex_body()).unwrap(); + let resp2 = reqwest::Client::new() + .post(raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .header("authorization", "Bearer some-client-supplied-key") + .body(body_bytes2) + .send() + .await + .unwrap(); + assert_eq!(resp2.status(), 200); + let auth2 = captured + .authorization + .lock() + .unwrap() + .clone() + .expect("auth header on second call"); + assert_eq!( + auth2, expected, + "proxy must overwrite client-supplied Authorization with the ADC bearer; \ + got {auth2}", + ); + + proxy.shutdown().await; +} + +// ─── TEST 3 ──────────────────────────────────────────────────────────── + +#[tokio::test] +async fn thinking_block_preserved() { + let upstream = MockServer::start().await; + let captured = mount_capture_json(&upstream).await; + + // Compression ON + LiveZone mode so the live-zone Anthropic + // dispatcher actually runs over the body. This test is the + // teeth of P4-37/P4-38: the Python LiteLLM converter dropped + // `thinking` and `redacted_thinking` blocks. The Rust path must + // preserve them byte-equal upstream. + let proxy = start_proxy_with_state( + &upstream.uri(), + |c| { + c.compression = true; + c.compression_mode = headroom_proxy::config::CompressionMode::LiveZone; + }, + |s| install_static_token_source(s, TEST_BEARER), + ) + .await; + + // Realistic Anthropic block content covering: + // - `thinking` block with a signature (cryptographically signed + // reasoning the model emitted in a prior turn). + // - `redacted_thinking` block (returned by Anthropic when + // thinking content is policy-redacted; carries an opaque + // `data` blob that MUST round-trip byte-equal). + // - text content alongside, so the live-zone walker has more + // than one block to consider. + let body_value = json!({ + "anthropic_version": "vertex-2023-10-16", + "max_tokens": 1024, + "messages": [ + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "I need to plan this carefully. Let me think step by step about the user's request and consider edge cases.", + "signature": "EuYBCkQYBCKMAQ.thinkingsig.example.base64.payload=" + }, + { + "type": "redacted_thinking", + "data": "EmkKAhgEEgwQ.redacted.opaque.bytes.must.roundtrip=" + }, + { + "type": "text", + "text": "Here is my answer." + } + ] + }, + { + "role": "user", + "content": "Can you elaborate?" + } + ], + "thinking": {"type": "enabled", "budget_tokens": 5000} + }); + let body_bytes = serde_json::to_vec(&body_value).unwrap(); + + let resp = reqwest::Client::new() + .post(raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .body(body_bytes.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + let got = captured + .body + .lock() + .unwrap() + .clone() + .expect("upstream got body"); + + // Strong assertion: byte-equal end-to-end. The live-zone + // dispatcher's RawValue-based surgery may rewrite live-zone + // messages but ours has only an assistant turn (frozen by + // definition) and a single short user turn that's below the + // compression-eligibility floor, so the body should be the + // same bytes. + assert_byte_equal_sha256(&body_bytes, &got); + + // Defensive: thinking + redacted_thinking + signature all + // present and unchanged. + let parsed: Value = serde_json::from_slice(&got).unwrap(); + let assistant_blocks = &parsed["messages"][0]["content"]; + assert_eq!(assistant_blocks[0]["type"], json!("thinking")); + assert!( + assistant_blocks[0]["signature"] + .as_str() + .unwrap() + .starts_with("EuYBCkQYBCKMAQ"), + "thinking.signature must round-trip byte-equal" + ); + assert_eq!(assistant_blocks[1]["type"], json!("redacted_thinking")); + assert_eq!( + assistant_blocks[1]["data"], + json!("EmkKAhgEEgwQ.redacted.opaque.bytes.must.roundtrip="), + "redacted_thinking.data must round-trip byte-equal" + ); + assert_eq!(assistant_blocks[2]["type"], json!("text")); + + proxy.shutdown().await; +} + +// ─── TEST 4 ──────────────────────────────────────────────────────────── + +#[tokio::test] +async fn stream_raw_predict_sse_handled() { + let upstream = MockServer::start().await; + let captured = mount_capture_sse(&upstream).await; + + let proxy = start_proxy_with_state( + &upstream.uri(), + |c| { + c.compression = false; + }, + |s| install_static_token_source(s, TEST_BEARER), + ) + .await; + + // Streaming envelope: same shape as :rawPredict, with `stream: + // true` (Vertex doesn't actually require the field — the verb + // disambiguates — but real clients send it for compatibility). + let body_value = json!({ + "anthropic_version": "vertex-2023-10-16", + "stream": true, + "max_tokens": 64, + "messages": [ + {"role": "user", "content": "Stream please."} + ] + }); + let body_bytes = serde_json::to_vec(&body_value).unwrap(); + + let resp = reqwest::Client::new() + .post(stream_raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .header("accept", "text/event-stream") + .body(body_bytes.clone()) + .send() + .await + .unwrap(); + assert_eq!(resp.status(), 200); + + // Confirm the response carries SSE content-type back to the + // client (the proxy MUST NOT translate to JSON or rewrap). + let ct = resp + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + assert!( + ct.eq_ignore_ascii_case("text/event-stream"), + "Vertex stream response must surface SSE content-type to client; got {ct}", + ); + + // Drain the stream and confirm we see the message_start + + // content_block_delta events the upstream emitted. Bytes pass + // through unchanged. + let body_text = resp.text().await.expect("read sse body"); + assert!( + body_text.contains("event: message_start"), + "sse body missing message_start: {body_text:?}", + ); + assert!( + body_text.contains("event: content_block_delta"), + "sse body missing content_block_delta: {body_text:?}", + ); + assert!( + body_text.contains("event: message_stop"), + "sse body missing message_stop: {body_text:?}", + ); + + // Body bytes upstream-received MUST match the inbound bytes (no + // request-side rewrite when compression is off). + let got = captured + .body + .lock() + .unwrap() + .clone() + .expect("upstream got body"); + assert_byte_equal_sha256(&body_bytes, &got); + + // Bearer was attached. + let auth = captured + .authorization + .lock() + .unwrap() + .clone() + .expect("upstream got Authorization"); + assert_eq!(auth, format!("Bearer {TEST_BEARER}")); + + proxy.shutdown().await; +} + +// ─── BONUS: ADC FAILURE PATH (no silent fallback) ────────────────────── + +#[tokio::test] +async fn adc_failure_returns_5xx_no_silent_forward() { + use async_trait::async_trait; + use headroom_proxy::vertex::{TokenSource, TokenSourceError}; + use std::sync::Arc as StdArc; + + // Token source that always fails — verifies the no-silent-fallback + // contract: the proxy must NOT forward a request to upstream + // without a bearer. + #[derive(Debug)] + struct AlwaysFail; + #[async_trait] + impl TokenSource for AlwaysFail { + async fn bearer(&self) -> Result { + Err(TokenSourceError::Fetch( + "synthetic test failure: ADC chain unreachable".into(), + )) + } + } + + let upstream = MockServer::start().await; + let captured = mount_capture_json(&upstream).await; + + let proxy = start_proxy_with_state( + &upstream.uri(), + |c| { + c.compression = false; + }, + |mut state| { + state.vertex_token_source = StdArc::new(AlwaysFail) as StdArc; + state + }, + ) + .await; + + let body_bytes = serde_json::to_vec(&minimal_vertex_body()).unwrap(); + let resp = reqwest::Client::new() + .post(raw_predict_url(&proxy.url())) + .header("content-type", "application/json") + .body(body_bytes) + .send() + .await + .unwrap(); + + assert!( + resp.status().is_server_error(), + "ADC failure must surface as 5xx, got {}", + resp.status() + ); + + // Critically: the upstream must NOT have been called at all. + assert!( + captured.body.lock().unwrap().is_none(), + "ADC failure must short-circuit; upstream must not be reached" + ); + + proxy.shutdown().await; +}