Kigi-CLI/crates/codegen/kigi-shell/src/sampling/error.rs

//! Sampling error types.
//!
//! The canonical error types live in `kigi_sampling_types::error`. This module
//! re-exports them and adds `map_sampling_err_to_acp`, which depends on
//! `agent_client_protocol::Error` (a kigi-shell dependency).

pub use kigi_sampling_types::error::*;

use agent_client_protocol as acp;

/// ACP error code for rate-limited requests (HTTP 429).
/// Uses the JSON-RPC implementation-defined server error range (-32000 to -32099).
///
/// Contract: this code must only be set for actual HTTP 429 responses from the
/// sampling client. Clients (desktop, pager) suppress error detail when they
/// see this code and show a user-friendly upgrade message instead.
pub const RATE_LIMITED_ERROR_CODE: i32 = -32003;

/// Subscription (OAuth) rate-limit copy. PRD Q3: the official Kimi CLI and
/// Kigi draw on the SAME subscription quota, so the message says so — a user
/// who also runs `kimi` should understand why the limit arrived early.
/// Deliberately promises no reset duration; the quota window is server-side.
pub static RATE_LIMITED_USER_MESSAGE_OAUTH: std::sync::LazyLock<String> =
    std::sync::LazyLock::new(|| {
        format!(
            "You\u{2019}ve hit the usage limit of your Kimi subscription. Note that Kigi and the \
             official Kimi CLI share the same subscription quota. Upgrade your plan at {} or try \
             again later.",
            kigi_env::upgrade_page_url()
        )
    });

/// Moonshot API-key rate-limit copy. Platform keys are tier-limited (RPM/TPM);
/// raising the tier happens in the Moonshot Open Platform console, not via a
/// Kimi subscription.
pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit the rate limit for your Moonshot API key. Check your tier\u{2019}s limits in the Moonshot Open Platform console (platform.moonshot.ai or platform.moonshot.cn), or try again later.";

/// Pick rate-limit copy from the *active* auth method.
///
/// Pass the real `is_api_key_auth` flag (pager `AppView`, `AuthMethodKind::is_api_key`
/// for the selected method). Do **not** decide from the key env var alone:
/// when both an env key and a cached OAuth session exist, auth prefers the
/// cached session over the API key.
pub fn rate_limited_user_message(is_api_key_auth: bool) -> &'static str {
    if is_api_key_auth {
        RATE_LIMITED_USER_MESSAGE_API_KEY
    } else {
        RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()
    }
}

/// Map a `SamplingError` to an ACP `Error` for client-facing responses.
pub fn map_sampling_err_to_acp(err: SamplingError) -> acp::Error {
    use reqwest::StatusCode;
    match err {
        SamplingError::Auth(msg) => acp::Error::auth_required().data(msg),
        SamplingError::InvalidConfiguration(msg) => acp::Error::invalid_params().data(msg),
        SamplingError::Http(e) => {
            acp::Error::internal_error().data(format!("http client init failed: {e}"))
        }
        SamplingError::Serialization(_) => acp::Error::invalid_params().data(err.to_string()),
        SamplingError::Api {
            status, message, ..
        } => match status {
            StatusCode::UNAUTHORIZED => acp::Error::auth_required().data(message),
            // 403 Forbidden is NOT an auth error — the request was
            // authenticated, but the action is not permitted (content-safety
            // blocks, ZDR-gated operations, remote-settings-blocked users).
            // Surfacing the proxy's message via internal_error keeps the
            // explanation visible to the user without triggering the client's
            // re-auth flow on -32000.
            StatusCode::FORBIDDEN => {
                let message = if message.contains("requires a Kigi subscription")
                    && crate::agent::auth_method::has_xai_api_key_env()
                {
                    format!(
                        "{message}\n\nYou have an API key set (KIGI_API_KEY). \
                         Your cached OAuth session is being used instead. \
                         To use your API key, run `kigi logout` or type /logout in the TUI."
                    )
                } else {
                    message
                };
                acp::Error::internal_error().data(message)
            }
            StatusCode::BAD_REQUEST => acp::Error::invalid_params().data(message),
            StatusCode::NOT_FOUND => acp::Error::resource_not_found(None).data(message),
            StatusCode::PAYLOAD_TOO_LARGE => acp::Error::invalid_params().data(message),
            StatusCode::TOO_MANY_REQUESTS => {
                acp::Error::new(RATE_LIMITED_ERROR_CODE, "Rate limited".to_string()).data(message)
            }
            _ => acp::Error::internal_error().data(message),
        },
        SamplingError::EventStreamError(message) => acp::Error::internal_error().data(message),
        SamplingError::StreamError {
            error_type,
            message,
        } => acp::Error::internal_error().data(format!("{error_type}: {message}")),
        SamplingError::EmptyResponse { context } => acp::Error::internal_error().data(format!(
            "empty response from model ({}): model={}, had_reasoning={}, finish_reason={}",
            context.reason,
            context.model,
            context.had_reasoning,
            context.finish_reason_str(),
        )),
        SamplingError::MaxTokensTruncation => {
            acp::Error::internal_error().data(terminal_error_data(
                err.to_string(),
                None,
                kigi_sampler::SamplingErrorKind::MaxTokensTruncation,
            ))
        }
        SamplingError::IdleTimeout { elapsed_secs } => acp::Error::internal_error().data(format!(
            "No response from model for {elapsed_secs}s — the model may be stuck"
        )),
        // Recovery consumes these inside the sampler's retry loop; a stray
        // terminal one still renders its labels.
        SamplingError::DoomLoopDetected { .. } => {
            acp::Error::internal_error().data(err.to_string())
        }
    }
}

pub fn error_data_with_status(message: String, http_status: Option<u16>) -> serde_json::Value {
    match http_status {
        Some(sc) => serde_json::json!({ "message": message, "http_status": sc }),
        None => serde_json::Value::String(message),
    }
}

/// Terminal-failure `acp::Error.data`: max-tokens truncation carries an `error_kind` marker (the kind's stable `as_str` name); other kinds keep the legacy shape.
pub fn terminal_error_data(
    message: String,
    http_status: Option<u16>,
    kind: kigi_sampler::SamplingErrorKind,
) -> serde_json::Value {
    if kind != kigi_sampler::SamplingErrorKind::MaxTokensTruncation {
        return error_data_with_status(message, http_status);
    }
    let mut data = serde_json::json!({ "message": message, "error_kind": kind.as_str() });
    if let Some(sc) = http_status {
        data["http_status"] = serde_json::json!(sc);
    }
    data
}

/// `turn_result.json` stop_reason for a failed turn: "MaxTokens" when the marker is present, else "Error" (matches the success path's `acp::StopReason` names).
pub fn stop_reason_for_turn_error(err: &acp::Error) -> &'static str {
    let is_max_tokens = err
        .data
        .as_ref()
        .and_then(|d| d.get("error_kind"))
        .and_then(|v| v.as_str())
        .is_some_and(|k| k == kigi_sampler::SamplingErrorKind::MaxTokensTruncation.as_str());
    if is_max_tokens { "MaxTokens" } else { "Error" }
}

fn error_message_from_data(data: &serde_json::Value) -> serde_json::Value {
    data.get("message").cloned().unwrap_or_else(|| data.clone())
}

pub fn error_detail_from_data(data: &serde_json::Value) -> Option<String> {
    if let Some(m) = data.get("message").and_then(|v| v.as_str()) {
        return Some(m.to_owned());
    }
    if let Some(s) = data.as_str() {
        return Some(s.to_owned());
    }
    data.get("detail")
        .and_then(|v| v.as_str())
        .map(str::to_owned)
}

/// The human-readable text of an ACP error: its `data` detail when present (Display would JSON-escape it), else the message.
pub fn acp_error_message(err: &acp::Error) -> String {
    err.data
        .as_ref()
        .and_then(error_detail_from_data)
        .unwrap_or_else(|| err.message.clone())
}

