magi-code 0.100.0

//! Request-scoped loopback relay: one admission, one official upstream generation.
//! Authorization headers remain in memory; errors never include request contents.
use crate::cancellation::{AgentCancellation, AgentCancellationHandle};
use anyhow::{Context, Result, bail};
use std::{
    io::{Read, Write},
    net::{TcpListener, TcpStream},
    sync::{
        Arc, Mutex,
        atomic::{AtomicUsize, Ordering},
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

const MAX_HEADERS: usize = 32 * 1024;
const MAX_REQUEST: usize = 8 * 1024 * 1024;
const MAX_RESPONSE: usize = 16 * 1024 * 1024;
// Relay polling is kept deliberately (issue #659): accept/finish add at most
// ~100ms each per generation, and blocking writes inside `block_on` stall
// cancellation by at most one write timeout because `write_all_cancellable`
// re-checks cancellation. Revisit if admission latency is measured as material.
const SOCKET_POLL: Duration = Duration::from_millis(100);
const ADMISSION_HEAD_TIMEOUT: Duration = Duration::from_secs(15);
// Concurrent local connections; extras are closed so stray clients cannot pile up threads.
const MAX_RELAY_CONNECTIONS: usize = 8;
pub(super) const CLAUDE_STARTUP_TIMEOUT: Duration = Duration::from_secs(300);
pub(super) const CLAUDE_MAX_REQUEST_DURATION: Duration = Duration::from_secs(30 * 60);

#[derive(Clone, Copy, Default, PartialEq, Eq)]
enum RelayPhase {
    #[default]
    AwaitingRequest,
    ReadingRequestBody,
    ConnectingUpstream,
    StreamingUpstream,
    ForwardingUpstream,
    FinishingLocalResponse,
    ResponseComplete,
}

impl std::fmt::Display for RelayPhase {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter.write_str(match self {
            Self::AwaitingRequest => "awaiting request",
            Self::ReadingRequestBody => "reading request body",
            Self::ConnectingUpstream => "connecting upstream",
            Self::StreamingUpstream => "streaming upstream",
            Self::ForwardingUpstream => "forwarding upstream",
            Self::FinishingLocalResponse => "finishing local response",
            Self::ResponseComplete => "response complete",
        })
    }
}

#[derive(Default)]
struct AdmissionState {
    claimed: bool,
    denied: u32,
    // Connections dropped before admission (wrong route, malformed or stalled head).
    // They never fail the generation; only the admitted request can.
    ignored_connections: u32,
    result: Option<RelayResult>,
    failure_reason: Option<&'static str>,
    phase: RelayPhase,
    last_upstream_bytes: Option<Instant>,
    // Payload bytes in fully forwarded chunks; a partial local write is not credited.
    forwarded_bytes: usize,
    last_forwarded_bytes: Option<Instant>,
    upstream_eof: bool,
}

impl AdmissionState {
    /// Keeps the first failure; later failures are usually its consequences.
    fn record_failure(&mut self, reason: &'static str) {
        self.failure_reason.get_or_insert(reason);
    }
}

pub(super) struct RelayResult {
    pub(super) status: u16,
    pub(super) body: Vec<u8>,
}

pub(super) struct ClaudeAdmission {
    pub(super) url: String,
    state: Arc<Mutex<AdmissionState>>,
    stop: AgentCancellationHandle,
    cancellation: AgentCancellation,
    worker: Option<JoinHandle<()>>,
}

impl ClaudeAdmission {
    /// `host_messages` is the complete ordered host conversation, never a CLI session replay.
    pub(super) fn start(
        cancellation: &AgentCancellation,
        host_messages: Vec<serde_json::Value>,
    ) -> Result<Self> {
        Self::start_at(
            "https://api.anthropic.com".to_string(),
            cancellation,
            host_messages,
        )
    }

    fn start_at(
        upstream: String,
        cancellation: &AgentCancellation,
        host_messages: Vec<serde_json::Value>,
    ) -> Result<Self> {
        if upstream != "https://api.anthropic.com" {
            bail!("Claude subscription relay requires official Anthropic endpoint");
        }
        let listener = TcpListener::bind(("127.0.0.1", 0)).context("binding Claude relay")?;
        listener.set_nonblocking(true)?;
        let route = format!("/admit/{}", uuid::Uuid::new_v4());
        let url = format!("http://127.0.0.1:{}{route}", listener.local_addr()?.port());
        let (worker_cancel, stop) = cancellation.child_token();
        let state = Arc::new(Mutex::new(AdmissionState::default()));
        let worker_state = Arc::clone(&state);
        let worker = thread::spawn(move || {
            let active_handlers = AtomicUsize::new(0);
            // Scoped handlers borrow request data; leaving the scope waits for every
            // handler, each of which re-checks cancellation within one socket poll.
            thread::scope(|scope| {
                while !worker_cancel.is_canceled() {
                    match listener.accept() {
                        Ok((socket, _)) => {
                            if worker_cancel.is_canceled() {
                                break;
                            }
                            if active_handlers.fetch_add(1, Ordering::AcqRel)
                                >= MAX_RELAY_CONNECTIONS
                            {
                                active_handlers.fetch_sub(1, Ordering::AcqRel);
                                let mut state =
                                    worker_state.lock().unwrap_or_else(|p| p.into_inner());
                                state.ignored_connections =
                                    state.ignored_connections.saturating_add(1);
                                continue;
                            }
                            let (route, upstream, host_messages, state, cancel, active) = (
                                &route,
                                &upstream,
                                &host_messages,
                                &*worker_state,
                                &worker_cancel,
                                &active_handlers,
                            );
                            scope.spawn(move || {
                                if let Err(error) =
                                    serve(socket, route, upstream, host_messages, state, cancel)
                                    && !cancel.is_canceled()
                                {
                                    state
                                        .lock()
                                        .unwrap_or_else(|p| p.into_inner())
                                        .record_failure(relay_failure_reason(&error));
                                }
                                active.fetch_sub(1, Ordering::AcqRel);
                            });
                        }
                        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                            thread::sleep(SOCKET_POLL);
                        }
                        Err(_) => {
                            worker_state
                                .lock()
                                .unwrap_or_else(|p| p.into_inner())
                                .record_failure("local connection accept failed");
                            break;
                        }
                    }
                }
            });
        });
        Ok(Self {
            url,
            state,
            stop,
            cancellation: cancellation.clone(),
            worker: Some(worker),
        })
    }

    /// Copy only new upstream bytes; never hold the relay lock while parsing or delivering.
    pub(super) fn stream_bytes(&self, offset: &mut usize) -> Result<Vec<u8>> {
        let state = self.state.lock().unwrap_or_else(|p| p.into_inner());
        if let Some(reason) = state.failure_reason {
            bail!("Claude relay failed: {reason} (phase={})", state.phase);
        }
        let Some(result) = state.result.as_ref().filter(|result| result.status == 200) else {
            return Ok(Vec::new());
        };
        let bytes = result.body[*offset..].to_vec();
        *offset = result.body.len();
        Ok(bytes)
    }

    /// Fixed fields only; capture before cleanup so a stalled local write stays visible.
    pub(super) fn diagnostic(&self) -> String {
        let state = self.state.lock().unwrap_or_else(|p| p.into_inner());
        let age = |time: Option<Instant>| {
            time.map(|time| time.elapsed().as_millis().to_string())
                .unwrap_or_else(|| "none".to_owned())
        };
        format!(
            "Claude relay state (phase={}, claimed={}, failed={}, additional_requests={}, ignored_connections={}, status={}, upstream_bytes={}, forwarded_bytes={}, upstream_byte_age_ms={}, forwarded_byte_age_ms={}, upstream_eof={})",
            state.phase,
            state.claimed,
            state.failure_reason.is_some(),
            state.denied,
            state.ignored_connections,
            state.result.as_ref().map_or(0, |result| result.status),
            state.result.as_ref().map_or(0, |result| result.body.len()),
            state.forwarded_bytes,
            age(state.last_upstream_bytes),
            age(state.last_forwarded_bytes),
            state.upstream_eof,
        )
    }

    pub(super) fn finish(mut self) -> Result<RelayResult> {
        let deadline = Instant::now() + Duration::from_secs(5);
        while Instant::now() < deadline {
            self.cancellation.check()?;
            let state = self.state.lock().unwrap_or_else(|p| p.into_inner());
            if state.failure_reason.is_some() || state.phase == RelayPhase::ResponseComplete {
                break;
            }
            drop(state);
            thread::sleep(SOCKET_POLL);
        }
        self.cancellation.check()?;
        self.shutdown();
        let mut state = self.state.lock().unwrap_or_else(|p| p.into_inner());
        // Upstream rejection first: Claude CLI may send a follow-up request after an
        // upstream error, and denying it must not hide the original cause.
        if let Some(result) = state.result.as_ref().filter(|result| result.status != 200) {
            return Err(upstream_error_diagnostic(result.status, &result.body));
        }
        if let Some(reason) = state.failure_reason {
            bail!(
                "Claude relay failed: {reason} (phase={}, additional_requests={})",
                state.phase,
                state.denied
            );
        }
        if state.denied != 0 {
            // HTTP 200 can still contain an explicit refusal. Validate the upstream
            // response before reporting the CLI's secondary retry attempts.
            if state.phase == RelayPhase::ResponseComplete
                && let Some(result) = state.result.as_ref()
            {
                super::anthropic::parse_captured_sse(&result.body)?;
            }
            bail!(
                "Claude CLI attempted {} additional model requests; magi-code admits one request per generation (claimed={}, phase={})",
                state.denied,
                state.claimed,
                state.phase
            );
        }
        if !state.claimed || state.phase != RelayPhase::ResponseComplete {
            bail!(
                "Claude relay incomplete (claimed={}, phase={}, ignored_connections={})",
                state.claimed,
                state.phase,
                state.ignored_connections
            );
        }
        state.result.take().context("Claude relay missing response")
    }

    fn shutdown(&mut self) {
        self.stop.cancel();
        if let Some(worker) = self.worker.take() {
            let _ = worker.join();
        }
    }
}

