mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
c10a2195af
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.
77 lines
3.8 KiB
TOML
77 lines
3.8 KiB
TOML
[workspace]
|
|
resolver = "2"
|
|
members = [
|
|
"crates/headroom-core",
|
|
"crates/headroom-proxy",
|
|
"crates/headroom-py",
|
|
"crates/headroom-parity",
|
|
]
|
|
# headroom-py is a Python extension module — it must be built via maturin, not
|
|
# plain cargo (the "extension-module" feature tells pyo3 not to link libpython,
|
|
# which is required for `import` to work). `cargo build --workspace` without
|
|
# explicit members skips it; `cargo test --workspace` still runs its tests
|
|
# because pyo3 can dynamically link here for the cdylib used by tests.
|
|
default-members = [
|
|
"crates/headroom-core",
|
|
"crates/headroom-proxy",
|
|
"crates/headroom-parity",
|
|
]
|
|
|
|
[workspace.package]
|
|
edition = "2021"
|
|
rust-version = "1.80"
|
|
license = "Apache-2.0"
|
|
repository = "https://github.com/chopratejas/headroom"
|
|
authors = ["Headroom Maintainers"]
|
|
|
|
[workspace.dependencies]
|
|
serde = { version = "1", features = ["derive"] }
|
|
# `preserve_order` makes `serde_json::Value::Object` use IndexMap so JSON
|
|
# parse order is preserved through Value→string→Value round-trips. The
|
|
# smart_crusher port relies on this to match Python's `str(dict)` output,
|
|
# which preserves insertion order; otherwise BTreeMap's sorted-key default
|
|
# would diverge from Python on every multi-key object.
|
|
#
|
|
# `arbitrary_precision` keeps the literal numeric token from the source
|
|
# JSON intact: `Value::Number` becomes a wrapper around the original
|
|
# digit string, so `1.0` does NOT collapse to `1`, and `12345678901234567`
|
|
# does NOT lose precision through f64. Required by Realignment invariant
|
|
# I1 (byte-faithful passthrough on unmutated bytes; see REALIGNMENT/02-
|
|
# architecture.md §2.2) and PR-A4 (see REALIGNMENT/03-phase-A-lockdown.md).
|
|
#
|
|
# `raw_value` exposes `serde_json::value::RawValue`, the unparsed JSON
|
|
# fragment type. Phase B PR-B2 uses this to forward unmodified
|
|
# `messages[*]` entries as exact byte copies — the parser captures the
|
|
# original byte slice, so byte-for-byte round-trips work even with
|
|
# whitespace, key order, or escape preferences the producer chose.
|
|
# Enabled here in Phase A so PR-B2 can land as a pure consumer change.
|
|
serde_json = { version = "1", features = ["preserve_order", "arbitrary_precision", "raw_value"] }
|
|
bytes = "1"
|
|
thiserror = "1"
|
|
tracing = "0.1"
|
|
anyhow = "1"
|
|
clap = { version = "4", features = ["derive"] }
|
|
tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal"] }
|
|
axum = "0.7"
|
|
tower = "0.5"
|
|
reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] }
|
|
pyo3 = "0.22"
|
|
# Phase D PR-D1: AWS SigV4 signing for native Bedrock InvokeModel route.
|
|
# `aws-sigv4` provides the canonical-request + signing-key implementation;
|
|
# `aws-config` resolves credentials from the standard provider chain
|
|
# (env vars, profiles, IMDS, ECS task role, etc); `aws-credential-types`
|
|
# exposes `Credentials` so the signer accepts whatever the chain returned.
|
|
aws-sigv4 = { version = "1", default-features = false, features = ["sign-http", "http1"] }
|
|
aws-config = { version = "1", default-features = false, features = ["behavior-version-latest", "rustls", "rt-tokio"] }
|
|
aws-credential-types = { version = "1", default-features = false }
|
|
# `Identity` lives in aws-smithy-runtime-api; the SigV4 builder
|
|
# 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"
|