/// First line of a compaction error without internal `compact failed:` prefixes, capped for a banner; diagnostic lines (request URLs) are dropped.
pub fn user_facing_compact_error(raw: &str) -> String {
    const INTERNAL_PREFIX: &str = "compact failed:";
    const MAX_CHARS: usize = 200;
    // Redact before line selection: a line break inside a URL must not hide its `@`.
    let raw = redact_urls_in_text(raw);
    let mut rest = raw.trim();
    while rest
        .get(..INTERNAL_PREFIX.len())
        .is_some_and(|head| head.eq_ignore_ascii_case(INTERNAL_PREFIX))
    {
        rest = rest[INTERNAL_PREFIX.len()..].trim_start();
    }
    rest.lines()
        .next()
        .unwrap_or_default()
        .trim()
        .chars()
        .take(MAX_CHARS)
        .collect()
}

/// Every canonical `scheme://` URL in `text` loses userinfo, sensitive query values and fragment; one that fails to parse is not echoed.
fn redact_urls_in_text(text: &str) -> String {
    let mut out = String::with_capacity(text.len());
    let mut rest = text;
    while let Some(sep) = rest.find("://") {
        // The delimiter before the scheme may be multi-byte (`“`, `：`): step past its UTF-8 length.
        let scheme_start = rest[..sep]
            .char_indices()
            .rev()
            .find(|(_, c)| !(c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.')))
            .map_or(0, |(at, c)| at + c.len_utf8());
        out.push_str(&rest[..scheme_start]);
        // The authority may hold whitespace (a configured password); the token ends at whitespace after it.
        let authority = &rest[sep + 3..];
        let authority_end = authority.find(['/', '?', '#']).unwrap_or(authority.len());
        let token_end = authority[authority_end..]
            .find(char::is_whitespace)
            .map_or(rest.len(), |at| sep + 3 + authority_end + at);
        let token = &rest[scheme_start..token_end];
        let url_text = token.trim_end_matches([')', ']', '}', ',', '.', ';', '\'', '"', '>']);
        match url::Url::parse(url_text) {
            Ok(mut url) => {
                kigi_secrets::redact_url(&mut url);
                out.push_str(url.as_str());
            }
            Err(_) => out.push_str("<unparseable url>"),
        }
        out.push_str(&token[url_text.len()..]);
        rest = &rest[token_end..];
    }
    out.push_str(rest);
    out
}

pub fn http_status_from_error(err: &acp::Error) -> Option<u16> {
    err.data
        .as_ref()?
        .get("http_status")?
        .as_u64()
        .map(|s| s as u16)
}

const PROMPT_USAGE_DATA_KEY: &str = "promptUsage";

pub fn attach_prompt_usage(
    err: acp::Error,
    usage: Option<crate::extensions::notification::PromptUsage>,
) -> acp::Error {
    let Some(usage) = usage else {
        return err;
    };
    let Ok(usage_val) = serde_json::to_value(&usage) else {
        tracing::warn!(
            "attach_prompt_usage: failed to serialize PromptUsage; leaving error unchanged"
        );
        return err;
    };
    let mut map = match err.data.clone() {
        Some(serde_json::Value::Object(map)) => map,
        Some(serde_json::Value::String(message)) => {
            let mut m = serde_json::Map::new();
            m.insert("message".into(), serde_json::Value::String(message));
            m
        }
        Some(other) => {
            let mut m = serde_json::Map::new();
            m.insert("message".into(), other);
            m
        }
        None => {
            let mut m = serde_json::Map::new();
            m.insert(
                "message".into(),
                serde_json::Value::String(err.message.clone()),
            );
            m
        }
    };
    map.insert(PROMPT_USAGE_DATA_KEY.into(), usage_val);
    err.data(serde_json::Value::Object(map))
}

pub fn prompt_usage_from_error(
    err: &acp::Error,
) -> Option<crate::extensions::notification::PromptUsage> {
    let data = err.data.as_ref()?;
    let raw = data.get(PROMPT_USAGE_DATA_KEY)?;
    serde_json::from_value(raw.clone()).ok()
}

/// Derive `(stopReason, agentResult)` JSON values for the `prompt_complete`
/// notification from a prompt result. Rate-limit errors produce
/// `("rate_limit", null)` so the client shows its own upgrade message;
/// other errors produce `("error", <detail>)`.
pub fn prompt_complete_fields(
    result: &std::result::Result<acp::StopReason, acp::Error>,
) -> (serde_json::Value, serde_json::Value) {
    match result {
        Ok(reason) => (serde_json::json!(*reason), serde_json::Value::Null),
        Err(err) => {
            let is_rate_limit = i32::from(err.code) == RATE_LIMITED_ERROR_CODE;
            let stop = if is_rate_limit { "rate_limit" } else { "error" };
            let result = if is_rate_limit {
                serde_json::Value::Null
            } else {
                err.data
                    .as_ref()
                    .map(error_message_from_data)
                    .unwrap_or_else(|| serde_json::Value::String(err.message.clone()))
            };
            (serde_json::json!(stop), result)
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn compact_error_text_keeps_the_first_decoded_line_only() {
        let err = super::acp::Error::internal_error().data(
            "compact failed: API error (status 500): overloaded\n\nRequest URL: http://u:TEST_PASSWORD@127.0.0.1/v1",
        );
        let detail = super::user_facing_compact_error(&super::acp_error_message(&err));
        assert_eq!(detail, "API error (status 500): overloaded");
        assert!(!detail.contains("TEST_PASSWORD"));
        assert_eq!(
            super::user_facing_compact_error(&"x".repeat(500))
                .chars()
                .count(),
            200
        );
        assert_eq!(
            super::acp_error_message(&super::acp::Error::internal_error()),
            "Internal error"
        );
    }

    /// A sampler 401 names the endpoint on the retained line; its userinfo must not survive.
    #[test]
    fn compact_error_text_redacts_url_credentials_on_the_first_line() {
        for password in [
            "TEST_PASSWORD",
            "TEST PASSWORD",
            "TEST%20PASSWORD",
            "TEST_PASSWORD\n",
            "TEST_PASSWORD\r\n",
        ] {
            let raw = format!(
                "compact failed: Unauthorized (401) from http://u:{password}@127.0.0.1:33263/401/v1/chat/completions: bad auth"
            );
            assert_eq!(
                super::user_facing_compact_error(&raw),
                "Unauthorized (401) from http://127.0.0.1:33263/401/v1/chat/completions: bad auth",
                "{password:?}"
            );
        }
        assert_eq!(
            super::redact_urls_in_text("see https://a:b@h.example/x and http://h2/"),
            "see https://h.example/x and http://h2/"
        );
    }

    /// A credential carried as a query value survives no failure shape: 401 message or transport error.
    #[test]
    fn compact_error_text_redacts_query_credentials_in_every_failure_shape() {
        let unauthorized = "compact failed: Unauthorized (401) from http://127.0.0.1:1/?access_token=TEST_QUERY_TOKEN&route=/v1/chat/completions: bad auth";
        let transport = "compact failed: request error: error sending request for url (http://127.0.0.1:1/?access_token=TEST_QUERY_TOKEN&route=/500/v1/chat/completions)";
        for raw in [unauthorized, transport] {
            let detail = super::user_facing_compact_error(raw);
            assert!(!detail.contains("TEST_QUERY_TOKEN"), "{detail}");
            assert!(detail.contains("route="), "{detail}");
        }
        assert!(super::user_facing_compact_error("see http://:://").contains("<unparseable url>"));
    }

    /// A multi-byte delimiter right before the URL must not break the scan.
    #[test]
    fn compact_error_text_handles_unicode_before_a_url() {
        for raw in [
            "compact failed: Service unavailable. See “https://u:TEST_PASSWORD@status.example/”.",
            "compact failed: 请参阅：https://u:TEST_PASSWORD@status.example/",
        ] {
            let detail = super::user_facing_compact_error(raw);
            assert!(!detail.contains("TEST_PASSWORD"), "{detail}");
            assert!(detail.contains("https://status.example/"), "{detail}");
        }
    }

    use super::*;
    use reqwest::StatusCode;

    #[test]
    fn attach_prompt_usage_preserves_error_kind_and_round_trips() {
        let mut ledger = kigi_chat_state::UsageLedger::default();
        ledger.record_main_loop_call(
            "m",
            &kigi_sampling_types::TokenUsage {
                prompt_tokens: 3,
                completion_tokens: 1,
                total_tokens: 4,
                reasoning_tokens: 0,
                cached_prompt_tokens: 0,
            },
            None,
            Some(10),
        );
        let usage = crate::extensions::notification::PromptUsage::from(&ledger);
        let err = attach_prompt_usage(
            acp::Error::internal_error().data(terminal_error_data(
                "truncated".into(),
                None,
                kigi_sampler::SamplingErrorKind::MaxTokensTruncation,
            )),
            Some(usage.clone()),
        );
        assert_eq!(stop_reason_for_turn_error(&err), "MaxTokens");
        let back = prompt_usage_from_error(&err).expect("usage attached");
        assert_eq!(back.totals.input_tokens, 3);
        assert_eq!(back.num_turns, 1);
    }

    #[test]
    fn attach_prompt_usage_keeps_string_message_readable() {
        let usage = crate::extensions::notification::PromptUsage {
            totals: Default::default(),
            model_usage: Default::default(),
            num_turns: 1,
            usage_is_incomplete: false,
        };
        let free = "subscription:free-usage-exhausted quota hit";
        let err = attach_prompt_usage(
            acp::Error::new(RATE_LIMITED_ERROR_CODE, "Rate limited").data(free),
            Some(usage),
        );
        let msg = err
            .data
            .as_ref()
            .and_then(|d| {
                d.as_str()
                    .or_else(|| d.get("message").and_then(|m| m.as_str()))
            })
            .unwrap_or("");
        assert!(msg.contains("subscription:free-usage-exhausted"));
        assert!(prompt_usage_from_error(&err).is_some());
        assert!(!err.data.as_ref().unwrap().is_string());
    }

    #[test]
    fn error_detail_from_data_reads_message_field() {
        let data = error_data_with_status("upstream unavailable".into(), Some(503));
        assert_eq!(
            error_detail_from_data(&data).as_deref(),
            Some("upstream unavailable")
        );
    }

    #[test]
    fn rate_limited_user_message_oauth_vs_api_key() {
        assert_eq!(
            rate_limited_user_message(false),
            RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()
        );
        assert_eq!(
            rate_limited_user_message(true),
            RATE_LIMITED_USER_MESSAGE_API_KEY
        );
        // PRD Q3: the subscription copy must state the shared quota with the
        // official Kimi CLI and point at the upgrade page.
        assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("official Kimi CLI"));
        assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("same subscription quota"));
        assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains(kigi_env::upgrade_page_url()));
        assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Moonshot"));
        assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("subscription quota"));
    }

    #[test]
    fn rate_limit_error_uses_dedicated_code() {
        let err = SamplingError::Api {
            status: StatusCode::TOO_MANY_REQUESTS,
            message: "Rate limit exceeded".into(),
            model_metadata: None,
            retry_after_secs: None,
        };
        let acp_err = map_sampling_err_to_acp(err);
        assert_eq!(acp_err.code, acp::ErrorCode::from(RATE_LIMITED_ERROR_CODE));
        assert_eq!(acp_err.message, "Rate limited");
        assert_eq!(
            acp_err.data,
            Some(serde_json::Value::String("Rate limit exceeded".into()))
        );
    }

    #[test]
    fn rate_limit_mapping_is_stable_with_retry_after() {
        let err = SamplingError::Api {
            status: StatusCode::TOO_MANY_REQUESTS,
            message: "Rate limit exceeded".into(),
            model_metadata: None,
            retry_after_secs: Some(60),
        };
        assert_eq!(err.retry_after(), Some(60));
        let acp_err = map_sampling_err_to_acp(err);
        assert_eq!(acp_err.code, acp::ErrorCode::from(RATE_LIMITED_ERROR_CODE));
        assert_eq!(acp_err.message, "Rate limited");
    }

    #[test]
    fn rate_limit_code_differs_from_internal_error() {
        let rate_err = SamplingError::Api {
            status: StatusCode::TOO_MANY_REQUESTS,
            message: "limited".into(),
            model_metadata: None,
            retry_after_secs: None,
        };
        let server_err = SamplingError::Api {
            status: StatusCode::INTERNAL_SERVER_ERROR,
            message: "oops".into(),
            model_metadata: None,
            retry_after_secs: None,
        };
        let rate_acp = map_sampling_err_to_acp(rate_err);
        let server_acp = map_sampling_err_to_acp(server_err);

        assert_eq!(rate_acp.code, acp::ErrorCode::from(RATE_LIMITED_ERROR_CODE));
        assert_ne!(rate_acp.code, server_acp.code);
        assert_eq!(server_acp.code, acp::Error::internal_error().code);
    }

    #[test]
    fn auth_errors_map_to_auth_required() {
        let err = SamplingError::Api {
            status: StatusCode::UNAUTHORIZED,
            message: "bad token".into(),
            model_metadata: None,
            retry_after_secs: None,
        };
        let acp_err = map_sampling_err_to_acp(err);
        assert_eq!(acp_err.code, acp::Error::auth_required().code);
    }

    /// Regression test: 403 Forbidden must NOT map to auth_required.
    /// The cli-chat-proxy returns 403 for policy denials that are unrelated to
    /// the caller's credentials (content-safety blocks like
    /// SAFETY_CHECK_TYPE_DATA_LEAKAGE, ZDR-gated operations, remote settings
    /// blocks). Mapping these to auth_required causes the desktop app to
    /// tear down the session and kick off silent re-auth on -32000, and
    /// can race with invalid_grant_threshold to wipe auth.json.
    #[test]
    fn forbidden_does_not_map_to_auth_required() {
        let err = SamplingError::Api {
            status: StatusCode::FORBIDDEN,
            message:
                "Content violates usage guidelines. Failed check: SAFETY_CHECK_TYPE_DATA_LEAKAGE"
                    .into(),
            model_metadata: None,
            retry_after_secs: None,
        };
        let acp_err = map_sampling_err_to_acp(err);
        assert_ne!(
            acp_err.code,
            acp::Error::auth_required().code,
            "403 Forbidden must not be surfaced as auth_required"
        );
        assert_eq!(
            acp_err.data,
            Some(serde_json::Value::String(
                "Content violates usage guidelines. Failed check: SAFETY_CHECK_TYPE_DATA_LEAKAGE"
                    .into()
            ))
        );
    }

    /// Helper: run a closure with the house BYOK key (KIGI_API_KEY) temporarily
    /// set (or cleared). Clears every house-key env (KIGI_API_KEY plus the
    /// back-compat XAI_API_KEY / KIGI_CODE_XAI_API_KEY) so the "no key" case is
    /// hermetic. Cleans up even if the closure panics.
    fn with_api_key_env<F: FnOnce()>(key: Option<&str>, f: F) {
        let prev_house = std::env::var("KIGI_API_KEY").ok();
        let prev = std::env::var("XAI_API_KEY").ok();
        let prev_legacy = std::env::var("KIGI_CODE_XAI_API_KEY").ok();
        // SAFETY: serial_test ensures no concurrent env mutation.
        unsafe {
            std::env::remove_var("KIGI_API_KEY");
            std::env::remove_var("XAI_API_KEY");
            std::env::remove_var("KIGI_CODE_XAI_API_KEY");
            if let Some(k) = key {
                std::env::set_var("KIGI_API_KEY", k);
            }
        }
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(f));
        unsafe {
            std::env::remove_var("KIGI_API_KEY");
            std::env::remove_var("XAI_API_KEY");
            std::env::remove_var("KIGI_CODE_XAI_API_KEY");
            if let Some(v) = prev_house {
                std::env::set_var("KIGI_API_KEY", v);
            }
            if let Some(v) = prev {
                std::env::set_var("XAI_API_KEY", v);
            }
            if let Some(v) = prev_legacy {
                std::env::set_var("KIGI_CODE_XAI_API_KEY", v);
            }
        }
        if let Err(e) = result {
            std::panic::resume_unwind(e);
        }
    }

    #[test]
    #[serial_test::serial]
    fn forbidden_subscription_error_includes_api_key_hint_when_env_set() {
        with_api_key_env(Some("xai-test"), || {
            let err = SamplingError::Api {
                status: StatusCode::FORBIDDEN,
                message: "The model 'kigi' requires a Kigi subscription.".into(),
                model_metadata: None,
                retry_after_secs: None,
            };
            let acp_err = map_sampling_err_to_acp(err);
            let data = acp_err.data.unwrap();
            let msg = data.as_str().unwrap();
            assert!(
                msg.contains("kigi logout"),
                "should suggest kigi logout when API key is available: {msg}"
            );
            assert!(
                msg.contains("/logout"),
                "should mention /logout TUI command: {msg}"
            );
        });
    }

    #[test]
    #[serial_test::serial]
    fn forbidden_subscription_error_no_hint_without_api_key() {
        with_api_key_env(None, || {
            let err = SamplingError::Api {
                status: StatusCode::FORBIDDEN,
                message: "The model 'kigi' requires a Kigi subscription.".into(),
                model_metadata: None,
                retry_after_secs: None,
            };
            let acp_err = map_sampling_err_to_acp(err);
            let data = acp_err.data.unwrap();
            let msg = data.as_str().unwrap();
            assert!(
                !msg.contains("kigi logout"),
                "should NOT suggest logout when no API key is available: {msg}"
            );
        });
    }

    #[test]
    #[serial_test::serial]
    fn forbidden_non_subscription_error_no_hint() {
        with_api_key_env(Some("xai-test"), || {
            let err = SamplingError::Api {
                status: StatusCode::FORBIDDEN,
                message: "Content violates usage guidelines.".into(),
                model_metadata: None,
                retry_after_secs: None,
            };
            let acp_err = map_sampling_err_to_acp(err);
            let data = acp_err.data.unwrap();
            let msg = data.as_str().unwrap();
            assert!(
                !msg.contains("kigi logout"),
                "should NOT suggest logout for non-subscription 403: {msg}"
            );
        });
    }

    #[test]
    fn prompt_complete_fields_ok_passes_through_stop_reason() {
        let result: std::result::Result<acp::StopReason, acp::Error> = Ok(acp::StopReason::EndTurn);
        let (stop, agent_result) = prompt_complete_fields(&result);
        assert_eq!(stop, serde_json::json!("end_turn"));
        assert_eq!(agent_result, serde_json::Value::Null);
    }

    #[test]
    fn prompt_complete_fields_rate_limit_omits_detail() {
        let err = acp::Error::new(RATE_LIMITED_ERROR_CODE, "Rate limited".to_string())
            .data("Rate limit exceeded");
        let (stop, agent_result) = prompt_complete_fields(&Err(err));
        assert_eq!(stop, serde_json::json!("rate_limit"));
        assert_eq!(agent_result, serde_json::Value::Null);
    }

    #[test]
    fn prompt_complete_fields_generic_error_includes_detail() {
        let err = acp::Error::internal_error().data("connection reset");
        let (stop, agent_result) = prompt_complete_fields(&Err(err));
        assert_eq!(stop, serde_json::json!("error"));
        assert_eq!(
            agent_result,
            serde_json::Value::String("connection reset".into())
        );
    }

    #[test]
    fn prompt_complete_fields_error_without_data_falls_back_to_message() {
        let err = acp::Error::new(-32000, "something broke".to_string());
        assert!(err.data.is_none());
        let (stop, agent_result) = prompt_complete_fields(&Err(err));
        assert_eq!(stop, serde_json::json!("error"));
        assert_eq!(
            agent_result,
            serde_json::Value::String("something broke".into())
        );
    }

    #[test]
    fn http_status_from_error_extracts_status() {
        let err = acp::Error::internal_error()
            .data(error_data_with_status("bad token".into(), Some(401)));
        assert_eq!(http_status_from_error(&err), Some(401));
    }

    /// The typed max-tokens kind round-trips through `acp::Error.data` to the uploaded stop_reason.
    #[test]
    fn stop_reason_for_turn_error_distinguishes_max_tokens() {
        let err = map_sampling_err_to_acp(SamplingError::MaxTokensTruncation);
        assert_eq!(stop_reason_for_turn_error(&err), "MaxTokens");
        assert_eq!(
            stop_reason_for_turn_error(&acp::Error::internal_error()),
            "Error"
        );
    }

    #[test]
    fn prompt_complete_fields_extracts_message_from_status_data() {
        let err = acp::Error::internal_error()
            .data(error_data_with_status("model not found".into(), Some(404)));
        let (stop, agent_result) = prompt_complete_fields(&Err(err));
        assert_eq!(stop, serde_json::json!("error"));
        assert_eq!(
            agent_result,
            serde_json::Value::String("model not found".into())
        );
    }
}