ZacharyZhang-NY/Kigi-CLI
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

F3: Kimi inference pipeline + full grok cloud-surface excision

Browse Source 
Sampler / inference (PRD F3):
- kimi_compat.rs: single adaptation point for the Kimi chat/completions
  dialect (thinking-field mapping, model_id stripping, empty-content
  tool-call message fix, stream_options.include_usage), with kimi-cli
  source citations
- Rate-limit handling reworked for Kimi/Moonshot semantics; UA kigi/{version}
- /models replaces the xAI models-v2 endpoint everywhere; idle model
  refresh carries X-Msh-* device headers only (X-XAI-Token-Auth and
  x-grok-client-mode/CLIENT_MODE_HEADER machinery deleted)

Cloud-surface excision (PRD §5, zero-egress):
- remote/ conversations lane, cli-chat-proxy-types crate, prod/ dir,
  share command, credit bar: deleted (single local session lane;
  paginate() replaces merge_and_paginate)
- Subscription/tier gate stack deleted end-to-end: AppView
  gate/tier/team/ZDR fields, app/subscription.rs watch loop,
  dispatch/billing.rs paywall + SuperGrok upsell, free-usage-exhausted
  chain, tier-restricted commands, GateInfo, RemoteSettings gate fields,
  SettingsUpdateNotification gate fields
- /privacy + coding-data-sharing setting deleted (backed by a dead xAI
  RPC; Kigi is zero-egress — nothing to share or retain remotely)

Auth UX correctness (user-reported):
- Device-flow fixtures now mirror the live Kimi payload shape
  (https://www.kimi.com/code/authorize_device?user_code=..., verified
  against auth.kimi.com); the fabricated auth.kimi.com/device?code=...
  URLs are gone
- open_browser_detached is a no-op under cfg(test): unit tests drove
  wiremock fixture URLs into the real browser (root cause of the
  "garbage mock link" ABCD-1234 tabs)
- Welcome/pager-minimal rebrand: Grok Build -> Kigi, grok.com ->
  kimi.com, "Sign in to Grok" -> "Sign in to Kimi"
