//! Sampler-turn pipeline for `SessionActor`: tool definitions, model auth
//! facts/gates and retry, sampler config reconstruction, sampling-failure
//! recovery, and per-response usage recording.
use super::*;
/// Auth-failure detector for tool errors. Matches strictly on HTTP 401
/// when the error carries a structured status code, mirroring
/// `SamplingError::is_auth_error` in kigi-sampling-types: 403 is
/// deliberately excluded because it means "authenticated but forbidden"
/// (content-safety blocks, ZDR-gated requests, remote settings gates), where
/// a token refresh would be a no-op and would surface to the client as
/// a spurious auth_required teardown.
///
/// String fallbacks remain for tools that surface auth failures without
/// going through the structured `HttpFailure` path (e.g. JSON-only
/// `invalid_token` payloads, BYOK key-validation messages).
pub(super) fn is_auth_tool_error(err: &kigi_tool_runtime::ToolError) -> bool {
    if let Some(details) = &err.details
        && let Some(status) = details
            .get(HTTP_STATUS_DETAILS_KEY)
            .and_then(|s| s.as_u64())
    {
        return status == 401;
    }
    let lower = err.to_string().to_ascii_lowercase();
    lower.contains("unauthorized")
        || lower.contains("invalid api key")
        || lower.contains("invalid_token")
}
/// Gate inputs bundled with the composed decision so the 401-recovery log can
/// report the components.
#[derive(Clone, Copy)]
struct SessionTokenAuthGate {
    is_session_based: bool,
    model_byok: crate::agent::auth_method::ModelByok,
    /// Whether the request targets a first-party host. Lets an `Unknown`
    /// BYOK status still refresh against the first-party cli-chat-proxy hosts without
    /// risking a session-token leak to a third-party BYOK endpoint.
    endpoint_is_first_party: bool,
}
impl SessionTokenAuthGate {
    /// Single place `is_session_based` / `endpoint_is_first_party` are derived,
    /// so all call sites assemble the gate identically.
    fn new(
        auth_method_id: Option<&acp::AuthMethodId>,
        model_byok: crate::agent::auth_method::ModelByok,
        base_url: &str,
    ) -> Self {
        Self {
            is_session_based: auth_method_id
                .is_some_and(crate::agent::auth_method::is_session_based_method),
            model_byok,
            endpoint_is_first_party: crate::util::is_first_party_url(base_url),
        }
    }
    fn active(self) -> bool {
        crate::agent::auth_method::session_token_auth_gate(
            self.is_session_based,
            self.model_byok,
            self.endpoint_is_first_party,
        )
    }
}
/// Run a tool call; on an auth-shaped failure, attempt recovery via
/// `AuthManager` and one retry. When `shared_recovery` is `Some`, concurrent
/// 401s in the same batch deduplicate via `OnceCell::get_or_init`.
pub(super) async fn call_with_auth_retry<F, Fut>(
    auth_manager: Option<&std::sync::Arc<crate::auth::AuthManager>>,
    shared_recovery: Option<&tokio::sync::OnceCell<bool>>,
    tool_name: &str,
    mut call: F,
) -> Result<kigi_tools::types::output::ToolRunResult, kigi_tool_runtime::ToolError>
where
    F: FnMut() -> Fut,
    Fut: std::future::Future<
            Output = Result<kigi_tools::types::output::ToolRunResult, kigi_tool_runtime::ToolError>,
        >,
{
    let result = call().await;
    let Err(ref err) = result else { return result };
    if !is_auth_tool_error(err) {
        return result;
    }
    let Some(am) = auth_manager else {
        return result;
    };
    let recovered = match shared_recovery {
        Some(cell) => *cell.get_or_init(|| am.try_recover_unauthorized()).await,
        None => am.try_recover_unauthorized().await,
    };
    if recovered {
        tracing::info!(
            tool = tool_name,
            "auth recovery: tool 401, recovered, retrying"
        );
        call().await
    } else {
        tracing::warn!(tool = tool_name, "auth recovery: tool 401, refresh failed");
        kigi_log::unified_log::warn(
            "auth recovery: tool 401, refresh failed",
            None,
            Some(serde_json::json!({ "tool" : tool_name })),
        );
        result
    }
}
/// Wraps an [`AuthManager`](crate::auth::AuthManager) as a sampler
/// [`BearerResolver`](kigi_sampler::BearerResolver), resolving the live
/// (current-or-expired) bearer at request time. Shared by
/// [`SessionActor::reconstruct_full_config`] (the session model) and the
/// aux-model bearer repoint ([`SessionActor::repoint_aux_bearer_resolver_for_oauth`])
/// so both wrap ONE definition. SECURITY: the bearer is resolved per request
/// and never logged.
pub(crate) struct AuthManagerBearerResolver(pub(crate) std::sync::Arc<crate::auth::AuthManager>);
impl std::fmt::Debug for AuthManagerBearerResolver {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AuthManagerBearerResolver").finish()
    }
}
impl kigi_sampler::BearerResolver for AuthManagerBearerResolver {
    fn current_bearer(&self) -> Option<String> {
        self.0.current_or_expired().map(|a| a.key)
    }
}
/// Wrap `am` as a shared sampler bearer resolver.
fn auth_manager_bearer_resolver(
    am: std::sync::Arc<crate::auth::AuthManager>,
) -> kigi_sampler::SharedBearerResolver {
    std::sync::Arc::new(AuthManagerBearerResolver(am))
}
impl SessionActor {
    pub(super) async fn prepare_tool_definitions_timed(&self) -> (Vec<ToolDefinition>, u64) {
        let mcp_wait_start = std::time::Instant::now();
        match self.mcp_strategy {
            McpInitStrategy::Blocking => {
                if !self.mcp_state.lock().await.is_initialized() {
                    tracing::info!(
                        "Blocking strategy: waiting for MCP initialization before first prompt..."
                    );
                    self.wait_for_mcp_initialized().await;
                }
            }
            McpInitStrategy::Progressive => {}
        }
        let mcp_wait_ms = mcp_wait_start.elapsed().as_millis() as u64;
        let defs = self.prepare_tool_definitions_inner().await;
        (defs, mcp_wait_ms)
    }
    pub(super) async fn prepare_tool_definitions(&self) -> Vec<ToolDefinition> {
        self.prepare_tool_definitions_timed().await.0
    }
    /// The exact tool specs a turn sends, BEFORE the turn-specific
    /// structured-output append. Single source of truth shared by the turn
    /// (`acp_session_impl/turn.rs`) and the `SnapshotToolDefinitions` handler, so
    /// a verbatim-fork child's tool prefix can never silently drift from what the
    /// parent turn actually sends. `defs` is the already-resolved tool list
    /// (`prepare_tool_definitions_*`); this applies only the `web_search` drop
    /// under backend search and the `ToolSpec::from` mapping.
    pub(crate) fn turn_base_tool_specs(&self, defs: &[ToolDefinition]) -> Vec<ToolSpec> {
        let use_backend_search =
            self.agent.borrow().backend_search_enabled() && self.supports_backend_search.get();
        defs.iter()
            .filter(|td| !use_backend_search || td.function.name != "web_search")
            .cloned()
            .map(ToolSpec::from)
            .collect()
    }
    pub(super) async fn prepare_tool_definitions_inner(&self) -> Vec<ToolDefinition> {
        let bridge = self.agent.borrow().tool_bridge().clone();
        let defs = bridge.tool_definitions_builtins_only().await;
        let plan_active = self.plan_mode.lock().is_active();
        filter_cursor_tools_by_plan_mode(defs, plan_active)
    }
    /// Memoized per-model [`ModelAuthFacts`](crate::agent::config::ModelAuthFacts),
    /// keyed by `model_id`.
    ///
    /// A fresh `Unknown` (config currently unparseable) falls back to the last
    /// definite value for the same `model_id` rather than demoting a live session
    /// to non-refreshable api-key mode. Because a config edit can turn the
    /// currently-selected model into a per-model BYOK model without changing
    /// `model_id`, keying on `model_id` alone is insufficient — each
    /// model/credential chokepoint must clear this memo (`replace(None)`).
    pub(super) fn model_auth_facts(&self, model_id: &str) -> crate::agent::config::ModelAuthFacts {
        use crate::agent::auth_method::ModelByok;
        if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
            && cached_id == model_id
            && facts.byok != ModelByok::Unknown
        {
            return *facts;
        }
        let fresh = crate::agent::config::resolve_model_auth_facts(model_id);
        if fresh.byok == ModelByok::Unknown {
            if let Some((cached_id, facts)) = self.model_auth_facts.borrow().as_ref()
                && cached_id == model_id
            {
                return *facts;
            }
            return fresh;
        }
        *self.model_auth_facts.borrow_mut() = Some((model_id.to_string(), fresh));
        fresh
    }
    /// Gate inputs for `model_id` routed to `base_url`. See
    /// [`crate::agent::auth_method::session_token_auth_gate`] for the rationale
    /// (`base_url` keeps an `Unknown` BYOK status refreshable only
    /// against first-party xAI hosts).
    fn auth_gate(&self, model_id: &str, base_url: &str) -> SessionTokenAuthGate {
        let byok = self.model_auth_facts(model_id).byok;
        let auth_method = self.auth_method_id.load();
        SessionTokenAuthGate::new(auth_method.as_deref(), byok, base_url)
    }
    /// The [`AuthManager`](crate::auth::AuthManager) that governs INFERENCE auth
    /// for the model whose routing slug is `model` (the sampling config's
    /// `model`). A generic device-code OAuth platform (xai-grok) routes to its
    /// OWN scope-keyed manager from the process-global OAuth pool
    /// ([`crate::auth::oauth_registry::manager_for_model`], built on demand from
    /// the on-disk token); every other model routes to the primary Kimi
    /// `auth_manager`.
    ///
    /// This is the single chokepoint that keeps a grok turn from ever sending
    /// the Kimi bearer (Facet B) and gives it its own proactive-refresh + 401
    /// recovery manager (Facet A). The pool is the single source of truth, so a
    /// grok login that lands AFTER this session spawned is resolved on the next
    /// grok turn (no frozen per-session snapshot). Cheap `Arc` clone. Kimi /
    /// first-party path is byte-identical: a non-oauth model always resolves to
    /// the primary.
    ///
    /// `None` when there is no governing manager: a BYOK / test session with no
    /// primary. A grok model always resolves to its pooled manager (never the
    /// Kimi primary); when the user has not logged into grok that manager simply
    /// holds no token, so no Kimi bearer can leak.
    pub(super) fn auth_manager_for_model(
        &self,
        model: &str,
    ) -> Option<std::sync::Arc<crate::auth::AuthManager>> {
        // `model` is the bare routing slug; recover the managed catalog key
        // (`{platform}/{model}`) so the platform — and thus its OAuth scope — is
        // unambiguous. A bare / config / unlisted model yields no managed key
        // and resolves to the primary.
        let managed_key = self.managed_key_for_model(model);
        crate::auth::oauth_registry::manager_for_model(
            &crate::util::kigi_home::kigi_home(),
            managed_key.as_deref().unwrap_or(model),
            self.auth_manager.as_ref(),
        )
    }
    /// Recover the managed catalog key (`{platform}/{model}`) for a routing slug
    /// from the live catalog. `None` for a bare / config / unlisted model.
    fn managed_key_for_model(&self, model: &str) -> Option<String> {
        let models = self.models_manager.models();
        crate::agent::config::find_model_by_id(&models, model).and_then(|e| e.info().id.clone())
    }
    /// Whether the aux/session model `model` routes to a generic device-code
    /// OAuth platform (xai-grok). The gate for re-pointing an aux model's
    /// bearer_resolver away from the session (Kimi) resolver — a first-party /
    /// non-oauth model returns `false` and keeps the stamped session resolver.
    fn model_is_oauth_platform(&self, model: &str) -> bool {
        let managed_key = self.managed_key_for_model(model);
        kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model))
            .and_then(|(platform, _)| platform.oauth())
            .is_some()
    }
    /// Whether `model` routes to the Claude Pro/Max OAuth-Messages platform
    /// (claude-pro-max) — the gate for the sampler's OAuth Messages adaptation
    /// (identity headers + "You are Claude Code" system prefix). A generic-OAuth
    /// platform speaking the Messages wire; every other model (incl. xai-grok,
    /// which is ChatCompletions) returns `false`, keeping the API-key Anthropic
    /// / MiniMax Messages requests byte-identical.
    fn model_is_anthropic_oauth(&self, model: &str) -> bool {
        let managed_key = self.managed_key_for_model(model);
        kigi_models::parse_managed_model_key(managed_key.as_deref().unwrap_or(model)).is_some_and(
            |(platform, _)| {
                platform.oauth().is_some()
                    && platform.wire_api() == kigi_models::PlatformWireApi::Messages
            },
        )
    }
    /// LEAK guard for the stamped aux paths (auto-mode classifier, image
    /// describe). After [`crate::agent::config::stamp_session_local_sampler_fields`]
    /// has copied the SESSION model's `bearer_resolver` onto an aux
    /// `SamplerConfig`, re-point it at the AUX model's OWN platform manager when
    /// the aux model is an oauth platform (xai-grok) — so a grok aux model never
    /// inherits the live Kimi session bearer (→ api.x.ai). No-op for a
    /// first-party / non-oauth aux model (keeps the stamped session resolver →
    /// byte-identical). SECURITY: no token is logged.
    pub(super) fn repoint_aux_bearer_resolver_for_oauth(
        &self,
        cfg: &mut kigi_sampler::SamplerConfig,
        slug: &str,
    ) {
        if self.model_is_oauth_platform(slug)
            && let Some(manager) = self.auth_manager_for_model(slug)
        {
            cfg.bearer_resolver = Some(auth_manager_bearer_resolver(manager));
        }
    }
    /// Emit a unified-log breadcrumb whenever the session-token refresh gate is
    /// evaluated with an **`Unknown`** per-model BYOK status on a session-based
    /// method — the condition that (pre-fix) silently demoted live sessions to
    /// stale-token 401s. The uploaded per-turn unified log then shows whether
    /// the first-party-endpoint fallback kept refresh active or withheld it, so
    /// we can confirm the fix works (or catch a residual demotion) per session
    /// even when server-side metrics only show the aggregate 401. No-op for a
    /// definite `Byok`/`NotByok`, so steady-state turns stay quiet — a burst of
    /// these is itself the signal that `Unknown` is being hit in the field.
    fn log_auth_gate_unknown(&self, site: &str, gate: SessionTokenAuthGate, base_url: &str) {
        use crate::agent::auth_method::ModelByok;
        if gate.model_byok != ModelByok::Unknown || !gate.is_session_based {
            return;
        }
        let refresh_active = gate.active();
        let ctx = serde_json::json!(
            { "site" : site, "model_byok" : gate.model_byok.as_str(), "is_session_based"
            : gate.is_session_based, "endpoint_is_first_party" : gate
            .endpoint_is_first_party, "refresh_active" : refresh_active, "base_url" :
            base_url, }
        );
        let sid = Some(self.session_info.id.0.as_ref());
        if refresh_active {
            kigi_log::unified_log::info(
                "auth gate: Unknown BYOK on first-party endpoint — session-token refresh kept active",
                sid,
                Some(ctx),
            );
        } else {
            kigi_log::unified_log::warn(
                "auth gate: Unknown BYOK on non-first-party endpoint — refresh withheld (may surface stale-token 401)",
                sid,
                Some(ctx),
            );
        }
    }
    /// Reconstruct a full `SamplerConfig` (with credentials) by combining
    /// the actor's `SamplingConfig` and `Credentials`. Folds in the
    /// URL-derived headers (cli-chat-proxy auth, the staging auth header)
    /// so the sampler crate stays URL-agnostic.
    pub(super) async fn reconstruct_full_config(&self) -> SamplingConfig {
        #[allow(clippy::items_after_statements)]
        #[derive(Debug)]
        struct TraceContextInjector;
        impl kigi_sampler::HeaderInjector for TraceContextInjector {
            fn inject(&self, headers: &mut reqwest::header::HeaderMap) {
                if let Some(tp) = kigi_file_utils::trace_context::current_traceparent()
                    && let Ok(v) = reqwest::header::HeaderValue::from_str(&tp)
                {
                    headers.insert("traceparent", v);
                }
            }
        }
        let cfg = self
            .chat_state_handle
            .get_sampling_config()
            .await
            .unwrap_or_else(|| kigi_sampling_types::SamplingConfig {
                base_url: String::new(),
                model: String::new(),
                max_completion_tokens: None,
                temperature: None,
                top_p: None,
                api_backend: Default::default(),
                chat_compat: Default::default(),
                extra_headers: Default::default(),
                context_window: std::num::NonZeroU64::new(256_000).unwrap(),
                reasoning_effort: None,
                stream_tool_calls: None,
            });
        let creds = self.chat_state_handle.get_credentials().await;
        let model_facts = self.model_auth_facts(cfg.model.as_str());
        let auth_method = self.auth_method_id.load();
        let gate =
            SessionTokenAuthGate::new(auth_method.as_deref(), model_facts.byok, &cfg.base_url);
        let use_bearer_resolver = gate.active();
        self.log_auth_gate_unknown("reconstruct_full_config", gate, &cfg.base_url);
        // Resolve the bearer from the ACTIVE model's OWN manager: a grok model
        // wraps the xai-grok manager, never the Kimi one (captured before
        // `cfg.model` is moved into the struct below). `None` when the gate is
        // inactive or the oauth provider has no manager (fail-fast, no Kimi
        // fallback).
        let inference_auth_manager = if use_bearer_resolver {
            self.auth_manager_for_model(&cfg.model)
        } else {
            None
        };
        // Claude Pro/Max OAuth Messages adaptation for THIS turn's model
        // (captured before `cfg.model` is moved into the struct below).
        let anthropic_oauth = self.model_is_anthropic_oauth(&cfg.model);
        let auth_scheme = model_facts.auth_scheme;
        let mut extra_headers = cfg.extra_headers;
        crate::agent::config::inject_url_derived_headers(
            &mut extra_headers,
            creds.alpha_test_key.as_deref(),
            &cfg.base_url,
        );
        let compaction_at_tokens = self.compaction_at_tokens.get();
        let compactions_remaining = self.compactions_remaining.get();
        if compactions_remaining.is_some() || compaction_at_tokens.is_some() {
            let has_compaction_summary = self
                .chat_state_handle
                .get_last_compaction_prompt_index()
                .await
                .is_some();
            if let Some(value) =
                compactions_remaining.and_then(|c| c.resolve(has_compaction_summary))
            {
                extra_headers.insert("x-compactions-remaining".to_string(), value.to_string());
            }
            if !has_compaction_summary
                && let Some(value) = compaction_at_tokens.and_then(|c| {
                    c.resolve(
                        cfg.context_window.get(),
                        self.compaction.threshold_percent.get(),
                    )
                })
            {
                extra_headers.insert("x-compaction-at".to_string(), value.to_string());
            }
        }
        SamplingConfig {
            api_key: creds.api_key,
            base_url: cfg.base_url,
            model: cfg.model,
            max_completion_tokens: cfg.max_completion_tokens,
            temperature: cfg.temperature,
            top_p: cfg.top_p,
            api_backend: cfg.api_backend,
            auth_scheme,
            anthropic_oauth,
            chat_compat: cfg.chat_compat,
            extra_headers,
            context_window: cfg.context_window.get(),
            reasoning_effort: cfg.reasoning_effort,
            force_http1: false,
            max_retries: Some(self.max_retries),
            stream_tool_calls: cfg.stream_tool_calls.unwrap_or(false),
            idle_timeout_secs: None,
            origin_client: self.origin_client.clone(),
            attribution_callback: self.attribution_callback.clone(),
            bearer_resolver: inference_auth_manager.map(auth_manager_bearer_resolver),
            supports_backend_search: self.supports_backend_search.get(),
            compactions_remaining: self.compactions_remaining.get(),
            compaction_at_tokens: self.compaction_at_tokens.get(),
            doom_loop_recovery: self.doom_loop_recovery,
            header_injector: Some(std::sync::Arc::new(TraceContextInjector)),
        }
    }
    /// Install auto-mode permission classifier with a live LLM side-query
    /// (laziness-classifier pattern: `prepare_chat_completion` +
    /// `conversation_collect` on a LocalSet task; channel bridges the
    /// `Send` permission actor). Heuristic runs only when the side-query
    /// errors or returns unparseable text.
    pub(crate) async fn wire_permission_auto_llm_classifier(self: &Arc<Self>) {
        if !self.permissions.is_auto_mode() {
            return;
        }
        if self.permissions.has_llm_side_query() {
            return;
        }
        let auto_cfg = crate::util::config::resolve_auto_mode_config_from_disk();
        let session_model = self
            .chat_state_handle
            .get_sampling_config()
            .await
            .map(|c| c.model)
            .unwrap_or_default();
        let aux_classifier_sampler = match auto_cfg.classifier_model.as_deref() {
            Some(slug) => self.resolve_auto_classifier_sampler(slug).await,
            None => None,
        };
        let models = self.models_manager.models();
        let effective_supports_re = crate::agent::config::effective_classifier_supports_re(
            aux_classifier_sampler
                .as_ref()
                .map(|(_, model)| model.as_str()),
            &session_model,
            &models,
        );
        let (prompt_type, classifier_reasoning_effort) =
            crate::util::config::auto_mode_classifier_defaults(&auto_cfg, effective_supports_re);
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<(
            Vec<kigi_workspace::permission::ClassifierMessage>,
            tokio::sync::oneshot::Sender<Result<String, String>>,
        )>();
        let session = Arc::clone(self);
        tokio::task::spawn_local(async move {
            const TIMEOUT_MS: u64 = 15_000;
            while let Some((messages, respond_to)) = rx.recv().await {
                let result = async {
                    let (sampling_client, model) = match &aux_classifier_sampler {
                        Some((client, model)) => (client.clone(), model.clone()),
                        None => {
                            let client = session
                                .prepare_chat_completion(false)
                                .await
                                .map_err(|e| e.to_string())?;
                            let model = session
                                .chat_state_handle
                                .get_sampling_config()
                                .await
                                .map(|c| c.model)
                                .unwrap_or_default();
                            (client, model)
                        }
                    };
                    let session_id = session.session_info.id.to_string();
                    let items = messages
                        .into_iter()
                        .map(|m| match m.role {
                            kigi_workspace::permission::ClassifierMessageRole::System => {
                                ConversationItem::system(m.text)
                            }
                            kigi_workspace::permission::ClassifierMessageRole::User => {
                                ConversationItem::user(m.text)
                            }
                        })
                        .collect::<Vec<_>>();
                    let request = ConversationRequest {
                        items,
                        tools: vec![],
                        hosted_tools: vec![],
                        tool_choice: None,
                        model: Some(model),
                        temperature: None,
                        max_output_tokens: None,
                        json_schema: Some(
                            kigi_workspace::permission::classifier_output_json_schema(),
                        ),
                        reasoning_effort: classifier_reasoning_effort,
                        x_kigi_conv_id: Some(session_id.clone()),
                        x_kigi_req_id: Some(format!("xai-perm-auto-{}", uuid::Uuid::new_v4())),
                        x_kigi_session_id: Some(session_id),
                        x_kigi_agent_id: Some(crate::util::agent_id::agent_id()),
                        ..ConversationRequest::default()
                    };
                    let fut = sampling_client.conversation_collect(request);
                    let response =
                        tokio::time::timeout(std::time::Duration::from_millis(TIMEOUT_MS), fut)
                            .await
                            .map_err(|_| "permission auto classifier timed out".to_string())?
                            .map_err(|e| e.to_string())?;
                    Ok(response.assistant_text())
                }
                .await;
                let _ = respond_to.send(result);
            }
        });
        let clf =
            kigi_workspace::permission::LlmPermissionClassifier::with_channel(tx, prompt_type);
        debug_assert!(
            clf.has_side_query(),
            "channel-wired classifier must report has_side_query"
        );
        self.permissions.set_classifier_with_side_query(clf, true);
        tracing::info!(
            session_id = % self.session_info.id,
            "Wired live LLM permission auto-mode classifier (session sampling channel)"
        );
    }
    /// Resolve a standalone aux-model `SamplerConfig` for `slug` via the shared
    /// catalog routing (Tier-1 catalog creds / Tier-2 xAI-proxy via session token
    /// / `XAI_API_KEY` / deployment key), gathering the session-local auth context
    /// once. Shared by image-describe and the classifier so the gather can't
    /// drift. `None` ⇒ caller falls back to the session model.
    pub(super) async fn resolve_aux_sampler_config(
        &self,
        slug: &str,
    ) -> Option<kigi_sampler::SamplerConfig> {
        let creds = self.chat_state_handle.get_credentials().await;
        // Resolve the aux token by the aux model's OWN platform: a grok
        // (oauth-platform) aux model draws its pooled grok token or `None` —
        // NEVER the primary Kimi session token (which `resolve_credentials`
        // would otherwise stamp onto an api.x.ai request). A first-party /
        // non-oauth aux model still gets the primary (byte-identical).
        let session_key = crate::auth::oauth_registry::session_key_for_model(
            &crate::util::kigi_home::kigi_home(),
            slug,
            self.auth_manager.as_ref(),
        );
        let models = self.models_manager.models();
        let endpoints = self.models_manager.endpoints();
        crate::agent::config::resolve_aux_model_sampling_config(
            slug,
            &models,
            &endpoints,
            session_key.as_deref(),
            creds.alpha_test_key.clone(),
        )
    }
    /// Resolve a dedicated sampler for the Auto-mode classifier model `slug`,
    /// stamping session-local auth/attribution like image-describe (which relies
    /// on the resolver, not a config override, for `base_url`/`api_backend` so
    /// credentials stay consistent). `None` ⇒ caller falls back to the session
    /// client + model.
    async fn resolve_auto_classifier_sampler(
        &self,
        slug: &str,
    ) -> Option<(kigi_sampler::SamplingClient, String)> {
        let active_session_config = self.reconstruct_full_config().await;
        let mut cfg = self.resolve_aux_sampler_config(slug).await?;
        crate::agent::config::stamp_session_local_sampler_fields(
            &mut cfg,
            &active_session_config,
            Some(self.max_retries),
        );
        // LEAK 1b: a grok aux classifier must not inherit the SESSION model's
        // (Kimi) bearer_resolver stamped above; re-point it at grok's own
        // manager (its pooled token, or None). No-op for a non-oauth aux.
        self.repoint_aux_bearer_resolver_for_oauth(&mut cfg, slug);
        let model = cfg.model.clone();
        let client = kigi_sampler::SamplingClient::new(cfg)
            .map_err(|e| {
                tracing::warn!(
                    error = % e,
                    "auto classifier aux sampler build failed; using session model"
                )
            })
            .ok()?;
        Some((client, model))
    }
    #[tracing::instrument(
        name = "session.prepare_chat_completion",
        skip_all,
        fields(force_http1)
    )]
    pub(super) async fn prepare_chat_completion(
        &self,
        force_http1: bool,
    ) -> Result<kigi_sampler::SamplingClient, acp::Error> {
        self.refresh_token_if_expired().await;
        let mut full_config = self.reconstruct_full_config().await;
        full_config.force_http1 = force_http1;
        let sampling_client =
            kigi_sampler::SamplingClient::new(full_config).map_err(|e| self.to_acp_error(e))?;
        Ok(sampling_client)
    }
    /// Push a fresh `SamplerConfig` into the per-session sampler actor
    /// before each turn. Mirrors `prepare_chat_completion`'s
    /// auth-refresh + config rebuild, but routes the result to the
    /// `kigi-sampler` instead of constructing a new
    /// `OaiCompatClient`.
    ///
    /// Behaviour parity: we run the same `refresh_token_if_expired()`
    /// and `reconstruct_full_config()` so the sampler picks up any
    /// newly issued session token. The previous client cache inside
    /// the sampler actor is invalidated automatically by
    /// `update_config`.
    pub(crate) async fn prepare_sampler_for_turn(&self) {
        self.refresh_token_if_expired().await;
        let mut sampler_config = self.reconstruct_full_config().await;
        sampler_config.idle_timeout_secs = Some(self.inference_idle_timeout.as_secs());
        self.sampler_handle.update_config(sampler_config);
    }
    fn log_terminal_failure(&self, error_type: &str, status_code: Option<u16>, message: &str) {
        let auth = self
            .auth_manager
            .as_ref()
            .and_then(|am| am.current_or_expired());
        let reauthable = is_reauthable_failure(Some(error_type), message);
        kigi_log::unified_log::warn(
            "turn.terminal_failure",
            Some(self.session_info.id.0.as_ref()),
            Some(serde_json::json!(
                { "error_type" : error_type, "status_code" : status_code,
                "reauthable" : reauthable, "auth_mode" : auth.as_ref().map(| a |
                format!("{:?}", a.auth_mode)), "key_prefix" : auth.as_ref().map(| a |
                crate ::auth::token_suffix(& a.key).to_owned()), "expires_at" : auth
                .as_ref().and_then(| a | a.expires_at.map(| e | e.to_rfc3339())),
                "message" : crate ::util::truncate(message, 300), }
            )),
        );
    }
    pub(crate) async fn handle_sampling_failure(
        self: &Arc<Self>,
        error: kigi_sampler::SamplingErrorInfo,
    ) -> Result<SamplerFailureRecovery, acp::Error> {
        use kigi_sampler::SamplingErrorKind;
        if self.should_compact_on_error(&error).await {
            let cw = error
                .model_metadata
                .as_ref()
                .and_then(|m| m.context_window)
                .expect("should_compact_on_error guarantees context_window");
            {
                let total_tokens = self.chat_state_handle.get_estimated_total_tokens().await;
                let percentage = kigi_token_estimation::usage_percentage_u8(total_tokens, cw);
                if let Some(mut cfg) = self.chat_state_handle.get_sampling_config().await
                    && let Some(new_cw) = std::num::NonZeroU64::new(cw)
                    && self.compaction.context_window_override.is_none()
                {
                    cfg.context_window = new_cw;
                    self.chat_state_handle.update_sampling_config(cfg);
                }
                let trigger_info = compaction::AutoCompactTriggerInfo {
                    tokens_used: total_tokens,
                    context_window: cw,
                    percentage,
                };
                self.run_compact_only(trigger_info).await?;
                return Ok(SamplerFailureRecovery::CompactAndResubmit);
            }
        }
        let detailed_message = error.message.clone();
        if matches!(error.kind, SamplingErrorKind::Api)
            && error.status_code == Some(400)
            && error.message.contains("encrypted_content")
        {
            self.signals_handle()
                .record_error_typed("encrypted_content_mismatch");
            let friendly = "This session's conversation history is incompatible \
                            with the current model. Please start a new session."
                .to_string();
            self.log_terminal_failure("encrypted_content_mismatch", error.status_code, &friendly);
            self.send_xai_notification(XaiSessionUpdate::RetryState(
                crate::extensions::notification::RetryState::Failed {
                    error_type: "encrypted_content_mismatch".to_string(),
                    message: friendly.clone(),
                },
            ))
            .await;
            return Err(acp::Error::invalid_params().data(friendly));
        }
        if matches!(error.kind, SamplingErrorKind::RateLimited) {
            self.log_terminal_failure("rate_limited", error.status_code, &detailed_message);
            self.send_xai_notification(XaiSessionUpdate::RetryState(
                crate::extensions::notification::RetryState::Exhausted {
                    attempts: 0,
                    reason: detailed_message.clone(),
                    is_rate_limited: true,
                },
            ))
            .await;
            let acp_err = acp::Error::new(
                crate::sampling::error::RATE_LIMITED_ERROR_CODE,
                "Rate limited".to_string(),
            )
            .data(detailed_message);
            return Err(acp_err);
        }
        let auth_recovery_eligible = matches!(error.kind, SamplingErrorKind::Auth) && {
            let (model_id, base_url) = self
                .chat_state_handle
                .get_sampling_config()
                .await
                .map(|c| (c.model, c.base_url))
                .unwrap_or_default();
            let gate = self.auth_gate(&model_id, &base_url);
            let eligible = gate.active();
            self.log_auth_gate_unknown("handle_sampling_failure", gate, &base_url);
            if !eligible {
                tracing::warn!(
                    session_id = % self.session_info.id.0, is_session_based = gate
                    .is_session_based, model_byok = gate.model_byok.as_str(),
                    endpoint_is_first_party = gate.endpoint_is_first_party,
                    "auth recovery: sampler 401 not refreshable (api-key auth) — surfacing 401",
                );
                kigi_log::unified_log::warn(
                    "auth recovery: sampler 401 not eligible (api-key auth)",
                    Some(self.session_info.id.0.as_ref()),
                    Some(serde_json::json!(
                        { "kind" : error.kind.as_str(), "status_code" : error
                        .status_code, "is_session_based" : gate.is_session_based,
                        "model_byok" : gate.model_byok.as_str(),
                        "endpoint_is_first_party" : gate.endpoint_is_first_party, }
                    )),
                );
            }
            eligible
        };
        if !matches!(error.kind, SamplingErrorKind::Auth) && error.status_code == Some(401) {
            kigi_log::unified_log::warn(
                "auth recovery: sampler 401 not eligible (non-auth error kind)",
                Some(self.session_info.id.0.as_ref()),
                Some(serde_json::json!(
                    { "kind" : error.kind.as_str(), "status_code" : error
                    .status_code, }
                )),
            );
        }
        // Recover via the ACTIVE model's OWN manager: a grok 401 recovers the
        // xai-grok session via the xai-grok manager, never the Kimi one. For a
        // Kimi / non-oauth model this resolves to the primary — byte-identical.
        if auth_recovery_eligible {
            let recovery_model = self
                .chat_state_handle
                .get_sampling_config()
                .await
                .map(|c| c.model)
                .unwrap_or_default();
            if let Some(am) = self.auth_manager_for_model(&recovery_model) {
                if am.try_recover_unauthorized().await {
                    tracing::info!(
                        session_id = % self.session_info.id.0,
                        "auth recovery: sampler 401, recovered, retrying"
                    );
                    kigi_log::unified_log::info(
                        "auth recovery: sampler 401, recovered, retrying",
                        Some(self.session_info.id.0.as_ref()),
                        None,
                    );
                    self.prepare_sampler_for_turn().await;
                    return Ok(SamplerFailureRecovery::RefreshAuthAndResubmit);
                }
                tracing::warn!(
                    session_id = % self.session_info.id.0,
                    "auth recovery: sampler 401, refresh failed"
                );
                kigi_log::unified_log::warn(
                    "auth recovery: sampler 401, refresh failed",
                    Some(self.session_info.id.0.as_ref()),
                    None,
                );
            }
        }
        if matches!(error.kind, SamplingErrorKind::IdleTimeout) {
            self.signals_handle().record_idle_timeout();
        }
        if matches!(error.kind, SamplingErrorKind::EmptyResponse) {
            if let Some(ref ctx) = error.empty_response_context {
                tracing::warn!(
                    empty_response = true, empty_reason = ctx.reason.as_str(),
                    had_reasoning = ctx.had_reasoning, content_len = ctx.content_len,
                    tool_call_count = ctx.tool_call_count, completion_tokens = ctx
                    .completion_tokens.unwrap_or(0), reasoning_tokens = ctx
                    .reasoning_tokens.unwrap_or(0), finish_reason = ctx
                    .finish_reason_str(), first_choice_seen = ctx.first_choice_seen,
                    model = % ctx.model,
                    "empty response after retries exhausted: {reason}", reason = ctx
                    .reason,
                );
                {
                    let mut cap = self.streaming_turn_capture.lock();
                    cap.reasoning_tokens = ctx.reasoning_tokens;
                    cap.completion_tokens = ctx.completion_tokens;
                    cap.finish_reason = ctx.finish_reason.clone();
                    cap.empty_reason = Some(ctx.reason.as_str().to_owned());
                }
            }
            self.signals_handle().record_error_typed("empty_response");
        }
        let auth_mode = self
            .auth_manager
            .as_ref()
            .and_then(|am| am.current())
            .map(|a| a.auth_mode)
            .unwrap_or(crate::auth::AuthMode::ApiKey);
        let auth_mode_str = format!("{auth_mode:?}");
        let client_version = kigi_version::VERSION;
        let is_model_404 =
            error.status_code == Some(404) && detailed_message.contains("does not exist");
        let is_auth_401 =
            error.status_code == Some(401) || matches!(error.kind, SamplingErrorKind::Auth);
        let detailed_message = if is_model_404 || is_auth_401 {
            let current_model = self
                .chat_state_handle
                .get_sampling_config()
                .await
                .map(|c| c.model)
                .unwrap_or_else(|| "unknown".to_string());
            let available: Vec<String> = self
                .models_manager
                .models()
                .values()
                .map(|m| m.model.clone())
                .collect();
            let mut msg = format!("{detailed_message}\n");
            msg.push_str(&format!("\n  Model:     {current_model}"));
            msg.push_str(&format!("\n  Auth:      {auth_mode_str}"));
            msg.push_str(&format!("\n  Version:   {client_version}"));
            if available.is_empty() {
                msg.push_str("\n  Available: (none)");
            } else {
                msg.push_str(&format!("\n  Available: {}", available.join(", ")));
            }
            if is_model_404 && !available.iter().any(|m| m == &current_model) {
                msg.push_str(&format!(
                    "\n\n  '{}' is not in your available models.",
                    current_model
                ));
                msg.push_str("\n  Switch models with /model or start a new session.");
            }
            msg
        } else {
            detailed_message
        };
        let error_type = if kigi_sampling_types::is_context_length_error(&error.message) {
            "context_length"
        } else {
            error.kind.as_str()
        };
        self.log_terminal_failure(error_type, error.status_code, &detailed_message);
        self.send_xai_notification(XaiSessionUpdate::RetryState(
            crate::extensions::notification::RetryState::Failed {
                error_type: error_type.to_string(),
                message: detailed_message.clone(),
            },
        ))
        .await;
        Err(
            acp::Error::internal_error().data(crate::sampling::error::terminal_error_data(
                detailed_message,
                error.status_code,
                error.kind,
            )),
        )
    }
    /// Drive a single turn through the sampler-based path.
    ///
    /// Calls `prepare_sampler_for_turn` first (auth refresh + config
    /// push), then submits via `SamplerHandle::submit_and_collect` and
    /// returns:
    /// * `Ok(SamplerTurnOutcome::Response(_))` - model responded.
    /// * `Ok(SamplerTurnOutcome::CompactAndResubmit)` - compaction
    ///    ran, the outer turn loop should `continue`.
    /// * `Ok(SamplerTurnOutcome::RefreshAuthAndResubmit)` - auth 401
    ///    recovery succeeded, credentials refreshed, retry once.
    /// * `Err(acp::Error)` - terminal failure already reported via
    ///    `send_xai_notification(RetryState::Failed)`.
    pub(crate) async fn run_turn_via_sampler(
        self: &Arc<Self>,
        request: ConversationRequest,
    ) -> Result<SamplerTurnOutcome, acp::Error> {
        self.prepare_sampler_for_turn().await;
        let stream_drained_rx = {
            let (tx, rx) = tokio::sync::oneshot::channel();
            *self.turn_stream_drained.lock() = Some(tx);
            rx
        };
        let request_id = kigi_sampler::RequestId::random();
        let request_id_str = request_id.as_str().to_string();
        match self
            .sampler_handle
            .submit_and_collect(request_id, request)
            .await
        {
            Ok((response, metrics)) => {
                let span = tracing::Span::current();
                span.record("request_id", request_id_str.as_str());
                if let Some(ttft) = metrics.time_to_first_token_ms {
                    span.record("ttft_ms", ttft as i64);
                }
                if metrics.attempts > 0 {
                    span.record("attempt", i64::from(metrics.attempts));
                }
                if tokio::time::timeout(std::time::Duration::from_secs(5), stream_drained_rx)
                    .await
                    .is_err()
                {
                    self.turn_stream_drained.lock().take();
                    tracing::warn!(
                        "stream-drain barrier timed out; proceeding to emit tool \
                         calls (eventId ordering may be imperfect this turn)"
                    );
                }
                Ok(SamplerTurnOutcome::Response(
                    Box::new(response),
                    Box::new(metrics),
                ))
            }
            Err(rich_err) => {
                self.turn_stream_drained.lock().take();
                let info = kigi_sampler::SamplingErrorInfo::from(&rich_err);
                match self.handle_sampling_failure(info).await? {
                    SamplerFailureRecovery::CompactAndResubmit => {
                        Ok(SamplerTurnOutcome::CompactAndResubmit)
                    }
                    SamplerFailureRecovery::RefreshAuthAndResubmit => {
                        Ok(SamplerTurnOutcome::RefreshAuthAndResubmit)
                    }
                }
            }
        }
    }
    /// Proactively refresh the auth token if near expiry.
    pub(super) async fn refresh_token_if_expired(&self) {
        let (model_id, base_url) = self
            .chat_state_handle
            .get_sampling_config()
            .await
            .map(|c| (c.model, c.base_url))
            .unwrap_or_default();
        // Refresh the ACTIVE model's OWN manager: a grok model refreshes the
        // xai-grok token via the xai-grok manager, never the Kimi one. For a
        // Kimi / non-oauth model this resolves to the primary — byte-identical.
        if let Some(am) = self.auth_manager_for_model(&model_id) {
            let creds = self.chat_state_handle.get_credentials().await;
            if self.auth_gate(&model_id, &base_url).active()
                && let Ok(key) = am.get_valid_token().await
            {
                if creds.api_key.as_deref() != Some(&key) {
                    let mut creds = creds;
                    creds.api_key = Some(key);
                    self.chat_state_handle.update_credentials(creds);
                }
                return;
            }
        } else {
            kigi_log::unified_log::debug(
                "token refresh skipped: no auth manager",
                Some(self.session_info.id.0.as_ref()),
                None,
            );
        }
        // BYOK path: pick up an externally rotated per-model key from
        // config.toml. Kimi bearers are opaque (no client-side expiry
        // probing); a changed on-disk key is adopted, an unchanged one is a
        // no-op.
        let creds = self.chat_state_handle.get_credentials().await;
        let current_key = creds.api_key;
        let current_model_id = self
            .chat_state_handle
            .get_sampling_config()
            .await
            .map(|c| c.model)
            .unwrap_or_default();
        let Some(ref key) = current_key else { return };
        let Some(new_key) = self.reload_api_key_from_config(&current_model_id) else {
            return;
        };
        if key == &new_key {
            return;
        }
        tracing::info!(
            model = % current_model_id, key_len = new_key.len(),
            "Refreshed API token from config.toml"
        );
        let mut creds = self.chat_state_handle.get_credentials().await;
        creds.api_key = Some(new_key);
        self.chat_state_handle.update_credentials(creds);
    }
    fn reload_api_key_from_config(&self, current_model_id: &str) -> Option<String> {
        let raw_config = crate::config::load_effective_config()
            .map_err(|e| tracing::warn!(error = % e, "Failed to reload config"))
            .ok()?;
        let config = crate::agent::config::Config::new_from_toml_cfg(&raw_config)
            .map_err(|e| tracing::warn!(error = % e, "Failed to parse reloaded config.toml"))
            .ok()?;
        let config_model = config
            .config_models
            .iter()
            .find(|(k, v)| v.model.as_deref().unwrap_or(k.as_str()) == current_model_id)
            .map(|(_, v)| v);
        let Some(model) = config_model else {
            tracing::warn!(
                model = % current_model_id, available = ? config.config_models.keys()
                .collect::< Vec < _ >> (), "Model not found in config.toml [model.*]"
            );
            return None;
        };
        let key = crate::agent::config::first_own_credential(
            model.api_key.as_deref(),
            model.env_key.as_ref(),
        );
        if key.is_none() {
            tracing::warn!(
                model = % current_model_id, env_key = ? model.env_key,
                "No api_key or env_key resolved for model"
            );
        }
        key
    }
    /// Propagate the model-reported token usage from a turn response into
    /// chat state, the per-prompt usage ledger, and per-turn signals.
    ///
    /// This is the only place per-turn `total_tokens` is refreshed in the
    /// post-sampler-refactor path; without it `state.total_tokens` would
    /// stay frozen at the `estimate_conversation_tokens` seed from
    /// `ChatState::new`, freezing `/context` and corrupting the resume
    /// restore that reads `meta.totalTokens` from `updates.jsonl`.
    /// Resetting `estimated_tokens_since_model = 0` here also keeps the
    /// preflight-overflow guard accurate against the next turn's
    /// tool-result deltas.
    pub(crate) fn record_response_token_usage(
        &self,
        response: &ConversationResponse,
        api_duration_ms: Option<u64>,
    ) {
        if let Some(ref u) = response.usage {
            self.chat_state_handle
                .record_token_usage(u64::from(u.total_tokens));
            self.chat_state_handle.record_last_turn_usage(u.clone());
            self.chat_state_handle.record_model_call_usage(
                response.assistant().and_then(|a| a.model_id.clone()),
                u.clone(),
                api_duration_ms,
                response.cost_usd_ticks,
            );
            self.signals_handle()
                .record_token_usage(u.completion_tokens, u.reasoning_tokens);
        }
    }
    pub(super) async fn record_assistant_response(&self, assistant_item: ConversationItem) {
        self.signals_handle().record_assistant_message();
        if let ConversationItem::Assistant(ref a) = assistant_item {
            tracing::info!(
                model_id = ? a.model_id, "DEBUG record_assistant_response model_id"
            );
        }
        if let ConversationItem::Assistant(ref a) = assistant_item
            && let Some(first_call) = a.tool_calls.first()
        {
            tracing::info!("Assistant requested tool call: {}", first_call.id);
        }
        self.chat_state_handle
            .push_assistant_response(assistant_item);
    }
}
#[cfg(test)]
mod bearer_resolver_tests {
    use super::AuthManagerBearerResolver;
    use kigi_sampler::BearerResolver;

    /// LEAK 1b: the shared `AuthManagerBearerResolver` resolves the LIVE bearer
    /// of the manager it wraps. So an aux bearer_resolver built over grok's OWN
    /// (oauth) pooled manager yields grok's token (or `None`) — NEVER the Kimi
    /// session token that a Kimi-manager resolver would. The
    /// `repoint_aux_bearer_resolver_for_oauth` fix wraps exactly this grok
    /// manager for a grok aux model.
    #[tokio::test]
    async fn resolver_resolves_the_wrapped_manager_never_kimi() {
        let dir = tempfile::tempdir().unwrap();
        let kimi = std::sync::Arc::new(crate::auth::AuthManager::new(
            dir.path(),
            crate::auth::KimiCodeConfig::default(),
        ));
        kimi.hot_swap(crate::auth::KimiAuth {
            key: "kimi-tok".to_string(),
            auth_mode: crate::auth::AuthMode::OAuth,
            ..crate::auth::KimiAuth::test_default()
        });
        // The Kimi-manager resolver yields the Kimi bearer.
        assert_eq!(
            AuthManagerBearerResolver(kimi.clone()).current_bearer(),
            Some("kimi-tok".to_string()),
        );
        // The grok (oauth) pooled manager is distinct — its resolver never
        // yields the Kimi bearer (grok's own token, or None).
        let oauth = kigi_models::PlatformId::XaiGrok
            .oauth()
            .expect("xai-grok carries an OAuthConfig");
        let grok = crate::auth::oauth_registry::global_manager_for(
            &crate::util::kigi_home::kigi_home(),
            oauth,
        );
        assert_ne!(
            AuthManagerBearerResolver(grok).current_bearer(),
            Some("kimi-tok".to_string()),
            "a grok aux bearer_resolver must never resolve the Kimi session token",
        );
    }
}
