amberfork_record/
proxy.rs

1//! The capture proxy: the boundary the record path records at.
2//!
3//! The agent is pointed at this proxy with a base-URL environment variable
4//! (`OPENAI_BASE_URL`, `ANTHROPIC_BASE_URL`, …) and needs no code change — the zero-code
5//! capture the design calls for. Every request is forwarded upstream verbatim and the round
6//! trip is appended to a [`Cassette`].
7//!
8//! It forwards over plain HTTP inbound and TLS outbound on purpose: the agent's SDK talks
9//! `http://127.0.0.1:PORT`, so there is no certificate to install and no TLS interception
10//! anywhere. This crate is an I/O edge — `tokio` lives here, never in the engine crates.
11//!
12//! **Known gap: streamed responses are buffered.** The relay reads the upstream body to
13//! completion before answering, so an agent using `stream: true` still receives every byte and
14//! still parses the SSE stream correctly, but receives it in one piece at the end instead of
15//! incrementally. Content is faithful; arrival timing is not. That is a real (if usually
16//! benign) change to the run being recorded, and it is stated here rather than left for
17//! someone to discover: an agent whose control flow keys on partial output — an early-abort on
18//! first token, a UI that streams — is not yet faithfully recordable. Passing the body through
19//! as a stream while teeing it to the tape is the fix, and it is a follow-up, not a rewrite.
20
21use crate::cassette::{
22    Body, CapturedRequest, CapturedResponse, Cassette, Exchange, retain_request_headers,
23    retain_response_headers,
24};
25use axum::Router;
26use axum::extract::{Request, State};
27use axum::http::StatusCode;
28use axum::response::{IntoResponse, Response};
29use std::fmt;
30use std::io;
31use std::net::{Ipv4Addr, SocketAddr};
32use std::sync::{Arc, Mutex};
33use tokio::net::TcpListener;
34
35/// Something that went wrong standing up or running a recording session.
36#[derive(Debug)]
37pub enum RecordError {
38    /// Binding the loopback listener failed — in practice: the port is already taken.
39    Bind { addr: SocketAddr, source: io::Error },
40    /// The accept loop died after a successful bind.
41    Serve(io::Error),
42    /// The upstream base URL is not a usable origin.
43    Upstream { url: String, source: reqwest::Error },
44}
45
46impl fmt::Display for RecordError {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            Self::Bind { addr, source } => write!(f, "cannot bind {addr}: {source}"),
50            Self::Serve(source) => write!(f, "capture proxy stopped: {source}"),
51            Self::Upstream { url, source } => {
52                write!(f, "unusable upstream base URL {url}: {source}")
53            }
54        }
55    }
56}
57
58impl std::error::Error for RecordError {
59    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
60        match self {
61            Self::Bind { source, .. } | Self::Serve(source) => Some(source),
62            Self::Upstream { source, .. } => Some(source),
63        }
64    }
65}
66
67/// The tape being written. Shared with every in-flight request handler.
68type Tape = Arc<Mutex<Vec<Exchange>>>;
69
70struct ProxyState {
71    /// Upstream origin, trailing slash trimmed, e.g. `https://api.openai.com`.
72    upstream: String,
73    client: reqwest::Client,
74    tape: Tape,
75}
76
77/// A bound, running capture proxy.
78pub struct CaptureProxy {
79    addr: SocketAddr,
80    tape: Tape,
81    id: String,
82}
83
84impl CaptureProxy {
85    /// Bind a capture proxy on loopback, forwarding to `upstream`.
86    ///
87    /// Loopback-only, like the serving edge: a proxy that relays an agent's traffic — and
88    /// holds the credential to do it — has no business being reachable off the machine.
89    /// Port 0 lets the OS pick, which is what keeps concurrent recordings from colliding.
90    pub async fn bind(id: impl Into<String>, upstream: &str) -> Result<Self, RecordError> {
91        let client =
92            reqwest::Client::builder()
93                .build()
94                .map_err(|source| RecordError::Upstream {
95                    url: upstream.to_string(),
96                    source,
97                })?;
98
99        let tape: Tape = Arc::new(Mutex::new(Vec::new()));
100        let state = Arc::new(ProxyState {
101            upstream: upstream.trim_end_matches('/').to_string(),
102            client,
103            tape: Arc::clone(&tape),
104        });
105
106        let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, 0));
107        let listener = TcpListener::bind(addr)
108            .await
109            .map_err(|source| RecordError::Bind { addr, source })?;
110        let addr = listener.local_addr().map_err(RecordError::Serve)?;
111
112        // `fallback` rather than routes: the proxy is path-agnostic by design. It must relay
113        // whatever the SDK asks for — today `/v1/chat/completions`, tomorrow an endpoint that
114        // did not exist when this shipped — so enumerating provider paths here would just be
115        // a list to keep stale.
116        let app = Router::new().fallback(capture).with_state(state);
117
118        tokio::spawn(async move {
119            let _ = axum::serve(listener, app).await;
120        });
121
122        Ok(Self {
123            addr,
124            tape,
125            id: id.into(),
126        })
127    }
128
129    /// The base URL to hand the agent, e.g. via `OPENAI_BASE_URL`.
130    #[must_use]
131    pub fn base_url(&self) -> String {
132        format!("http://{}", self.addr)
133    }
134
135    /// The address the proxy is bound to.
136    #[must_use]
137    pub fn addr(&self) -> SocketAddr {
138        self.addr
139    }
140
141    /// Snapshot the tape as a cassette. Takes `&self` so a recording can be inspected while
142    /// the session is still open.
143    #[must_use]
144    pub fn cassette(&self) -> Cassette {
145        let exchanges = self
146            .tape
147            .lock()
148            .unwrap_or_else(std::sync::PoisonError::into_inner)
149            .clone();
150        Cassette {
151            exchanges,
152            ..Cassette::new(self.id.clone())
153        }
154    }
155}
156
157/// Relay one request upstream and record the round trip.
158async fn capture(State(state): State<Arc<ProxyState>>, req: Request) -> Response {
159    let method = req.method().clone();
160    let path = req
161        .uri()
162        .path_and_query()
163        .map_or_else(|| req.uri().path().to_string(), ToString::to_string);
164
165    let (parts, body) = req.into_parts();
166    let body = match axum::body::to_bytes(body, usize::MAX).await {
167        Ok(bytes) => bytes,
168        Err(err) => {
169            return (
170                StatusCode::BAD_REQUEST,
171                format!("cannot read request body: {err}"),
172            )
173                .into_response();
174        }
175    };
176
177    // Forward every inbound header, credentials included: the proxy is a relay, and the
178    // agent's key is how upstream authenticates it. The allowlist governs what reaches the
179    // *cassette*, which is a different question from what reaches the provider.
180    let url = format!("{}{}", state.upstream, path);
181    let mut outbound = state.client.request(method.clone(), &url);
182    for (name, value) in &parts.headers {
183        // `host` belongs to the proxy's own origin; forwarding it would point upstream at
184        // 127.0.0.1 and vhost-routed providers would reject the request.
185        if name.as_str().eq_ignore_ascii_case("host") {
186            continue;
187        }
188        outbound = outbound.header(name, value);
189    }
190
191    let upstream_response = match outbound.body(body.clone()).send().await {
192        Ok(response) => response,
193        Err(err) => {
194            // Relay failure is reported to the agent as a gateway error rather than a panic:
195            // the agent under recording should see a normal HTTP failure it can handle.
196            return (
197                StatusCode::BAD_GATEWAY,
198                format!("upstream request failed: {err}"),
199            )
200                .into_response();
201        }
202    };
203
204    let status = upstream_response.status();
205    let response_headers = upstream_response.headers().clone();
206    let response_body = match upstream_response.bytes().await {
207        Ok(bytes) => bytes,
208        Err(err) => {
209            return (
210                StatusCode::BAD_GATEWAY,
211                format!("upstream response failed: {err}"),
212            )
213                .into_response();
214        }
215    };
216
217    push(
218        &state.tape,
219        CapturedRequest {
220            method: method.to_string(),
221            path,
222            headers: retain_request_headers(
223                parts
224                    .headers
225                    .iter()
226                    .map(|(n, v)| (n.as_str(), v.as_bytes())),
227            ),
228            body: Body::from_bytes(&body),
229        },
230        CapturedResponse {
231            status: status.as_u16(),
232            headers: retain_response_headers(
233                response_headers
234                    .iter()
235                    .map(|(n, v)| (n.as_str(), v.as_bytes())),
236            ),
237            body: Body::from_bytes(&response_body),
238        },
239    );
240
241    // The caller gets the upstream's answer, not ours. A recording that changes the run it
242    // records is not a recording.
243    let mut relayed = Response::builder().status(status);
244    for (name, value) in &response_headers {
245        // Hop-by-hop framing is the relay's own concern: we buffered the body, so upstream's
246        // transfer-encoding/content-length no longer describe what we are about to send.
247        if is_hop_by_hop(name.as_str()) {
248            continue;
249        }
250        relayed = relayed.header(name, value);
251    }
252    relayed
253        .body(axum::body::Body::from(response_body))
254        .map_or_else(
255            |err| {
256                (
257                    StatusCode::INTERNAL_SERVER_ERROR,
258                    format!("cannot relay response: {err}"),
259                )
260                    .into_response()
261            },
262            IntoResponse::into_response,
263        )
264}
265
266/// Headers that describe a single hop's framing and must not be copied across a relay.
267fn is_hop_by_hop(name: &str) -> bool {
268    const HOP_BY_HOP: &[&str] = &[
269        "connection",
270        "content-length",
271        "keep-alive",
272        "proxy-authenticate",
273        "proxy-authorization",
274        "te",
275        "trailer",
276        "transfer-encoding",
277        "upgrade",
278    ];
279    HOP_BY_HOP.iter().any(|h| name.eq_ignore_ascii_case(h))
280}
281
282/// Append one exchange to the tape, stamping its capture order.
283fn push(tape: &Tape, request: CapturedRequest, response: CapturedResponse) {
284    let mut tape = tape
285        .lock()
286        .unwrap_or_else(std::sync::PoisonError::into_inner);
287    let idx = tape.len();
288    tape.push(Exchange {
289        idx,
290        request,
291        response,
292    });
293}
amberfork_record/proxy.rs

amberfork_record/
proxy.rs