This commit is contained in:
ZacharyZhang-NY
2026-07-17 16:05:51 -04:00
parent fe1f885bb3
commit ea0ce9d15f
231 changed files with 4730 additions and 26358 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/codegen/kigi-shell/src/agent/app.rs
+4 -21
View File
@@ -621,10 +621,8 @@ pub async fn run_leader(
let fetch_auth_for_prefetch = ModelFetchAuth::resolve(&endpoints_for_prefetch);
let platform_keys_for_prefetch =
crate::agent::models::PlatformApiKeys::resolve(&agent_config.platforms);
// The shared pair helper owns the remote_fetch gate for both halves, so a
// disabled knob cannot block leader readiness on settings retries.
let (prefetched_models, remote_settings) = tokio::task::spawn_blocking(move || {
crate::agent::models::prefetch_models_and_settings_blocking(
let prefetched_models = tokio::task::spawn_blocking(move || {
crate::agent::models::prefetch_models_blocking(
&endpoints_for_prefetch,
auth_for_prefetch.as_ref(),
fetch_auth_for_prefetch,
@@ -632,20 +630,7 @@ pub async fn run_leader(
)
})
.await
.unwrap_or((None, None));
// Process-wide image normalize cache: off by default, toggled here from
// `RemoteSettings.image_normalize_cache_enabled` once at startup.
let image_normalize_cache_enabled = remote_settings
.as_ref()
.and_then(|r| r.image_normalize_cache_enabled)
.unwrap_or(false);
crate::session::normalize_cache::NormalizeCache::global()
.set_enabled(image_normalize_cache_enabled);
tracing::debug!(
enabled = image_normalize_cache_enabled,
"image normalize cache toggle resolved from remote settings"
);
.unwrap_or(None);
// ── Phase 7: Signal readiness ─────────────────────────────────────────────
//
@@ -657,9 +642,7 @@ pub async fn run_leader(
// ── Phase 8: LocalSet — agent, bridges, config watcher ───────────────────
let local_set = tokio::task::LocalSet::new();
let remote_settings_for_reloader = remote_settings.clone();
let mut agent_config_for_spawn = agent_config.clone();
agent_config_for_spawn.remote_settings = remote_settings;
crate::util::config::sync_campaign_fields(&mut agent_config_for_spawn);
let agent_to_ipc_tx_clone = agent_to_ipc_tx.clone();
let cancel_clone = cancel.clone();
@@ -891,7 +874,7 @@ pub async fn run_leader(
initial_auth_key_hash,
initial_config,
auth_scope,
remote_settings_for_reloader,
None,
config_update_tx,
agent_config.cli_experimental_memory,
agent_config.cli_no_memory,
crates/codegen/kigi-shell/src/agent/chat_modes.rs
+12 -330
View File
@@ -1,334 +1,16 @@
//! grok.com chat-product model catalog: caches `/rest/modes` and maps modes to
//! the `SessionModelState` returned by `load_chat_session` (the chat analogue of
//! [`crate::agent::models::ModelsManager`]). NB: these "modes" populate the
//! desktop MODEL picker, not the ACP session plan-modes in `LoadSessionResponse.modes`.
use crate::auth::AuthManager;
use crate::remote::chat_models_client::{
ChatModelsClient, ChatModelsError, ListModesResponse, Mode,
};
use agent_client_protocol as acp;
use parking_lot::RwLock;
use std::sync::Arc;
use std::time::{Duration, Instant};
/// ~54 min, matching grok-web's refetch cadence.
const CACHE_TTL: Duration = Duration::from_secs(54 * 60);
/// Cold-miss budget on the `session/load` critical path (warm/stale served instantly).
const COLD_FETCH_TIMEOUT: Duration = Duration::from_secs(2);
const DEFAULT_LOCALE: &str = "en";
/// Process-wide flag set by the pager when started with `--chat` so initialize
/// and early UI seed the chat `/rest/modes` catalog instead of build models.
//! Legacy `--chat` gateway gate.
//!
//! The grok.com chat-product model picker (`/rest/modes`, `ChatModesManager`)
//! was removed with the xAI proxy: those "modes" came from a grok backend with
//! no Kimi counterpart. Only the process-mode gate survives so the `--chat`
//! frontend path stays a compile-time-off no-op across crates without a
//! cross-crate churn to delete every reference.
/// Process-wide flag set by the pager when started with `--chat`.
pub const KIGI_CHAT_MODE_ENV: &str = "KIGI_CHAT_MODE";
/// True when the process is a gateway light-frontend (`--chat`) agent.
/// Hard-off in release builds so it can't be enabled via env.
/// Hard-off: the grok chat-modes backend is gone, so this is always `false`.
pub fn process_chat_mode_enabled() -> bool {
if true {
return false;
}
match std::env::var(KIGI_CHAT_MODE_ENV) {
Ok(v) => {
let v = v.trim();
!v.is_empty() && v != "0" && !v.eq_ignore_ascii_case("false")
}
Err(_) => false,
}
}
#[derive(Clone)]
struct CachedModes {
/// Keyed by identity; a mismatch is a miss so one user's modes never leak to another.
user_id: String,
locale: String,
fetched_at: Instant,
response: ListModesResponse,
}
/// Thread-safe, cheaply-cloneable manager. Cloning bumps the inner `Arc`.
#[derive(Clone)]
pub struct ChatModesManager {
inner: Arc<Inner>,
}
struct Inner {
auth: Arc<AuthManager>,
cache: RwLock<Option<CachedModes>>,
/// Single-flight guard so concurrent fetches coalesce.
fetch_lock: tokio::sync::Mutex<()>,
}
impl ChatModesManager {
pub fn new(auth: Arc<AuthManager>) -> Self {
Self {
inner: Arc::new(Inner {
auth,
cache: RwLock::new(None),
fetch_lock: tokio::sync::Mutex::new(()),
}),
}
}
/// The active grok.com identity, or `None` when unauthenticated. Modes are
/// per-identity (tier/ACL), so every cache key and store is gated on it.
fn current_user_id(&self) -> Option<String> {
self.inner.auth.current_or_expired().map(|a| a.user_id)
}
/// Chat model state for a `session/load` response. On missing auth or fetch
/// failure, serves last-good cache else empty — never the build catalog.
pub async fn model_state(&self) -> acp::SessionModelState {
let Some(user_id) = self.current_user_id() else {
return empty_state();
};
let locale = DEFAULT_LOCALE;
{
let guard = self.inner.cache.read();
if let Some(c) = guard.as_ref()
&& c.user_id == user_id
&& c.locale == locale
{
if c.fetched_at.elapsed() < CACHE_TTL {
return modes_to_model_state(&c.response);
}
let stale = c.response.clone();
drop(guard);
self.spawn_refresh(user_id, locale);
return modes_to_model_state(&stale);
}
}
let _flight = self.inner.fetch_lock.lock().await;
{
let guard = self.inner.cache.read();
if let Some(c) = guard.as_ref()
&& c.user_id == user_id
&& c.locale == locale
&& c.fetched_at.elapsed() < CACHE_TTL
{
return modes_to_model_state(&c.response);
}
}
match self.fetch(locale).await {
Ok(resp) if !resp.modes.is_empty() => {
if self.current_user_id().as_deref() != Some(user_id.as_str()) {
return empty_state();
}
let mapped = modes_to_model_state(&resp);
if mapped.available_models.is_empty() {
tracing::warn!(
raw_modes = resp.modes.len(),
"chat modes: fetch returned modes but none selectable after availability filter"
);
}
self.store(user_id, locale.to_owned(), resp);
mapped
}
Ok(_) => empty_state(),
Err(err) => {
tracing::warn!(
error = % err, "chat modes fetch failed; serving cache/empty"
);
let guard = self.inner.cache.read();
match guard.as_ref() {
Some(c) if c.user_id == user_id => modes_to_model_state(&c.response),
_ => empty_state(),
}
}
}
}
async fn fetch(&self, locale: &str) -> Result<ListModesResponse, ChatModelsError> {
let client = ChatModelsClient::new(self.inner.auth.clone());
match tokio::time::timeout(COLD_FETCH_TIMEOUT, client.list_modes(locale)).await {
Ok(result) => result,
Err(_elapsed) => Err(ChatModelsError::Timeout),
}
}
fn store(&self, user_id: String, locale: String, response: ListModesResponse) {
*self.inner.cache.write() = Some(CachedModes {
user_id,
locale,
fetched_at: Instant::now(),
response,
});
}
/// Best-effort stale refresh; skips if a fetch is already in flight.
fn spawn_refresh(&self, user_id: String, locale: &'static str) {
let me = self.clone();
tokio::spawn(async move {
let Ok(_flight) = me.inner.fetch_lock.try_lock() else {
return;
};
if me.current_user_id().as_deref() != Some(user_id.as_str()) {
return;
}
if let Ok(resp) = me.fetch(locale).await
&& !resp.modes.is_empty()
&& me.current_user_id().as_deref() == Some(user_id.as_str())
{
me.store(user_id, locale.to_owned(), resp);
}
});
}
/// Kick a background `/rest/modes` fill when auth is already present so
/// `--chat` initialize / first `session/new` hit a warm cache.
pub fn warm_in_background(&self) {
let Some(user_id) = self.current_user_id() else {
return;
};
self.spawn_refresh(user_id, DEFAULT_LOCALE);
}
}
fn empty_state() -> acp::SessionModelState {
acp::SessionModelState::new(acp::ModelId::from(String::new()), Vec::new())
}
/// Maps grok.com modes → `SessionModelState`: keeps only `available` modes,
/// reconciles `current_model_id` (default → first available → empty, never
/// out-of-set), and stashes `badgeText`/`iconHint`/`tags` in `_meta`.
pub fn modes_to_model_state(resp: &ListModesResponse) -> acp::SessionModelState {
let available_models: Vec<acp::ModelInfo> = resp
.modes
.iter()
.filter(|m| m.is_available())
.map(mode_to_model_info)
.collect();
let current_model_id = reconcile_current(&resp.default_mode_id, &available_models);
acp::SessionModelState::new(current_model_id, available_models)
}
fn mode_to_model_info(m: &Mode) -> acp::ModelInfo {
let name = if m.title.trim().is_empty() {
m.id.clone()
} else {
m.title.clone()
};
acp::ModelInfo::new(acp::ModelId::from(m.id.clone()), name)
.description(if m.description.is_empty() {
None
} else {
Some(m.description.clone())
})
.meta(build_meta(m))
}
fn build_meta(m: &Mode) -> Option<acp::Meta> {
let mut map = serde_json::Map::new();
if let Some(badge) = m.badge_text.as_deref().filter(|s| !s.is_empty()) {
map.insert("badgeText".to_owned(), serde_json::json!(badge));
}
if !m.icon_hint.is_empty() {
map.insert("iconHint".to_owned(), serde_json::json!(m.icon_hint));
}
if !m.tags.is_empty() {
map.insert("tags".to_owned(), serde_json::json!(m.tags));
}
if map.is_empty() { None } else { Some(map) }
}
fn reconcile_current(default_mode_id: &str, available: &[acp::ModelInfo]) -> acp::ModelId {
let in_set = |id: &str| available.iter().any(|m| m.model_id.0.as_ref() == id);
if !default_mode_id.is_empty() && in_set(default_mode_id) {
acp::ModelId::from(default_mode_id.to_owned())
} else if let Some(first) = available.first() {
first.model_id.clone()
} else {
acp::ModelId::from(String::new())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::remote::chat_models_client::ModeAvailability;
fn available(id: &str, title: &str) -> Mode {
Mode {
id: id.to_owned(),
title: title.to_owned(),
availability: ModeAvailability {
available: Some(serde_json::json!({})),
..Default::default()
},
..Default::default()
}
}
fn requires_upgrade(id: &str) -> Mode {
Mode {
id: id.to_owned(),
availability: ModeAvailability {
requires_upgrade: Some(serde_json::json!({ "message" : "Upgrade" })),
..Default::default()
},
..Default::default()
}
}
#[test]
fn filters_to_available_modes() {
let resp = ListModesResponse {
modes: vec![
available("auto", "Auto"),
requires_upgrade("heavy"),
available("fast", "Fast"),
],
default_mode_id: "auto".to_owned(),
};
let state = modes_to_model_state(&resp);
let ids: Vec<String> = state
.available_models
.iter()
.map(|m| m.model_id.0.to_string())
.collect();
assert_eq!(ids, vec!["auto".to_string(), "fast".to_string()]);
assert_eq!(state.current_model_id.0.as_ref(), "auto");
}
#[test]
fn default_outside_filtered_set_falls_back_to_first_available() {
let resp = ListModesResponse {
modes: vec![requires_upgrade("heavy"), available("fast", "Fast")],
default_mode_id: "heavy".to_owned(),
};
let state = modes_to_model_state(&resp);
assert_eq!(state.current_model_id.0.as_ref(), "fast");
assert!(
state
.available_models
.iter()
.any(|m| m.model_id == state.current_model_id)
);
}
#[test]
fn empty_default_falls_back_to_first() {
let resp = ListModesResponse {
modes: vec![available("a", "A"), available("b", "B")],
default_mode_id: String::new(),
};
let state = modes_to_model_state(&resp);
assert_eq!(state.current_model_id.0.as_ref(), "a");
}
#[test]
fn no_available_modes_yields_empty_current() {
let resp = ListModesResponse {
modes: vec![requires_upgrade("heavy")],
default_mode_id: "heavy".to_owned(),
};
let state = modes_to_model_state(&resp);
assert!(state.available_models.is_empty());
assert_eq!(state.current_model_id.0.as_ref(), "");
}
#[test]
fn maps_fields_and_meta() {
let mut m = available("auto", "Auto");
m.description = "Picks the best model".to_owned();
m.badge_text = Some("New".to_owned());
m.icon_hint = "rocket".to_owned();
m.tags = vec!["TAG_PRIMARY".to_owned()];
let resp = ListModesResponse {
modes: vec![m],
default_mode_id: "auto".to_owned(),
};
let state = modes_to_model_state(&resp);
let info = &state.available_models[0];
assert_eq!(info.name, "Auto");
assert_eq!(info.description.as_deref(), Some("Picks the best model"));
let meta = info.meta.as_ref().unwrap();
assert_eq!(meta["badgeText"], serde_json::json!("New"));
assert_eq!(meta["iconHint"], serde_json::json!("rocket"));
assert_eq!(meta["tags"], serde_json::json!(["TAG_PRIMARY"]));
}
#[test]
fn name_falls_back_to_id_when_title_blank() {
let mut m = available("grok-4.5", "");
m.title = " ".to_owned();
let resp = ListModesResponse {
modes: vec![m],
default_mode_id: String::new(),
};
let state = modes_to_model_state(&resp);
assert_eq!(state.available_models[0].name, "grok-4.5");
}
false
}
crates/codegen/kigi-shell/src/agent/config.rs
+67 -128
View File
@@ -1,6 +1,6 @@
use crate::agent::auth_method::ModelByok;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
use crate::auth::{AuthManager, KimiCodeConfig};
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::{config::StorageMode, sampling::ApiBackend, tools::config::ShellToolsetConfig};
use agent_client_protocol as acp;
use indexmap::IndexMap;
@@ -141,7 +141,7 @@ pub struct EndpointsConfig {
/// `Some` = explicitly configured. Tracking explicitness (vs comparing to the
/// default value) lets an org pin the proxy to the default on purpose.
#[serde(skip_serializing_if = "Option::is_none")]
pub cli_chat_proxy_base_url: Option<String>,
pub coding_api_base_url: Option<String>,
/// Base URL for the public xAI API.
pub xai_api_base_url: String,
/// Optional extra access-header value (applied only with the optional
@@ -271,11 +271,11 @@ impl EndpointsConfig {
resolved
}
/// The subscription proxy base URL through which all auxiliary services (and
/// OAuth/session inference) resolve: explicit `cli_chat_proxy_base_url`, else
/// OAuth/session inference) resolve: explicit `coding_api_base_url`, else
/// [`kigi_env::coding_api_base_url`]. NEVER falls back to `xai_api_base_url` —
/// that is the inference endpoint (API-key auth) only.
pub fn proxy_url(&self) -> String {
blank_as_unset(&self.cli_chat_proxy_base_url).unwrap_or_else(kigi_env::coding_api_base_url)
blank_as_unset(&self.coding_api_base_url).unwrap_or_else(kigi_env::coding_api_base_url)
}
pub fn resolve_inference_base_url(&self) -> String {
self.models_base_url
@@ -406,7 +406,7 @@ impl EndpointsConfig {
impl Default for EndpointsConfig {
fn default() -> Self {
Self {
cli_chat_proxy_base_url: std::env::var("KIGI_CLI_CHAT_PROXY_BASE_URL").ok(),
coding_api_base_url: std::env::var("KIGI_CODE_BASE_URL").ok(),
xai_api_base_url: std::env::var("KIGI_XAI_API_BASE_URL")
.unwrap_or_else(|_| XAI_API_BASE_URL_DEFAULT.to_owned()),
alpha_test_key: None,
@@ -4169,19 +4169,11 @@ pub fn resolve_aux_model_sampling_config(
endpoints: &EndpointsConfig,
session_key: Option<&str>,
alpha_test_key: Option<String>,
client_version: Option<String>,
) -> Option<SamplerConfig> {
let catalog_entry = find_model_by_id(models, model_id).cloned();
if let Some(entry) = &catalog_entry {
let credentials = resolve_credentials(entry, session_key);
let sampler = sampling_config_for_model(
entry,
credentials,
alpha_test_key.clone(),
client_version.clone(),
None,
None,
);
let sampler = sampling_config_for_model(entry, credentials, alpha_test_key.clone());
if sampler.api_key.is_some() {
return Some(sampler);
}
@@ -4232,14 +4224,7 @@ pub fn resolve_aux_model_sampling_config(
api_base_url: None,
};
let credentials = resolve_credentials(&entry, session_key);
let sampler = sampling_config_for_model(
&entry,
credentials,
alpha_test_key,
client_version,
None,
None,
);
let sampler = sampling_config_for_model(&entry, credentials, alpha_test_key);
return Some(sampler);
}
tracing::warn!(
@@ -4252,21 +4237,19 @@ pub fn resolve_aux_model_sampling_config(
/// Shared so the aux resolve happy path and the
/// `None` fallback cannot diverge between those entry points.
///
/// On aux resolve `Some`, stamp session-local fields (client id, attribution, bearer,
/// On aux resolve `Some`, stamp session-local fields (attribution, bearer,
/// retries) onto the helper config. On `None`, fall back to the active session model and
/// full config (not forcing `image_description_model` onto the agent endpoint, which 404s
/// on BYOK / non-proxy routes for internal slugs like `grok-build`).
/// Stamp the session-local fields (client id, attribution, bearer resolver,
/// retries) from the active session onto a routed aux `SamplerConfig` so a
/// on BYOK / non-proxy routes for internal slugs).
/// Stamp the session-local fields (attribution, bearer resolver, retries)
/// from the active session onto a routed aux `SamplerConfig` so a
/// helper model keeps the session's auth/attribution. Shared by image-describe
/// and the auto-mode classifier so the two can't drift.
pub fn stamp_session_local_sampler_fields(
cfg: &mut SamplerConfig,
active_session_config: &SamplerConfig,
client_identifier: Option<String>,
max_retries: Option<u32>,
) {
cfg.client_identifier = client_identifier;
cfg.attribution_callback = active_session_config.attribution_callback.clone();
cfg.bearer_resolver = active_session_config.bearer_resolver.clone();
cfg.max_retries = max_retries;
@@ -4274,7 +4257,6 @@ pub fn stamp_session_local_sampler_fields(
pub fn finalize_image_describe_sampler_config(
resolved_aux: Option<SamplerConfig>,
active_session_config: &SamplerConfig,
client_identifier: Option<String>,
max_retries: Option<u32>,
) -> (String, SamplerConfig) {
match resolved_aux {
@@ -4282,7 +4264,6 @@ pub fn finalize_image_describe_sampler_config(
stamp_session_local_sampler_fields(
&mut describe_cfg,
active_session_config,
client_identifier,
max_retries,
);
let model = describe_cfg.model.clone();
@@ -4310,9 +4291,6 @@ pub fn sampling_config_for_model(
model: &ModelEntry,
credentials: ResolvedCredentials,
alpha_test_key: Option<String>,
client_version: Option<String>,
deployment_id: Option<String>,
user_id: Option<String>,
) -> SamplerConfig {
let info = model.info();
let model_name = info.model.clone();
@@ -4337,15 +4315,11 @@ pub fn sampling_config_for_model(
auth_scheme: credentials.auth_scheme,
extra_headers,
context_window: info.context_window.get(),
client_version,
reasoning_effort: info.reasoning_effort,
force_http1: false,
max_retries: info.max_retries,
stream_tool_calls: info.stream_tool_calls.unwrap_or(false),
idle_timeout_secs: None,
client_identifier: None,
deployment_id,
user_id,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
@@ -4359,11 +4333,19 @@ pub fn sampling_config_for_model(
/// Fold URL-derived headers into `extra_headers`.
///
/// The sampler crate is intentionally URL-agnostic: it does not inspect
/// `base_url` to decide which auth or staging headers to add. Replicate the
/// `base_url` to decide which auth or identity headers to add. Replicate the
/// URL-derived header logic at the shell boundary so callers downstream see a
/// single homogenous header bag.
///
/// * First-party bases get the client-mode header.
/// * First-party (Kimi subscription) bases get the `X-Msh-Device-*` identity
/// headers, mirroring the official client, which sends its OAuth device
/// headers on every inference request (kimi-cli src/kimi_cli/llm.py:317-323
/// `_kimi_default_headers` merges `oauth.common_headers()`). Third-party /
/// Moonshot-open-platform bases get none — only the bearer and User-Agent.
///
/// A device-id failure only skips the headers (with a warning): inference
/// must not hard-fail because `~/.kigi/device_id` is unwritable — unlike
/// OAuth login, where the id is mandatory.
///
/// Existing entries are never overwritten so callers can pre-set a value.
pub fn inject_url_derived_headers(
@@ -4371,10 +4353,20 @@ pub fn inject_url_derived_headers(
alpha_test_key: Option<&str>,
base_url: &str,
) {
if crate::util::is_cli_chat_proxy_url(base_url) {
headers
.entry(crate::http::CLIENT_MODE_HEADER.to_string())
.or_insert_with(|| crate::http::process_client_mode().to_string());
if crate::util::is_production_coding_api_url(base_url) {
match crate::auth::device_headers() {
Ok(device_headers) => {
for (name, value) in device_headers {
headers.entry(name.to_string()).or_insert(value);
}
}
Err(e) => {
tracing::warn!(
error = %e,
"device identity headers unavailable; sending inference request without them"
);
}
}
}
let _ = (alpha_test_key, base_url);
}
@@ -4383,7 +4375,6 @@ pub fn resolve_model_to_sampling_config(
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
alpha_test_key: Option<String>,
client_version: Option<String>,
fallback_entry: Option<ModelEntry>,
) -> Option<SamplerConfig> {
let entry = find_model_by_id(models, model_id)
@@ -4394,16 +4385,12 @@ pub fn resolve_model_to_sampling_config(
&entry,
credentials,
alpha_test_key,
client_version,
None,
None,
))
}
fn resolve_hidden_default_web_search_sampling_config(
model_id: &str,
session_key: Option<&str>,
alpha_test_key: Option<String>,
client_version: Option<String>,
endpoints: &EndpointsConfig,
) -> SamplerConfig {
let entry = ModelEntry {
@@ -4445,21 +4432,13 @@ fn resolve_hidden_default_web_search_sampling_config(
api_base_url: None,
};
let credentials = resolve_credentials(&entry, session_key);
sampling_config_for_model(
&entry,
credentials,
alpha_test_key,
client_version,
None,
None,
)
sampling_config_for_model(&entry, credentials, alpha_test_key)
}
pub fn resolve_web_search_sampling_config(
model_id: &str,
models: &IndexMap<String, ModelEntry>,
session_key: Option<&str>,
alpha_test_key: Option<String>,
client_version: Option<String>,
endpoints: &EndpointsConfig,
) -> Option<SamplerConfig> {
let resolved = if let Some(entry) = find_model_by_id(models, model_id).cloned() {
@@ -4468,16 +4447,12 @@ pub fn resolve_web_search_sampling_config(
&entry,
credentials,
alpha_test_key,
client_version,
None,
None,
))
} else if model_id == crate::models::default_web_search_model() {
Some(resolve_hidden_default_web_search_sampling_config(
model_id,
session_key,
alpha_test_key,
client_version,
endpoints,
))
} else {
@@ -4757,21 +4732,23 @@ reasoning_effort = "low"
}
}
#[test]
fn inject_url_derived_headers_adds_client_mode_for_first_party_url() {
fn inject_url_derived_headers_adds_device_identity_for_first_party_url() {
let mut headers = IndexMap::new();
inject_url_derived_headers(
&mut headers,
None,
kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url,
);
assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_some());
assert!(headers.get("X-Msh-Device-Id").is_some());
assert!(headers.get("X-Msh-Device-Name").is_some());
assert!(headers.get("X-XAI-Token-Auth").is_none());
}
#[test]
fn inject_url_derived_headers_skips_headers_for_external_url() {
let mut headers = IndexMap::new();
inject_url_derived_headers(&mut headers, None, "https://api.example.com/v1");
assert!(headers.get(crate::http::CLIENT_MODE_HEADER).is_none());
assert!(headers.get("X-Msh-Device-Id").is_none());
assert!(headers.get("X-Msh-Device-Name").is_none());
}
#[test]
fn inject_url_derived_headers_preserves_caller_extra_headers() {
@@ -4924,7 +4901,6 @@ reasoning_effort = "low"
&IndexMap::new(),
Some("session-token"),
None,
None,
&endpoints,
)
.expect("hidden default web search model should resolve");
@@ -4943,7 +4919,7 @@ reasoning_effort = "low"
model: "composer-session-model".into(),
..Default::default()
};
let (model, cfg) = finalize_image_describe_sampler_config(None, &active, None, Some(3));
let (model, cfg) = finalize_image_describe_sampler_config(None, &active, Some(3));
assert_eq!(model, "composer-session-model");
assert_eq!(cfg.model, "composer-session-model");
assert_ne!(cfg.model, "grok-build");
@@ -4958,11 +4934,9 @@ reasoning_effort = "low"
model: "grok-build".into(),
..Default::default()
};
let (model, cfg) =
finalize_image_describe_sampler_config(Some(aux), &active, Some("cli".into()), Some(7));
let (model, cfg) = finalize_image_describe_sampler_config(Some(aux), &active, Some(7));
assert_eq!(model, "grok-build");
assert_eq!(cfg.model, "grok-build");
assert_eq!(cfg.client_identifier.as_deref(), Some("cli"));
assert_eq!(cfg.max_retries, Some(7));
}
#[test]
@@ -4980,7 +4954,7 @@ reasoning_effort = "low"
),
);
let resolved =
resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None, None)
resolve_aux_model_sampling_config("grok-build", &catalog, &endpoints, None, None)
.expect("override entry has an API key, so resolution succeeds");
assert_eq!(resolved.model, "v9m-rl-learnability-tp8");
assert_eq!(resolved.base_url, "https://vendor.example/v1");
@@ -5085,14 +5059,8 @@ reasoning_effort = "low"
None,
None,
);
let sampling_config = sampling_config_for_model(
&model,
resolve_credentials(&model, None),
None,
None,
None,
None,
);
let sampling_config =
sampling_config_for_model(&model, resolve_credentials(&model, None), None);
assert_eq!(
sampling_config.api_key,
Some("model-specific-key".to_string())
@@ -5111,9 +5079,6 @@ reasoning_effort = "low"
auth_scheme: AuthScheme::Bearer,
},
None,
None,
None,
None,
);
assert_eq!(sampling_config.api_key, Some("fallback-key".to_string()));
}
@@ -5388,14 +5353,8 @@ reasoning_effort = "low"
None,
);
model.info.api_backend = ApiBackend::Messages;
let config = sampling_config_for_model(
&model,
resolve_credentials(&model, Some("tok")),
None,
None,
None,
None,
);
let config =
sampling_config_for_model(&model, resolve_credentials(&model, Some("tok")), None);
assert_eq!(config.api_backend, ApiBackend::Messages);
assert_eq!(config.auth_scheme, AuthScheme::Bearer);
assert_eq!(config.api_key, Some("tok".to_string()));
@@ -5440,7 +5399,7 @@ reasoning_effort = "low"
assert_eq!(creds.auth_scheme, AuthScheme::XApiKey);
assert_eq!(creds.auth_type, kigi_chat_state::AuthType::ApiKey);
assert_eq!(creds.api_key, Some("sk-ant-test-key".to_string()));
let config = sampling_config_for_model(&model, creds, None, None, None, None);
let config = sampling_config_for_model(&model, creds, None);
assert_eq!(config.auth_scheme, AuthScheme::XApiKey);
assert_eq!(config.api_backend, ApiBackend::Messages);
let client = kigi_sampler::SamplingClient::new(config).expect("client should build");
@@ -5459,7 +5418,7 @@ reasoning_effort = "low"
assert_eq!(model.info.auth_scheme, AuthScheme::Bearer);
let creds = resolve_credentials(&model, None);
assert_eq!(creds.auth_scheme, AuthScheme::Bearer);
let config = sampling_config_for_model(&model, creds, None, None, None, None);
let config = sampling_config_for_model(&model, creds, None);
assert_eq!(config.auth_scheme, AuthScheme::Bearer);
let client = kigi_sampler::SamplingClient::new(config).expect("client should build");
let info = client.auth_info();
@@ -5771,25 +5730,11 @@ reasoning_effort = "low"
#[test]
fn sampling_config_context_window_from_entry_or_default() {
let model = test_model_entry("any-model", "https://api.x.ai/v1", None, None, None);
let config = sampling_config_for_model(
&model,
resolve_credentials(&model, None),
None,
None,
None,
None,
);
let config = sampling_config_for_model(&model, resolve_credentials(&model, None), None);
assert_eq!(config.context_window, 200_000);
let mut model = test_model_entry("any-model", "https://api.x.ai/v1", None, None, None);
model.info.context_window = NonZeroU64::new(256_000).unwrap();
let config = sampling_config_for_model(
&model,
resolve_credentials(&model, None),
None,
None,
None,
None,
);
let config = sampling_config_for_model(&model, resolve_credentials(&model, None), None);
assert_eq!(config.context_window, 256_000);
}
#[test]
@@ -5917,14 +5862,8 @@ reasoning_effort = "low"
let mut model =
test_model_entry("test-model", "https://api.example.com/v1", None, None, None);
model.info.api_backend = ApiBackend::Responses;
let sampling_config = sampling_config_for_model(
&model,
resolve_credentials(&model, None),
None,
None,
None,
None,
);
let sampling_config =
sampling_config_for_model(&model, resolve_credentials(&model, None), None);
assert_eq!(sampling_config.api_backend, ApiBackend::Responses);
}
#[test]
@@ -6636,7 +6575,7 @@ reasoning_effort = "low"
}
fn resolve_sampling(model: &ModelEntry, session_key: Option<&str>) -> SamplerConfig {
let credentials = resolve_credentials(model, session_key);
sampling_config_for_model(model, credentials, None, None, None, None)
sampling_config_for_model(model, credentials, None)
}
#[test]
#[serial]
@@ -7022,7 +6961,7 @@ reasoning_effort = "low"
&format!(
r#"
[endpoints]
cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1"
coding_api_base_url = "https://enterprise-proxy.acme.com/v1"
[model."{BUNDLED_DEFAULT_KEY}"]
api_key = "acme-api-key"
@@ -7052,14 +6991,14 @@ reasoning_effort = "low"
let (_, models) = resolve_models_from_toml(
r#"
[endpoints]
cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1"
coding_api_base_url = "https://enterprise-proxy.acme.com/v1"
"#,
None,
);
let model = models.get(BUNDLED_DEFAULT_KEY).expect("model should exist");
assert_eq!(
model.info.base_url, "https://enterprise-proxy.acme.com/v1",
"default model should use enterprise cli_chat_proxy_base_url"
"default model should use enterprise coding_api_base_url"
);
// The open-platform fallback entries keep their fixed moonshot bases;
// only the subscription entry follows the proxy override.
@@ -7073,7 +7012,7 @@ reasoning_effort = "low"
/// the ambient environment. Gated behind `#[serial]`.
fn unset_endpoint_env_vars() {
for k in [
"KIGI_CLI_CHAT_PROXY_BASE_URL",
"KIGI_CODE_BASE_URL",
kigi_env::CODE_BASE_URL_ENV,
"KIGI_XAI_API_BASE_URL",
"KIGI_FEEDBACK_BASE_URL",
@@ -7101,7 +7040,7 @@ reasoning_effort = "low"
let inference = "https://inference.acme-corp.example/xai/v1";
let cfg = EndpointsConfig {
xai_api_base_url: inference.to_string(),
cli_chat_proxy_base_url: None,
coding_api_base_url: None,
..Default::default()
};
let proxy = kigi_env::PRODUCTION_ENDPOINTS.coding_api_base_url;
@@ -7119,7 +7058,7 @@ reasoning_effort = "low"
);
assert_eq!(cfg.xai_api_base_url, inference);
let overridden = EndpointsConfig {
cli_chat_proxy_base_url: Some("https://proxy.enterprise.example/v1".to_string()),
coding_api_base_url: Some("https://proxy.enterprise.example/v1".to_string()),
managed_config_url: Some(
"https://control.enterprise.example/deployment/config".to_string(),
),
@@ -7159,7 +7098,7 @@ reasoning_effort = "low"
.unwrap(),
)
.expect("config should parse");
assert!(cfg.endpoints.cli_chat_proxy_base_url.is_none());
assert!(cfg.endpoints.coding_api_base_url.is_none());
assert_eq!(
cfg.endpoints.resolve_managed_config_url(),
format!(
@@ -7181,7 +7120,7 @@ reasoning_effort = "low"
&format!(
r#"
[endpoints]
cli_chat_proxy_base_url = "https://enterprise-proxy.acme.com/v1"
coding_api_base_url = "https://enterprise-proxy.acme.com/v1"
[model."{dm}"]
base_url = "https://my-special-proxy.example.com/v1"
@@ -8656,7 +8595,7 @@ agent_type = "cursor"
fn otlp_traces_endpoint_precedence() {
let proxy = "https://inference.acme.com/v1".to_string();
let derived = EndpointsConfig {
cli_chat_proxy_base_url: Some(proxy.clone()),
coding_api_base_url: Some(proxy.clone()),
..Default::default()
};
assert_eq!(
@@ -8664,7 +8603,7 @@ agent_type = "cursor"
"https://inference.acme.com/v1/traces"
);
let base = EndpointsConfig {
cli_chat_proxy_base_url: Some(proxy.clone()),
coding_api_base_url: Some(proxy.clone()),
otel_exporter_otlp_endpoint: Some("https://otel.acme.com".to_string()),
..Default::default()
};
@@ -8673,7 +8612,7 @@ agent_type = "cursor"
"https://otel.acme.com/v1/traces"
);
let full = EndpointsConfig {
cli_chat_proxy_base_url: Some(proxy),
coding_api_base_url: Some(proxy),
otel_exporter_otlp_endpoint: Some("https://ignored.example".to_string()),
otel_exporter_otlp_traces_endpoint: Some("https://otel.acme.com/v1/traces".to_string()),
..Default::default()
@@ -8702,7 +8641,7 @@ agent_type = "cursor"
/// explicitly unset so ambient env (via `Default`) can't leak in.
fn internal_otlp_test_config() -> EndpointsConfig {
EndpointsConfig {
cli_chat_proxy_base_url: Some("https://proxy.example/v1".to_string()),
coding_api_base_url: Some("https://proxy.example/v1".to_string()),
otel_exporter_otlp_endpoint: None,
otel_exporter_otlp_traces_endpoint: None,
otel_exporter_otlp_headers: None,
crates/codegen/kigi-shell/src/agent/feedback_client.rs
+181 -1277
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/agent/handlers/mod.rs
-1
View File
@@ -1,3 +1,2 @@
pub(crate) mod model_switch;
pub(crate) mod session;
pub(crate) mod workspaces;
crates/codegen/kigi-shell/src/agent/handlers/session.rs
+2 -11
View File
@@ -271,19 +271,10 @@ async fn handle_session_list(
// (never union) so every list surface is conversations-only.
let req = unified_list::parse_list_req(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
tracing::debug!(
chat_mode_forced_kind = crate::agent::chat_modes::process_chat_mode_enabled(),
"session/list"
);
tracing::debug!("session/list");
let registry_client = agent.session_registry_client();
let conversations_client = agent.conversations_client();
let result = unified_list::build_unified_list(
registry_client.as_ref(),
conversations_client.as_ref(),
req,
)
.await;
let result = unified_list::build_unified_list(registry_client.as_ref(), req).await;
ExtMethodResult::success(unified_list::ext_list_response(result))
.to_ext_response()
crates/codegen/kigi-shell/src/agent/handlers/workspaces.rs
-174
View File
@@ -1,174 +0,0 @@
use agent_client_protocol::{self as acp};
use serde::{Deserialize, Serialize};
use super::super::mvp_agent::MvpAgent;
use crate::remote::{ListWorkspacesPage, WsError, WsQuery};
use crate::session::ExtMethodResult;
const DEFAULT_PAGE_SIZE: i64 = 50;
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
struct WorkspacesListRequest {
#[serde(default)]
page_size: Option<i64>,
#[serde(default)]
page_token: Option<String>,
#[serde(default)]
query: Option<String>,
#[serde(default)]
kind: Option<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct WorkspaceRow {
id: String,
name: String,
#[serde(skip_serializing_if = "Option::is_none")]
kind: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
create_time: Option<String>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
struct WorkspacesListResponse {
workspaces: Vec<WorkspaceRow>,
#[serde(skip_serializing_if = "Option::is_none")]
next_page_token: Option<String>,
#[serde(rename = "_meta", skip_serializing_if = "Option::is_none")]
meta: Option<WorkspacesMeta>,
}
#[derive(Debug, Serialize)]
struct WorkspacesMeta {
#[serde(rename = "x.ai/partial")]
partial: PartialInfo,
}
#[derive(Debug, Serialize)]
struct PartialInfo {
workspaces: bool,
reason: &'static str,
}
pub async fn handle(
agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
let req: WorkspacesListRequest = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
let q = WsQuery {
// Clamp to a sane positive page size: a missing, zero, or negative
// `pageSize` falls back to the default rather than being forwarded
// verbatim to `/rest/workspaces`.
page_size: match req.page_size {
Some(n) if n > 0 => n,
_ => DEFAULT_PAGE_SIZE,
},
page_token: req.page_token,
query: req.query,
kind: req.kind,
};
let response = match agent.workspaces_client().list_workspaces(&q).await {
Ok(page) => success_response(page),
Err(WsError::NoOauth) => degraded_response("no_oauth"),
Err(e) => {
// Degrade to a partial result, but don't silently swallow the
// cause — log it so field failures are diagnosable.
tracing::warn!("workspaces/list fetch failed: {e}");
degraded_response("error")
}
};
ExtMethodResult::success(response)
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
fn success_response(page: ListWorkspacesPage) -> WorkspacesListResponse {
WorkspacesListResponse {
workspaces: page
.workspaces
.into_iter()
.map(|w| WorkspaceRow {
id: w.workspace_id,
name: w.name,
kind: w.kind,
create_time: w.create_time,
})
.collect(),
next_page_token: page.next_page_token,
meta: None,
}
}
fn degraded_response(reason: &'static str) -> WorkspacesListResponse {
WorkspacesListResponse {
workspaces: Vec::new(),
next_page_token: None,
meta: Some(WorkspacesMeta {
partial: PartialInfo {
workspaces: true,
reason,
},
}),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::remote::Workspace;
#[test]
fn request_parses_camelcase_and_defaults_page_size() {
let req: WorkspacesListRequest =
serde_json::from_value(serde_json::json!({})).expect("empty params parse");
assert!(req.page_size.is_none());
let req: WorkspacesListRequest = serde_json::from_value(serde_json::json!({
"pageSize": 10,
"pageToken": "tok",
"query": "gpu",
"kind": "WORKSPACE_KIND_IMAGINE"
}))
.expect("full params parse");
assert_eq!(req.page_size, Some(10));
assert_eq!(req.page_token.as_deref(), Some("tok"));
assert_eq!(req.query.as_deref(), Some("gpu"));
assert_eq!(req.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE"));
}
#[test]
fn success_response_projects_grok_workspace_fields() {
let page = ListWorkspacesPage {
workspaces: vec![Workspace {
workspace_id: "ws_1".into(),
name: "Research".into(),
create_time: Some("2026-06-18T17:30:00Z".into()),
kind: Some("WORKSPACE_KIND_IMAGINE".into()),
}],
next_page_token: Some("tok2".into()),
};
let value = serde_json::to_value(success_response(page)).unwrap();
assert_eq!(value["workspaces"][0]["id"], "ws_1");
assert_eq!(value["workspaces"][0]["name"], "Research");
assert_eq!(value["workspaces"][0]["kind"], "WORKSPACE_KIND_IMAGINE");
assert_eq!(value["workspaces"][0]["createTime"], "2026-06-18T17:30:00Z");
assert_eq!(value["nextPageToken"], "tok2");
assert!(value.get("_meta").is_none());
}
#[test]
fn degraded_response_carries_partial_reason() {
let value = serde_json::to_value(degraded_response("no_oauth")).unwrap();
assert_eq!(value["workspaces"].as_array().unwrap().len(), 0);
assert!(value.get("nextPageToken").is_none());
assert_eq!(value["_meta"]["x.ai/partial"]["workspaces"], true);
assert_eq!(value["_meta"]["x.ai/partial"]["reason"], "no_oauth");
}
}
crates/codegen/kigi-shell/src/agent/init.rs
-21
View File
@@ -73,27 +73,6 @@ fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig
tracing::info!(field = %e.path, value = %e.value, source = %e.source, "policy override");
}
// Fallback: if the client didn't pre-supply remote settings, fetch them
// now so remote-settings-gated features work regardless of which client
// spawned us. Clients that already call `start_early_prefetch()` and
// thread the result into `cfg.remote_settings` skip this entirely.
if cfg.remote_settings.is_none()
&& let Some(handle) =
crate::agent::models::start_early_prefetch(Some(cfg.kimi_code_config.clone()))
{
match handle.join() {
Ok(result) => {
cfg.remote_settings = result.settings;
crate::util::config::set_remote_campaigns_from_settings(
cfg.remote_settings.as_ref(),
);
tracing::info!("remote_settings fetched as shell-level fallback");
}
Err(_) => {
tracing::warn!("remote_settings fallback prefetch thread panicked");
}
}
}
crate::util::config::sync_campaign_fields(&mut cfg);
crate::agent::config::apply_remote_settings_side_effects(cfg.remote_settings.as_ref());
crates/codegen/kigi-shell/src/agent/mod.rs
+2 -1
View File
@@ -5,11 +5,12 @@ pub mod chat_modes;
pub mod config;
pub mod config_model_override_parse;
mod ext_parsers;
pub mod feedback_client;
pub(crate) mod feedback_client;
pub mod folder_trust;
pub(crate) mod handlers;
pub mod init;
pub mod models;
pub(crate) mod models_fetch;
pub mod mvp_agent;
pub(crate) mod proxy;
pub(crate) mod restore_code;
crates/codegen/kigi-shell/src/agent/models.rs
+38 -63
View File
@@ -10,8 +10,8 @@ use chrono::{DateTime, Duration as ChronoDuration, Utc};
use indexmap::IndexMap;
use crate::agent::config::{self, ModelEntry, resolve_credentials, sampling_config_for_model};
use crate::agent::models_fetch::{FetchModelsResult, fetch_models_blocking};
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use crate::remote::{FetchModelsResult, fetch_models_blocking};
use crate::sampling::SamplerConfig as SamplingConfig;
use globset::{Glob, GlobSet, GlobSetBuilder};
use kigi_sampling_types::{ReasoningEffort, ReasoningEffortOption};
@@ -290,7 +290,7 @@ impl ModelsManager {
cache
.load_fresh(
&fetch_auth.cache_auth_method(),
&crate::remote::models_fetch_origin(
&crate::agent::models_fetch::models_fetch_origin(
&cfg.endpoints,
fetch_auth,
has_session,
@@ -1035,11 +1035,6 @@ impl ModelsManager {
current_model,
credentials,
config.endpoints.alpha_test_key.clone(),
config.client_version.clone(),
crate::managed_config::resolve_deployment_id(
config.endpoints.deployment_key.as_deref(),
),
None,
)
}
@@ -1053,7 +1048,12 @@ impl ModelsManager {
let fetch_auth = *self.inner.fetch_auth.read();
let has_oauth = self.inner.auth_manager.current_or_expired().is_some();
let platform_keys = PlatformApiKeys::resolve(&platforms);
crate::remote::models_fetch_origin(&endpoints, fetch_auth, has_oauth, &platform_keys)
crate::agent::models_fetch::models_fetch_origin(
&endpoints,
fetch_auth,
has_oauth,
&platform_keys,
)
}
fn try_load_cache(&self) -> bool {
@@ -1327,7 +1327,7 @@ struct ModelsCache {
#[serde(default, skip_serializing_if = "Option::is_none")]
auth_method: Option<CacheAuthMethod>,
/// Models-list URL this catalog was fetched from
/// ([`crate::remote::models_list_url`]). Compared on load so a cache
/// ([`crate::agent::models_fetch::models_fetch_origin`]). Compared on load so a cache
/// written against one backend is a miss for another: entries embed
/// absolute `base_url`s, so adopting a foreign-origin cache silently
/// re-points inference (the windows lifecycle e2e failed exactly this
@@ -1575,43 +1575,6 @@ pub(crate) fn prefetch_models_blocking(
.models
}
/// Blocking models + `/v1/settings` prefetch pair, shared by the early
/// prefetch thread and the leader's startup phase so the settings gate lives
/// once. The remote_fetch knob is resolved a single time so the two fetch
/// decisions cannot disagree mid-startup.
pub(crate) fn prefetch_models_and_settings_blocking(
endpoints: &config::EndpointsConfig,
auth: Option<&KimiAuth>,
fetch_auth: ModelFetchAuth,
platform_keys: &PlatformApiKeys,
) -> (
Option<IndexMap<String, ModelEntry>>,
Option<crate::util::config::RemoteSettings>,
) {
let remote_fetch_enabled = crate::util::config::resolve_remote_fetch_enabled();
let models = prefetch_models_blocking_gated(
endpoints,
auth,
fetch_auth,
platform_keys,
remote_fetch_enabled,
)
.models;
// Settings need a subscription session; skip for API-key-only setups.
let settings = match auth {
Some(auth) if remote_fetch_enabled => {
let _timer = crate::instrumentation_timer!("startup.early_settings_fetch");
crate::remote::fetch_settings_blocking(
&endpoints.proxy_url(),
auth,
endpoints.alpha_test_key.as_deref(),
)
}
_ => None,
};
(models, settings)
}
/// `remote_fetch_enabled` is a parameter so the pair helper above resolves the
/// knob once for both halves.
fn prefetch_models_blocking_gated(
@@ -1624,8 +1587,12 @@ fn prefetch_models_blocking_gated(
let cache_auth = fetch_auth.cache_auth_method();
// Same fetch plan the network path below executes — the cache is only
// valid for it.
let cache_origin =
crate::remote::models_fetch_origin(endpoints, fetch_auth, auth.is_some(), platform_keys);
let cache_origin = crate::agent::models_fetch::models_fetch_origin(
endpoints,
fetch_auth,
auth.is_some(),
platform_keys,
);
let cache = ModelsCacheManager::new();
if let Some(cached) = cache.load_fresh(&cache_auth, &cache_origin) {
tracing::info!(
@@ -1703,10 +1670,9 @@ fn stale_cache_or_failure(
ModelsFetchOutcome::failed(oauth_unauthorized)
}
/// Startup prefetch result: models + remote settings.
/// Startup prefetch result: the model catalog, when a fetch plan existed.
pub struct EarlyPrefetchResult {
pub models: Option<IndexMap<String, ModelEntry>>,
pub settings: Option<crate::util::config::RemoteSettings>,
}
/// Handle for a startup prefetch thread.
@@ -1743,7 +1709,7 @@ fn resolve_prefetch_env_with_auth(auth: Option<KimiAuth>) -> Option<PrefetchEnv>
/// `has_custom_endpoint()` (which otherwise forces the prefetch to run): the
/// explicit off switch must hold even when a stray login, a platform API key,
/// or a `deployment_key` would re-arm the prefetch — and with it the
/// `/v1/settings` fetch and the deployment-config sync on the prefetch thread.
/// deployment-config sync on the prefetch thread.
///
/// PRD F2 acceptance: a moonshot API key alone (no subscription login) must
/// arm the prefetch so the catalog syncs on startup.
@@ -1754,7 +1720,7 @@ fn resolve_prefetch_env_from_parts(
remote_fetch_enabled: bool,
) -> Option<PrefetchEnv> {
if !remote_fetch_enabled {
tracing::info!("startup model/settings prefetch skipped: remote_fetch disabled");
tracing::info!("startup model prefetch skipped: remote_fetch disabled");
return None;
}
@@ -1779,7 +1745,7 @@ fn resolve_prefetch_env(kimi_code_config: Option<KimiCodeConfig>) -> Option<Pref
resolve_prefetch_env_with_auth(auth)
}
/// Start model + settings prefetch on a background thread using pre-resolved auth.
/// Start the model-catalog prefetch on a background thread using pre-resolved auth.
///
/// When the caller has already obtained valid credentials (e.g. via
/// `try_ensure_fresh_auth`), pass them here to avoid re-reading stale cached
@@ -1789,7 +1755,7 @@ pub fn start_early_prefetch_with_auth(auth: Option<KimiAuth>) -> Option<EarlyPre
Some(spawn_prefetch_thread(env))
}
/// Start model + settings prefetch on a background thread.
/// Start the model-catalog prefetch on a background thread.
///
/// Convenience wrapper that reads cached auth from disk. Prefer
/// `start_early_prefetch_with_auth` when you have pre-resolved credentials.
@@ -1805,7 +1771,7 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle {
let mut timer = crate::instrumentation_timer!("startup.early_prefetch");
let proxy_endpoint = env.endpoints.proxy_url();
timer.with_field("endpoint", proxy_endpoint.as_str());
let (models, settings) = prefetch_models_and_settings_blocking(
let models = prefetch_models_blocking(
&env.endpoints,
env.auth.as_ref(),
env.model_fetch_auth,
@@ -1824,7 +1790,7 @@ fn spawn_prefetch_thread(env: PrefetchEnv) -> EarlyPrefetchHandle {
let _ = rt.block_on(crate::managed_config::sync());
}
EarlyPrefetchResult { models, settings }
EarlyPrefetchResult { models }
})
}
@@ -3868,7 +3834,7 @@ mod tests {
fn proxied_endpoints(server_uri: &str) -> config::EndpointsConfig {
config::EndpointsConfig {
cli_chat_proxy_base_url: Some(server_uri.to_string()),
coding_api_base_url: Some(server_uri.to_string()),
models_base_url: None,
models_list_url: None,
..config::EndpointsConfig::default()
@@ -3900,7 +3866,7 @@ mod tests {
..KimiAuth::test_default()
};
let result = tokio::task::spawn_blocking(move || {
crate::remote::fetch_models_blocking(
crate::agent::models_fetch::fetch_models_blocking(
&endpoints,
Some(&auth),
ModelFetchAuth::Platforms,
@@ -3969,7 +3935,12 @@ mod tests {
let endpoints = config::EndpointsConfig::default();
let keys = PlatformApiKeys::test_keys(Some("sk-cn-secret"), None);
let result = tokio::task::spawn_blocking(move || {
crate::remote::fetch_models_blocking(&endpoints, None, ModelFetchAuth::Platforms, &keys)
crate::agent::models_fetch::fetch_models_blocking(
&endpoints,
None,
ModelFetchAuth::Platforms,
&keys,
)
})
.await
.unwrap()
@@ -4062,7 +4033,7 @@ mod tests {
auth_manager.set_refresher(Arc::new(SwapRefresher));
let mut cfg = config::Config::default();
cfg.endpoints.cli_chat_proxy_base_url = Some(server.uri());
cfg.endpoints.coding_api_base_url = Some(server.uri());
let mgr = ModelsManager::new(
None,
IndexMap::new(),
@@ -4120,8 +4091,12 @@ mod tests {
assert!(bundled.contains_key("moonshot-ai/kimi-k2-turbo-preview"));
// 2. A STALE cache for the same fetch plan is served on sync failure.
let origin =
crate::remote::models_fetch_origin(&endpoints, ModelFetchAuth::Platforms, true, &keys);
let origin = crate::agent::models_fetch::models_fetch_origin(
&endpoints,
ModelFetchAuth::Platforms,
true,
&keys,
);
let cache = ModelsCacheManager::new();
let stale = ModelsCache {
fetched_at: Utc::now() - ChronoDuration::seconds(86_400),
@@ -4203,7 +4178,7 @@ mod tests {
let _cache = EnvGuard::set("KIGI_MODELS_CACHE_DIR", cache_dir.path().to_str().unwrap());
let endpoints = proxied_endpoints("http://127.0.0.1:9");
// Cache written when a moonshot key was ALSO configured...
let with_key_origin = crate::remote::models_fetch_origin(
let with_key_origin = crate::agent::models_fetch::models_fetch_origin(
&endpoints,
ModelFetchAuth::Platforms,
true,
crates/codegen/kigi-shell/src/agent/models_fetch.rs
+1084
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/agent/mvp_agent/acp_agent.rs
+8 -244
View File
@@ -48,7 +48,6 @@ impl acp::Agent for MvpAgent {
);
});
kigi_workspace::trust::migrate_legacy_hook_trust();
self.maybe_sync_bundle_in_background(false);
let mut client_type = arguments
.meta
.as_ref()
@@ -278,11 +277,7 @@ impl acp::Agent for MvpAgent {
}
self.spawn_initialize_launch_mcp_setup(fetch_managed_mcps);
self.spawn_managed_gateway_tool_catalog_fetch();
let init_model_state = if crate::agent::chat_modes::process_chat_mode_enabled() {
self.chat_modes.model_state().await
} else {
self.model_state(None)
};
let init_model_state = self.model_state(None);
Ok(
acp::InitializeResponse::new(acp::ProtocolVersion::V1)
.agent_capabilities(
@@ -374,9 +369,6 @@ impl acp::Agent for MvpAgent {
}
}
self.set_auth_method(arguments.method_id.clone());
if crate::agent::chat_modes::process_chat_mode_enabled() {
self.chat_modes.warm_in_background();
}
emit_login_span(true, "api_key", None, None);
Ok(Default::default())
}
@@ -421,10 +413,8 @@ impl acp::Agent for MvpAgent {
.authenticate_after_cached_token_unavailable(arguments)
.await;
};
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.maybe_sync_bundle_in_background(false);
{
{
let mut sampling_config = self.sampling_config.borrow_mut();
sampling_config.api_key = Some(auth.key);
tracing::debug!(
@@ -437,12 +427,8 @@ impl acp::Agent for MvpAgent {
);
}
self.set_auth_method(arguments.method_id.clone());
if crate::agent::chat_modes::process_chat_mode_enabled() {
self.chat_modes.warm_in_background();
}
let uid = self.auth_manager.current().map(|a| a.user_id);
emit_login_span(true, "cached_token", uid.as_deref(), None);
self.maybe_fetch_post_auth_settings().await;
Ok(self.auth_response_with_meta())
}
auth_method::KIGI_COM_METHOD_ID => {
@@ -517,21 +503,15 @@ impl acp::Agent for MvpAgent {
);
}
self.auth_manager.hot_swap(auth.clone());
self.refresh_remote_settings(&auth).await;
self.emit_settings_update_notification();
self.maybe_sync_bundle_in_background(false);
self.set_auth_method(arguments.method_id.clone());
self.set_auth_method(arguments.method_id.clone());
self.models_manager.on_auth_changed().await;
if crate::agent::chat_modes::process_chat_mode_enabled() {
self.chat_modes.warm_in_background();
}
emit_login_span(
true,
arguments.method_id.0.as_ref(),
Some(auth.user_id.as_str()),
None,
);
self.maybe_fetch_post_auth_settings().await;
Ok(self.auth_response_with_meta())
}
_ => {
@@ -559,9 +539,7 @@ impl acp::Agent for MvpAgent {
.data("initialize must be called before new_session")
})?;
self.seed_client_config_auth_if_available();
if let Ok(auth) = self.auth_manager.auth().await {
self.refresh_settings_and_reapply(&auth).await;
}
self.refresh_settings_and_reapply().await;
let cwd = AbsPathBuf::new(arguments.cwd.clone())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let remote_settings = self.cfg.borrow().remote_settings.clone();
@@ -858,8 +836,10 @@ impl acp::Agent for MvpAgent {
Some(serde_json::json!({ "cwd" : cwd.as_str() })),
);
let models = if is_chat_kind {
// The grok.com chat-mode model picker was removed with the xAI
// proxy; a chat-kind session has no managed catalog to offer.
chat_new_session_model_state(
self.chat_modes.model_state().await,
acp::SessionModelState::new(acp::ModelId::from(String::new()), Vec::new()),
session_initial_model
.filter(|_| matches!(bridge_attach, BridgeAttach::Spawned)),
)
@@ -982,14 +962,6 @@ impl acp::Agent for MvpAgent {
.build_summary_client(&load_session_sampling)?;
let mut persistence_timer = crate::instrumentation_timer!("session.load_light");
persistence_timer.with_field("session_id", session_id.0.as_ref());
let backend = if self.build_registry_config().is_some() {
Some(
crate::remote::BackendClient::new()
.with_auth_manager(self.auth_manager.clone()),
)
} else {
None
};
let registry_title_sync = self
.session_registry_client()
.map(|client| crate::session::persistence::RegistryGeneratedTitleSync {
@@ -999,9 +971,6 @@ impl acp::Agent for MvpAgent {
let (persistence_info, persistence) = crate::session::persistence::load_light(
&session_info,
summary_client,
self.storage_mode,
Some(self.auth_manager.clone()),
backend.as_ref(),
Some(self.gateway.clone()),
summary_model,
registry_title_sync,
@@ -2095,9 +2064,6 @@ impl acp::Agent for MvpAgent {
| "x.ai/sessions/list" => {
crate::agent::handlers::session::handle(self, &args).await
}
"x.ai/workspaces/list" => {
crate::agent::handlers::workspaces::handle(self, &args).await
}
"x.ai/session/updates" => {
crate::extensions::session_updates::handle(&args, &self.gateway).await
}
@@ -2122,6 +2088,7 @@ impl acp::Agent for MvpAgent {
crate::extensions::session_admin::handle(self, &args).await
}
"x.ai/session/repair" => crate::extensions::repair::handle(self, &args).await,
"x.ai/billing" => crate::extensions::billing::handle(self, &args).await,
"x.ai/memory/flush" | "x.ai/memory/rewrite" => {
crate::extensions::memory::handle(self, &args).await
}
@@ -2136,206 +2103,6 @@ impl acp::Agent for MvpAgent {
crate::extensions::feedback::handle(self, &args).await
}
"x.ai/recap" => crate::extensions::recap::handle(self, &args).await,
"x.ai/cloud/terminate" => {
crate::extensions::auth_gate::require_xai_auth(
&self.auth_manager,
"Authentication required",
"Run `grok login` to authenticate.",
)?;
let params: serde_json::Value = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let sandbox_id = params
.get("sandbox_id")
.and_then(|v| v.as_str())
.ok_or_else(|| {
acp::Error::invalid_params().data("missing sandbox_id")
})?;
let sandbox_client = crate::remote::SandboxClient::new(
self.cli_chat_proxy_base_url(),
self.auth_manager.clone(),
);
sandbox_client
.terminate_session(
sandbox_id,
&crate::remote::SandboxTerminateRequest {
environment_id: None,
},
)
.await
.map_err(|e| {
acp::Error::internal_error()
.data(format!("Failed to terminate sandbox: {e}"))
})?;
crate::extensions::to_raw_response(&serde_json::json!({ "ok" : true }))
}
"x.ai/cloud/env/list" => {
crate::extensions::auth_gate::require_xai_auth(
&self.auth_manager,
"Authentication required",
"Run `grok login` to authenticate.",
)?;
let sandbox_client = crate::remote::SandboxClient::new(
self.cli_chat_proxy_base_url(),
self.auth_manager.clone(),
);
let resp = sandbox_client
.list_environments(
&crate::remote::SandboxListEnvironmentsRequest::default(),
)
.await
.map_err(|e| {
acp::Error::internal_error()
.data(format!("Failed to list environments: {e}"))
})?;
crate::extensions::to_raw_response(
&serde_json::json!({ "environments" : resp.environments, }),
)
}
"x.ai/cloud/env/create" => {
crate::extensions::auth_gate::require_xai_auth(
&self.auth_manager,
"Authentication required",
"Run `grok login` to authenticate.",
)?;
let params: serde_json::Value = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let sandbox_client = crate::remote::SandboxClient::new(
self.cli_chat_proxy_base_url(),
self.auth_manager.clone(),
);
let resp = sandbox_client
.create_environment(
&crate::remote::SandboxCreateEnvironmentRequest {
name: params
.get("name")
.and_then(|v| v.as_str())
.map(String::from),
description: params
.get("description")
.and_then(|v| v.as_str())
.map(String::from),
repository: params
.get("repository")
.and_then(|v| v.as_str())
.map(String::from),
default_branch: params
.get("default_branch")
.and_then(|v| v.as_str())
.map(String::from),
container_image: params
.get("container_image")
.and_then(|v| v.as_str())
.map(String::from),
setup_script: params
.get("setup_script")
.and_then(|v| v.as_str())
.map(String::from),
workspace_directory: Some("/workspace".to_string()),
internet_enabled: Some(true),
domain_allowlist_preset: Some("common".to_string()),
allowed_http_methods: Some("all".to_string()),
..Default::default()
},
)
.await
.map_err(|e| {
acp::Error::internal_error()
.data(format!("Failed to create environment: {e}"))
})?;
crate::extensions::to_raw_response(
&serde_json::json!({ "environment" : resp.environment, }),
)
}
"x.ai/cloud/env/update" => {
crate::extensions::auth_gate::require_xai_auth(
&self.auth_manager,
"Authentication required",
"Run `grok login` to authenticate.",
)?;
let params: serde_json::Value = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let environment_id = params
.get("environment_id")
.and_then(|v| v.as_str())
.ok_or_else(|| {
acp::Error::invalid_params().data("missing environment_id")
})?;
let sandbox_client = crate::remote::SandboxClient::new(
self.cli_chat_proxy_base_url(),
self.auth_manager.clone(),
);
let resp = sandbox_client
.update_environment(
environment_id,
&crate::remote::SandboxUpdateEnvironmentRequest {
name: params
.get("name")
.and_then(|v| v.as_str())
.map(String::from),
description: params
.get("description")
.and_then(|v| v.as_str())
.map(String::from),
repository: params
.get("repository")
.and_then(|v| v.as_str())
.map(String::from),
default_branch: params
.get("default_branch")
.and_then(|v| v.as_str())
.map(String::from),
container_image: params
.get("container_image")
.and_then(|v| v.as_str())
.map(String::from),
setup_script: params
.get("setup_script")
.and_then(|v| v.as_str())
.map(String::from),
..Default::default()
},
)
.await
.map_err(|e| {
acp::Error::internal_error()
.data(format!("Failed to update environment: {e}"))
})?;
crate::extensions::to_raw_response(
&serde_json::json!({ "environment" : resp.environment, }),
)
}
"x.ai/cloud/env/delete" => {
crate::extensions::auth_gate::require_xai_auth(
&self.auth_manager,
"Authentication required",
"Run `grok login` to authenticate.",
)?;
let params: serde_json::Value = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(e.to_string()))?;
let environment_id = params
.get("environment_id")
.and_then(|v| v.as_str())
.ok_or_else(|| {
acp::Error::invalid_params().data("missing environment_id")
})?;
let sandbox_client = crate::remote::SandboxClient::new(
self.cli_chat_proxy_base_url(),
self.auth_manager.clone(),
);
sandbox_client
.delete_environment(environment_id)
.await
.map_err(|e| {
acp::Error::internal_error()
.data(format!("Failed to delete environment: {e}"))
})?;
crate::extensions::to_raw_response(&serde_json::json!({ "ok" : true }))
}
"x.ai/billing" => crate::extensions::billing::handle(self, &args).await,
"x.ai/auto-topup-rule" => {
crate::extensions::billing::handle(self, &args).await
}
"x.ai/share_session" => crate::extensions::share::handle(self, &args).await,
"x.ai/rollout/survey" => {
crate::extensions::rollout::handle(self, &args).await
}
@@ -2395,9 +2162,6 @@ impl acp::Agent for MvpAgent {
s if s.starts_with("x.ai/search/") => {
crate::extensions::search::handle(self, &args).await
}
s if s.starts_with("x.ai/bundle/") => {
crate::extensions::bundle::handle(self, &args).await
}
s if s.starts_with("x.ai/code/") => {
let ops = self.resolve_workspace_ops()?;
crate::extensions::code_nav::handle(self, &ops, &args).await
crates/codegen/kigi-shell/src/agent/mvp_agent/agent_ops.rs
+27 -201
View File
@@ -28,23 +28,15 @@ impl MvpAgent {
let session_key = self.auth_manager.current_or_expired().map(|a| a.key.clone());
let models = self.models_manager.models();
let endpoints = self.models_manager.endpoints();
let (alpha_test_key, client_version) = {
let cfg = self.cfg.borrow();
(
cfg.endpoints.alpha_test_key.clone(),
cfg.client_version.clone(),
)
};
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
let config = match crate::agent::config::resolve_aux_model_sampling_config(
&slug,
&models,
&endpoints,
session_key.as_deref(),
alpha_test_key,
client_version,
) {
Some(mut cfg) => {
cfg.client_identifier = primary.client_identifier.clone();
cfg.attribution_callback = primary.attribution_callback.clone();
cfg.bearer_resolver = primary.bearer_resolver.clone();
cfg.max_retries = primary.max_retries;
@@ -60,10 +52,6 @@ impl MvpAgent {
let client = OaiCompatClient::new(config).map_err(map_sampling_err_to_acp)?;
Ok((client, model))
}
fn has_proxy_credentials(&self) -> bool {
self.cfg.borrow().endpoints.deployment_key.is_some()
|| self.auth_manager.current_or_expired().is_some_and(|a| a.is_session_auth())
}
/// `true` for session-based ACP auth methods.
fn is_session_based_auth(&self) -> bool {
self.auth_method_id
@@ -384,39 +372,23 @@ impl MvpAgent {
);
}
}
/// Extract feedback credentials when proxy credentials are available.
///
/// Returns `(base_url, user_token, optional_extra_access_key, deployment_key)`.
/// Used by both [`feedback_client`] and session spawning to avoid
/// duplicating the credential assembly logic.
#[allow(clippy::type_complexity)]
fn feedback_credentials(
&self,
) -> Option<(String, Option<String>, Option<String>, Option<String>)> {
if !self.has_proxy_credentials() {
return None;
}
let user_token = self
/// Feedback endpoint base when this is a subscription (OAuth) session —
/// the Kimi Code feedback endpoint only takes the OAuth Bearer, so
/// API-key-only setups get `None` (they are pointed at the issue
/// tracker instead; kimi-cli slash.py parity).
fn feedback_base_url(&self) -> Option<String> {
let has_session = self
.auth_manager
.current_or_expired()
.filter(|a| a.is_session_auth())
.map(|a| a.key.clone());
let cfg = self.cfg.borrow();
let base_url = cfg.endpoints.resolve_feedback_base_url();
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let deployment_key = cfg.endpoints.deployment_key.clone();
Some((base_url, user_token, alpha_test_key, deployment_key))
.is_some_and(|a| a.is_session_auth());
has_session.then(|| self.cfg.borrow().endpoints.resolve_feedback_base_url())
}
/// Build a `FeedbackClient` with resolved feedback URL and credentials.
/// Build a `FeedbackClient` for subscription sessions.
pub(crate) fn feedback_client(&self) -> Option<FeedbackClient> {
let (base_url, user_token, alpha_test_key, deployment_key) = self
.feedback_credentials()?;
Some(
FeedbackClient::new(base_url, user_token)
.with_alpha_test_key(alpha_test_key)
.with_deployment_key(deployment_key)
.with_auth_manager(self.auth_manager.clone()),
)
Some(FeedbackClient::new(
self.feedback_base_url()?,
self.auth_manager.clone(),
))
}
/// Build a `RegistryConfig` if the feature is enabled (for passing to persistence actor).
pub(super) fn build_registry_config(
@@ -460,17 +432,6 @@ impl MvpAgent {
.with_auth(self.auth_manager.clone()),
)
}
pub(crate) fn conversations_client(
&self,
) -> Option<crate::remote::ConversationsClient> {
if !crate::session::unified_list::conversations_lane_active() {
return None;
}
Some(crate::remote::ConversationsClient::new(self.auth_manager.clone()))
}
pub(crate) fn workspaces_client(&self) -> crate::remote::WorkspacesClient {
crate::remote::WorkspacesClient::new(self.auth_manager.clone())
}
/// Pre-session command availability snapshot.
///
/// Used by the `x.ai/commands/list` ext method and the
@@ -515,13 +476,9 @@ impl MvpAgent {
) -> &kigi_agent::plugins::SharedPluginRegistryHandle {
&self.plugin_registry_handle
}
/// `true` when the agent runs in writeback storage mode.
pub(crate) fn is_writeback_storage(&self) -> bool {
matches!(self.storage_mode, StorageMode::Writeback)
}
/// Resolved cli-chat-proxy base for session features (via
/// `proxy_url`). Not for the deployment-config fetch.
pub(crate) fn cli_chat_proxy_base_url(&self) -> String {
pub(crate) fn coding_api_base_url(&self) -> String {
self.cfg.borrow().endpoints.proxy_url()
}
pub(crate) fn alpha_test_key(&self) -> Option<String> {
@@ -635,54 +592,14 @@ impl MvpAgent {
pub(crate) fn deployment_key(&self) -> Option<String> {
self.cfg.borrow().endpoints.deployment_key.clone()
}
/// Re-fetch remote settings and re-init the telemetry client.
///
/// Called unconditionally from both auth handlers so that:
/// - First install / expired OIDC token: settings are fetched for
/// the first time (the early prefetch had no auth to use).
/// - Reauth / account switch: settings are refreshed to reflect
/// the new user's remote settings targeting attributes.
///
/// This only refreshes `cfg.remote_settings` and re-inits the
/// telemetry client (the only global static). Other settings
/// derived from `remote_settings` (`web_fetch_enabled`, etc.) are
/// resolved lazily per-turn from `cfg` and pick up the new values
/// automatically.
/// Agent-level fields materialised at startup (`worktree_type`,
/// `restore_code`) are NOT re-resolved here; that requires a
/// broader refactor of the init path.
pub(super) async fn refresh_remote_settings(&self, auth: &crate::auth::KimiAuth) {
if !crate::util::config::resolve_remote_fetch_enabled() {
tracing::debug!("post-auth settings refresh skipped: remote_fetch disabled");
return;
}
let Some(settings) = self.fetch_remote_settings(auth.clone()).await else {
tracing::warn!("post-auth settings refresh failed (HTTP or parse error)");
return;
};
tracing::info!("post-auth settings refreshed");
{
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::util::config::sync_campaign_fields(&mut cfg);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
);
}
}
/// Refresh remote settings settings and re-resolve eagerly-resolved config fields.
/// Re-resolve eagerly-resolved config fields from the local config.
///
/// Called on `/new` session creation so feature flags reflect the latest
/// remote settings state without requiring a TUI restart. Extends
/// [`refresh_remote_settings`] by also re-running [`resolve_runtime_fields`]
/// with the fresh settings.
/// on-disk config without requiring a TUI restart. (Formerly this also
/// re-fetched the xAI proxy's remote settings; that endpoint is gone.)
///
/// In-flight sessions are unaffected — they snapshot config at creation.
pub(super) async fn refresh_settings_and_reapply(
&self,
auth: &crate::auth::KimiAuth,
) {
self.refresh_remote_settings(auth).await;
pub(super) async fn refresh_settings_and_reapply(&self) {
let cwd = std::env::current_dir().ok();
{
let mut cfg = self.cfg.borrow_mut();
@@ -698,36 +615,6 @@ impl MvpAgent {
}
self.emit_settings_update_notification();
}
/// Shared fetch half of every settings refresh: endpoint fields from a
/// scoped `cfg` borrow, `fetch_settings_blocking` off-executor (it already
/// retries transient errors internally), failures normalized to `None`.
/// Callers own their miss logging.
pub(super) async fn fetch_remote_settings(
&self,
auth: crate::auth::KimiAuth,
) -> Option<crate::util::config::RemoteSettings> {
if !crate::util::config::resolve_remote_fetch_enabled() {
tracing::debug!("settings fetch skipped: remote_fetch disabled");
return None;
}
let (base_url, alpha_test_key) = {
let cfg = self.cfg.borrow();
(cfg.endpoints.proxy_url(), cfg.endpoints.alpha_test_key.clone())
};
match tokio::task::spawn_blocking(move || crate::remote::fetch_settings_blocking(
&base_url,
&auth,
alpha_test_key.as_deref(),
))
.await
{
Ok(settings) => settings,
Err(e) => {
tracing::warn!(error = % e, "settings fetch task panicked");
None
}
}
}
pub(super) async fn send_model_auto_switched(
&self,
session_id: &acp::SessionId,
@@ -828,26 +715,9 @@ impl MvpAgent {
),
);
}
let cfg = self.cfg.borrow();
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let client_version = cfg.client_version.clone();
let deployment_id = crate::managed_config::resolve_deployment_id(
cfg.endpoints.deployment_key.as_deref(),
);
drop(cfg);
let user_id = self
.auth_manager
.current_or_expired()
.filter(|a| a.is_session_auth())
.map(|a| a.user_id);
let mut config = crate::agent::config::sampling_config_for_model(
model,
credentials,
alpha_test_key,
client_version,
deployment_id,
user_id,
);
let alpha_test_key = self.cfg.borrow().endpoints.alpha_test_key.clone();
let mut config =
crate::agent::config::sampling_config_for_model(model, credentials, alpha_test_key);
config.origin_client = origin_client;
config
}
@@ -912,13 +782,7 @@ impl MvpAgent {
.unwrap_or_else(|| kigi_version::VERSION.to_string());
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let mut headers = indexmap::IndexMap::new();
headers.insert("user-agent".to_string(), format!("xai-grok-build/{version}"));
inject_proxy_headers(
&mut headers,
cfg.client_version.as_deref(),
alpha_test_key.as_deref(),
&base_url,
);
headers.insert("user-agent".to_string(), format!("kigi/{version}"));
ImageGenConfig::Enabled {
api_key: api_key.clone(),
base_url,
@@ -961,13 +825,7 @@ impl MvpAgent {
.unwrap_or_else(|| kigi_version::VERSION.to_string());
let alpha_test_key = cfg.endpoints.alpha_test_key.clone();
let mut headers = indexmap::IndexMap::new();
headers.insert("user-agent".to_string(), format!("xai-grok-build/{version}"));
inject_proxy_headers(
&mut headers,
cfg.client_version.as_deref(),
alpha_test_key.as_deref(),
&base_url,
);
headers.insert("user-agent".to_string(), format!("kigi/{version}"));
VideoGenConfig::Enabled {
api_key,
base_url,
@@ -987,15 +845,8 @@ impl MvpAgent {
&models,
session.as_ref().map(|a| a.key.as_str()),
alpha_test_key.clone(),
client_version,
&self.cfg.borrow().endpoints,
)?;
inject_proxy_headers(
&mut cfg.extra_headers,
cfg.client_version.as_deref(),
alpha_test_key.as_deref(),
&cfg.base_url,
);
Some(cfg)
}
/// Returns `Err` with a user-facing message on invalid config; the caller at
@@ -1112,15 +963,6 @@ impl MvpAgent {
.map(|(name, p)| p.render_io_summary(name))
.collect(),
models_manager,
chat_modes: {
let chat_modes = crate::agent::chat_modes::ChatModesManager::new(
auth_manager.clone(),
);
if crate::agent::chat_modes::process_chat_mode_enabled() {
chat_modes.warm_in_background();
}
chat_modes
},
cfg: RefCell::new(cfg.clone()),
auth_method_id: crate::agent::auth_method::new_shared_auth_method_id(None),
sampling_config: RefCell::new(sampling_config),
@@ -1159,7 +1001,6 @@ impl MvpAgent {
subagent_event_rx: RefCell::new(Some(subagent_event_rx)),
subagent_coordinator: RefCell::new(subagent_coordinator),
monitor_event_buffer: kigi_tools::implementations::grok_build::task::types::MonitorEventBuffer::default(),
bundle_sync_in_flight: Arc::new(std::sync::atomic::AtomicBool::new(false)),
workspace_ops: RefCell::new(None),
require_gateway_sessions: Rc::new(
RefCell::new(std::collections::HashSet::new()),
@@ -2221,19 +2062,9 @@ impl MvpAgent {
let auto_update = self.cfg.borrow().cli.auto_update;
let client_type = *self.client_type.borrow();
let buffering_settings = self.buffering_settings.borrow().clone();
let (
feedback_proxy_url,
feedback_user_token,
feedback_alpha_test_key,
deployment_key,
) = if let Some((url, token, alpha, deploy)) = self.feedback_credentials() {
(Some(url), token, alpha, deploy)
} else {
(None, None, None, None)
};
let feedback_base_url = self.feedback_base_url();
tracing::info!(
session_id = % session_info.id.0, feedback_url = ? feedback_proxy_url,
authenticated = feedback_user_token.is_some(),
session_id = % session_info.id.0, feedback_url = ? feedback_base_url,
"Initializing feedback manager for session"
);
let skills = self.cfg.borrow().skills.clone();
@@ -2489,7 +2320,6 @@ impl MvpAgent {
self.auth_type(),
),
alpha_test_key: self.alpha_test_key(),
client_version: sampling_config.client_version.clone(),
};
let attribution_callback: Option<
kigi_sampler::SharedAttributionCallback,
@@ -2599,10 +2429,7 @@ impl MvpAgent {
self.codebase_indexes.clone(),
client_code_nav_enabled,
fs_watch_caps,
feedback_proxy_url,
feedback_user_token,
feedback_alpha_test_key,
deployment_key,
feedback_base_url,
client_terminal,
client_fs_read && client_fs_write,
gateway_enabled,
@@ -2617,7 +2444,6 @@ impl MvpAgent {
persisted_goal_mode,
persisted_announcement_state,
self.memory_config.clone(),
loc_tracking_enabled,
feedback_flags,
self.managed_mcp_cache.clone(),
managed_mcp_expires_at,
crates/codegen/kigi-shell/src/agent/mvp_agent/mod.rs
-188
View File
@@ -406,17 +406,11 @@ struct SettingsUpdateNotification {
sharing_enabled: Option<bool>,
session_picker_grouped: Option<bool>,
tips: Option<Vec<String>>,
gate_message: Option<String>,
gate_url: Option<String>,
gate_label: Option<String>,
allow_access: Option<bool>,
subscription_tier_display: Option<String>,
auto_permission_mode_enabled: Option<bool>,
/// Soft-default permission mode for the pager (post-auth / `/new` refresh).
permission_mode: Option<String>,
group_tool_verbs: Option<bool>,
collapsed_edit_blocks: Option<bool>,
subscription_watch_interval_secs: Option<u64>,
}
/// Reason why a client is not eligible to use codebase indexing.
///
@@ -509,9 +503,6 @@ pub struct MvpAgent {
pub(crate) sampling_config: RefCell<SamplingConfig>,
pub(crate) auth_manager: Arc<AuthManager>,
pub(crate) models_manager: crate::agent::models::ModelsManager,
/// grok.com chat-product catalog (`/rest/modes`) for chat sessions; distinct
/// from `models_manager` (the build `/v1/models` catalog).
pub(crate) chat_modes: crate::agent::chat_modes::ChatModesManager,
/// Forwards pasted codes from `handle_auth_submit_code` to the auth flow.
pub(crate) auth_code_tx: RefCell<Option<tokio::sync::mpsc::Sender<String>>>,
/// Receives the auth URL from the auth flow; read by `handle_auth_get_url`.
@@ -672,20 +663,6 @@ pub struct MvpAgent {
/// this flag keeps that to a single discovery walk.
plugin_registry_initialized: std::cell::Cell<bool>,
persona_io_summaries: Vec<String>,
/// Single-flight guard for the proactive bundle sync background task.
///
/// `maybe_sync_bundle_in_background` is invoked from each post-auth path
/// (initialize, cached-token reauth, oidc) and a rapid reconnect can fire
/// all three within the TTL window, giving us multiple concurrent
/// `tokio::task::spawn_local` tasks racing to extract the tar archive,
/// rewrite `manifest.json`, and prune stale files. The non-atomic
/// per-file write/prune semantics in `bundle::extract_bundle_archive`
/// make that race observable as a partially-written cache.
///
/// We use an `Arc<AtomicBool>` so the spawned task can clear the flag
/// on completion without re-borrowing `&self`. `Send` is required
/// because the inner `sync_bundle_to_root` now uses `spawn_blocking`.
bundle_sync_in_flight: Arc<std::sync::atomic::AtomicBool>,
/// Local workspace ops, built lazily via [`Self::ensure_local_workspace_ops`].
/// The agent never opens Computer Hub as a harness/client; remote cloud
/// sandboxes are gateway-owned (`gateway_bridge` / `computer_sessions`).
@@ -944,50 +921,6 @@ impl AuthRequestMeta {
.unwrap_or_default()
}
}
/// Inject standard proxy headers into an `extra_headers` map.
///
/// Every authenticated request to cli-chat-proxy (web search, image gen, and
/// any future tools that go through the proxy) must carry these headers.
/// Centralising them here means new tool code paths only need one call instead
/// of remembering which headers the proxy expects.
///
/// Headers injected:
/// - `x-grok-client-version` -- required by the proxy's version-gate check.
/// Uses `client_version` when provided, otherwise falls back to cli-chat-proxy
/// compile-time `CARGO_PKG_VERSION`.
/// - `X-XAI-Token-Auth` / `x-authenticateresponse` -- required by the
/// cli-chat-proxy auth middleware when the `base_url` is a known proxy URL.
/// - optional extra access header -- only set when the corresponding key is
/// `Some` *and* the `base_url` points at a matching non-production host
/// (requires the optional non-production feature).
///
/// Existing entries are never overwritten so callers can pre-set a value.
fn inject_proxy_headers(
headers: &mut indexmap::IndexMap<String, String>,
client_version: Option<&str>,
alpha_test_key: Option<&str>,
base_url: &str,
) {
headers
.entry("x-grok-client-version".to_string())
.or_insert_with(|| {
client_version
.map(String::from)
.unwrap_or_else(|| kigi_version::VERSION.to_string())
});
if crate::util::is_cli_chat_proxy_url(base_url) {
headers
.entry("X-XAI-Token-Auth".to_string())
.or_insert_with(|| "xai-grok-cli".to_string());
headers
.entry("x-authenticateresponse".to_string())
.or_insert_with(|| "authenticate-response".to_string());
headers
.entry(crate::http::CLIENT_MODE_HEADER.to_string())
.or_insert_with(|| crate::http::process_client_mode().to_string());
}
let _ = (alpha_test_key, base_url);
}
fn resolve_inference_idle_timeout_secs(
models: &indexmap::IndexMap<String, crate::agent::config::ModelEntry>,
model: &str,
@@ -1580,47 +1513,6 @@ impl MvpAgent {
});
AuthenticateResponse::new().meta(meta)
}
/// Fetch remote settings after authentication when early prefetch had none.
/// Notifies the pager so soft-default permission_mode applies post-login.
pub(super) async fn maybe_fetch_post_auth_settings(&self) {
if self.cfg.borrow().remote_settings.is_some() {
return;
}
let Some(auth) = self.auth_manager.current() else {
return;
};
let is_session_auth = auth.is_session_auth();
let Some(settings) = self.fetch_remote_settings(auth).await else {
return;
};
tracing::info!("post-auth remote_settings fetch succeeded");
{
let mut cfg = self.cfg.borrow_mut();
cfg.remote_settings = Some(settings);
crate::agent::config::apply_remote_settings_side_effects(
cfg.remote_settings.as_ref(),
);
if cfg.storage_mode == StorageMode::Local
&& cfg.mode != crate::agent::config::AgentMode::Generic
{
cfg.storage_mode = StorageMode::resolve(
None,
cfg.remote_settings.as_ref(),
);
if cfg.storage_mode == StorageMode::Writeback && !is_session_auth {
cfg.storage_mode = StorageMode::Local;
}
}
if let Some(v) = cfg
.remote_settings
.as_ref()
.and_then(|s| s.path_not_found_hints)
{
cfg.path_not_found_hints = v;
}
}
self.emit_settings_update_notification();
}
/// Fire-and-forget `x.ai/settings/update` from the current remote snapshot.
pub(super) fn emit_settings_update_notification(&self) {
let payload = {
@@ -1631,20 +1523,12 @@ impl MvpAgent {
sharing_enabled: rs.and_then(|s| s.sharing_enabled),
session_picker_grouped: rs.and_then(|s| s.session_picker_grouped),
tips: rs.and_then(|s| s.tips.clone()),
gate_message: rs.and_then(|s| s.gate_message.clone()),
gate_url: rs.and_then(|s| s.gate_url.clone()),
gate_label: rs.and_then(|s| s.gate_label.clone()),
allow_access: rs.and_then(|s| s.allow_access),
subscription_tier_display: rs
.and_then(|s| s.subscription_tier_display.clone()),
auto_permission_mode_enabled: crate::util::config::remote_auto_mode_enabled(
rs,
),
permission_mode: rs.and_then(|s| s.permission_mode.clone()),
group_tool_verbs: rs.and_then(|s| s.group_tool_verbs),
collapsed_edit_blocks: rs.and_then(|s| s.collapsed_edit_blocks),
subscription_watch_interval_secs: rs
.and_then(|s| s.subscription_watch_interval_secs),
}
};
if let Ok(params) = serde_json::value::to_raw_value(&payload) {
@@ -1719,78 +1603,6 @@ impl MvpAgent {
});
}
}
/// Spawn a best-effort bundle sync. Re-fires on every call site (init,
/// cached_token, grok.com/oidc); the cheap pre-checks below absorb repeats
/// so reconnects are cheap.
///
/// Pre-spawn gating order (cheapest first, all synchronous):
/// 1. Auth gate — avoid spawning a no-op task on every init.
/// 2. Freshness check — skip the sender snapshot + spawn entirely on
/// cache hits, which is the steady-state on every reconnect.
/// 3. Single-flight guard — if a previous sync is still in flight (e.g.,
/// initialize + cached_token + oidc fired in quick succession before
/// the first sync's tar extract finished), drop this call to avoid
/// racing concurrent extracts that would interleave per-file writes
/// against `~/.kigi/bundled/` and the manifest.
pub(crate) fn maybe_sync_bundle_in_background(&self, force: bool) {
use crate::extensions::bundle::{
BUNDLE_SYNC_TTL, bundle_cache_is_fresh, has_bundle_credentials,
maybe_sync_bundle_to_root,
};
use std::sync::atomic::Ordering;
let am = self.auth_manager.clone();
let deployment_key = self.deployment_key();
if !has_bundle_credentials(Some(&am), deployment_key.as_deref()) {
return;
}
let root = crate::bundle::bundled_root();
if !force && bundle_cache_is_fresh(&root, BUNDLE_SYNC_TTL) {
tracing::debug!("proactive bundle sync skipped pre-spawn: cache is fresh");
return;
}
let in_flight = self.bundle_sync_in_flight.clone();
if in_flight
.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
.is_err()
{
tracing::debug!(
"proactive bundle sync skipped: another sync is already in flight"
);
return;
}
let proxy_base_url = self.cli_chat_proxy_base_url();
let alpha_test_key = self.alpha_test_key();
let senders: Vec<
tokio::sync::mpsc::UnboundedSender<crate::session::SessionCommand>,
> = self.sessions.borrow().values().map(|h| h.cmd_tx.clone()).collect();
tokio::task::spawn_local(async move {
let result = maybe_sync_bundle_to_root(
&root,
&proxy_base_url,
Some(&am),
deployment_key.as_deref(),
alpha_test_key.as_deref(),
force,
BUNDLE_SYNC_TTL,
)
.await;
in_flight.store(false, Ordering::Release);
match result {
Ok(Some(res)) => {
tracing::info!(
version = % res.version, personas = res.personas_count, roles =
res.roles_count, agents = res.agents_count, skills = res
.skills_count, "proactive bundle sync complete"
);
Self::broadcast_refresh_skill_baseline(senders);
}
Ok(None) => {}
Err(err) => {
tracing::warn!(error = % err, "proactive bundle sync failed");
}
}
});
}
}
/// Parse `_meta.agentProfile` as a JSON object or string name.
/// Returns `None` if absent or invalid.
crates/codegen/kigi-shell/src/agent/mvp_agent/subagent_coordinator.rs
+1 -1
View File
@@ -402,7 +402,7 @@ impl MvpAgent {
client_hooks: Default::default(),
sampling_config: self.sampling_config.borrow().clone(),
managed_mcp_proxy_base_url: parent_managed_mcp_proxy_base_url
.unwrap_or_else(|| self.cli_chat_proxy_base_url()),
.unwrap_or_else(|| self.coding_api_base_url()),
alpha_test_key: self.alpha_test_key(),
auth_method_id: self
.auth_method_id
crates/codegen/kigi-shell/src/agent/mvp_agent/tests.rs
-12
View File
@@ -1822,18 +1822,6 @@ fn orphaned_tasks_filters_rewind_dead_branches() {
);
}
#[test]
fn allow_access_from_remote_settings() {
let json = serde_json::json!({ "allow_access" : true });
let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap();
assert_eq!(rs.allow_access, Some(true));
let json = serde_json::json!({ "allow_access" : false });
let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap();
assert_eq!(rs.allow_access, Some(false));
let json = serde_json::json!({});
let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap();
assert_eq!(rs.allow_access, None);
}
#[test]
fn on_demand_enabled_from_remote_settings() {
let json = serde_json::json!({ "on_demand_enabled" : false });
let rs: crate::util::config::RemoteSettings = serde_json::from_value(json).unwrap();
crates/codegen/kigi-shell/src/agent/subagent/handle_request.rs
-5
View File
@@ -708,7 +708,6 @@ pub(crate) async fn handle_subagent_request(
api_key: effective_sampling_config.api_key.clone(),
auth_type: inherited_auth_type,
alpha_test_key: ctx.alpha_test_key.clone(),
client_version: effective_sampling_config.client_version.clone(),
};
kigi_log::unified_log::info(
"subagent spawn credentials",
@@ -1020,9 +1019,6 @@ pub(crate) async fn handle_subagent_request(
false,
subagent_fs_watch,
None,
None,
None,
None,
false,
false,
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true)),
@@ -1057,7 +1053,6 @@ pub(crate) async fn handle_subagent_request(
} else {
ctx.memory_config.clone()
},
false,
Default::default(),
ctx.managed_mcp_state.clone(),
None,
crates/codegen/kigi-shell/src/agent/subagent/mod.rs
+1 -12
View File
@@ -898,15 +898,11 @@ async fn read_parent_sampling_config(
auth_scheme,
extra_headers,
context_window: cfg.context_window.get(),
client_version: creds.client_version,
reasoning_effort: cfg.reasoning_effort,
force_http1: false,
max_retries: None,
stream_tool_calls: cfg.stream_tool_calls.unwrap_or(false),
idle_timeout_secs: None,
client_identifier: ctx.sampling_config.client_identifier.clone(),
deployment_id: ctx.sampling_config.deployment_id.clone(),
user_id: ctx.sampling_config.user_id.clone(),
origin_client: ctx.sampling_config.origin_client.clone(),
attribution_callback: ctx.attribution_callback.clone(),
bearer_resolver: None,
@@ -997,14 +993,7 @@ fn resolve_model_override_to_config(
let mut credentials = resolve_credentials(&entry, session_key);
credentials.auth_type = subagent_auth_type(Some(&entry), &ctx.auth_method_id);
let resolved_auth_type = credentials.auth_type;
let config = sampling_config_for_model(
&entry,
credentials,
ctx.alpha_test_key.clone(),
ctx.sampling_config.client_version.clone(),
ctx.sampling_config.deployment_id.clone(),
ctx.sampling_config.user_id.clone(),
);
let config = sampling_config_for_model(&entry, credentials, ctx.alpha_test_key.clone());
kigi_log::unified_log::debug(
"subagent resolve_model_override_to_config",
None,