mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(proxy): PR-D4 native Vertex publisher path + ADC bearer auth
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
`:<verb>` 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 <token>` (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<dyn TokenSource>`
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
<static-test-token>` 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.
This commit is contained in:
chopratejas
parent
41ca0544c7
commit
c10a2195af
15 changed files with 2227 additions and 0 deletions
158
Cargo.lock
generated
158
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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 <token>` 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"
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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<Duration, String> {
|
||||
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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<dyn TokenSource>` 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<Config>,
|
||||
|
|
@ -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<dyn crate::vertex::TokenSource>,
|
||||
}
|
||||
|
||||
/// 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<dyn crate::vertex::TokenSource> =
|
||||
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<dyn crate::vertex::TokenSource>,
|
||||
) -> Result<Self, ProxyError> {
|
||||
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 `:<verb>` 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
|
||||
|
|
|
|||
315
crates/headroom-proxy/src/vertex/adc.rs
Normal file
315
crates/headroom-proxy/src/vertex/adc.rs
Normal file
|
|
@ -0,0 +1,315 @@
|
|||
//! GCP Application Default Credentials → bearer token resolution.
|
||||
//!
|
||||
//! Vertex AI's `:rawPredict` / `:streamRawPredict` endpoints expect
|
||||
//! `Authorization: Bearer <jwt>` 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 <token>` header. The token is cached
|
||||
/// internally; concurrent calls de-dup to a single fetch.
|
||||
async fn bearer(&self) -> Result<String, TokenSourceError>;
|
||||
}
|
||||
|
||||
/// 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<Option<...>>`
|
||||
/// (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<Option<Arc<dyn ::gcp_auth::TokenProvider>>>,
|
||||
/// 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<Option<CachedToken>>,
|
||||
}
|
||||
|
||||
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<String>) -> 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<Arc<dyn ::gcp_auth::TokenProvider>, 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<String, TokenSourceError> {
|
||||
// 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<Utc>`.
|
||||
// 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<String>) -> Self {
|
||||
Self {
|
||||
token: token.into(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl TokenSource for StaticTokenSource {
|
||||
async fn bearer(&self) -> Result<String, TokenSourceError> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
167
crates/headroom-proxy/src/vertex/envelope.rs
Normal file
167
crates/headroom-proxy/src/vertex/envelope.rs
Normal file
|
|
@ -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::<Value>` 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<ParsedEnvelope, VertexEnvelopeError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
277
crates/headroom-proxy/src/vertex/mod.rs
Normal file
277
crates/headroom-proxy/src/vertex/mod.rs
Normal file
|
|
@ -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 <jwt>`. 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 `<model>:<verb>`
|
||||
/// (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<AppState>,
|
||||
ConnectInfo(client_addr): ConnectInfo<SocketAddr>,
|
||||
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<Self> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
562
crates/headroom-proxy/src/vertex/raw_predict.rs
Normal file
562
crates/headroom-proxy/src/vertex/raw_predict.rs
Normal file
|
|
@ -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 <token>`. 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/<path>`).
|
||||
//! 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 `<model_id>:<verb>` 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<AppState>,
|
||||
ConnectInfo(client_addr): ConnectInfo<SocketAddr>,
|
||||
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/<path-and-query>`.
|
||||
// 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::<bytes::Bytes>(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<bytes::Bytes>,
|
||||
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"
|
||||
);
|
||||
}
|
||||
49
crates/headroom-proxy/src/vertex/stream_raw_predict.rs
Normal file
49
crates/headroom-proxy/src/vertex/stream_raw_predict.rs
Normal file
|
|
@ -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;
|
||||
|
|
@ -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<F, G>(
|
||||
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<dyn TokenSource>;
|
||||
state
|
||||
}
|
||||
|
||||
/// Hold a reference to the config so dead_code doesn't strip its use.
|
||||
#[allow(dead_code)]
|
||||
pub fn _config_ref() -> Arc<Config> {
|
||||
|
|
|
|||
556
crates/headroom-proxy/tests/integration_vertex_raw_predict.rs
Normal file
556
crates/headroom-proxy/tests/integration_vertex_raw_predict.rs
Normal file
|
|
@ -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 <static-test-token>`, 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<Option<Vec<u8>>>,
|
||||
authorization: Mutex<Option<String>>,
|
||||
content_type_response: Mutex<Option<String>>,
|
||||
}
|
||||
|
||||
impl CapturedUpstream {
|
||||
fn new() -> Arc<Self> {
|
||||
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<CapturedUpstream> {
|
||||
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<CapturedUpstream> {
|
||||
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 <token>` 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<String, TokenSourceError> {
|
||||
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<dyn TokenSource>;
|
||||
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;
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue