dirge-agent 0.12.2

use bytes::Bytes;
use rig::http_client::{
    self, HttpClientExt, LazyBody, MultipartForm, Request, Response, StreamingResponse,
};

/// Normalizes Anthropic OAuth requests at the transport boundary: swaps
/// `x-api-key` for `Authorization: Bearer`, adds the Claude Code identity
/// headers, and shapes the body. rig 0.37 exposes no per-request header seam.
//
// `bearer_token` is `Option` only to satisfy the `HttpClientExt: Default`
// bound; a default instance never sends.
#[derive(Clone, Default)]
pub(crate) struct AnthropicHttpClient {
    inner: reqwest::Client,
    bearer_token: Option<String>,
}

// Redacts the token so it can't leak via `{:?}`.
impl std::fmt::Debug for AnthropicHttpClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AnthropicHttpClient")
            .field("bearer_token", &"<redacted>")
            .finish()
    }
}

/// Anthropic requires this exact text as the first system block when
/// authenticating with a Claude Code OAuth token.
const CLAUDE_CODE_SYSTEM_PROMPT: &str = "You are Claude Code, Anthropic's official CLI for Claude.";

/// Beta flags Anthropic's API requires for the Claude Code OAuth wire path.
const ANTHROPIC_OAUTH_BETA: &str = "claude-code-20250219,oauth-2025-04-20";

/// `User-Agent` that identifies the request as first-party Claude Code.
/// Without it the subscription path returns a third-party "extra usage" 400.
const CLAUDE_CODE_USER_AGENT: &str = "claude-cli/2.1.75";

impl AnthropicHttpClient {
    pub(crate) fn new(bearer_token: String) -> Self {
        Self {
            inner: reqwest::Client::new(),
            bearer_token: Some(bearer_token),
        }
    }

    fn normalized_request<T>(&self, req: Request<T>) -> http_client::Result<Request<Bytes>>
    where
        T: Into<Bytes>,
    {
        let (mut parts, body) = req.into_parts();
        parts.headers.remove("x-api-key");
        if let Some(token) = self.bearer_token.as_deref()
            && let Ok(value) = http::HeaderValue::from_str(&format!("Bearer {token}"))
        {
            parts.headers.insert(http::header::AUTHORIZATION, value);
        }
        parts.headers.insert(
            http::HeaderName::from_static("anthropic-beta"),
            http::HeaderValue::from_static(ANTHROPIC_OAUTH_BETA),
        );
        parts.headers.insert(
            http::HeaderName::from_static("anthropic-dangerous-direct-browser-access"),
            http::HeaderValue::from_static("true"),
        );
        parts.headers.insert(
            http::HeaderName::from_static("x-app"),
            http::HeaderValue::from_static("cli"),
        );
        parts.headers.insert(
            http::header::USER_AGENT,
            http::HeaderValue::from_static(CLAUDE_CODE_USER_AGENT),
        );

        let body = body.into();
        let body = if is_messages_path(parts.uri.path()) && self.is_oauth_token() {
            shape_oauth_messages_payload(body)
        } else {
            body
        };

        let mut builder = Request::builder()
            .method(parts.method)
            .uri(parts.uri)
            .version(parts.version);
        if let Some(headers) = builder.headers_mut() {
            *headers = parts.headers;
        }
        builder.body(body).map_err(http_client::Error::Protocol)
    }

    fn is_oauth_token(&self) -> bool {
        self.bearer_token
            .as_deref()
            .is_some_and(|token| token.contains("sk-ant-oat"))
    }
}

impl HttpClientExt for AnthropicHttpClient {
    fn send<T, U>(
        &self,
        req: Request<T>,
    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + Send + 'static
    where
        T: Into<Bytes>,
        T: Send,
        U: From<Bytes>,
        U: Send + 'static,
    {
        let inner = self.inner.clone();
        let req = self.normalized_request(req);
        async move {
            let req = req?;
            inner.send(req).await
        }
    }

    fn send_multipart<U>(
        &self,
        req: Request<MultipartForm>,
    ) -> impl Future<Output = http_client::Result<Response<LazyBody<U>>>> + Send + 'static
    where
        U: From<Bytes> + Send + 'static,
    {
        self.inner.send_multipart(req)
    }

    fn send_streaming<T>(
        &self,
        req: Request<T>,
    ) -> impl Future<Output = http_client::Result<StreamingResponse>> + Send
    where
        T: Into<Bytes> + Send,
    {
        let inner = self.inner.clone();
        let req = self.normalized_request(req);
        async move {
            let req = req?;
            inner.send_streaming(req).await
        }
    }
}

fn is_messages_path(path: &str) -> bool {
    path.ends_with("/messages")
}

const CLAUDE_CODE_VERSION: &str = "2.1.169";
const BILLING_HEADER_SALT: &str = "59cf53e54c78";
const BILLING_HEADER_POSITIONS: [usize; 3] = [4, 7, 20];
const TEXT_REPLACEMENTS: [(&str, &str); 1] = [(
    "Here is some useful information about the environment you are running in:",
    "Environment context you are running in:",
)];

fn shape_oauth_messages_payload(body: Bytes) -> Bytes {
    let Ok(mut value) = serde_json::from_slice::<serde_json::Value>(&body) else {
        return body;
    };
    if !is_anthropic_messages_payload(&value) {
        return body;
    }

    if let Some(messages) = value
        .get_mut("messages")
        .and_then(serde_json::Value::as_array_mut)
    {
        split_assistant_tool_use_trailing_content(messages);
    }

    // Anthropic's Messages API only permits user/assistant in `messages[]`.
    // Dirge injects `role:"system"` entries there for compaction summaries
    // (compression.rs) and mid-session memory re-injection (run.rs). Under the
    // OAuth/Claude-Code classifier a stray system turn is read as non-Claude-Code
    // traffic and rejected as a third-party app ("extra usage" 400), so fold any
    // such entries into the top-level `system` array before the prepends run.
    hoist_system_messages(&mut value);

    shape_system_blocks(&mut value);
    // Order matters: the billing header must be the FIRST system block. With
    // the Claude Code identity block ahead of it, Anthropic's OAuth classifier
    // rejects the request as a third-party app ("extra usage" 400). Prepend the
    // identity first, then the billing header, so billing lands at index 0.
    prepend_claude_code_system_value(&mut value);
    prepend_billing_header(&mut value);

    serde_json::to_vec(&value).map(Bytes::from).unwrap_or(body)
}

fn hoist_system_messages(value: &mut serde_json::Value) {
    let Some(messages) = value
        .get_mut("messages")
        .and_then(serde_json::Value::as_array_mut)
    else {
        return;
    };

    let mut hoisted: Vec<String> = Vec::new();
    messages.retain(|m| {
        if m.get("role").and_then(serde_json::Value::as_str) != Some("system") {
            return true;
        }
        if let Some(text) = system_message_text(m) {
            if !text.is_empty() {
                hoisted.push(text);
            }
        }
        false
    });

    if hoisted.is_empty() {
        return;
    }

    match value.get_mut("system") {
        Some(serde_json::Value::Array(_)) => {}
        Some(existing @ serde_json::Value::String(_)) => {
            let prior = existing.take();
            *existing = serde_json::Value::Array(vec![prior]);
        }
        _ => {
            if let Some(obj) = value.as_object_mut() {
                obj.insert("system".to_string(), serde_json::Value::Array(Vec::new()));
            }
        }
    }
    let blocks = value
        .get_mut("system")
        .and_then(serde_json::Value::as_array_mut)
        .expect("system normalized to an array above");
    for text in hoisted {
        blocks.push(serde_json::json!({ "type": "text", "text": text }));
    }
}

fn system_message_text(message: &serde_json::Value) -> Option<String> {
    match message.get("content") {
        Some(serde_json::Value::String(s)) => Some(s.clone()),
        Some(serde_json::Value::Array(blocks)) => {
            let joined = blocks
                .iter()
                .filter(|b| b.get("type").and_then(serde_json::Value::as_str) == Some("text"))
                .filter_map(|b| b.get("text").and_then(serde_json::Value::as_str))
                .collect::<Vec<_>>()
                .join("\n");
            Some(joined)
        }
        _ => None,
    }
}

fn is_anthropic_messages_payload(value: &serde_json::Value) -> bool {
    value.get("model").is_some_and(serde_json::Value::is_string)
        && value
            .get("messages")
            .is_some_and(serde_json::Value::is_array)
        && value
            .get("stream")
            .is_some_and(serde_json::Value::is_boolean)
}

#[cfg_attr(not(test), allow(dead_code))]
fn prepend_claude_code_system(body: Bytes) -> Bytes {
    let Ok(mut value) = serde_json::from_slice::<serde_json::Value>(&body) else {
        return body;
    };
    let Some(obj) = value.as_object_mut() else {
        return body;
    };

    let claude_block = serde_json::json!({
        "type": "text",
        "text": CLAUDE_CODE_SYSTEM_PROMPT,
    });

    match obj.get_mut("system") {
        // Already an array of content blocks: prepend unless it's already first.
        Some(serde_json::Value::Array(items)) => {
            if first_system_block_is_claude_code(items) {
                return body;
            }
            items.insert(0, claude_block);
        }
        // A bare string system prompt: lift it into the array form behind the
        // required Claude Code block.
        Some(serde_json::Value::String(text)) => {
            let existing = std::mem::take(text);
            obj.insert(
                "system".to_string(),
                serde_json::json!([
                    claude_block,
                    { "type": "text", "text": existing },
                ]),
            );
        }
        // No system prompt at all.
        _ => {
            obj.insert("system".to_string(), serde_json::json!([claude_block]));
        }
    }

    serde_json::to_vec(&value).map(Bytes::from).unwrap_or(body)
}

fn prepend_claude_code_system_value(value: &mut serde_json::Value) {
    let Some(obj) = value.as_object_mut() else {
        return;
    };
    let claude_block = serde_json::json!({ "type": "text", "text": CLAUDE_CODE_SYSTEM_PROMPT });
    match obj.get_mut("system") {
        Some(serde_json::Value::Array(items)) => {
            if !first_system_block_is_claude_code(items) {
                items.insert(0, claude_block);
            }
        }
        Some(serde_json::Value::String(text)) => {
            let existing = std::mem::take(text);
            obj.insert(
                "system".to_string(),
                serde_json::json!([claude_block, { "type": "text", "text": existing }]),
            );
        }
        _ => {
            obj.insert("system".to_string(), serde_json::json!([claude_block]));
        }
    }
}

fn first_system_block_is_claude_code(items: &[serde_json::Value]) -> bool {
    items
        .first()
        .and_then(|item| item.get("text"))
        .and_then(serde_json::Value::as_str)
        == Some(CLAUDE_CODE_SYSTEM_PROMPT)
}

fn prepend_billing_header(value: &mut serde_json::Value) {
    let Some(messages) = value.get("messages").and_then(serde_json::Value::as_array) else {
        return;
    };
    let Some(first_user_text) = first_user_text(messages) else {
        return;
    };
    let billing = build_billing_header(first_user_text);
    let Some(obj) = value.as_object_mut() else {
        return;
    };
    let blocks = obj.entry("system").or_insert_with(|| serde_json::json!([]));
    if !blocks.is_array() {
        let old = blocks.take();
        *blocks = serde_json::json!([normalize_system_block(old)]);
    }
    let Some(items) = blocks.as_array_mut() else {
        return;
    };
    if items.iter().any(|b| {
        b.get("text")
            .and_then(serde_json::Value::as_str)
            .is_some_and(|t| t.contains("x-anthropic-billing-header:"))
    }) {
        return;
    }
    items.insert(0, serde_json::json!({ "type": "text", "text": billing }));
}

fn first_user_text(messages: &[serde_json::Value]) -> Option<&str> {
    messages
        .iter()
        .find(|m| m.get("role").and_then(serde_json::Value::as_str) == Some("user"))
        .and_then(|m| match m.get("content") {
            Some(serde_json::Value::String(s)) => Some(s.as_str()),
            Some(serde_json::Value::Array(blocks)) => blocks
                .iter()
                .find(|b| b.get("type").and_then(serde_json::Value::as_str) == Some("text"))
                .and_then(|b| b.get("text"))
                .and_then(serde_json::Value::as_str),
            _ => None,
        })
        .filter(|s| !s.is_empty())
}

fn build_billing_header(message_text: &str) -> String {
    let cch = hex_sha256(message_text.as_bytes())[0..5].to_string();
    let utf16: Vec<u16> = message_text.encode_utf16().collect();
    let sampled: String = BILLING_HEADER_POSITIONS
        .iter()
        .map(|i| {
            utf16
                .get(*i)
                .copied()
                .map(utf16_code_unit_to_js_char)
                .unwrap_or('0')
        })
        .collect();
    let suffix_input = format!("{BILLING_HEADER_SALT}{sampled}{CLAUDE_CODE_VERSION}");
    let suffix = hex_sha256(suffix_input.as_bytes())[0..3].to_string();
    format!(
        "x-anthropic-billing-header: cc_version={CLAUDE_CODE_VERSION}.{suffix}; cc_entrypoint=sdk-cli; cch={cch};"
    )
}

fn utf16_code_unit_to_js_char(unit: u16) -> char {
    char::from_u32(unit as u32).unwrap_or(char::REPLACEMENT_CHARACTER)
}

fn hex_sha256(bytes: &[u8]) -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(bytes);
    format!("{:x}", hasher.finalize())
}

fn normalize_system_block(value: serde_json::Value) -> serde_json::Value {
    match value {
        serde_json::Value::String(text) => serde_json::json!({ "type": "text", "text": text }),
        serde_json::Value::Object(mut map) => {
            map.insert(
                "type".to_string(),
                serde_json::Value::String("text".to_string()),
            );
            serde_json::Value::Object(map)
        }
        _ => serde_json::json!({ "type": "text", "text": "" }),
    }
}

fn shape_system_blocks(value: &mut serde_json::Value) {
    let Some(items) = value
        .get_mut("system")
        .and_then(serde_json::Value::as_array_mut)
    else {
        return;
    };
    for block in items {
        let Some(text) = block.get("text").and_then(serde_json::Value::as_str) else {
            continue;
        };
        let shaped = shape_system_text(text);
        if let Some(obj) = block.as_object_mut() {
            obj.insert("text".to_string(), serde_json::Value::String(shaped));
        }
    }
}

fn shape_system_text(text: &str) -> String {
    // dirge's full system prompt is retained on OAuth requests: with the
    // billing header as the first system block, Anthropic's classifier accepts
    // the verbose prompt and tool set, so only the known classifier-trigger
    // phrase is rewritten.
    text.replace(TEXT_REPLACEMENTS[0].0, TEXT_REPLACEMENTS[0].1)
}

fn split_assistant_tool_use_trailing_content(messages: &mut Vec<serde_json::Value>) {
    let mut out = Vec::with_capacity(messages.len());
    for message in messages.drain(..) {
        if message.get("role").and_then(serde_json::Value::as_str) != Some("assistant") {
            out.push(message);
            continue;
        }
        let Some(content) = message.get("content").and_then(serde_json::Value::as_array) else {
            out.push(message);
            continue;
        };
        let Some(first_tool_idx) = content
            .iter()
            .position(|b| b.get("type").and_then(serde_json::Value::as_str) == Some("tool_use"))
        else {
            out.push(message);
            continue;
        };
        if !content[first_tool_idx..]
            .iter()
            .any(|b| b.get("type").and_then(serde_json::Value::as_str) != Some("tool_use"))
        {
            out.push(message);
            continue;
        }
        let mut non_tools = Vec::new();
        let mut tools = Vec::new();
        for block in content {
            if block.get("type").and_then(serde_json::Value::as_str) == Some("tool_use") {
                tools.push(block.clone());
            } else {
                non_tools.push(block.clone());
            }
        }
        let mut first = message.clone();
        first
            .as_object_mut()
            .unwrap()
            .insert("content".to_string(), serde_json::Value::Array(non_tools));
        let mut second = message;
        second
            .as_object_mut()
            .unwrap()
            .insert("content".to_string(), serde_json::Value::Array(tools));
        out.push(first);
        out.push(second);
    }
    *messages = out;
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn prepends_claude_code_block_to_system_array() {
        let body = Bytes::from(
            serde_json::json!({
                "system": [{ "type": "text", "text": "Real prompt." }],
                "messages": []
            })
            .to_string(),
        );

        let value: serde_json::Value =
            serde_json::from_slice(&prepend_claude_code_system(body)).unwrap();

        let system = value["system"].as_array().unwrap();
        assert_eq!(system.len(), 2);
        assert_eq!(system[0]["text"], CLAUDE_CODE_SYSTEM_PROMPT);
        assert_eq!(system[1]["text"], "Real prompt.");
    }

    #[test]
    fn lifts_string_system_into_array_behind_claude_code_block() {
        let body = Bytes::from(
            serde_json::json!({ "system": "Real prompt.", "messages": [] }).to_string(),
        );

        let value: serde_json::Value =
            serde_json::from_slice(&prepend_claude_code_system(body)).unwrap();

        let system = value["system"].as_array().unwrap();
        assert_eq!(system.len(), 2);
        assert_eq!(system[0]["text"], CLAUDE_CODE_SYSTEM_PROMPT);
        assert_eq!(system[1]["text"], "Real prompt.");
    }

    #[test]
    fn adds_system_when_absent() {
        let body = Bytes::from(serde_json::json!({ "messages": [] }).to_string());

        let value: serde_json::Value =
            serde_json::from_slice(&prepend_claude_code_system(body)).unwrap();

        let system = value["system"].as_array().unwrap();
        assert_eq!(system.len(), 1);
        assert_eq!(system[0]["text"], CLAUDE_CODE_SYSTEM_PROMPT);
    }

    #[test]
    fn does_not_double_prepend_claude_code_block() {
        let body = Bytes::from(
            serde_json::json!({
                "system": [
                    { "type": "text", "text": CLAUDE_CODE_SYSTEM_PROMPT },
                    { "type": "text", "text": "Real prompt." }
                ]
            })
            .to_string(),
        );

        let value: serde_json::Value =
            serde_json::from_slice(&prepend_claude_code_system(body)).unwrap();

        let system = value["system"].as_array().unwrap();
        assert_eq!(system.len(), 2);
    }

    #[test]
    fn oauth_shaper_preserves_full_system_prompt() {
        use crate::agent::prompt::SYSTEM_PROMPT;
        let body = Bytes::from(
            serde_json::json!({
                "model": "claude-sonnet-4-5",
                "stream": true,
                "system": [{"type": "text", "text": SYSTEM_PROMPT}],
                "messages": [{"role": "user", "content": "hi"}]
            })
            .to_string(),
        );
        let value: serde_json::Value =
            serde_json::from_slice(&shape_oauth_messages_payload(body)).unwrap();
        let system = value["system"].as_array().unwrap();
        // [billing, identity, full dirge prompt] — prompt retained verbatim.
        assert!(
            system[0]["text"]
                .as_str()
                .unwrap()
                .contains("x-anthropic-billing-header:")
        );
        assert_eq!(system[1]["text"], CLAUDE_CODE_SYSTEM_PROMPT);
        assert_eq!(system[2]["text"], SYSTEM_PROMPT);
    }

    #[test]
    fn oauth_shaper_adds_billing_and_classifier_rewrite() {
        let body = Bytes::from(serde_json::json!({
            "model": "claude-sonnet-4-5",
            "stream": true,
            "system": [{"type":"text", "text":"Here is some useful information about the environment you are running in:"}],
            "messages": [{"role":"user", "content":"hello from dirge oauth"}]
        }).to_string());

        let value: serde_json::Value =
            serde_json::from_slice(&shape_oauth_messages_payload(body)).unwrap();
        let system = value["system"].as_array().unwrap();
        // Billing header must be first; Claude Code identity second.
        assert!(
            system[0]["text"]
                .as_str()
                .unwrap_or("")
                .contains("x-anthropic-billing-header: cc_version=2.1.169.")
        );
        assert_eq!(system[1]["text"], CLAUDE_CODE_SYSTEM_PROMPT);
        assert!(system.iter().any(|block| {
            block["text"]
                .as_str()
                .unwrap_or("")
                .contains("Environment context you are running in:")
        }));
    }

    #[test]
    fn oauth_shaper_hoists_system_messages_out_of_messages_array() {
        // A compaction summary (or mid-session memory re-injection) lands in
        // `messages[]` as a `role:"system"` entry. Anthropic's Messages API
        // only allows user/assistant there; under the OAuth classifier a stray
        // system turn trips the "third-party app / extra usage 400". The shaper
        // must hoist it into the top-level `system` array.
        let body = Bytes::from(
            serde_json::json!({
                "model": "claude-sonnet-4-5",
                "stream": true,
                "system": [{"type": "text", "text": "Real prompt."}],
                "messages": [
                    {"role": "user", "content": "first"},
                    {"role": "system", "content": "[Previous conversation summary]\n## Active Task\nfix"},
                    {"role": "assistant", "content": "ok"}
                ]
            })
            .to_string(),
        );

        let value: serde_json::Value =
            serde_json::from_slice(&shape_oauth_messages_payload(body)).unwrap();

        let messages = value["messages"].as_array().unwrap();
        assert!(
            messages
                .iter()
                .all(|m| m.get("role").and_then(serde_json::Value::as_str) != Some("system")),
            "no system-role message may survive in messages[]"
        );
        assert_eq!(messages.len(), 2);

        let system = value["system"].as_array().unwrap();
        assert!(
            system
                .iter()
                .any(|b| b["text"].as_str().unwrap_or("").contains("## Active Task")),
            "hoisted summary text must land in the system array"
        );
    }

    #[test]
    fn oauth_shaper_splits_assistant_text_after_tool_use() {
        let body = Bytes::from(
            serde_json::json!({
                "model":"claude-sonnet-4-5", "stream": true, "messages":[
                    {"role":"user", "content":"please"},
                    {"role":"assistant", "content":[
                        {"type":"tool_use", "id":"t1", "name":"read", "input":{}},
                        {"type":"text", "text":"after"}
                    ]}
                ]
            })
            .to_string(),
        );
        let value: serde_json::Value =
            serde_json::from_slice(&shape_oauth_messages_payload(body)).unwrap();
        let messages = value["messages"].as_array().unwrap();
        assert_eq!(messages.len(), 3);
        assert_eq!(messages[1]["content"][0]["type"], "text");
        assert_eq!(messages[2]["content"][0]["type"], "tool_use");
    }

    #[test]
    fn non_messages_payload_passes_through_unchanged() {
        let body = Bytes::from(r#"{"messages":[]}"#);
        assert_eq!(shape_oauth_messages_payload(body.clone()), body);
    }

    #[test]
    fn billing_header_samples_utf16_code_units_like_reference() {
        let header = build_billing_header("abcd😀fg0123456789z");
        let cch = hex_sha256("abcd😀fg0123456789z".as_bytes())[0..5].to_string();
        let suffix =
            hex_sha256(format!("{BILLING_HEADER_SALT}\u{fffd}g0{CLAUDE_CODE_VERSION}").as_bytes())
                [0..3]
                .to_string();
        assert_eq!(
            header,
            format!(
                "x-anthropic-billing-header: cc_version={CLAUDE_CODE_VERSION}.{suffix}; cc_entrypoint=sdk-cli; cch={cch};"
            )
        );
    }

    #[test]
    fn non_oauth_messages_payload_passes_through_unchanged() {
        let client = AnthropicHttpClient::new("sk-ant-api03-test".to_string());
        let body = Bytes::from(
            r#"{"model":"claude","stream":true,"system":[{"type":"text","text":"Real prompt"}],"messages":[{"role":"user","content":"hello"}]}"#,
        );
        let req = Request::builder()
            .method("POST")
            .uri("https://api.anthropic.com/v1/messages")
            .body(body.clone())
            .unwrap();

        let normalized = client.normalized_request(req).unwrap();
        assert_eq!(normalized.into_body(), body);
    }
}