impl Drop for ClaudeAdmission {
    fn drop(&mut self) {
        self.shutdown();
    }
}

/// Fixed categories only: reqwest errors can contain credential-bearing URLs.
fn relay_failure_reason(error: &anyhow::Error) -> &'static str {
    if let Some(error) = error.downcast_ref::<reqwest::Error>() {
        if error.is_timeout() {
            "upstream request timed out"
        } else if error.is_connect() {
            "upstream connection failed"
        } else if error.is_body() || error.is_decode() {
            "upstream response read failed"
        } else {
            "upstream transport failed"
        }
    } else if error
        .downcast_ref::<std::io::Error>()
        .is_some_and(|error| error.kind() == std::io::ErrorKind::TimedOut)
    {
        "local response write timed out"
    } else {
        "local relay I/O or protocol failed"
    }
}

const PROMPT_TOO_LONG_HINT: &str = "Claude CLI rejected prompt as too long; run /compact or start /new; session history was not truncated";

pub(super) fn claude_error_with_hint(error: anyhow::Error, hint: Option<&str>) -> anyhow::Error {
    match hint {
        Some(PROMPT_TOO_LONG_HINT) if !crate::cancellation::is_run_canceled(&error) => {
            super::error::ProviderError::prompt_too_long(format!("{error}; {PROMPT_TOO_LONG_HINT}"))
                .into()
        }
        Some(hint) => {
            let message = format!("{error}; {hint}");
            error.context(message)
        }
        None => error,
    }
}

/// Return only a fixed diagnostic category, never an upstream body or CLI stderr.
pub(super) fn claude_error_hint(message: &[u8]) -> Option<&'static str> {
    let text = String::from_utf8_lossy(&message[..message.len().min(8192)]).to_ascii_lowercase();
    if text.contains("prompt") && text.contains("too long") {
        Some(PROMPT_TOO_LONG_HINT)
    } else if text.contains("model")
        && [
            "not available",
            "does not exist",
            "not found",
            "not allowed",
            "no access",
            "don't have access",
            "do not have access",
            "not authorized",
            "invalid model",
            "unknown model",
        ]
        .iter()
        .any(|phrase| text.contains(phrase))
    {
        Some("model unavailable for this Claude account or CLI; choose a supported model")
    } else if text.contains("rate_limit") || text.contains("rate limit") {
        Some("Claude rate limit reached; retry later")
    } else if text.contains("overloaded_error") {
        Some("Claude upstream overloaded; retry later")
    } else if text.contains("api_error") {
        Some("Claude upstream API failed; retry later")
    } else if text.contains("authentication_error") || text.contains("unauthorized") {
        Some("Claude authentication rejected; check claude auth status")
    } else {
        None
    }
}

fn upstream_error_diagnostic(status: u16, body: &[u8]) -> anyhow::Error {
    let hint = (body.len() <= 8192)
        .then(|| serde_json::from_slice::<serde_json::Value>(body).ok())
        .flatten()
        .and_then(|error| {
            error
                .pointer("/error/message")
                .and_then(serde_json::Value::as_str)
                .and_then(|message| claude_error_hint(message.as_bytes()))
        })
        .or(match status {
            401 | 403 => {
                Some("Claude authentication or model access rejected; check claude auth status")
            }
            429 => Some("Claude rate limit reached; retry later"),
            _ => None,
        });
    claude_error_with_hint(
        anyhow::anyhow!("Claude upstream returned HTTP {status}"),
        hint,
    )
}

/// Replace the CLI placeholder with ordered host history. Retain native reminders
/// and cache boundaries; reject unexpected shapes before any upstream request.
fn restore_host_conversation(
    payload: &[u8],
    host_messages: &[serde_json::Value],
) -> Option<Vec<u8>> {
    use serde_json::{Value, json};
    const MAX_CACHE_BREAKPOINTS: usize = 4;
    if host_messages.is_empty() {
        return None;
    }
    let mut body: Value = serde_json::from_slice(payload).ok()?;
    let cli_messages = body.get("messages")?.as_array()?;
    let mut reminders = Vec::new();
    let mut cli_system_messages = Vec::new();
    for message in cli_messages {
        match message.get("role").and_then(Value::as_str)? {
            "system" => cli_system_messages.push(message.clone()),
            "user" => reminders.extend(
                message
                    .get("content")
                    .and_then(Value::as_array)
                    .into_iter()
                    .flatten()
                    .filter(|block| {
                        block
                            .get("text")
                            .and_then(Value::as_str)
                            .is_some_and(|text| text.starts_with("<system-reminder>"))
                    })
                    .cloned(),
            ),
            _ => {}
        }
    }
    let mut messages = host_messages.to_vec();
    let used_breakpoints = serde_json::to_string(&body["system"])
        .ok()?
        .matches("\"cache_control\"")
        .count()
        + serde_json::to_string(&body["tools"])
            .ok()?
            .matches("\"cache_control\"")
            .count()
        + serde_json::to_string(&cli_system_messages)
            .ok()?
            .matches("\"cache_control\"")
            .count();
    if used_breakpoints < MAX_CACHE_BREAKPOINTS
        && let Some(last_block) = messages
            .last_mut()
            .and_then(|message| message.get_mut("content"))
            .and_then(Value::as_array_mut)
            .and_then(|content| content.last_mut())
    {
        last_block["cache_control"] = json!({"type":"ephemeral","ttl":"1h"});
    }
    if !reminders.is_empty() {
        let first_user_content = messages
            .iter_mut()
            .find(|message| message["role"] == "user")
            .and_then(|message| message.get_mut("content"))
            .and_then(Value::as_array_mut)?;
        first_user_content.splice(0..0, reminders);
    }
    messages.extend(cli_system_messages);
    body["messages"] = Value::Array(messages);
    serde_json::to_vec(&body).ok()
}

/// Admission candidate: route matched and headers parsed, body not yet read.
struct AdmissionRequest {
    query: String,
    headers: reqwest::header::HeaderMap,
    content_length: Option<usize>,
    body_prefix: Vec<u8>,
}

fn serve(
    mut socket: TcpStream,
    route: &str,
    upstream: &str,
    host_messages: &[serde_json::Value],
    state: &Mutex<AdmissionState>,
    cancel: &AgentCancellation,
) -> Result<()> {
    // Accepted sockets can inherit the listener's nonblocking mode (BSD/macOS); reads and
    // writes rely on blocking timeouts so each handler waits instead of spinning.
    socket.set_nonblocking(false)?;
    socket.set_read_timeout(Some(SOCKET_POLL))?;
    socket.set_write_timeout(Some(SOCKET_POLL))?;
    // Stray, malformed or misrouted connections are dropped without failing the
    // generation; only the admitted request can fail it.
    let request = match read_admission_request(&mut socket, route, cancel) {
        Ok(Some(request)) => request,
        Ok(None) | Err(_) => {
            cancel.check()?;
            let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
            state.ignored_connections = state.ignored_connections.saturating_add(1);
            return Ok(());
        }
    };
    let AdmissionRequest {
        query,
        mut headers,
        content_length,
        body_prefix: mut payload,
    } = request;
    // Claim before reading request body or attempting any upstream I/O.
    {
        let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
        if state.claimed {
            state.denied = state.denied.saturating_add(1);
            drop(state);
            write_response(&mut socket, 400, br#"{"type":"error","error":{"type":"invalid_request_error","message":"MODEL_ADMISSION_CONSUMED"}}"#, cancel)?;
            return Ok(());
        }
        state.claimed = true;
        state.phase = RelayPhase::ReadingRequestBody;
    }
    let length = content_length.context("missing Claude relay content length")?;
    if length > MAX_REQUEST {
        bail!("Claude relay request too large");
    }
    if payload.len() > length {
        bail!("Claude relay request exceeds content length");
    }
    let deadline = Instant::now() + CLAUDE_STARTUP_TIMEOUT;
    let mut buffer = [0u8; 8192];
    while payload.len() < length {
        cancel.check()?;
        if Instant::now() >= deadline {
            bail!("Claude relay request timed out");
        }
        let count = buffer.len().min(length - payload.len());
        match socket.read(&mut buffer[..count]) {
            Ok(0) => bail!("Claude relay request body ended early"),
            Ok(count) => payload.extend_from_slice(&buffer[..count]),
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
                ) =>
            {
                continue;
            }
            Err(_) => bail!("Claude relay request body read failed"),
        }
    }
    let payload = restore_host_conversation(&payload, host_messages)
        .context("Claude request could not restore host conversation; no upstream request sent")?;
    if payload.len() > MAX_REQUEST {
        bail!("Claude restored request too large");
    }
    state.lock().unwrap_or_else(|p| p.into_inner()).phase = RelayPhase::ConnectingUpstream;
    headers.insert(
        reqwest::header::ACCEPT_ENCODING,
        reqwest::header::HeaderValue::from_static("identity"),
    );
    // Blocking only on local socket reads with short timeouts; upstream request can be
    // dropped immediately by select! when cancellation is observed.
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let canceled = async {
        while !cancel.is_canceled() {
            tokio::time::sleep(SOCKET_POLL).await;
        }
    };
    let relay = runtime.block_on(async {
        let client = reqwest::Client::builder()
            .redirect(reqwest::redirect::Policy::none())
            .connect_timeout(CLAUDE_STARTUP_TIMEOUT)
            // The native reader owns the progress deadline and drops this relay on timeout.
            .timeout(CLAUDE_MAX_REQUEST_DURATION)
            .build()?;
        let request = client
            .post(format!("{upstream}/v1/messages{query}"))
            .headers(headers)
            .body(payload);
        let upstream = async {
            let mut response = request.send().await?;
            let status = response.status().as_u16();
            {
                let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
                state.phase = RelayPhase::StreamingUpstream;
                state.result = Some(RelayResult {
                    status,
                    body: Vec::new(),
                });
            }
            let mut socket_started = false;
            while let Some(chunk) = response.chunk().await? {
                {
                    let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
                    let result = state
                        .result
                        .as_mut()
                        .context("Claude relay missing response")?;
                    if result.body.len().saturating_add(chunk.len()) > MAX_RESPONSE {
                        state.record_failure("upstream response exceeded size limit");
                        bail!("Claude upstream response too large");
                    }
                    result.body.extend_from_slice(&chunk);
                    state.last_upstream_bytes = Some(Instant::now());
                    state.phase = RelayPhase::ForwardingUpstream;
                }
                if !socket_started {
                    write_stream_headers(&mut socket, status, cancel)?;
                    socket_started = true;
                }
                write_chunk(&mut socket, &chunk, cancel)?;
                let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
                state.forwarded_bytes += chunk.len();
                state.last_forwarded_bytes = Some(Instant::now());
                state.phase = RelayPhase::StreamingUpstream;
            }
            {
                let mut state = state.lock().unwrap_or_else(|p| p.into_inner());
                state.upstream_eof = true;
                state.phase = RelayPhase::FinishingLocalResponse;
            }
            if !socket_started {
                write_stream_headers(&mut socket, status, cancel)?;
            }
            write_all_cancellable(&mut socket, b"0\r\n\r\n", cancel)?;
            state.lock().unwrap_or_else(|p| p.into_inner()).phase = RelayPhase::ResponseComplete;
            Ok::<_, anyhow::Error>(())
        };
        tokio::select! {
            result = upstream => result,
            _ = canceled => bail!("Claude relay canceled"),
        }
    });
    relay?;
    Ok(())
}

