mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
4 commits
| Author | SHA1 | Message | Date | |
|---|---|---|---|---|
chopratejas
|
c10a2195af |
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.
|
||
|
chopratejas
|
f2d4fe39cb |
fix(proxy): PR-D1 native Bedrock InvokeModel route + SigV4
Phase D PR-D1 lands the first native Rust path for AWS Bedrock,
replacing the lossy LiteLLM Python converter for Anthropic-on-Bedrock
non-streaming requests. Eliminates part of P4-37 and P4-39.
What landed
-----------
- New crates/headroom-proxy/src/bedrock/ module:
- envelope.rs: parses the {"anthropic_version": "...", ...}
Bedrock body shape; re-emits with anthropic_version preserved
as the first key (relies on serde_json preserve_order).
- sigv4.rs: AWS SigV4 signer wrapping the aws-sigv4 crate.
Forces PayloadChecksumKind::XAmzSha256 so x-amz-content-sha256
is in the canonical request, hashed over the post-compression
body bytes (the bytes that actually hit Bedrock). No silent
fallback: signing failures return 5xx with
event=bedrock_sigv4_failed.
- invoke.rs: POST handler for /model/{model_id}/invoke
(and /converse - same wire shape for anthropic.claude-*).
Detects Anthropic vendor via literal starts_with("anthropic.")
(no regex per project rule), routes Anthropic-shape bodies
through the existing compress_anthropic_request live-zone
dispatcher, then signs and forwards to the configured Bedrock
endpoint.
- Modified:
- proxy.rs: routes /model/:model_id/invoke and
/model/:model_id/converse when enable_bedrock_native is on
(default). Adds bedrock_credentials: Option<Arc<Credentials>>
to AppState.
- config.rs: new flags --bedrock-region (default us-east-1,
env HEADROOM_PROXY_BEDROCK_REGION), --bedrock-endpoint
(operator override for FIPS/VPC/test setups),
--enable-bedrock-native (default true), --aws-profile.
- main.rs: resolves AWS credentials at startup via
aws_config::defaults(BehaviorVersion::latest()). Failure logs
event=bedrock_credentials_unavailable at WARN; the handler
refuses to forward unsigned (event=bedrock_credentials_missing).
- Cargo.toml: workspace deps aws-sigv4, aws-config,
aws-credential-types, aws-smithy-runtime-api.
Tests
-----
8 integration tests under
crates/headroom-proxy/tests/integration_bedrock_invoke.rs:
1. native_envelope_round_trip_byte_equal
2. sigv4_signed_correctly_after_compression - confirms
authorization is SigV4-shape and x-amz-content-sha256
matches sha256(body received by upstream).
3. thinking_block_preserved_through_bedrock
4. redacted_thinking_preserved
5. document_block_preserved
6. tool_result_array_with_image_preserved
7. stop_sequence_null_only_when_present - pins that the proxy
does NOT inject stop_sequence: null (P4-37 hardcode).
8. tool_use_input_byte_equal_preserves_key_order
All eight pass. Full workspace test run is green; clippy + fmt
clean. make ci-precheck (rust + python + commitlint) passes
locally.
Build constraints honoured
--------------------------
- No silent fallbacks: missing creds / signing failures return
5xx with structured event=... log; no path ever forwards
unsigned.
- No hardcodes: region, endpoint, profile, enable-flag all
configurable via CLI + env.
- No regexes: vendor detection is str::starts_with.
- Comprehensive structured logs: event=bedrock_invoke_received,
bedrock_envelope_parsed, bedrock_compression_skipped,
bedrock_credentials_missing, sigv4_signed,
bedrock_invoke_forwarded, etc.
- Performant: body buffered once, passed by &[u8] to signer
(zero-copy), Bytes::clone only for ownership transfer to
reqwest. Sign exactly once per request.
- Elegant: 4 small focused modules mirror handlers/ + sse/.
- Tests use realistic Anthropic block content (real thinking,
redacted_thinking, document, base64 image fixtures).
Streaming (PR-D2) and observability (PR-D3) follow.
|
||
|
chopratejas
|
fa5fbfabf4 |
fix(rust): wire ICM compressor into Rust proxy on /v1/messages
Adds an opt-in compression interceptor that buffers Anthropic
/v1/messages requests, runs IntelligentContextManager over the
messages array, and forwards the (possibly trimmed) body upstream.
All other paths, methods, and content-types stay on the original
streaming passthrough — so existing operators see zero change.
Behaviour gates ALL must be true to buffer + compress:
- --compression flag (or HEADROOM_PROXY_COMPRESSION=1)
- method == POST
- path == /v1/messages
- Content-Type: application/json
- ICM constructed successfully at startup
Falls through to streaming on any failure: parse, missing fields,
unknown model, body-too-large. Compression must never break a
request — that's the safety contract.
Model context windows come from a vendored LiteLLM snapshot at
crates/headroom-proxy/data/model_prices_and_context_window.json
parsed once into an OnceLock<HashMap>. Refresh via
scripts/refresh_model_limits.sh. Rationale documented inline:
hardcoded tables silently rot; LiteLLM is the canonical source
the entire LLM-tooling ecosystem relies on.
New tests:
- 16 unit tests across compression::{anthropic, icm, model_limits}
- 5 integration tests: off-passthrough, on-short-passthrough,
on-oversized-trim, on-non-json-skip, on-non-llm-path-skip
Verification:
- cargo test --workspace -> 884 passed, 0 failed
- cargo clippy --workspace -- -D warnings -> clean
- cargo fmt --check -> clean
|
||
|
chopratejas
|
bcc2ad810a |
test(rust): integration tests + wiremock harness (phase-1)
15 integration tests across five suites that spin up the proxy on an ephemeral port pointed at a per-test mock upstream: - integration_http: all 7 methods round-trip with body, status passthrough for 404/500/502, query strings preserved, 1MB POST streams through. - integration_sse: a 10-event in-process hyper SSE upstream emits at 50ms cadence; chunks reach the client with max gap < 500ms (loose CI bound) and a client disconnect propagates to the upstream within 2s. - integration_ws: 5 text + 5 binary messages echo through a tungstenite upstream byte-equal; client-initiated close propagates. - integration_headers: hop-by-hop strip both directions, X-Forwarded-* injection, X-Forwarded-For appends to existing value, multi-valued response headers preserved. - integration_body: 5MB POST round-trips byte-equal; streaming response yields first byte before the upstream finishes sending. - integration_health: own /healthz always 200; /healthz/upstream is 200 when upstream healthy and 503 when down. The Sec-WebSocket-Protocol forwarding is exercised implicitly by the WS tests via tungstenite handshake. The harness lives at tests/common/mod.rs and is shared by every integration suite. |