/// Reads one request head. `Ok(None)` means the connection is not an admission
/// candidate (wrong method or route) and has been answered with 404.
fn read_admission_request(
    socket: &mut TcpStream,
    route: &str,
    cancel: &AgentCancellation,
) -> Result<Option<AdmissionRequest>> {
    // Claude CLI writes its request head immediately after connecting; a short
    // bound keeps an idle stray connection from holding the serial accept loop.
    let deadline = Instant::now() + ADMISSION_HEAD_TIMEOUT;
    let mut data = Vec::new();
    let mut buffer = [0u8; 8192];
    let end = loop {
        cancel.check()?;
        if Instant::now() >= deadline {
            bail!("Claude relay request head timed out");
        }
        if data.len() > MAX_HEADERS {
            bail!("Claude relay request headers too large");
        }
        match socket.read(&mut buffer) {
            Ok(0) => bail!("Claude relay request ended early"),
            Ok(count) => data.extend_from_slice(&buffer[..count]),
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
                ) =>
            {
                continue;
            }
            Err(_) => bail!("Claude relay request read failed"),
        }
        if let Some(offset) = data.windows(4).position(|part| part == b"\r\n\r\n") {
            break offset + 4;
        }
    };
    if end > MAX_HEADERS {
        bail!("Claude relay request headers too large");
    }
    let header_text = std::str::from_utf8(&data[..end]).context("invalid Claude relay headers")?;
    let mut lines = header_text.split("\r\n");
    let requested = lines.next().context("missing Claude relay request line")?;
    let expected = format!("{route}/v1/messages");
    let Some((_, query)) = requested
        .strip_prefix("POST ")
        .and_then(|rest| rest.strip_suffix(" HTTP/1.1"))
        .map(|target| target.split_once('?').unwrap_or((target, "")))
        .filter(|(path, _)| *path == expected)
    else {
        write_response(socket, 404, b"", cancel)?;
        return Ok(None);
    };
    let query = if query.is_empty() {
        String::new()
    } else {
        format!("?{query}")
    };
    let mut headers = reqwest::header::HeaderMap::new();
    let mut content_length = None;
    for line in lines.filter(|line| !line.is_empty()) {
        let (key, value) = line
            .split_once(':')
            .context("invalid Claude relay header")?;
        let key = reqwest::header::HeaderName::from_bytes(key.as_bytes())
            .context("invalid Claude relay header name")?;
        let value = reqwest::header::HeaderValue::from_str(value.trim())
            .context("invalid Claude relay header value")?;
        if key == reqwest::header::CONTENT_LENGTH {
            if content_length.is_some() {
                bail!("duplicate Claude relay content length");
            }
            content_length = Some(
                value
                    .to_str()?
                    .parse::<usize>()
                    .context("invalid Claude relay content length")?,
            );
        }
        if key == reqwest::header::ORIGIN || key == reqwest::header::TRANSFER_ENCODING {
            bail!("unsupported Claude relay request header");
        }
        if key.as_str() != "proxy-connection"
            && !matches!(
                key,
                reqwest::header::HOST
                    | reqwest::header::CONNECTION
                    | reqwest::header::CONTENT_LENGTH
                    | reqwest::header::ACCEPT_ENCODING
                    | reqwest::header::PROXY_AUTHORIZATION
            )
        {
            headers.append(key, value);
        }
    }
    Ok(Some(AdmissionRequest {
        query,
        headers,
        content_length,
        body_prefix: data[end..].to_vec(),
    }))
}

fn write_all_cancellable(
    socket: &mut TcpStream,
    mut bytes: &[u8],
    cancel: &AgentCancellation,
) -> Result<()> {
    let deadline = Instant::now() + CLAUDE_STARTUP_TIMEOUT;
    while !bytes.is_empty() {
        cancel.check()?;
        if Instant::now() >= deadline {
            return Err(std::io::Error::new(
                std::io::ErrorKind::TimedOut,
                "Claude relay socket write timed out",
            )
            .into());
        }
        match socket.write(bytes) {
            Ok(0) => bail!("Claude relay socket closed"),
            Ok(count) => bytes = &bytes[count..],
            Err(error)
                if matches!(
                    error.kind(),
                    std::io::ErrorKind::TimedOut | std::io::ErrorKind::WouldBlock
                ) =>
            {
                continue;
            }
            Err(_) => bail!("Claude relay socket write failed"),
        }
    }
    Ok(())
}

fn write_stream_headers(
    socket: &mut TcpStream,
    status: u16,
    cancel: &AgentCancellation,
) -> Result<()> {
    let content_type = if status == 200 {
        "text/event-stream"
    } else {
        "application/json"
    };
    write_all_cancellable(socket, format!("HTTP/1.1 {status} Upstream\r\nTransfer-Encoding: chunked\r\nContent-Type: {content_type}\r\nConnection: close\r\n\r\n").as_bytes(), cancel)
}

fn write_chunk(socket: &mut TcpStream, chunk: &[u8], cancel: &AgentCancellation) -> Result<()> {
    write_all_cancellable(socket, format!("{:X}\r\n", chunk.len()).as_bytes(), cancel)?;
    write_all_cancellable(socket, chunk, cancel)?;
    write_all_cancellable(socket, b"\r\n", cancel)
}

fn write_response(
    socket: &mut TcpStream,
    status: u16,
    body: &[u8],
    cancel: &AgentCancellation,
) -> Result<()> {
    write_all_cancellable(socket, format!("HTTP/1.1 {status} Denied\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", body.len()).as_bytes(), cancel)?;
    write_all_cancellable(socket, body, cancel)
}