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/Cargo.toml
-1
View File
@@ -147,7 +147,6 @@ kigi-auth = { workspace = true, features = ["middleware"] }
kigi-log = { workspace = true }
kigi-http = { workspace = true }
kigi-models = { workspace = true }
prod-mc-cli-chat-proxy-types = { path = "../../../prod/mc/cli-chat-proxy-types" }
flate2 = { workspace = true }
fs2 = { workspace = true }
zstd = { workspace = true }
crates/codegen/kigi-shell/benches/session_list.rs
-2
View File
@@ -536,7 +536,6 @@ fn summary_ids(summaries: &[Summary]) -> Vec<String> {
async fn build_local_list_with_delayed_peer(cwd: String) -> UnifiedListResult {
let local = build_unified_list(
None,
None,
ListReq {
cwd: Some(cwd),
@@ -642,7 +641,6 @@ fn bench_session_list(c: &mut Criterion) {
|b| {
b.iter_with_large_drop(|| {
black_box(runtime.block_on(build_unified_list(
None,
None,
ListReq {
cwd: Some(black_box(fixture.picker_cwd.clone())),
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,
crates/codegen/kigi-shell/src/auth/device.rs
+5 -3
View File
@@ -29,12 +29,14 @@ pub(crate) fn ascii_header_value(value: &str) -> String {
}
}
/// The three device-identity headers sent on every OAuth call.
/// The three device-identity headers sent on every OAuth call and, via
/// `agent::config::inject_url_derived_headers`, on every first-party
/// inference request (mirroring kimi-cli src/kimi_cli/llm.py:317-323).
///
/// Errors when the persistent device id cannot be created (e.g. read-only
/// `~/.kigi`): the OAuth endpoints require `X-Msh-Device-Id`, so login cannot
/// proceed without it.
pub(crate) fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> {
/// proceed without it. Inference callers treat the error as skip-with-warning.
pub fn device_headers() -> anyhow::Result<[(&'static str, String); 3]> {
Ok([
("X-Msh-Device-Name", ascii_header_value(&device_name())),
("X-Msh-Device-Model", ascii_header_value(device_model())),
crates/codegen/kigi-shell/src/auth/device_code.rs
+12 -4
View File
@@ -150,6 +150,11 @@ async fn complete_device_code_login(
/// caller can decide how to notify the user (eprintln on CLI, nothing on TUI
/// where the URL is already rendered in the widget).
async fn open_browser_detached(url: &str) -> bool {
// Unit tests drive the full login flow against mock servers — their
// fixture URLs must never reach a real browser.
if cfg!(test) {
return false;
}
let url = url.to_owned();
match tokio::task::spawn_blocking(move || webbrowser::open(&url)).await {
Ok(Ok(())) => true,
@@ -171,13 +176,16 @@ mod tests {
use wiremock::matchers::{body_string_contains, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
/// Fixture mirroring the live `device_authorization` payload (verified
/// against auth.kimi.com): verification URLs are passed through verbatim
/// by the login flow, so they use the real shape.
fn device_auth_json(code: &str) -> serde_json::Value {
serde_json::json!({
"user_code": "ABCD-1234",
"user_code": "WXYZ-6789",
"device_code": code,
"verification_uri": "https://auth.kimi.com/device",
"verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234",
"expires_in": 600,
"verification_uri": "https://www.kimi.com/code/authorize_device",
"verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789",
"expires_in": 1800,
"interval": 0, // floored to 1s by the poll loop
})
}
crates/codegen/kigi-shell/src/auth/kimi_oauth.rs
+9 -9
View File
@@ -359,11 +359,11 @@ mod tests {
"client_id={KIMI_CODE_CLIENT_ID}"
)))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "ABCD-1234",
"user_code": "WXYZ-6789",
"device_code": "dev-code-1",
"verification_uri": "https://auth.kimi.com/device",
"verification_uri_complete": "https://auth.kimi.com/device?code=ABCD-1234",
"expires_in": 600,
"verification_uri": "https://www.kimi.com/code/authorize_device",
"verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789",
"expires_in": 1800,
"interval": 7,
})))
.expect(1)
@@ -371,13 +371,13 @@ mod tests {
.await;
let auth = request_device_authorization(&server.uri()).await.unwrap();
assert_eq!(auth.user_code, "ABCD-1234");
assert_eq!(auth.user_code, "WXYZ-6789");
assert_eq!(auth.device_code, "dev-code-1");
assert_eq!(auth.interval, 7);
assert_eq!(auth.expires_in, Some(600));
assert_eq!(auth.expires_in, Some(1800));
assert_eq!(
auth.verification_uri_complete,
"https://auth.kimi.com/device?code=ABCD-1234"
"https://www.kimi.com/code/authorize_device?user_code=WXYZ-6789"
);
}
@@ -392,7 +392,7 @@ mod tests {
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
"verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=AAAA",
"interval": 5,
})))
.expect(1)
@@ -409,7 +409,7 @@ mod tests {
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"user_code": "AAAA",
"device_code": "d",
"verification_uri_complete": "https://auth.kimi.com/device?code=AAAA",
"verification_uri_complete": "https://www.kimi.com/code/authorize_device?user_code=AAAA",
})))
.mount(&server)
.await;
crates/codegen/kigi-shell/src/auth/meta.rs
-12
View File
@@ -1,17 +1,5 @@
use serde::{Deserialize, Serialize};
/// Access-gate copy resolved from remote settings (message + optional CTA).
/// Auth no longer produces gates (tier gating was an xAI concept); the pager
/// still renders one when remote settings carry a gate message.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GateInfo {
pub message: String,
#[serde(default)]
pub url: Option<String>,
#[serde(default)]
pub label: Option<String>,
}
/// Typed auth metadata passed from the shell to the pager via ACP.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AuthMeta {
crates/codegen/kigi-shell/src/auth/mod.rs
+2 -1
View File
@@ -20,9 +20,10 @@ pub use flow::{
run_auth_flow_with_stderr_bridge, run_cli_login, run_cli_logout, try_ensure_fresh_auth,
};
mod meta;
pub use device::device_headers;
pub use error::{AuthError, RefreshTokenError, RefreshTokenFailedReason};
pub use manager::{AuthManager, shared_api_key_provider};
pub use meta::{AuthMeta, GateInfo};
pub use meta::AuthMeta;
pub use model::{AuthMode, KimiAuth, lookup_auth};
pub(crate) use model::{TOKEN_TTL, is_expired, token_suffix};
pub use storage::{
crates/codegen/kigi-shell/src/bundle.rs
+10 -1502
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/config/mod.rs
+4 -4
View File
@@ -1094,11 +1094,11 @@ fn apply_requirements_inner(
config.endpoints.xai_api_base_url = val.to_owned();
push("endpoints.xai_api_base_url", val.to_owned());
}
if let Some(val) = req_str(req, "endpoints", "cli_chat_proxy_base_url")
&& config.endpoints.cli_chat_proxy_base_url.as_deref() != Some(val)
if let Some(val) = req_str(req, "endpoints", "coding_api_base_url")
&& config.endpoints.coding_api_base_url.as_deref() != Some(val)
{
config.endpoints.cli_chat_proxy_base_url = Some(val.to_owned());
push("endpoints.cli_chat_proxy_base_url", val.to_owned());
config.endpoints.coding_api_base_url = Some(val.to_owned());
push("endpoints.coding_api_base_url", val.to_owned());
}
enforce_str!(
"endpoints",
crates/codegen/kigi-shell/src/config/tests.rs
+5 -5
View File
@@ -2564,7 +2564,7 @@ fn config_layers_user_overrides_managed() {
fn enterprise_two_file_merge_routes_deployment_key_to_proxy() {
for k in [
"KIGI_MANAGED_CONFIG_URL",
"KIGI_CLI_CHAT_PROXY_BASE_URL",
"KIGI_CODE_BASE_URL",
"KIGI_TRACE_UPLOAD_ENDPOINT_URL",
] {
unsafe { std::env::remove_var(k) };
@@ -2573,7 +2573,7 @@ fn enterprise_two_file_merge_routes_deployment_key_to_proxy() {
r#"
[endpoints]
xai_api_base_url = "https://inference.acme-corp.example/xai/v1"
cli_chat_proxy_base_url = "https://cli-chat-proxy.kigi.com/v1"
coding_api_base_url = "https://cli-chat-proxy.kigi.com/v1"
[model.kigi-build]
base_url = "https://inference.acme-corp.example/xai/v1"
@@ -2668,13 +2668,13 @@ fn config_layers_system_managed_lowest_priority() {
#[test]
fn apply_requirements_value_overrides_user_settings() {
let raw_config: toml::Value = toml::from_str(
"[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\nweb_search = \"user-ws-model\"\n\n[endpoints]\ncli_chat_proxy_base_url = \"https://user-proxy.example/v1\"\nxai_api_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n",
"[cli]\nauto_update = true\nchannel = \"beta\"\n\n[features]\nfeedback = true\nlsp_tools = true\nweb_fetch = true\nwrite_file = true\n\n[ui]\nyolo = true\n\n[models]\ndefault = \"user-model\"\nweb_search = \"user-ws-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://user-proxy.example/v1\"\nxai_api_base_url = \"https://user-api.example/v1\"\nmodels_base_url = \"https://user-models.example/v1\"\nmodels_list_url = \"https://user-models.example/v1/models\"\n",
)
.unwrap();
let mut cfg = crate::agent::config::Config::new_from_toml_cfg(&raw_config).unwrap();
cfg.default_yolo_mode = true;
let requirements: toml::Value = toml::from_str(
"[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\nweb_search = \"managed-ws-model\"\n\n[endpoints]\ncli_chat_proxy_base_url = \"https://managed-proxy.example/v1\"\nxai_api_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n",
"[cli]\nauto_update = false\nchannel = \"stable\"\n\n[features]\nfeedback = false\nlsp_tools = false\nweb_fetch = false\nwrite_file = false\nremote_fetch = false\n\n[ui]\nyolo = false\n\n[models]\ndefault = \"managed-model\"\nweb_search = \"managed-ws-model\"\n\n[endpoints]\ncoding_api_base_url = \"https://managed-proxy.example/v1\"\nxai_api_base_url = \"https://managed-api.example/v1\"\nmodels_base_url = \"https://managed-models.example/v1\"\nmodels_list_url = \"https://managed-models.example/v1/models\"\ndeployment_key = \"enterprise-deploy-key-should-not-log\"\n",
)
.unwrap();
let source = RequirementSource::Requirements {
@@ -2697,7 +2697,7 @@ fn apply_requirements_value_overrides_user_settings() {
assert_eq!(Some("managed-ws-model"), cfg.models.web_search.as_deref());
assert_eq!(Some("stable"), cfg.cli.channel.as_deref());
assert_eq!(
Some("https://managed-proxy.example/v1"), cfg.endpoints.cli_chat_proxy_base_url
Some("https://managed-proxy.example/v1"), cfg.endpoints.coding_api_base_url
.as_deref()
);
assert_eq!("https://managed-api.example/v1", cfg.endpoints.xai_api_base_url);
crates/codegen/kigi-shell/src/extensions/billing.rs
+344 -538
View File
@@ -1,8 +1,11 @@
//! `x.ai/billing` extension handler.
//! `x.ai/billing` extension handler — Kimi Code usage/quota.
//!
//! Fetches the authenticated user's Grok Build billing configuration
//! (credit limit, usage, on-demand cap, billing period, history) from
//! the backend. Used by the pager/desktop to display credits and usage.
//! Port of kimi-cli's `/usage` command (kimi-cli `src/kimi_cli/ui/shell/usage.py`):
//! `GET {coding_api_base_url}/usages` with the OAuth Bearer token, parsed into
//! display rows (`{usage: {...}, limits: [{detail, window, ...}]}` payload
//! shape). The TUI renders the rows as label + remaining-quota bar +
//! reset hint. The xAI credits/auto-topup surface this file used to serve is
//! gone with the xAI proxy.
use agent_client_protocol as acp;
use serde::{Deserialize, Serialize};
@@ -10,593 +13,396 @@ use serde::{Deserialize, Serialize};
use super::{ExtResult, to_raw_response};
use crate::agent::MvpAgent;
/// Billing period cycle identifier.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BillingCycle {
pub year: i32,
pub month: i32,
}
/// Cent value from the billing API (USD cents).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Cent {
/// proto3 JSON omits zero-valued scalars, so a `$0` Cent arrives as `{}`;
/// default to 0 rather than failing the whole parse.
#[serde(default)]
pub val: i64,
}
/// A usage period (weekly or monthly) from the newer credits config.
/// One usage row: a named quota with `used`/`limit` and an optional
/// human-readable reset hint (e.g. "resets in 2h 5m").
///
/// `start`/`end` are RFC 3339 timestamps. `period_type` is the proto enum name
/// (e.g. `USAGE_PERIOD_TYPE_WEEKLY`); kept so callers can distinguish weekly
/// vs monthly cycles.
/// `Deserialize` because the TUI parses this back out of the
/// `x.ai/billing` ext response.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UsageRow {
pub label: String,
pub used: i64,
pub limit: i64,
#[serde(skip_serializing_if = "Option::is_none")]
pub reset_hint: Option<String>,
}
/// Response for `x.ai/billing`: the parsed usage rows, in display order
/// (summary row first when the payload carries one).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct UsagePeriod {
#[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
pub period_type: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub start: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub end: Option<String>,
pub struct UsageResponse {
pub rows: Vec<UsageRow>,
}
/// Usage summary for one past billing period.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BillingPeriodUsage {
#[serde(skip_serializing_if = "Option::is_none")]
pub billing_cycle: Option<BillingCycle>,
#[serde(skip_serializing_if = "Option::is_none")]
pub included_used: Option<Cent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub on_demand_used: Option<Cent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub total_used: Option<Cent>,
}
/// Current billing configuration for Grok Build coding credits.
///
/// Carries both the newer credits-config fields (`credit_usage_percent`,
/// `current_period`) and the deprecated `GrokBuildBillingConfig` fields
/// (`monthly_limit`, `used`, `billing_period_*`). Consumers should prefer the
/// new fields and fall back to the deprecated ones, so the same struct works
/// against both the new `GetGrokCreditsConfig` and the legacy
/// `GetGrokBuildBillingConfig` backend responses.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct BillingConfig {
/// Included credit usage as a percentage of the allowance (0.0–100.0).
/// Preferred over deriving from `monthly_limit`/`used`.
#[serde(skip_serializing_if = "Option::is_none")]
pub credit_usage_percent: Option<f64>,
/// Current usage period (weekly or monthly). Preferred over
/// `billing_period_start`/`billing_period_end`.
#[serde(skip_serializing_if = "Option::is_none")]
pub current_period: Option<UsagePeriod>,
/// Deprecated: included monthly credit budget. Use `credit_usage_percent`.
#[serde(skip_serializing_if = "Option::is_none")]
pub monthly_limit: Option<Cent>,
/// Deprecated: credits used this period. Use `credit_usage_percent`.
#[serde(skip_serializing_if = "Option::is_none")]
pub used: Option<Cent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub on_demand_cap: Option<Cent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub on_demand_used: Option<Cent>,
/// Remaining prepaid (purchased) credit balance, positive — the "bought
/// credits" the user has topped up. Populated from the credits config
/// (`GetGrokCreditsConfig.prepaid_balance`); absent in the legacy billing
/// shape.
#[serde(skip_serializing_if = "Option::is_none")]
pub prepaid_balance: Option<Cent>,
/// Whether this user is on unified usage billing (shared weekly/monthly
/// pool). From `GrokCreditsConfig.is_unified_billing_user`, which billing
/// sets from remote settings `unified_consumer_billing_enabled`. `None` when
/// absent (legacy `GetGrokBuildBillingConfig` shape or older servers).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub is_unified_billing_user: Option<bool>,
/// Deprecated: use `current_period.start`.
#[serde(skip_serializing_if = "Option::is_none")]
pub billing_period_start: Option<String>,
/// Deprecated: use `current_period.end`.
#[serde(skip_serializing_if = "Option::is_none")]
pub billing_period_end: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub history: Vec<BillingPeriodUsage>,
}
/// Top-level response (primarily from `GET /rest/grok/credits` + auto-topup-rule).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BillingConfigResponse {
pub config: Option<BillingConfig>,
/// Whether on-demand credit usage is enabled. When `false`, the pager
/// should hide on-demand controls. Populated from `RemoteSettings`.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub on_demand_enabled: Option<bool>,
/// User-friendly subscription tier name (e.g. "SuperGrok Heavy").
/// Populated from `RemoteSettings` so the pager can update its cached
/// tier on every billing fetch without an extra request.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub subscription_tier: Option<String>,
}
/// Auto top-up configuration (from GetAutoTopupRule).
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct AutoTopupRule {
/// proto3 JSON omits `false`, so a disabled rule arrives without this field;
/// default to `false` rather than failing the parse (which would otherwise
/// keep a stale cached rule in the pager).
#[serde(default)]
pub enabled: bool,
pub min_before_hitting_sl: Option<Cent>,
pub topup_amount: Option<Cent>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_amount_per_month: Option<Cent>,
}
/// Wrapper for the auto top-up rule response.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GetAutoTopupRuleResponse {
#[serde(default)]
pub rule: Option<AutoTopupRule>,
/// Error from the usages fetch, mapped to the same user-facing messages
/// kimi-cli shows (usage.py error handling).
#[derive(Debug, thiserror::Error)]
pub enum UsageError {
#[error("Authorization failed. Please check your credentials.")]
Unauthorized,
#[error("Usage endpoint not available. Try Kimi for Coding.")]
NotFound,
#[error("Failed to fetch usage (HTTP {status}).")]
Http { status: u16 },
#[error("Failed to fetch usage: {0}")]
Network(#[from] reqwest::Error),
#[error("Failed to parse usage response: {0}")]
Parse(#[from] serde_json::Error),
}
#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/billing" => {
tracing::info!("handling billing config request");
handle_get_billing(agent).await
}
"x.ai/auto-topup-rule" => {
tracing::info!("handling auto top-up rule request");
handle_get_auto_topup_rule(agent).await
tracing::info!("handling usage request");
handle_get_usage(agent).await
}
_ => Err(acp::Error::method_not_found()),
}
}
/// Structured context for unified-log entries from a successful billing fetch.
///
/// Keeps history to a count + the most recent period so `~/.kigi/logs/unified.jsonl`
/// stays useful without dumping unbounded period arrays.
fn billing_unified_log_ctx(billing: &BillingConfigResponse) -> serde_json::Value {
let history_len = billing
.config
.as_ref()
.map(|c| c.history.len())
.unwrap_or(0);
let latest_history = billing
.config
.as_ref()
.and_then(|c| c.history.last())
.and_then(|p| serde_json::to_value(p).ok());
let mut config_value = billing
.config
.as_ref()
.and_then(|c| serde_json::to_value(c).ok())
.unwrap_or(serde_json::Value::Null);
if let Some(obj) = config_value.as_object_mut() {
// Drop full history array; surface length + latest entry instead.
obj.remove("history");
obj.insert("historyLen".into(), serde_json::json!(history_len));
if let Some(latest) = latest_history {
obj.insert("latestHistory".into(), latest);
}
}
serde_json::json!({
"config": config_value,
"onDemandEnabled": billing.on_demand_enabled,
"subscriptionTier": billing.subscription_tier,
})
}
async fn handle_get_billing(agent: &MvpAgent) -> ExtResult {
async fn handle_get_usage(agent: &MvpAgent) -> ExtResult {
let auth = super::auth_gate::require_xai_auth(
&agent.auth_manager,
"Authentication required to fetch billing data",
"Billing data requires auth with grok.com. Run `grok login` to authenticate.",
"Authentication required to fetch usage data",
"Usage data requires a Kimi Code subscription session. Run `kigi login` to authenticate.",
)?;
let proxy_base = agent.cli_chat_proxy_base_url();
let base = proxy_base.trim_end_matches('/');
// Credits balance / usage (new billing system) via the CLI proxy, which
// forwards to the backend `GetGrokCreditsConfig`.
let credits_url = format!("{}/billing?format=credits", base);
let credits_resp = crate::http::shared_client()
.get(&credits_url)
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.timeout(std::time::Duration::from_secs(15))
.send()
let base = agent.cfg.borrow().endpoints.proxy_url();
let usage = fetch_usage(&crate::http::shared_client(), &base, &auth.key)
.await
.map_err(|e| {
tracing::error!(error = %e, "billing: upstream request failed");
tracing::warn!(error = %e, "usage fetch failed");
kigi_log::unified_log::warn(
"billing: upstream request failed",
"usage: fetch failed",
None,
Some(serde_json::json!({ "error": e.to_string() })),
);
acp::Error::internal_error().data(format!("Failed to fetch billing data: {e}"))
acp::Error::internal_error().data(e.to_string())
})?;
if !credits_resp.status().is_success() {
let status = credits_resp.status().as_u16();
let body = credits_resp.text().await.unwrap_or_default();
tracing::warn!(status, url = %credits_url, "billing: upstream error");
let detail = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from))
.unwrap_or_else(|| format!("HTTP {status}"));
kigi_log::unified_log::warn(
"billing: upstream error",
None,
Some(serde_json::json!({
"status": status,
"detail": detail,
})),
);
return Err(acp::Error::internal_error().data(format!("Billing service error: {detail}")));
}
let mut billing: BillingConfigResponse = credits_resp.json().await.map_err(|e| {
tracing::error!(error = %e, "billing: failed to parse response");
kigi_log::unified_log::warn(
"billing: failed to parse response",
None,
Some(serde_json::json!({ "error": e.to_string() })),
);
acp::Error::internal_error().data(format!("Failed to parse billing data: {e}"))
})?;
// Enrich with fields from remote settings.
let rs = agent.cfg.borrow().remote_settings.clone();
billing.on_demand_enabled = rs.as_ref().and_then(|rs| rs.on_demand_enabled);
billing.subscription_tier = rs.as_ref().and_then(|rs| {
rs.subscription_tier_display
.clone()
.or_else(|| rs.subscription_tier.clone())
});
// Every prompt / /usage / poll path hits `x.ai/billing`; log the fetched
// credits snapshot so support can correlate limit UX with real balances.
kigi_log::unified_log::info(
"billing: fetched credits config",
"usage: fetched quota rows",
None,
Some(billing_unified_log_ctx(&billing)),
serde_json::to_value(&usage).ok(),
);
to_raw_response(&billing)
to_raw_response(&usage)
}
async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult {
let auth = super::auth_gate::require_xai_auth(
&agent.auth_manager,
"Authentication required to fetch auto top-up rule",
"Auto top-up data requires auth with grok.com. Run `grok login` to authenticate.",
)?;
let proxy_base = agent.cli_chat_proxy_base_url();
let base = proxy_base.trim_end_matches('/');
// Auto top-up rule via the CLI proxy, which forwards to the backend
// `GetAutoTopupRule`.
let url = format!("{}/auto-topup-rule", base);
let response = crate::http::shared_client()
/// `GET {base}/usages` with a Bearer token, parsed per kimi-cli usage.py.
pub(crate) async fn fetch_usage(
http: &reqwest::Client,
base_url: &str,
token: &str,
) -> Result<UsageResponse, UsageError> {
let url = format!("{}/usages", base_url.trim_end_matches('/'));
let response = http
.get(&url)
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.timeout(std::time::Duration::from_secs(10))
.bearer_auth(token)
.timeout(std::time::Duration::from_secs(15))
.send()
.await
.map_err(|e| {
tracing::error!(error = %e, "auto-topup: upstream request failed");
acp::Error::internal_error().data(format!("Failed to fetch auto top-up rule: {e}"))
})?;
.await?;
match response.status().as_u16() {
200..=299 => {}
401 => return Err(UsageError::Unauthorized),
404 => return Err(UsageError::NotFound),
status => return Err(UsageError::Http { status }),
}
let payload: serde_json::Value = serde_json::from_str(&response.text().await?)?;
Ok(parse_usage_payload(&payload))
}
if !response.status().is_success() {
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
tracing::warn!(status, url = %url, "auto-topup: upstream error");
/// Port of usage.py `_parse_usage_payload`: `usage` (summary) + `limits[]`.
fn parse_usage_payload(payload: &serde_json::Value) -> UsageResponse {
let mut rows = Vec::new();
let detail = serde_json::from_str::<serde_json::Value>(&body)
.ok()
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from))
.unwrap_or_else(|| format!("HTTP {status}"));
return Err(
acp::Error::internal_error().data(format!("Auto top-up service error: {detail}"))
);
if let Some(usage) = payload.get("usage").filter(|v| v.is_object())
&& let Some(row) = to_usage_row(usage, "Weekly limit")
{
rows.push(row);
}
// Return the upstream response body verbatim (as a JSON value) so /usage
// can print the exact data from this request unformatted.
let body_text = response.text().await.unwrap_or_default();
let value: serde_json::Value =
serde_json::from_str(&body_text).unwrap_or(serde_json::json!({"raw": body_text}));
to_raw_response(&value)
if let Some(limits) = payload.get("limits").and_then(|v| v.as_array()) {
for (idx, item) in limits.iter().enumerate() {
if !item.is_object() {
continue;
}
let detail = match item.get("detail") {
Some(d) if d.is_object() => d,
_ => item,
};
let empty = serde_json::json!({});
let window = match item.get("window") {
Some(w) if w.is_object() => w,
_ => &empty,
};
let label = limit_label(item, detail, window, idx);
if let Some(row) = to_usage_row(detail, &label) {
rows.push(row);
}
}
}
UsageResponse { rows }
}
/// Port of usage.py `_to_usage_row`: `used`/`limit`, with
/// `used = limit - remaining` fallback; row dropped when both absent.
fn to_usage_row(data: &serde_json::Value, default_label: &str) -> Option<UsageRow> {
let limit = to_int(data.get("limit"));
let used = to_int(data.get("used")).or_else(|| match (to_int(data.get("remaining")), limit) {
(Some(remaining), Some(limit)) => Some(limit - remaining),
_ => None,
});
if used.is_none() && limit.is_none() {
return None;
}
let label = data
.get("name")
.and_then(non_empty_str)
.or_else(|| data.get("title").and_then(non_empty_str))
.map(str::to_owned)
.unwrap_or_else(|| default_label.to_owned());
Some(UsageRow {
label,
used: used.unwrap_or(0),
limit: limit.unwrap_or(0),
reset_hint: reset_hint(data),
})
}
/// Port of usage.py `_limit_label`: name/title/scope, else the window
/// duration ("5h limit"), else "Limit #N".
fn limit_label(
item: &serde_json::Value,
detail: &serde_json::Value,
window: &serde_json::Value,
idx: usize,
) -> String {
for key in ["name", "title", "scope"] {
if let Some(val) = item
.get(key)
.and_then(non_empty_str)
.or_else(|| detail.get(key).and_then(non_empty_str))
{
return val.to_owned();
}
}
let duration = to_int(window.get("duration"))
.or_else(|| to_int(item.get("duration")))
.or_else(|| to_int(detail.get("duration")));
let time_unit = window
.get("timeUnit")
.and_then(non_empty_str)
.or_else(|| item.get("timeUnit").and_then(non_empty_str))
.or_else(|| detail.get("timeUnit").and_then(non_empty_str))
.unwrap_or("");
if let Some(duration) = duration.filter(|&d| d != 0) {
if time_unit.contains("MINUTE") {
if duration >= 60 && duration % 60 == 0 {
return format!("{}h limit", duration / 60);
}
return format!("{duration}m limit");
}
if time_unit.contains("HOUR") {
return format!("{duration}h limit");
}
if time_unit.contains("DAY") {
return format!("{duration}d limit");
}
return format!("{duration}s limit");
}
format!("Limit #{}", idx + 1)
}
/// Port of usage.py `_reset_hint`: absolute reset keys first, then
/// seconds-until keys.
fn reset_hint(data: &serde_json::Value) -> Option<String> {
for key in ["reset_at", "resetAt", "reset_time", "resetTime"] {
if let Some(val) = data.get(key).and_then(non_empty_str) {
return Some(format_reset_time(val));
}
}
for key in ["reset_in", "resetIn", "ttl", "window"] {
if let Some(seconds) = to_int(data.get(key)).filter(|&s| s != 0) {
return Some(format!(
"resets in {}",
format_duration(seconds.max(0) as u64)
));
}
}
None
}
/// Port of usage.py `_format_reset_time`: ISO timestamp → "resets in …" /
/// "reset" (already past) / "resets at <raw>" when unparseable.
fn format_reset_time(val: &str) -> String {
match chrono::DateTime::parse_from_rfc3339(val) {
Ok(dt) => {
let delta = dt.with_timezone(&chrono::Utc) - chrono::Utc::now();
let seconds = delta.num_seconds();
if seconds <= 0 {
"reset".to_owned()
} else {
format!("resets in {}", format_duration(seconds as u64))
}
}
Err(_) => format!("resets at {val}"),
}
}
/// Port of kimi-cli `utils/datetime.py` `format_duration`: short units,
/// seconds shown only for sub-minute durations.
fn format_duration(seconds: u64) -> String {
let days = seconds / 86_400;
let hours = (seconds % 86_400) / 3_600;
let minutes = (seconds % 3_600) / 60;
let secs = seconds % 60;
let mut parts = Vec::new();
if days > 0 {
parts.push(format!("{days}d"));
}
if hours > 0 {
parts.push(format!("{hours}h"));
}
if minutes > 0 {
parts.push(format!("{minutes}m"));
}
if secs > 0 && parts.is_empty() {
parts.push(format!("{secs}s"));
}
if parts.is_empty() {
"0s".to_owned()
} else {
parts.join(" ")
}
}
fn non_empty_str(v: &serde_json::Value) -> Option<&str> {
v.as_str().filter(|s| !s.is_empty())
}
/// Port of usage.py `_to_int`: ints and int-shaped floats/strings; anything
/// else is `None`.
fn to_int(value: Option<&serde_json::Value>) -> Option<i64> {
let value = value?;
if let Some(i) = value.as_i64() {
return Some(i);
}
if let Some(f) = value.as_f64() {
return Some(f as i64);
}
value.as_str()?.trim().parse::<i64>().ok()
}
#[cfg(test)]
mod tests {
use super::*;
use wiremock::matchers::{bearer_token, method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[test]
fn auto_topup_disabled_rule_omits_enabled_field() {
// proto3 JSON omits `false` / `0`, so a disabled rule arrives without
// `enabled` (and zero Cents as `{}`). It must still deserialize (as
// disabled) rather than erroring — otherwise the pager keeps a stale
// cached rule.
let json = serde_json::json!({
"rule": { "topupAmount": {"val": 500}, "minBeforeHittingSl": {} }
});
let resp: GetAutoTopupRuleResponse = serde_json::from_value(json).unwrap();
let rule = resp.rule.expect("rule present");
assert!(!rule.enabled);
assert_eq!(rule.topup_amount.unwrap().val, 500);
assert_eq!(rule.min_before_hitting_sl.unwrap().val, 0);
}
#[test]
fn billing_config_response_deserializes_from_backend_json() {
let json = serde_json::json!({
"config": {
"monthlyLimit": {"val": 2000},
"used": {"val": 1234},
"onDemandCap": {"val": 500},
"billingPeriodStart": "2025-04-01T00:00:00Z",
"billingPeriodEnd": "2025-05-01T00:00:00Z",
"history": [
/// Happy path: GET /usages with Bearer, kimi payload shape → rows with
/// summary first, remaining-derived `used`, and window-derived labels.
#[tokio::test]
async fn fetch_usage_parses_kimi_payload() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/usages"))
.and(bearer_token("tok-42"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"usage": { "limit": 1000, "used": 250, "reset_at": "2099-01-01T00:00:00Z" },
"limits": [
{
"billingCycle": {"year": 2025, "month": 3},
"includedUsed": {"val": 1800},
"onDemandUsed": {"val": 0},
"totalUsed": {"val": 1800}
}
"window": { "duration": 300, "timeUnit": "TIME_UNIT_MINUTE" },
"detail": { "limit": 50, "remaining": 30, "resetIn": 1800 }
},
{ "name": "RPM", "limit": 60, "used": 12 }
]
}
});
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
let config = resp.config.unwrap();
assert_eq!(config.monthly_limit.unwrap().val, 2000);
assert_eq!(config.used.unwrap().val, 1234);
assert_eq!(config.on_demand_cap.unwrap().val, 500);
assert_eq!(
config.billing_period_start.as_deref(),
Some("2025-04-01T00:00:00Z")
);
assert_eq!(config.history.len(), 1);
let period = &config.history[0];
let cycle = period.billing_cycle.as_ref().unwrap();
assert_eq!(cycle.year, 2025);
assert_eq!(cycle.month, 3);
assert_eq!(period.included_used.as_ref().unwrap().val, 1800);
assert_eq!(period.total_used.as_ref().unwrap().val, 1800);
}
})))
.expect(1)
.mount(&server)
.await;
#[test]
fn billing_unified_log_ctx_includes_credits_and_collapses_history() {
let resp = BillingConfigResponse {
config: Some(BillingConfig {
credit_usage_percent: Some(42.5),
current_period: Some(UsagePeriod {
period_type: Some("USAGE_PERIOD_TYPE_WEEKLY".into()),
start: Some("2025-04-01T00:00:00Z".into()),
end: Some("2025-04-08T00:00:00Z".into()),
}),
monthly_limit: Some(Cent { val: 2000 }),
used: Some(Cent { val: 850 }),
on_demand_cap: Some(Cent { val: 500 }),
on_demand_used: Some(Cent { val: 0 }),
prepaid_balance: Some(Cent { val: 100 }),
is_unified_billing_user: Some(true),
billing_period_start: None,
billing_period_end: None,
history: vec![
BillingPeriodUsage {
billing_cycle: Some(BillingCycle {
year: 2025,
month: 2,
}),
included_used: Some(Cent { val: 1000 }),
on_demand_used: Some(Cent { val: 0 }),
total_used: Some(Cent { val: 1000 }),
},
BillingPeriodUsage {
billing_cycle: Some(BillingCycle {
year: 2025,
month: 3,
}),
included_used: Some(Cent { val: 1800 }),
on_demand_used: Some(Cent { val: 0 }),
total_used: Some(Cent { val: 1800 }),
},
],
}),
on_demand_enabled: Some(true),
subscription_tier: Some("SuperGrok".into()),
};
let ctx = billing_unified_log_ctx(&resp);
assert_eq!(ctx["onDemandEnabled"], true);
assert_eq!(ctx["subscriptionTier"], "SuperGrok");
let config = ctx["config"].as_object().expect("config object");
let usage = fetch_usage(&reqwest::Client::new(), &server.uri(), "tok-42")
.await
.expect("usage fetch should succeed");
assert_eq!(usage.rows.len(), 3);
assert_eq!(usage.rows[0].label, "Weekly limit");
assert_eq!(usage.rows[0].used, 250);
assert_eq!(usage.rows[0].limit, 1000);
assert!(
config.get("history").is_none(),
"full history must be collapsed"
usage.rows[0]
.reset_hint
.as_deref()
.is_some_and(|h| h.starts_with("resets in")),
"absolute reset_at renders a relative hint: {:?}",
usage.rows[0].reset_hint
);
assert_eq!(config["historyLen"], 2);
assert_eq!(
config["latestHistory"]["billingCycle"]["month"], 3,
"latest history period retained"
);
assert_eq!(config["creditUsagePercent"], 42.5);
assert_eq!(config["prepaidBalance"]["val"], 100);
// 300 minutes → "5h limit"; used derived from remaining (50-30=20).
assert_eq!(usage.rows[1].label, "5h limit");
assert_eq!(usage.rows[1].used, 20);
assert_eq!(usage.rows[1].limit, 50);
assert_eq!(usage.rows[1].reset_hint.as_deref(), Some("resets in 30m"));
// Item-level fields when there is no `detail` object.
assert_eq!(usage.rows[2].label, "RPM");
assert_eq!(usage.rows[2].used, 12);
assert_eq!(usage.rows[2].limit, 60);
}
/// Auth failure: 401 maps to the typed `Unauthorized` error (kimi-cli's
/// "Authorization failed" path).
#[tokio::test]
async fn fetch_usage_maps_401_to_unauthorized() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/usages"))
.respond_with(ResponseTemplate::new(401))
.expect(1)
.mount(&server)
.await;
let err = fetch_usage(&reqwest::Client::new(), &server.uri(), "bad")
.await
.expect_err("401 must fail");
assert!(matches!(err, UsageError::Unauthorized));
}
/// 404 maps to the "endpoint not available" error (kimi-cli parity).
#[tokio::test]
async fn fetch_usage_maps_404_to_not_found() {
let server = MockServer::start().await;
Mock::given(method("GET"))
.and(path("/usages"))
.respond_with(ResponseTemplate::new(404))
.expect(1)
.mount(&server)
.await;
let err = fetch_usage(&reqwest::Client::new(), &server.uri(), "tok")
.await
.expect_err("404 must fail");
assert!(matches!(err, UsageError::NotFound));
}
/// Empty payload parses to zero rows (TUI shows "No usage data").
#[test]
fn parse_usage_payload_empty_object_yields_no_rows() {
let usage = parse_usage_payload(&serde_json::json!({}));
assert!(usage.rows.is_empty());
}
#[test]
fn billing_config_response_roundtrips_through_json() {
let config = BillingConfig {
credit_usage_percent: None,
current_period: None,
monthly_limit: Some(Cent { val: 5000 }),
used: Some(Cent { val: 123 }),
on_demand_cap: Some(Cent { val: 0 }),
on_demand_used: Some(Cent { val: 50 }),
prepaid_balance: Some(Cent { val: 750 }),
is_unified_billing_user: None,
billing_period_start: Some("2025-04-01T00:00:00Z".to_string()),
billing_period_end: Some("2025-05-01T00:00:00Z".to_string()),
history: vec![BillingPeriodUsage {
billing_cycle: Some(BillingCycle {
year: 2025,
month: 3,
}),
included_used: Some(Cent { val: 4500 }),
on_demand_used: Some(Cent { val: 100 }),
total_used: Some(Cent { val: 4600 }),
}],
};
let resp = BillingConfigResponse {
config: Some(config),
on_demand_enabled: None,
subscription_tier: None,
};
let json = serde_json::to_value(&resp).unwrap();
let roundtripped: BillingConfigResponse = serde_json::from_value(json).unwrap();
let rt_config = roundtripped.config.unwrap();
assert_eq!(rt_config.monthly_limit.unwrap().val, 5000);
assert_eq!(rt_config.used.unwrap().val, 123);
assert_eq!(rt_config.prepaid_balance.unwrap().val, 750);
assert_eq!(rt_config.history.len(), 1);
}
#[test]
fn billing_config_response_handles_null_config() {
let json = serde_json::json!({"config": null});
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
assert!(resp.config.is_none());
}
#[test]
fn billing_config_response_handles_empty_history() {
let json = serde_json::json!({
"config": {
"monthlyLimit": {"val": 1000},
"used": {"val": 0}
}
});
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
let config = resp.config.unwrap();
assert_eq!(config.monthly_limit.unwrap().val, 1000);
assert!(config.history.is_empty());
}
#[test]
fn billing_config_serializes_camel_case() {
let config = BillingConfig {
credit_usage_percent: None,
current_period: None,
monthly_limit: Some(Cent { val: 100 }),
used: None,
on_demand_cap: None,
on_demand_used: None,
prepaid_balance: None,
is_unified_billing_user: None,
billing_period_start: None,
billing_period_end: None,
history: vec![],
};
let json = serde_json::to_value(&config).unwrap();
assert!(json.get("monthlyLimit").is_some());
// Fields with None are skipped
assert!(json.get("creditUsagePercent").is_none());
assert!(json.get("currentPeriod").is_none());
assert!(json.get("used").is_none());
assert!(json.get("onDemandCap").is_none());
assert!(json.get("onDemandUsed").is_none());
assert!(json.get("prepaidBalance").is_none());
assert!(json.get("billingPeriodStart").is_none());
// Empty history is skipped
assert!(json.get("history").is_none());
}
#[test]
fn billing_config_deserializes_credits_config_shape() {
// Newer `GetGrokCreditsConfig` response: percentage-based usage,
// a typed current period, and history keyed by `period`.
let json = serde_json::json!({
"config": {
"creditUsagePercent": 42.5,
"currentPeriod": {
"type": "USAGE_PERIOD_TYPE_WEEKLY",
"start": "2026-06-01T00:00:00Z",
"end": "2026-06-08T00:00:00Z"
},
"onDemandCap": {"val": 5000},
"onDemandUsed": {"val": 300},
"prepaidBalance": {"val": 1250},
"isUnifiedBillingUser": true,
"productUsage": [
{"product": "PRODUCT_GROK_BUILD", "usagePercent": 61.2}
],
"history": [
{
"period": {
"type": "USAGE_PERIOD_TYPE_WEEKLY",
"start": "2026-05-25T00:00:00Z",
"end": "2026-06-01T00:00:00Z"
},
"onDemandUsed": {"val": 120}
}
]
}
});
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
let config = resp.config.unwrap();
assert_eq!(config.credit_usage_percent, Some(42.5));
let period = config.current_period.as_ref().unwrap();
assert_eq!(
period.period_type.as_deref(),
Some("USAGE_PERIOD_TYPE_WEEKLY")
);
assert_eq!(period.end.as_deref(), Some("2026-06-08T00:00:00Z"));
// Deprecated fields are absent in the credits shape.
assert!(config.monthly_limit.is_none());
assert!(config.billing_period_end.is_none());
assert_eq!(config.on_demand_cap.unwrap().val, 5000);
assert_eq!(config.on_demand_used.unwrap().val, 300);
// Bought (prepaid) credit balance is parsed from the credits config.
assert_eq!(config.prepaid_balance.unwrap().val, 1250);
assert_eq!(config.is_unified_billing_user, Some(true));
// productUsage is still unused by the CLI billing surface.
assert_eq!(config.history.len(), 1);
assert_eq!(config.history[0].on_demand_used.as_ref().unwrap().val, 120);
}
#[test]
fn cent_serializes_as_val_field() {
let c = Cent { val: 4299 };
let json = serde_json::to_value(&c).unwrap();
assert_eq!(json, serde_json::json!({"val": 4299}));
fn format_duration_matches_kimi_semantics() {
assert_eq!(format_duration(0), "0s");
assert_eq!(format_duration(45), "45s");
assert_eq!(format_duration(90), "1m");
assert_eq!(format_duration(3_661), "1h 1m");
assert_eq!(format_duration(90_000), "1d 1h");
}
}
crates/codegen/kigi-shell/src/extensions/bundle.rs
-1204
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/extensions/feedback.rs
+30 -47
View File
@@ -1,12 +1,15 @@
//! `x.ai/feedback`, `x.ai/feedback/dismiss`, `x.ai/btw`, and `x.ai/review/*`
//! extension handlers.
//!
//! - `feedback`/`feedback/dismiss`: persist user ratings/text locally and
//! forward to cli-chat-proxy.
//! - `feedback`/`feedback/dismiss`: persist user ratings/text locally; text
//! feedback from subscription (OAuth) sessions is forwarded to the Kimi
//! Code feedback endpoint (`POST {base}/feedback`, kimi-cli slash.py
//! parity). Without a subscription session the record stays local and the
//! response points at the GitHub issue tracker.
//! - `btw`: dispatch a side question to the active session via
//! `SessionCommand::SideQuestion` and return the answer.
//! - `review/comment` and `review/comment/delete`: record inline code review
//! events to cloud storage.
//! events locally.
use std::sync::Arc;
@@ -15,6 +18,7 @@ use tokio::sync::oneshot;
use super::{ExtResult, parse_params};
use crate::agent::MvpAgent;
use crate::agent::feedback_client::FEEDBACK_ISSUES_URL;
use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry};
use crate::session::{
ClientFeedbackInput, CommentDeleteRequest, CommentDeleteResponse, CommentRequest,
@@ -34,7 +38,7 @@ pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
}
m if m.starts_with("x.ai/review") => {
tracing::info!("handling review comment");
handle_review(agent, args).await
handle_review(args).await
}
_ => Err(acp::Error::method_not_found()),
}
@@ -97,8 +101,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
let simple: crate::session::FeedbackRequest = parse_params(args)?;
ClientFeedbackInput {
session_id: simple.session_id,
client_type:
prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui,
client_type: crate::session::feedback_types::ClientType::Tui,
rating_type: None,
rating_value: None,
feedback_text: Some(simple.feedback_text),
@@ -151,7 +154,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
submission.merge_metadata(user_meta);
}
// Enrich with session context for Slack notifications (best-effort).
// Enrich with session context (persisted alongside the record).
if let Some(ref session_handle) = session_handle {
let (tx, rx) = tokio::sync::oneshot::channel();
let _ = session_handle
@@ -174,7 +177,7 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
if let (Some(session_handle), Some(rating_value)) =
(&session_handle, feedback_input.rating_value)
{
use prod_mc_cli_chat_proxy_types::feedback_types::RatingType;
use crate::session::feedback_types::RatingType;
let (is_positive, is_negative) = match feedback_input.rating_type {
// Thumbs: -1 = down, 0 = neutral, 1 = up
Some(RatingType::Thumbs) | None => (rating_value > 0, rating_value < 0),
@@ -208,8 +211,8 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
let client = agent.feedback_client();
if client.is_none() {
tracing::warn!(
"no feedback client available (missing proxy credentials); feedback saved locally only"
tracing::info!(
"no subscription session; feedback saved locally — submit at {FEEDBACK_ISSUES_URL}"
);
}
let outcome = crate::session::feedback_manager::submit_feedback_workflow(
@@ -222,15 +225,17 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
match &outcome {
crate::session::feedback_manager::SubmitOutcome::Submitted => {
tracing::info!("feedback submitted to proxy successfully");
tracing::info!("feedback submitted to the Kimi Code endpoint");
}
crate::session::feedback_manager::SubmitOutcome::LocalOnly => {
tracing::warn!("feedback saved locally only (no proxy client)");
tracing::info!("feedback saved locally only");
}
crate::session::feedback_manager::SubmitOutcome::Failed(e) => {
tracing::error!(error = %e, "feedback submission to proxy failed");
return Err(acp::Error::internal_error()
.data(format!("Feedback submission failed: {e}")));
tracing::error!(error = %e, "feedback submission failed");
return Err(acp::Error::internal_error().data(format!(
"Feedback submission failed: {e}. \
Please submit feedback at {FEEDBACK_ISSUES_URL}"
)));
}
}
@@ -281,36 +286,14 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
}
}
let request_id = dismiss_input.request_id.clone();
let client = agent
.feedback_client()
.ok_or_else(|| acp::Error::internal_error().data("No credentials for feedback"))?;
let feedback_base_url = agent.cfg.borrow().endpoints.resolve_feedback_base_url();
match client.dismiss_request(&request_id).await {
Ok(response) => {
tracing::info!(
request_id = %response.request_id,
status = %response.status,
feedback_url = %feedback_base_url,
"Feedback request dismissed"
);
let value = serde_json::to_value(&response)
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
.expect("to work")
.expect("to work");
Ok(acp::ExtResponse::new(value))
}
Err(e) => {
tracing::warn!(
error = %e,
request_id = %request_id,
feedback_url = %feedback_base_url,
"Failed to dismiss feedback request"
);
Err(acp::Error::internal_error()
.data(format!("Failed to dismiss feedback request: {e}")))
}
}
let value = serde_json::to_value(serde_json::json!({
"requestId": dismiss_input.request_id,
"status": "dismissed",
}))
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
.expect("to work")
.expect("to work");
Ok(acp::ExtResponse::new(value))
}
_ => Err(acp::Error::method_not_found()),
}
@@ -319,9 +302,9 @@ async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult
/// Record inline code review events.
///
/// Methods:
/// - `x.ai/review/comment`: record a new inline code comment to cloud storage
/// - `x.ai/review/comment`: record a new inline code comment
/// - `x.ai/review/comment/delete`: record a tombstone event for a deleted comment
async fn handle_review(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
async fn handle_review(args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/review/comment" => {
let request: CommentRequest = parse_params(args)?;
crates/codegen/kigi-shell/src/extensions/mod.rs
-2
View File
@@ -1,7 +1,6 @@
pub mod auth;
pub(crate) mod auth_gate;
pub mod billing;
pub mod bundle;
pub mod chat_conversation_history;
pub mod code_nav;
pub mod debug;
@@ -28,7 +27,6 @@ pub mod search;
pub mod session_admin;
pub mod session_search;
pub mod session_updates;
pub mod share;
pub mod skills;
pub mod suggest;
pub mod task;
crates/codegen/kigi-shell/src/extensions/session_admin.rs
+12 -114
View File
@@ -4,8 +4,8 @@
//! persistent or shared agent state but are not part of the per-turn prompt
//! lifecycle:
//!
//! - `x.ai/session/rename` rename a session locally + remote
//! - `x.ai/session/delete` delete a session locally + remote
//! - `x.ai/session/rename` rename a session locally
//! - `x.ai/session/delete` delete a session locally
//! - `x.ai/session/update_mcp_servers` mid-session MCP server swap
//! - `x.ai/session/fork` fork a session into a new one
//! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions
@@ -77,7 +77,8 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
}
if req.kind == SessionKind::Chat {
return rename_chat_conversation(agent, &req.session_id, &req.title).await;
return Err(acp::Error::invalid_request()
.data("chat conversations are not available in kigi (local sessions only)"));
}
let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));
@@ -111,22 +112,6 @@ async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
// Send a SessionSummaryGenerated notification so the TUI updates its title
notify_session_title(agent, session_id, &req.title).await;
if agent.is_writeback_storage() && agent.current_auth().is_some() {
use crate::remote::client::BackendClient;
use crate::session::export::ExportedMetadata;
let mut metadata = ExportedMetadata::from_summary(summary);
metadata.title = Some(req.title.clone());
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
if let Err(e) = BackendClient::new()
.with_auth_manager(agent.auth_manager.clone())
.save_session_data(&req.session_id, &[], Some(&metadata))
.await
{
tracing::warn!(?e, session_id = %req.session_id, "failed to sync renamed title to backend");
}
}
// Hook 2: update session replica with summary (fire-and-forget)
if let Some(client) = agent.session_registry_client() {
let sid = req.session_id.to_string();
@@ -169,49 +154,6 @@ async fn notify_session_title(agent: &MvpAgent, session_id: acp::SessionId, titl
}
}
async fn rename_chat_conversation(
agent: &MvpAgent,
conversation_id: &str,
title: &str,
) -> ExtResult {
use crate::remote::{ConvError, UpdateConversationBody};
let Some(client) = agent.conversations_client() else {
return Err(acp::Error::invalid_request()
.data("chat session rename requires the conversations lane (OIDC + chat feature)"));
};
let body = UpdateConversationBody {
title: Some(title.to_owned()),
starred: None,
};
client
.update_conversation(conversation_id, &body)
.await
.map_err(|e| match e {
ConvError::NoOauth => acp::Error::invalid_request()
.data("chat session rename requires xAI OAuth credentials"),
ConvError::Http { status: 404 } => acp::Error::invalid_request()
.data(format!("conversation not found: {conversation_id}")),
other => acp::Error::internal_error()
.data(format!("chat conversation rename failed: {other}")),
})?;
// If this conversation is open live, notify clients of the new title.
let session_id = acp::SessionId::new(Arc::from(conversation_id));
if agent.sessions.borrow().contains_key(&session_id) {
notify_session_title(agent, session_id, title).await;
}
tracing::info!(
session_id = %conversation_id,
title = %title,
"Chat conversation renamed"
);
to_raw_response(&serde_json::json!({ "success": true }))
}
// session/delete
/// Delete a session from history.
@@ -229,31 +171,17 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
let req: DeleteRequest = parse_params(args)?;
if req.kind == SessionKind::Chat {
return soft_delete_chat_conversation(agent, &req.session_id).await;
return Err(acp::Error::invalid_request()
.data("chat conversations are not available in kigi (local sessions only)"));
}
let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));
// For writeback storage (non-ZDR): remote delete is authoritative for
// the cloud history and runs first; on failure no local bits are
// touched so the pager does not remove the row or toast success.
let needs_remote = agent.is_writeback_storage() && agent.current_auth().is_some();
// Shared delete: remote-first, then local disk + FTS eviction.
// Mirrored by the `grok sessions delete <id>` CLI path.
crate::session::persistence::delete_session_history(
&req.session_id,
req.cwd.as_deref(),
needs_remote,
agent.auth_manager.clone(),
)
.await
.map_err(|e| {
if let crate::session::persistence::DeleteSessionError::Remote(_) = &e {
tracing::warn!(?e, session_id = %req.session_id, "failed to delete remote session data");
}
acp::Error::internal_error().data(e.to_string())
})?;
// Local disk + FTS eviction. Mirrored by the `kigi sessions delete <id>`
// CLI path.
crate::session::persistence::delete_session_history(&req.session_id, req.cwd.as_deref())
.await
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
// If an in-memory live session exists for this id (e.g. the user
// deleted history for a session that is still open in another agent
@@ -269,35 +197,6 @@ async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtR
to_raw_response(&serde_json::json!({ "success": true }))
}
async fn soft_delete_chat_conversation(agent: &MvpAgent, conversation_id: &str) -> ExtResult {
use crate::remote::ConvError;
let Some(client) = agent.conversations_client() else {
return Err(acp::Error::invalid_request()
.data("chat session delete requires the conversations lane (OIDC + chat feature)"));
};
client
.soft_delete_conversation(conversation_id)
.await
.map_err(|e| match e {
ConvError::NoOauth => acp::Error::invalid_request()
.data("chat session delete requires xAI OAuth credentials"),
other => acp::Error::internal_error()
.data(format!("chat conversation soft-delete failed: {other}")),
})?;
let session_id = acp::SessionId::new(Arc::from(conversation_id));
if agent.sessions.borrow().contains_key(&session_id) {
agent.request_session_shutdown(&session_id);
agent.remove_session(&session_id);
}
tracing::info!(session_id = %conversation_id, "Chat conversation soft-deleted");
to_raw_response(&serde_json::json!({ "success": true }))
}
// session/update_mcp_servers
async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
@@ -708,8 +607,7 @@ async fn handle_session_fork(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtRes
let request: ForkSessionRequest = parse_params(args)?;
let agent_id = crate::util::agent_id::agent_id();
let response = fork_session(request, &agent_id, Some(agent.auth_manager.clone()))
let response = fork_session(request)
.await
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
crates/codegen/kigi-shell/src/extensions/share.rs
-204
View File
@@ -1,204 +0,0 @@
//! `x.ai/share_session` extension handler.
//!
//! Loads a local session, exports it, and asks the backend for a public
//! share URL.
use agent_client_protocol as acp;
use super::{ExtResult, parse_params, to_raw_response};
use crate::agent::MvpAgent;
use crate::remote::client::BackendClient;
use crate::session::export::ExportedSession;
use crate::session::info::Info as SessionInfo;
use crate::session::persistence::list_summaries;
use crate::session::share::{ShareSessionRequest, ShareSessionResponse};
#[tracing::instrument(skip_all, fields(method = %args.method))]
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
match args.method.as_ref() {
"x.ai/share_session" => {
tracing::info!("handling share session request");
handle_share_session(agent, args).await
}
_ => Err(acp::Error::method_not_found()),
}
}
async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
let request: ShareSessionRequest = parse_params(args)?;
// Get auth - required for sharing.
let auth = require_xai_auth_for_share(&agent.auth_manager)?;
// Remote settings / feature-flag gate: sharing_enabled defaults to false
// and is only enabled for eligible accounts.
let sharing_enabled = agent
.cfg
.borrow()
.remote_settings
.as_ref()
.and_then(|rs| rs.sharing_enabled)
.unwrap_or(false);
if !sharing_enabled {
return Err(
acp::Error::invalid_params().data("Session sharing is not available for your account.")
);
}
// Find session info by searching through summaries
let summaries = list_summaries(None).await.map_err(|e| {
acp::Error::internal_error().data(format!("Failed to list sessions: {}", e))
})?;
let summary = summaries
.iter()
.find(|s| s.info.id.0.as_ref() == request.session_id.as_str())
.ok_or_else(|| acp::Error::resource_not_found(Some("Session not found".into())))?;
let info = SessionInfo {
id: acp::SessionId::new(request.session_id.clone()),
cwd: summary.info.cwd.clone(),
};
// Load and export session
let exported = ExportedSession::from_local_session(&info)
.await
.map_err(|e| acp::Error::internal_error().data(format!("Failed to load session: {}", e)))?;
// Check for empty session
if exported.messages.is_empty() {
return Err(acp::Error::invalid_params().data("No messages to share yet"));
}
// Upload to backend and get share URL.
let client = BackendClient::new().with_auth_manager(agent.auth_manager.clone());
let agent_id = crate::util::agent_id::agent_id();
let share_url = client
.share_session(&exported, &agent_id)
.await
.map_err(|e| {
tracing::error!(error = %e, "Failed to share session with backend");
acp::Error::internal_error().data(format!("Failed to share session: {}", e))
})?;
let response = ShareSessionResponse { share_url };
to_raw_response(&response)
}
fn require_xai_auth_for_share(
auth_manager: &crate::auth::AuthManager,
) -> Result<crate::auth::KimiAuth, acp::Error> {
super::auth_gate::require_xai_auth(
auth_manager,
"Authentication required to share session",
"Share session is disabled. Run `grok login` to authenticate.",
)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::auth::KimiCodeConfig;
use crate::auth::{AuthMode, KimiAuth};
use chrono::{Duration, Utc};
use std::sync::Arc;
use tempfile::tempdir;
fn make_auth_manager_with_token_expiring_in(
ttl: Duration,
) -> (Arc<crate::auth::AuthManager>, tempfile::TempDir) {
let dir = tempdir().expect("tempdir for share auth test");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
KimiCodeConfig::default(),
));
let expires_at = Utc::now() + ttl;
// We must explicitly set oidc_issuer to a first-party xAI issuer.
// Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent)
// return true from is_xai_auth(). This is required for the share tests to
// exercise the happy path through require_xai_auth_for_share.
let auth = KimiAuth {
auth_mode: AuthMode::OAuth,
key: "test-key".into(),
expires_at: Some(expires_at),
create_time: Utc::now() - Duration::hours(1),
..Default::default()
};
mgr.hot_swap(auth);
(mgr, dir)
}
#[test]
fn share_works_outside_the_5m_early_invalidation_window() {
let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::minutes(10));
assert!(mgr.current().is_some());
assert!(require_xai_auth_for_share(&mgr).is_ok());
}
#[test]
fn share_succeeds_inside_the_5m_early_invalidation_window() {
let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::seconds(1));
// This is exactly the state that triggered the user bug:
assert!(
mgr.current().is_none(),
"current() drops the token inside the buffer"
);
assert!(mgr.expired_auth().is_some());
// Now that we use current_or_expired(), this passes.
let res = require_xai_auth_for_share(&mgr);
assert!(
res.is_ok(),
"require_xai_auth_for_share must succeed for a still-valid buffered xAI token"
);
}
#[test]
fn share_fails_with_no_auth_at_all() {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
KimiCodeConfig::default(),
));
assert!(require_xai_auth_for_share(&mgr).is_err());
}
#[test]
fn share_rejects_non_xai_auth_with_actionable_grok_login_message() {
let dir = tempdir().expect("tempdir");
let mgr = Arc::new(crate::auth::AuthManager::new(
dir.path(),
KimiCodeConfig::default(),
));
// API key is the simplest non-xAI credential (External and enterprise OIDC
// are also rejected the same way).
let non_xai = KimiAuth {
auth_mode: AuthMode::ApiKey,
key: "xai-test-key".into(),
create_time: Utc::now(),
..Default::default()
};
mgr.hot_swap(non_xai);
let err = require_xai_auth_for_share(&mgr)
.expect_err("non-xAI accounts (API key, External, enterprise IdP) must be rejected");
// This is the key assertion the review asked for: we must test the *exact*
// actionable data string for the non-xAI path (distinct from the generic
// "Authentication required to share session" path).
let serialized =
serde_json::to_value(&err).expect("acp::Error serializes to JSON-RPC shape");
let data = serialized
.get("data")
.and_then(|v| v.as_str())
.expect("auth_required error carries a data string");
assert_eq!(
data,
"Share session is disabled. Run `grok login` to authenticate."
);
}
}
crates/codegen/kigi-shell/src/lib.rs
-1
View File
@@ -28,7 +28,6 @@ pub mod managed_config;
pub mod mcp_doctor;
pub use kigi_models as models;
pub mod plugin;
pub mod remote;
pub mod sampling;
pub mod session;
pub mod terminal;
crates/codegen/kigi-shell/src/remote/agent.rs
-326
View File
@@ -1,326 +0,0 @@
//! Remote sandbox client for cli-chat-proxy.
//!
//! This module provides an HTTP client to interact with cli-chat-proxy
//! for managing sandbox sessions and environments via REST API.
use std::sync::Arc;
use crate::auth::{AuthManager, KimiCodeConfig};
use anyhow::{Context, Result, bail};
use serde::de::DeserializeOwned;
// Re-export sandbox API types from cli-chat-proxy-types for convenience.
// Sorted alphabetically; see sandbox_types.rs for logical grouping.
pub use prod_mc_cli_chat_proxy_types::{
SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse,
SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest,
SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse,
SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse,
SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse,
SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse,
SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse,
SandboxTerminateRequest, SandboxUpdateEnvironmentRequest,
};
// ============================================================================
// Sandbox Client
// ============================================================================
/// HTTP client for interacting with the sandbox API via cli-chat-proxy.
///
/// Path parameters (`session_id`, `environment_id`) are interpolated directly
/// into URLs without percent-encoding. This is safe because these IDs are
/// UUIDs in practice. If ID formats ever change to include URL-unsafe
/// characters, the `format!()` calls should be updated to use percent-encoding.
pub struct SandboxClient {
client: reqwest::Client,
base_url: String,
auth_manager: Arc<AuthManager>,
}
impl SandboxClient {
pub fn new(base_url: impl Into<String>, auth_manager: Arc<AuthManager>) -> Self {
Self {
client: crate::http::shared_client(),
base_url: base_url.into(),
auth_manager,
}
}
/// Returns the base URL.
pub fn base_url(&self) -> &str {
&self.base_url
}
// Do not set Content-Type — callers use .json() and reqwest .header() appends.
async fn auth_headers(
&self,
builder: reqwest::RequestBuilder,
) -> Result<reqwest::RequestBuilder> {
let auth = self
.auth_manager
.auth()
.await
.context("failed to resolve sandbox auth")?;
let mut builder = builder
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION);
if let Some(email) = &auth.email {
builder = builder.header("x-email", email);
}
builder = builder
.header(
"x-grok-client-identifier",
crate::http::process_client_identifier(),
)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
);
Ok(kigi_file_utils::trace_context::inject_trace_context_into_request(builder))
}
/// Check an HTTP response for errors, then deserialize the JSON body.
async fn parse_response<T: DeserializeOwned>(
response: reqwest::Response,
operation: &str,
) -> Result<T> {
if !response.status().is_success() {
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
bail!("{operation} failed: {status} - {body}");
}
response
.json()
.await
.with_context(|| format!("failed to parse {operation} response"))
}
/// Check an HTTP response for errors, discarding the body.
async fn check_response(response: reqwest::Response, operation: &str) -> Result<()> {
if !response.status().is_success() {
let status = response.status().as_u16();
let body = response.text().await.unwrap_or_default();
bail!("{operation} failed: {status} - {body}");
}
Ok(())
}
/// Fork an existing sandbox session.
pub async fn fork_session(&self, request: &SandboxForkRequest) -> Result<SandboxForkResponse> {
let url = format!("{}/sandbox/sessions/fork", self.base_url);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send fork session request")?;
Self::parse_response(response, "fork session").await
}
/// Terminate a sandbox session.
pub async fn terminate_session(
&self,
session_id: &str,
request: &SandboxTerminateRequest,
) -> Result<()> {
let mut url = format!("{}/sandbox/sessions/{}", self.base_url, session_id);
if let Some(env_id) = &request.environment_id {
url = format!("{}?environmentId={}", url, env_id);
}
let response = self
.auth_headers(self.client.delete(&url))
.await?
.send()
.await
.context("failed to send terminate session request")?;
if response.status().as_u16() == 404 {
bail!("session not found: {session_id}");
}
Self::check_response(response, "terminate session").await
}
// ========================================================================
// Session Lifecycle
// ========================================================================
/// Start a sandbox session (non-TUI).
pub async fn start_session(
&self,
request: &SandboxStartRequest,
) -> Result<SandboxStartResponse> {
let url = format!("{}/sandbox/sessions/start", self.base_url);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send start session request")?;
Self::parse_response(response, "start session").await
}
/// Get sandbox session status.
pub async fn get_session_status(&self, session_id: &str) -> Result<SandboxStatusResponse> {
let url = format!("{}/sandbox/sessions/{}/status", self.base_url, session_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get session status request")?;
Self::parse_response(response, "get session status").await
}
/// Get sandbox session logs.
pub async fn get_session_logs(&self, session_id: &str) -> Result<SandboxLogsResponse> {
let url = format!("{}/sandbox/sessions/{}/logs", self.base_url, session_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get session logs request")?;
Self::parse_response(response, "get session logs").await
}
/// Hibernate a sandbox session (snapshot rootfs to GCS and terminate).
pub async fn hibernate_session(&self, session_id: &str) -> Result<SandboxHibernateResponse> {
let url = format!(
"{}/sandbox/sessions/{}/hibernate",
self.base_url, session_id
);
let response = self
.auth_headers(self.client.post(&url))
.await?
.send()
.await
.context("failed to send hibernate session request")?;
Self::parse_response(response, "hibernate session").await
}
/// Restore a previously hibernated sandbox session from its snapshot.
pub async fn restore_session(
&self,
session_id: &str,
request: &SandboxRestoreRequest,
) -> Result<SandboxRestoreResponse> {
let url = format!("{}/sandbox/sessions/{}/restore", self.base_url, session_id);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send restore session request")?;
Self::parse_response(response, "restore session").await
}
// ========================================================================
// Environment CRUD
// ========================================================================
/// List sandbox environments.
pub async fn list_environments(
&self,
request: &SandboxListEnvironmentsRequest,
) -> Result<SandboxListEnvironmentsResponse> {
let url = format!("{}/sandbox/environments", self.base_url);
let mut builder = self.auth_headers(self.client.get(&url)).await?;
if let Some(page) = request.page {
builder = builder.query(&[("page", page)]);
}
if let Some(page_size) = request.page_size {
builder = builder.query(&[("pageSize", page_size)]);
}
let response = builder
.send()
.await
.context("failed to send list environments request")?;
Self::parse_response(response, "list environments").await
}
/// Create a new sandbox environment.
pub async fn create_environment(
&self,
request: &SandboxCreateEnvironmentRequest,
) -> Result<SandboxEnvironmentResponse> {
let url = format!("{}/sandbox/environments", self.base_url);
let response = self
.auth_headers(self.client.post(&url))
.await?
.json(request)
.send()
.await
.context("failed to send create environment request")?;
Self::parse_response(response, "create environment").await
}
/// Get a sandbox environment by ID.
pub async fn get_environment(
&self,
environment_id: &str,
) -> Result<SandboxEnvironmentResponse> {
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send get environment request")?;
Self::parse_response(response, "get environment").await
}
/// Update a sandbox environment.
pub async fn update_environment(
&self,
environment_id: &str,
request: &SandboxUpdateEnvironmentRequest,
) -> Result<SandboxEnvironmentResponse> {
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
let response = self
.auth_headers(self.client.put(&url))
.await?
.json(request)
.send()
.await
.context("failed to send update environment request")?;
Self::parse_response(response, "update environment").await
}
/// Delete a sandbox environment.
pub async fn delete_environment(&self, environment_id: &str) -> Result<()> {
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
let response = self
.auth_headers(self.client.delete(&url))
.await?
.send()
.await
.context("failed to send delete environment request")?;
Self::check_response(response, "delete environment").await
}
/// List preinstalled packages available for sandbox environments.
pub async fn list_preinstalled_packages(
&self,
) -> Result<SandboxListPreinstalledPackagesResponse> {
let url = format!(
"{}/sandbox/environments/preinstalled-packages",
self.base_url
);
let response = self
.auth_headers(self.client.get(&url))
.await?
.send()
.await
.context("failed to send list preinstalled packages request")?;
Self::parse_response(response, "list preinstalled packages").await
}
}
crates/codegen/kigi-shell/src/remote/chat_models_client.rs
-202
View File
@@ -1,202 +0,0 @@
//! grok.com chat-product model catalog (`POST /rest/modes`) — the models
//! grok-web's chat picker shows, distinct from the CLI `/v1/models` build
//! catalog. Transport only; cache + ACP mapping live in
//! [`crate::agent::chat_modes`].
use std::sync::Arc;
use serde::Deserialize;
use crate::auth::AuthManager;
const KIGI_WEB_URL: &str = "https://grok.com";
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Mode {
#[serde(default)]
pub id: String,
#[serde(default)]
pub title: String,
#[serde(default)]
pub description: String,
#[serde(default)]
pub badge_text: Option<String>,
#[serde(default)]
pub availability: ModeAvailability,
#[serde(default)]
pub icon_hint: String,
#[serde(default)]
pub tags: Vec<String>,
}
impl Mode {
pub fn is_available(&self) -> bool {
self.availability.available.is_some()
}
}
/// proto3-JSON oneof: exactly one field is present.
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ModeAvailability {
#[serde(default)]
pub available: Option<serde_json::Value>,
#[serde(default)]
pub unavailable: Option<serde_json::Value>,
#[serde(default)]
pub requires_upgrade: Option<serde_json::Value>,
#[serde(default)]
pub coming_soon: Option<serde_json::Value>,
}
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ListModesResponse {
#[serde(default)]
pub modes: Vec<Mode>,
#[serde(default)]
pub default_mode_id: String,
}
#[derive(Debug, thiserror::Error)]
pub enum ChatModelsError {
#[error("no grok.com credentials")]
NoAuth,
#[error("request timed out")]
Timeout,
#[error("network error: {0}")]
Network(#[from] reqwest::Error),
#[error("request failed: {status}")]
Http { status: u16 },
#[error("parse error: {0}")]
Parse(#[from] serde_json::Error),
}
/// Stateless transport for `POST /rest/modes`; caching lives in
/// [`crate::agent::chat_modes::ChatModesManager`].
pub struct ChatModelsClient {
http: reqwest::Client,
base_url: String,
auth: Arc<AuthManager>,
}
impl ChatModelsClient {
pub fn new(auth: Arc<AuthManager>) -> Self {
let base_url = std::env::var("KIGI_MODES_BASE_URL")
.ok()
.filter(|s| !s.is_empty())
.or_else(|| {
std::env::var("KIGI_CONVERSATIONS_BASE_URL")
.ok()
.filter(|s| !s.is_empty())
})
.or_else(|| {
std::env::var("KIGI_CODE_WEB_URL")
.ok()
.filter(|s| !s.is_empty())
})
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
Self {
http: crate::http::shared_client(),
base_url,
auth,
}
}
/// Gated only on a valid grok.com bearer — deliberately NOT `is_xai_auth()`
/// (unlike workspaces/conversations), since `/rest/modes` is the public chat
/// endpoint and that gate would exclude API-key / cached-token chat users.
pub async fn list_modes(&self, locale: &str) -> Result<ListModesResponse, ChatModelsError> {
let auth = self
.auth
.auth()
.await
.map_err(|_| ChatModelsError::NoAuth)?;
let url = format!("{}/rest/modes", self.base_url);
let body = serde_json::json!({ "locale": locale });
let mut builder = self
.http
.post(&url)
.json(&body)
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
"x-grok-client-identifier",
crate::http::process_client_identifier(),
)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.header(reqwest::header::ACCEPT, "application/json");
if let Some(email) = &auth.email {
builder = builder.header("x-email", email);
}
let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder);
let response = builder.send().await?;
let status = response.status();
if !status.is_success() {
return Err(ChatModelsError::Http {
status: status.as_u16(),
});
}
let bytes = response.bytes().await?;
let resp: ListModesResponse = serde_json::from_slice(&bytes)?;
Ok(resp)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn modes_parse_camelcase_wire() {
let json = serde_json::json!({
"modes": [{
"id": "auto",
"title": "Auto",
"description": "Picks the best model",
"badgeText": "New",
"availability": { "available": {} },
"iconHint": "rocket",
"tags": ["TAG_PRIMARY"]
}, {
"id": "heavy",
"title": "Heavy",
"availability": { "requiresUpgrade": { "message": "Upgrade" } }
}],
"defaultModeId": "auto"
});
let resp: ListModesResponse = serde_json::from_value(json).unwrap();
assert_eq!(resp.modes.len(), 2);
assert_eq!(resp.default_mode_id, "auto");
let auto = &resp.modes[0];
assert_eq!(auto.id, "auto");
assert_eq!(auto.title, "Auto");
assert_eq!(auto.badge_text.as_deref(), Some("New"));
assert_eq!(auto.icon_hint, "rocket");
assert_eq!(auto.tags, vec!["TAG_PRIMARY".to_string()]);
assert!(auto.is_available());
assert!(!resp.modes[1].is_available());
}
#[test]
fn missing_fields_default_gracefully() {
let json = serde_json::json!({ "modes": [{ "id": "m1" }] });
let resp: ListModesResponse = serde_json::from_value(json).unwrap();
let m = &resp.modes[0];
assert_eq!(m.id, "m1");
assert!(m.title.is_empty());
assert!(m.description.is_empty());
assert!(m.badge_text.is_none());
// No availability field on the wire → not selectable.
assert!(!m.is_available());
assert!(resp.default_mode_id.is_empty());
}
}
crates/codegen/kigi-shell/src/remote/client.rs
-2268
View File
File diff suppressed because it is too large Load Diff
crates/codegen/kigi-shell/src/remote/conversations_client.rs
-305
View File
@@ -1,305 +0,0 @@
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use crate::auth::{AuthManager, KimiAuth};
const KIGI_WEB_URL: &str = "https://grok.com";
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Conversation {
#[serde(default)]
pub conversation_id: String,
#[serde(default)]
pub title: String,
#[serde(default)]
pub starred: bool,
#[serde(default)]
pub create_time: Option<String>,
#[serde(default)]
pub modify_time: Option<String>,
#[serde(default)]
pub workspaces: Vec<Workspace>,
}
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Workspace {
#[serde(default)]
pub workspace_id: String,
}
#[derive(Debug, Clone, Default)]
pub struct ConvQuery {
pub page_size: i64,
pub page_token: Option<String>,
pub search_query: Option<String>,
pub workspace_id: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct ListConversationsPage {
pub conversations: Vec<Conversation>,
pub next_page_token: Option<String>,
}
/// Body for `PUT /rest/app-chat/conversations/{id}` (grok-web `chatUpdateConversation`).
#[derive(Debug, Clone, Default, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpdateConversationBody {
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub starred: Option<bool>,
}
#[derive(Debug, thiserror::Error)]
pub enum ConvError {
#[error("no OAuth credentials for conversations:read")]
NoOauth,
#[error("network error: {0}")]
Network(#[from] reqwest::Error),
#[error("request failed: {status}")]
Http { status: u16 },
#[error("parse error: {0}")]
Parse(#[from] serde_json::Error),
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListConversationsResponseWire {
#[serde(default)]
conversations: Vec<Conversation>,
#[serde(default)]
next_page_token: Option<String>,
#[serde(default)]
text_search_matches: Vec<ListConversationsMatchWire>,
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListConversationsMatchWire {
#[serde(default)]
conversation: Option<Conversation>,
}
pub struct ConversationsClient {
http: reqwest::Client,
base_url: String,
auth: Arc<AuthManager>,
}
impl ConversationsClient {
pub fn new(auth: Arc<AuthManager>) -> Self {
let base_url = std::env::var("KIGI_CONVERSATIONS_BASE_URL")
.ok()
.filter(|s| !s.is_empty())
.or_else(|| {
std::env::var("KIGI_CODE_WEB_URL")
.ok()
.filter(|s| !s.is_empty())
})
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
Self {
http: crate::http::shared_client(),
base_url,
auth,
}
}
async fn require_xai_auth(&self) -> Result<KimiAuth, ConvError> {
let auth = self.auth.auth().await.map_err(|_| ConvError::NoOauth)?;
if !auth.is_session_auth() {
return Err(ConvError::NoOauth);
}
Ok(auth)
}
fn apply_auth_headers(
&self,
builder: reqwest::RequestBuilder,
auth: &KimiAuth,
) -> reqwest::RequestBuilder {
let mut builder = builder
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
"x-grok-client-identifier",
crate::http::process_client_identifier(),
)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.header(reqwest::header::ACCEPT, "application/json");
if let Some(email) = &auth.email {
builder = builder.header("x-email", email);
}
kigi_file_utils::trace_context::inject_trace_context_into_request(builder)
}
pub async fn list_conversations(
&self,
q: &ConvQuery,
) -> Result<ListConversationsPage, ConvError> {
let auth = self.require_xai_auth().await?;
let url = format!("{}/rest/app-chat/conversations", self.base_url);
let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())];
if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) {
query.push(("pageToken", token.to_owned()));
}
if let Some(search) = q.search_query.as_deref().filter(|s| !s.is_empty()) {
query.push(("searchQuery", search.to_owned()));
}
if let Some(workspace) = q.workspace_id.as_deref().filter(|s| !s.is_empty()) {
query.push(("workspaceId", workspace.to_owned()));
}
let builder = self.apply_auth_headers(self.http.get(&url).query(&query), &auth);
let response = builder.send().await?;
let status = response.status();
if !status.is_success() {
return Err(ConvError::Http {
status: status.as_u16(),
});
}
let bytes = response.bytes().await?;
let wire: ListConversationsResponseWire = serde_json::from_slice(&bytes)?;
let searching = q.search_query.as_deref().is_some_and(|s| !s.is_empty());
// During an active search, results come exclusively from
// `text_search_matches`. Never fall back to `wire.conversations` here:
// an empty match set means "no hits", and the server may return
// recent/unfiltered conversations in `conversations` that are NOT search
// matches — surfacing those would be wrong.
let conversations = if searching {
wire.text_search_matches
.into_iter()
.filter_map(|m| m.conversation)
.collect()
} else {
wire.conversations
};
Ok(ListConversationsPage {
conversations,
next_page_token: wire.next_page_token.filter(|t| !t.is_empty()),
})
}
/// `PUT /rest/app-chat/conversations/{conversation_id}` — rename and/or star.
pub async fn update_conversation(
&self,
conversation_id: &str,
body: &UpdateConversationBody,
) -> Result<(), ConvError> {
let auth = self.require_xai_auth().await?;
let url = format!(
"{}/rest/app-chat/conversations/{}",
self.base_url,
urlencoding::encode(conversation_id)
);
let builder = self
.apply_auth_headers(self.http.put(&url), &auth)
.json(body);
let response = builder.send().await?;
let status = response.status();
if !status.is_success() {
return Err(ConvError::Http {
status: status.as_u16(),
});
}
Ok(())
}
/// `DELETE /rest/app-chat/conversations/soft/{conversation_id}` — soft-delete.
pub async fn soft_delete_conversation(&self, conversation_id: &str) -> Result<(), ConvError> {
let auth = self.require_xai_auth().await?;
let url = format!(
"{}/rest/app-chat/conversations/soft/{}",
self.base_url,
urlencoding::encode(conversation_id)
);
let builder = self.apply_auth_headers(self.http.delete(&url), &auth);
let response = builder.send().await?;
let status = response.status();
// 404 = already soft-deleted; keep deletion idempotent like the
// build path's `classify_remote_delete`.
if !status.is_success() && status.as_u16() != 404 {
return Err(ConvError::Http {
status: status.as_u16(),
});
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn conversation_parses_camelcase_wire() {
let json = serde_json::json!({
"conversations": [{
"conversationId": "conv_abc",
"title": "Compare GPU vendors",
"starred": true,
"createTime": "2026-06-18T17:30:00Z",
"modifyTime": "2026-06-18T18:02:00Z",
"workspaces": [{ "workspaceId": "ws_9f3a" }]
}],
"nextPageToken": "tok2"
});
let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap();
assert_eq!(wire.conversations.len(), 1);
let c = &wire.conversations[0];
assert_eq!(c.conversation_id, "conv_abc");
assert_eq!(c.title, "Compare GPU vendors");
assert!(c.starred);
assert_eq!(c.modify_time.as_deref(), Some("2026-06-18T18:02:00Z"));
assert_eq!(c.workspaces[0].workspace_id, "ws_9f3a");
assert_eq!(wire.next_page_token.as_deref(), Some("tok2"));
}
#[test]
fn missing_fields_default_gracefully() {
let json = serde_json::json!({ "conversations": [{ "conversationId": "c1" }] });
let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap();
let c = &wire.conversations[0];
assert_eq!(c.conversation_id, "c1");
assert!(c.title.is_empty());
assert!(c.modify_time.is_none());
assert!(c.create_time.is_none());
assert!(c.workspaces.is_empty());
assert!(wire.next_page_token.is_none());
}
#[test]
fn update_body_serializes_only_set_fields() {
let title_only = UpdateConversationBody {
title: Some("New title".into()),
starred: None,
};
assert_eq!(
serde_json::to_value(&title_only).unwrap(),
serde_json::json!({ "title": "New title" })
);
let both = UpdateConversationBody {
title: Some("T".into()),
starred: Some(true),
};
assert_eq!(
serde_json::to_value(&both).unwrap(),
serde_json::json!({ "title": "T", "starred": true })
);
}
}
crates/codegen/kigi-shell/src/remote/mod.rs
-37
View File
@@ -1,37 +0,0 @@
//! Remote storage client for the backend.
pub mod agent;
pub mod chat_models_client;
pub mod client;
pub mod conversations_client;
pub mod pull;
#[cfg(test)]
mod pull_smoke_test;
pub mod sync;
pub mod workspaces_client;
pub use agent::{
SandboxClient, SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse,
SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest,
SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse,
SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse,
SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse,
SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse,
SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse,
SandboxTerminateRequest, SandboxUpdateEnvironmentRequest,
};
pub use chat_models_client::{
ChatModelsClient, ChatModelsError, ListModesResponse, Mode, ModeAvailability,
};
pub use client::{
BackendClient, BackendError, FetchModelsResult, FetchedBundle, fetch_bundle,
fetch_login_device_flow, fetch_settings_blocking, fetch_subagent_bundle, share_url,
};
pub(crate) use client::{DEFAULT_CONTEXT_WINDOW, fetch_models_blocking, models_fetch_origin};
pub use conversations_client::{
ConvError, ConvQuery, Conversation, ConversationsClient, ListConversationsPage,
UpdateConversationBody,
};
pub use pull::{PullResult, pull_session_to_local};
pub use sync::RemoteSync;
pub use workspaces_client::{ListWorkspacesPage, Workspace, WorkspacesClient, WsError, WsQuery};
crates/codegen/kigi-shell/src/remote/pull.rs
-772
View File
@@ -1,772 +0,0 @@
//! Pull-on-miss: fetch a session from the backend and hydrate local JSONL storage.
use crate::remote::client::{BackendClient, BackendError};
#[derive(Debug)]
pub enum PullResult {
/// Written to local storage. The [`Info`] cwd comes from the backend (may differ from caller's).
Hydrated(crate::session::info::Info),
/// Not found on the backend.
NotFound,
}
/// Fetch a session from the backend and hydrate local JSONL storage.
pub async fn pull_session_to_local(
session_id: &str,
client: &BackendClient,
) -> Result<PullResult, BackendError> {
let loaded = match client.load_session_data(session_id).await {
Ok(resp) => resp,
Err(BackendError::SessionNotFound { .. }) => return Ok(PullResult::NotFound),
Err(e) => return Err(e),
};
let remote = match loaded.session.as_ref() {
Some(s) => s,
None => return Ok(PullResult::NotFound),
};
// cwd required for local dir placement; null means pre-writeback session.
let cwd = match remote.cwd.as_ref() {
Some(cwd) => cwd,
None => {
tracing::warn!(session_id, "Cannot pull session: backend has cwd=null");
return Ok(PullResult::NotFound);
}
};
let info = crate::session::info::Info {
id: agent_client_protocol::SessionId::new(std::sync::Arc::from(session_id)),
cwd: cwd.clone(),
};
let dir = crate::session::persistence::session_dir(&info);
let num_messages = hydrate::write_to_dir(&dir, &loaded)?;
tracing::info!(session_id, %cwd, num_messages, "Pulled session from backend");
Ok(PullResult::Hydrated(info))
}
pub(crate) mod hydrate {
use std::path::Path;
use std::sync::Arc;
use crate::remote::client::{BackendError, LoadDataResponse, LoadedMessage, SessionInfo};
use crate::session::info::Info;
use crate::session::persistence::{CHAT_FORMAT_VERSION, Summary, default_model_id};
fn io_err(path: &Path, source: std::io::Error) -> BackendError {
BackendError::Hydration {
path: path.to_path_buf(),
source,
}
}
/// Write all session files to `dir`.
pub(super) fn write_to_dir(
dir: &Path,
loaded: &LoadDataResponse,
) -> Result<usize, BackendError> {
let remote = loaded
.session
.as_ref()
.expect("caller checked session.is_some()");
let info = Info {
id: agent_client_protocol::SessionId::new(Arc::from(remote.session_id.as_str())),
cwd: remote.cwd.clone().expect("caller verified cwd is Some"),
};
std::fs::create_dir_all(dir).map_err(|e| io_err(dir, e))?;
let num_messages = loaded.messages.as_ref().map_or(0, |m| m.len());
let mut num_chat_messages = 0;
if let Some(ref messages) = loaded.messages {
write_updates(dir, messages)?;
num_chat_messages = rebuild_chat_history(dir)?;
}
write_summary(dir, &info, remote, num_messages, num_chat_messages)?;
write_remote_origin_marker(dir);
Ok(num_messages)
}
fn write_summary(
dir: &Path,
info: &Info,
remote: &SessionInfo,
num_messages: usize,
num_chat_messages: usize,
) -> Result<(), BackendError> {
let meta = remote.metadata.as_ref();
let model_id = meta
.and_then(|m| m.get("modelId"))
.and_then(|v| v.as_str())
.map(agent_client_protocol::ModelId::new)
.unwrap_or_else(default_model_id);
let parent_session_id = meta
.and_then(|m| m.get("parentSessionId"))
.and_then(|v| v.as_str())
.map(String::from);
let summary = Summary {
info: info.clone(),
session_summary: remote.title.clone().unwrap_or_default(),
created_at: parse_rfc3339_or_now(remote.created_at.as_deref()),
updated_at: parse_rfc3339_or_now(remote.updated_at.as_deref()),
num_messages,
num_chat_messages,
current_model_id: model_id,
parent_session_id,
forked_at: None,
collection_id: None,
next_trace_turn: 0,
chat_format_version: CHAT_FORMAT_VERSION,
prompt_display_cwd: None,
session_kind: None,
fork_context_source: None,
fork_parent_prompt_id: None,
inherited_prefix_len: None,
hidden: None,
source_workspace_dir: None,
git_root_dir: None,
git_remotes: Vec::new(),
head_commit: None,
head_branch: None,
request_id: None,
// Record the *local* kigi_home (where this hydrated copy lives),
// not the original remote session's, since reconstruction runs locally.
kigi_home: crate::session::persistence::kigi_home_string(),
last_active_at: None,
generated_title: None,
title_is_manual: false,
worktree_label: None,
agent_name: None,
// Hydrated locally — record the profile this process runs under.
sandbox_profile: kigi_sandbox::configured_profile_name().map(String::from),
reasoning_effort: None,
};
let json = serde_json::to_string_pretty(&summary)?;
write_file(&dir.join("summary.json"), json.as_bytes())
}
/// Convert backend JSON-RPC messages to local updates.jsonl (replayable methods only).
pub(super) fn write_updates(
dir: &Path,
messages: &[LoadedMessage],
) -> Result<(), BackendError> {
use std::io::Write;
let path = dir.join("updates.jsonl");
let file = std::fs::File::create(&path).map_err(|e| io_err(&path, e))?;
let mut w = std::io::BufWriter::new(file);
for msg in messages {
let parsed = match serde_json::from_str::<serde_json::Value>(&msg.content) {
Ok(v) => v,
Err(_) => continue,
};
if !is_session_update(&parsed) {
continue;
}
if let Some(line) = to_envelope_line(&parsed) {
let _ = w.write_all(line.as_bytes());
let _ = w.write_all(b"\n");
}
}
w.flush().map_err(|e| io_err(&path, e))
}
/// Rebuild `chat_history.jsonl` from `updates.jsonl` so pulled sessions are continuable.
fn rebuild_chat_history(dir: &Path) -> Result<usize, BackendError> {
use crate::session::storage::UpdatesIterator;
use std::io::{Seek, Write};
let updates_path = dir.join("updates.jsonl");
let Some(iter) =
UpdatesIterator::open(&updates_path).map_err(|e| io_err(&updates_path, e))?
else {
return Ok(0);
};
let chat_path = dir.join("chat_history.jsonl");
let file = std::fs::File::create(&chat_path).map_err(|e| io_err(&chat_path, e))?;
let mut writer = std::io::BufWriter::new(file);
let mut reducer = ChatReducer::new();
for result in iter {
let update = match result {
Ok(u) => u,
Err(_) => continue,
};
for item in reducer.process(&update) {
if let Ok(line) = serde_json::to_string(&item) {
let _ = writer.write_all(line.as_bytes());
let _ = writer.write_all(b"\n");
}
}
// CompactionCheckpoint: truncate file and reset
if reducer.should_truncate() {
reducer.clear_truncate_flag();
let _ = writer.seek(std::io::SeekFrom::Start(0));
let _ = writer.get_mut().set_len(0);
}
}
// Flush trailing state
for item in reducer.flush() {
if let Ok(line) = serde_json::to_string(&item) {
let _ = writer.write_all(line.as_bytes());
let _ = writer.write_all(b"\n");
}
}
writer.flush().map_err(|e| io_err(&chat_path, e))?;
Ok(reducer.count())
}
use crate::sampling::{AssistantItem, ContentPart, ConversationItem, ToolCall};
use agent_client_protocol as acp;
use std::collections::{HashMap, HashSet};
/// Reduces ACP session updates into conversation items.
///
/// Turn boundaries: User→Agent flushes user, Agent→User flushes agent,
/// tool completion flushes agent before emitting result.
struct ChatReducer {
user_parts: Vec<ContentPart>,
agent_text: String,
agent_tool_calls: Vec<ToolCall>,
in_user_turn: bool,
has_agent_content: bool,
needs_truncate: bool,
tool_args: HashMap<String, String>,
emitted_tool_results: HashSet<String>,
item_count: usize,
}
impl ChatReducer {
fn new() -> Self {
Self {
user_parts: Vec::new(),
agent_text: String::new(),
agent_tool_calls: Vec::new(),
in_user_turn: false,
has_agent_content: false,
needs_truncate: false,
tool_args: HashMap::new(),
emitted_tool_results: HashSet::new(),
item_count: 0,
}
}
fn process(
&mut self,
update: &crate::session::storage::SessionUpdate,
) -> Vec<ConversationItem> {
use crate::session::storage::SessionUpdate;
match update {
SessionUpdate::Acp(n) => self.handle_acp(&n.update),
SessionUpdate::Xai(n) => self.handle_xai(&n.update),
}
}
fn handle_acp(&mut self, update: &acp::SessionUpdate) -> Vec<ConversationItem> {
match update {
acp::SessionUpdate::UserMessageChunk(chunk) => self.on_user_chunk(chunk),
acp::SessionUpdate::AgentMessageChunk(chunk) => self.on_agent_chunk(chunk),
acp::SessionUpdate::ToolCall(tc) => self.on_tool_call(tc),
acp::SessionUpdate::ToolCallUpdate(tc) => self.on_tool_call_update(tc),
_ => Vec::new(), // AgentThoughtChunk, Retry, Plan not needed
}
}
fn handle_xai(
&mut self,
update: &crate::extensions::notification::SessionUpdate,
) -> Vec<ConversationItem> {
use crate::extensions::notification::SessionUpdate as XaiUpdate;
match update {
XaiUpdate::CompactionCheckpoint(_) => {
self.reset();
self.needs_truncate = true;
Vec::new()
}
_ => Vec::new(), // DiffReview, MemoryFlush, etc. not needed
}
}
fn on_user_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec<ConversationItem> {
let mut out = Vec::new();
if !self.in_user_turn {
out.extend(self.flush_agent());
self.in_user_turn = true;
}
match &chunk.content {
acp::ContentBlock::Text(t) => {
self.user_parts.push(ContentPart::Text {
text: std::sync::Arc::<str>::from(t.text.clone()),
});
}
acp::ContentBlock::Image(img) => {
if let Some(uri) = &img.uri {
self.user_parts.push(ContentPart::Image {
url: std::sync::Arc::<str>::from(uri.clone()),
});
}
}
_ => {} // Audio, Resource, etc. not needed for chat replay
}
out
}
fn on_agent_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec<ConversationItem> {
let mut out = Vec::new();
if self.in_user_turn {
out.extend(self.flush_user());
self.in_user_turn = false;
}
if let acp::ContentBlock::Text(t) = &chunk.content {
self.agent_text.push_str(&t.text);
self.has_agent_content = true;
}
out
}
fn on_tool_call(&mut self, tc: &acp::ToolCall) -> Vec<ConversationItem> {
let id = tc.tool_call_id.0.to_string();
let args = tc
.raw_input
.as_ref()
.map(|v| v.to_string())
.unwrap_or_default();
self.tool_args.insert(id.clone(), args.clone());
self.agent_tool_calls.push(ToolCall {
id: std::sync::Arc::<str>::from(id),
name: tc.title.clone(),
arguments: std::sync::Arc::<str>::from(args),
});
Vec::new()
}
fn on_tool_call_update(&mut self, tc: &acp::ToolCallUpdate) -> Vec<ConversationItem> {
let id = tc.tool_call_id.0.to_string();
self.maybe_backfill_args(&id, &tc.fields);
if Self::is_completed(&tc.fields) && self.emitted_tool_results.insert(id.clone()) {
return self.emit_tool_result(&id, &tc.fields);
}
Vec::new()
}
/// Backfill tool arguments from ToolCallUpdate if ToolCall didn't have them.
fn maybe_backfill_args(&mut self, id: &str, fields: &acp::ToolCallUpdateFields) {
let Some(raw) = &fields.raw_input else { return };
let needs_backfill = self.tool_args.get(id).is_none_or(String::is_empty);
if !needs_backfill {
return;
}
let args = raw.to_string();
self.tool_args.insert(id.to_string(), args.clone());
if let Some(call) = self
.agent_tool_calls
.iter_mut()
.find(|c| c.id.as_ref() == id)
{
call.arguments = std::sync::Arc::<str>::from(args);
}
}
fn is_completed(fields: &acp::ToolCallUpdateFields) -> bool {
matches!(
fields.status,
Some(acp::ToolCallStatus::Completed | acp::ToolCallStatus::Failed)
)
}
fn emit_tool_result(
&mut self,
id: &str,
fields: &acp::ToolCallUpdateFields,
) -> Vec<ConversationItem> {
let mut out = Vec::new();
out.extend(self.flush_agent());
let content = extract_tool_result_text(fields);
let item = ConversationItem::tool_result(id.to_string(), content);
self.item_count += 1;
out.push(item);
out
}
fn flush_user(&mut self) -> Option<ConversationItem> {
if self.user_parts.is_empty() {
return None;
}
let item = ConversationItem::user_with_parts(std::mem::take(&mut self.user_parts));
self.item_count += 1;
Some(item)
}
fn flush_agent(&mut self) -> Option<ConversationItem> {
if !self.has_agent_content && self.agent_tool_calls.is_empty() {
return None;
}
let item = ConversationItem::Assistant(AssistantItem {
content: std::sync::Arc::<str>::from(std::mem::take(&mut self.agent_text)),
tool_calls: std::mem::take(&mut self.agent_tool_calls),
model_id: None,
model_fingerprint: None,
reasoning_effort: None,
});
self.has_agent_content = false;
self.item_count += 1;
Some(item)
}
fn flush(&mut self) -> Vec<ConversationItem> {
let mut out = Vec::new();
out.extend(self.flush_user());
out.extend(self.flush_agent());
out
}
fn reset(&mut self) {
self.user_parts.clear();
self.agent_text.clear();
self.agent_tool_calls.clear();
self.tool_args.clear();
self.emitted_tool_results.clear();
self.in_user_turn = false;
self.has_agent_content = false;
self.item_count = 0;
}
fn should_truncate(&self) -> bool {
self.needs_truncate
}
fn clear_truncate_flag(&mut self) {
self.needs_truncate = false;
}
fn count(&self) -> usize {
self.item_count
}
}
/// Extract displayable text from a completed ToolCallUpdate.
fn extract_tool_result_text(fields: &agent_client_protocol::ToolCallUpdateFields) -> String {
if let Some(content) = &fields.content {
let text: String = content
.iter()
.filter_map(|c| match c {
agent_client_protocol::ToolCallContent::Content(
agent_client_protocol::Content {
content: agent_client_protocol::ContentBlock::Text(t),
..
},
) => Some(t.text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("");
if !text.is_empty() {
return text;
}
}
if let Some(raw) = &fields.raw_output {
return raw.to_string();
}
String::new()
}
fn write_remote_origin_marker(dir: &Path) {
let _ = std::fs::write(
dir.join(".remote_origin"),
format!("pulled_at={}\n", chrono::Utc::now().to_rfc3339()),
);
}
/// Replayable JSON-RPC methods (excludes metadata like `prompt_complete`).
const REPLAYABLE_METHODS: &[&str] = &["session/update", "_x.ai/session/update"];
fn is_session_update(json_rpc: &serde_json::Value) -> bool {
json_rpc
.get("method")
.and_then(|v| v.as_str())
.is_some_and(|m| REPLAYABLE_METHODS.contains(&m))
}
fn to_envelope_line(json_rpc: &serde_json::Value) -> Option<String> {
let method = json_rpc.get("method").and_then(|v| v.as_str())?;
let params = json_rpc.get("params").cloned().unwrap_or_default();
serde_json::to_string(&serde_json::json!({
"timestamp": 0u64,
"method": method,
"params": params,
}))
.ok()
}
fn parse_rfc3339_or_now(s: Option<&str>) -> chrono::DateTime<chrono::Utc> {
s.and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
.map(|dt| dt.with_timezone(&chrono::Utc))
.unwrap_or_else(chrono::Utc::now)
}
fn write_file(path: &Path, data: &[u8]) -> Result<(), BackendError> {
std::fs::write(path, data).map_err(|e| io_err(path, e))
}
}
#[cfg(test)]
mod tests {
use crate::remote::client::LoadedMessage;
#[test]
fn hydrate_writes_valid_updates_jsonl() {
let tmp = tempfile::TempDir::new().unwrap();
let messages = vec![
LoadedMessage {
id: "1".into(),
content: r#"{"method":"session/update","params":{"update":"hello"}}"#.into(),
timestamp: None,
},
LoadedMessage {
id: "2".into(),
content: r#"{"method":"session/update","params":{"update":"world"}}"#.into(),
timestamp: None,
},
];
super::hydrate::write_updates(tmp.path(), &messages).unwrap();
let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap();
let lines: Vec<&str> = content.lines().collect();
assert_eq!(lines.len(), 2);
for line in &lines {
let v: serde_json::Value = serde_json::from_str(line).unwrap();
assert_eq!(v["timestamp"], 0);
assert_eq!(v["method"], "session/update");
assert!(v["params"].is_object());
}
}
#[test]
fn rebuild_chat_history_merges_chunks() {
use crate::session::export::ExportedMessage;
use agent_client_protocol::{ContentBlock, ContentChunk, SessionUpdate, TextContent};
use std::sync::Arc;
// Build ACP notifications matching the RemoteSync path
let sid = agent_client_protocol::SessionId::new(Arc::from("test"));
let notifications = [
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("hello "),
))),
),
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("world"),
))),
),
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("hi back"),
))),
),
];
// Serialize through ExportedMessage (writeback path)
let messages: Vec<LoadedMessage> = notifications
.iter()
.map(|n| {
let exported = ExportedMessage::from_notification(n);
LoadedMessage {
id: "x".into(),
content: exported.content,
timestamp: None,
}
})
.collect();
let data = crate::remote::client::LoadDataResponse {
messages: Some(messages),
session: Some(crate::remote::client::SessionInfo {
session_id: "test".into(),
title: None,
cwd: Some("/tmp".into()),
status: None,
created_at: None,
updated_at: None,
metadata: None,
}),
};
let tmp = tempfile::TempDir::new().unwrap();
super::hydrate::write_to_dir(tmp.path(), &data).unwrap();
let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap();
let items: Vec<crate::sampling::ConversationItem> = chat
.lines()
.filter(|l| !l.is_empty())
.filter_map(|l| serde_json::from_str(l).ok())
.collect();
assert_eq!(items.len(), 2, "should have 1 user + 1 agent item");
assert!(matches!(
&items[0],
crate::sampling::ConversationItem::User(_)
));
assert!(matches!(
&items[1],
crate::sampling::ConversationItem::Assistant(_)
));
if let crate::sampling::ConversationItem::User(u) = &items[0] {
let text: String = u
.content
.iter()
.filter_map(|p| match p {
crate::sampling::ContentPart::Text { text } => Some(text.as_ref()),
_ => None,
})
.collect();
assert_eq!(text, "hello world");
}
}
#[test]
fn rebuild_chat_history_preserves_user_images() {
use crate::session::export::ExportedMessage;
use agent_client_protocol::{
ContentBlock, ContentChunk, ImageContent, SessionUpdate, TextContent,
};
use std::sync::Arc;
let sid = agent_client_protocol::SessionId::new(Arc::from("test"));
let notifications = [
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("look at this"),
))),
),
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Image(
ImageContent::new(String::new(), String::new())
.uri(Some("data:image/png;base64,abc".into())),
))),
),
agent_client_protocol::SessionNotification::new(
sid.clone(),
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("I see an image"),
))),
),
];
let messages: Vec<LoadedMessage> = notifications
.iter()
.map(|n| LoadedMessage {
id: "x".into(),
content: ExportedMessage::from_notification(n).content,
timestamp: None,
})
.collect();
let data = crate::remote::client::LoadDataResponse {
messages: Some(messages),
session: Some(crate::remote::client::SessionInfo {
session_id: "test".into(),
title: None,
cwd: Some("/tmp".into()),
status: None,
created_at: None,
updated_at: None,
metadata: None,
}),
};
let tmp = tempfile::TempDir::new().unwrap();
super::hydrate::write_to_dir(tmp.path(), &data).unwrap();
let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap();
let items: Vec<crate::sampling::ConversationItem> = chat
.lines()
.filter(|l| !l.is_empty())
.filter_map(|l| serde_json::from_str(l).ok())
.collect();
assert_eq!(items.len(), 2);
if let crate::sampling::ConversationItem::User(u) = &items[0] {
assert_eq!(u.content.len(), 2, "should have text + image parts");
assert!(matches!(
&u.content[0],
crate::sampling::ContentPart::Text { .. }
));
assert!(matches!(
&u.content[1],
crate::sampling::ContentPart::Image { .. }
));
} else {
panic!("expected User item");
}
}
#[test]
fn hydrate_skips_invalid_messages() {
let tmp = tempfile::TempDir::new().unwrap();
let messages = vec![
LoadedMessage {
id: "1".into(),
content: r#"{"method":"session/update","params":{}}"#.into(),
timestamp: None,
},
LoadedMessage {
id: "bad".into(),
content: "not valid json".into(),
timestamp: None,
},
LoadedMessage {
id: "3".into(),
content: r#"{"method":"session/update","params":{"x":1}}"#.into(),
timestamp: None,
},
];
super::hydrate::write_updates(tmp.path(), &messages).unwrap();
let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap();
let lines: Vec<&str> = content.lines().collect();
assert_eq!(lines.len(), 2, "invalid message should be skipped");
}
}
crates/codegen/kigi-shell/src/remote/pull_smoke_test.rs
-127
View File
@@ -1,127 +0,0 @@
//! Push → pull round-trip smoke test against the live backend.
//!
//! Run with: `cargo test -p kigi-shell -- pull_smoke --ignored --nocapture`
#[cfg(test)]
mod tests {
use crate::auth::KimiAuth;
use crate::remote::client::BackendClient;
use crate::session::storage::{JsonlStorageAdapter, StorageAdapter};
use std::collections::BTreeMap;
use std::sync::Arc;
fn load_prod_auth() -> Option<KimiAuth> {
let path = crate::util::kigi_home::kigi_home().join("auth.json");
let contents = std::fs::read_to_string(&path).ok()?;
let store: BTreeMap<String, KimiAuth> = serde_json::from_str(&contents).ok()?;
let scope = crate::auth::KimiCodeConfig::default().auth_scope();
crate::auth::lookup_auth(&store, &scope)
}
/// Full round-trip using the real RemoteSync production code path:
/// create RemoteSync → queue ACP notifications → flush → verify on
/// backend → pull back → verify local hydration + storage adapter load.
#[tokio::test]
#[ignore]
async fn smoke_push_pull_round_trip() {
use crate::remote::sync::RemoteSync;
use crate::session::export::ExportedMetadata;
use agent_client_protocol::{
ContentBlock, ContentChunk, SessionNotification, SessionUpdate, TextContent,
};
let auth = load_prod_auth().expect("No auth.json — run `grok login`");
let am = Arc::new(crate::auth::AuthManager::new(
&crate::util::kigi_home::kigi_home(),
crate::auth::KimiCodeConfig::default(),
));
am.hot_swap(auth);
let client = BackendClient::new().with_auth_manager(am.clone());
let session_id = format!("test-rt-{}", uuid::Uuid::new_v4());
let test_cwd = "/tmp/smoke-test".to_string();
let test_title = "Push-Pull Round Trip Test";
// PUSH via RemoteSync (real production path)
let metadata = ExportedMetadata {
title: Some(test_title.into()),
cwd: test_cwd.clone(),
model_id: Some("grok-3".into()),
created_at: Some(chrono::Utc::now().to_rfc3339()),
updated_at: Some(chrono::Utc::now().to_rfc3339()),
total_messages: None,
parent_session_id: None,
session_kind: None,
subagent_type: None,
subagent_persona: None,
subagent_role: None,
fork_context_source: None,
subagent_depth: None,
};
let sync = RemoteSync::new(
session_id.clone(),
metadata,
BackendClient::new().with_auth_manager(am.clone()),
);
let sid = agent_client_protocol::SessionId::new(Arc::from(session_id.as_str()));
sync.queue(SessionNotification::new(
sid.clone(),
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("Hello from smoke test — user".to_string()),
))),
));
sync.queue(SessionNotification::new(
sid.clone(),
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
TextContent::new("Hello from smoke test — agent".to_string()),
))),
));
sync.flush();
tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
// Verify backend has cwd, title, messages
let loaded = client
.load_session_data(&session_id)
.await
.expect("load after push failed");
let remote = loaded.session.as_ref().expect("no session row");
assert_eq!(remote.cwd.as_deref(), Some(test_cwd.as_str()));
assert_eq!(remote.title.as_deref(), Some(test_title));
assert!(loaded.messages.as_ref().map_or(0, |m| m.len()) >= 2);
// PULL back to local
let result = crate::remote::pull_session_to_local(&session_id, &client)
.await
.expect("pull failed");
let pulled = match result {
crate::remote::PullResult::Hydrated(info) => info,
crate::remote::PullResult::NotFound => panic!("pull returned NotFound"),
};
assert_eq!(pulled.cwd, test_cwd);
// Verify local storage loads
let local_dir = crate::session::persistence::session_dir(&pulled);
assert!(local_dir.join("summary.json").exists());
assert!(local_dir.join("updates.jsonl").exists());
let storage = JsonlStorageAdapter::default();
let data = storage
.load_session_without_updates(&pulled)
.await
.expect("storage load failed");
assert_eq!(data.summary.session_summary, test_title);
// Verify chat_history has both turns
let chat =
std::fs::read_to_string(local_dir.join("chat_history.jsonl")).unwrap_or_default();
assert!(chat.contains("user"), "chat_history missing user turn");
assert!(chat.contains("agent"), "chat_history missing agent turn");
// Cleanup
drop(sync);
let _ = client.delete_session_data(&session_id).await;
let _ = std::fs::remove_dir_all(&local_dir);
}
}
crates/codegen/kigi-shell/src/remote/sync.rs
-167
View File
@@ -1,167 +0,0 @@
//! Writeback push: async queue that flushes session updates to the backend.
//!
//! `RemoteSync` runs a background tokio task that buffers ACP notifications
//! and flushes them to the backend via [`BackendClient::save_session_data()`].
//!
//! ## Backpressure
//!
//! When the buffer exceeds [`MAX_PENDING`], the task attempts an emergency
//! flush. If that also fails (network down), the oldest messages are dropped
//! to prevent unbounded memory growth.
//!
//! ## Drop behavior
//!
//! When `RemoteSync` is dropped, the sender half of the channel closes and
//! the background task exits. **Pending buffered messages are lost.** This
//! is acceptable because the local JSONL files are the source of truth —
//! writeback is best-effort.
use crate::remote::BackendClient;
use crate::session::export::{ExportedMessage, ExportedMetadata};
use agent_client_protocol as acp;
use tokio::sync::mpsc;
/// Max buffered notifications before triggering an emergency flush.
/// Sized to keep memory under ~50MB even with large notifications.
const MAX_PENDING: usize = 512;
/// How many oldest messages to drop when an emergency flush fails.
/// Dropping a batch (not one-by-one) avoids repeated failed flushes.
const DROP_BATCH_SIZE: usize = 64;
enum SyncMsg {
Queue(Box<acp::SessionNotification>),
Flush,
SetTitle(String),
SetModelId(String),
}
#[derive(Clone)]
pub struct RemoteSync {
tx: mpsc::UnboundedSender<SyncMsg>,
}
impl RemoteSync {
/// Metadata is included on every flush to keep the backend session row current.
pub(crate) fn new(
session_id: String,
metadata: ExportedMetadata,
client: BackendClient,
) -> Self {
let (tx, rx) = mpsc::unbounded_channel();
tokio::spawn(sync_task(session_id, metadata, client, rx));
Self { tx }
}
pub fn queue(&self, notification: acp::SessionNotification) {
let _ = self.tx.send(SyncMsg::Queue(Box::new(notification)));
}
pub fn flush(&self) {
let _ = self.tx.send(SyncMsg::Flush);
}
pub fn set_title(&self, title: String) {
let _ = self.tx.send(SyncMsg::SetTitle(title));
}
pub fn set_model_id(&self, model_id: String) {
let _ = self.tx.send(SyncMsg::SetModelId(model_id));
}
}
async fn do_flush(
client: &BackendClient,
session_id: &str,
metadata: &ExportedMetadata,
pending: &mut Vec<acp::SessionNotification>,
) -> bool {
if pending.is_empty() {
return true;
}
let messages: Vec<ExportedMessage> = pending
.iter()
.map(ExportedMessage::from_notification)
.collect();
match client
.save_session_data(session_id, &messages, Some(metadata))
.await
{
Ok(()) => {
tracing::debug!(count = pending.len(), "Writeback: synced");
pending.clear();
// Link session to agent so the relay can route requests to it.
if let Err(e) = client
.upsert_session(session_id, metadata, &crate::util::agent_id::agent_id())
.await
{
tracing::warn!(error = %e, "Writeback: failed to upsert session");
}
true
}
Err(e) => {
tracing::warn!(error = %e, pending = pending.len(), "Writeback: flush failed");
false
}
}
}
async fn sync_task(
session_id: String,
mut metadata: ExportedMetadata,
client: BackendClient,
mut rx: mpsc::UnboundedReceiver<SyncMsg>,
) {
let mut pending: Vec<acp::SessionNotification> = Vec::new();
while let Some(msg) = rx.recv().await {
match msg {
SyncMsg::Queue(n) => {
if pending.len() >= MAX_PENDING {
tracing::warn!(
pending = pending.len(),
"Writeback: buffer full, attempting emergency flush"
);
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
if !do_flush(&client, &session_id, &metadata, &mut pending).await {
let dropped = pending.drain(0..DROP_BATCH_SIZE.min(pending.len())).count();
tracing::error!(
dropped = dropped,
"Writeback: emergency flush failed, dropping oldest messages"
);
}
}
pending.push(*n);
}
SyncMsg::Flush => {
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
do_flush(&client, &session_id, &metadata, &mut pending).await;
}
SyncMsg::SetTitle(title) => {
metadata.title = Some(title);
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
if let Err(e) = client
.save_session_data(&session_id, &[], Some(&metadata))
.await
{
tracing::warn!(?e, "Writeback: failed to sync title to backend");
}
}
SyncMsg::SetModelId(id) => {
metadata.model_id = Some(id);
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
if let Err(e) = client
.save_session_data(&session_id, &[], Some(&metadata))
.await
{
tracing::warn!(?e, "Writeback: failed to sync model_id to backend");
}
}
}
}
}
crates/codegen/kigi-shell/src/remote/workspaces_client.rs
-176
View File
@@ -1,176 +0,0 @@
use std::sync::Arc;
use serde::Deserialize;
use crate::auth::AuthManager;
const KIGI_WEB_URL: &str = "https://grok.com";
#[derive(Debug, Clone, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Workspace {
#[serde(default)]
pub workspace_id: String,
#[serde(default)]
pub name: String,
#[serde(default)]
pub create_time: Option<String>,
#[serde(default)]
pub kind: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct WsQuery {
pub page_size: i64,
pub page_token: Option<String>,
pub query: Option<String>,
pub kind: Option<String>,
}
#[derive(Debug, Clone, Default)]
pub struct ListWorkspacesPage {
pub workspaces: Vec<Workspace>,
pub next_page_token: Option<String>,
}
#[derive(Debug, thiserror::Error)]
pub enum WsError {
#[error("no OAuth credentials for workspaces:read")]
NoOauth,
#[error("network error: {0}")]
Network(#[from] reqwest::Error),
#[error("request failed: {status}")]
Http { status: u16 },
#[error("parse error: {0}")]
Parse(#[from] serde_json::Error),
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
struct ListWorkspacesResponseWire {
#[serde(default)]
workspaces: Vec<Workspace>,
#[serde(default)]
next_page_token: Option<String>,
}
pub struct WorkspacesClient {
http: reqwest::Client,
base_url: String,
auth: Arc<AuthManager>,
}
impl WorkspacesClient {
pub fn new(auth: Arc<AuthManager>) -> Self {
let base_url = first_nonempty_env(&[
"KIGI_WORKSPACES_BASE_URL",
"KIGI_CONVERSATIONS_BASE_URL",
"KIGI_CODE_WEB_URL",
])
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
Self {
http: crate::http::shared_client(),
base_url,
auth,
}
}
pub async fn list_workspaces(&self, q: &WsQuery) -> Result<ListWorkspacesPage, WsError> {
let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?;
if !auth.is_session_auth() {
return Err(WsError::NoOauth);
}
let url = format!("{}/rest/workspaces", self.base_url);
let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())];
if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) {
query.push(("pageToken", token.to_owned()));
}
if let Some(search) = q.query.as_deref().filter(|s| !s.is_empty()) {
query.push(("query", search.to_owned()));
}
if let Some(kind) = q.kind.as_deref().filter(|s| !s.is_empty()) {
query.push(("kind", kind.to_owned()));
}
let mut builder = self
.http
.get(&url)
.query(&query)
.header("Authorization", format!("Bearer {}", auth.key))
.header("x-userid", &auth.user_id)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
"x-grok-client-identifier",
crate::http::process_client_identifier(),
)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.header(reqwest::header::ACCEPT, "application/json");
if let Some(email) = &auth.email {
builder = builder.header("x-email", email);
}
let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder);
let response = builder.send().await?;
let status = response.status();
if !status.is_success() {
return Err(WsError::Http {
status: status.as_u16(),
});
}
let bytes = response.bytes().await?;
let wire: ListWorkspacesResponseWire = serde_json::from_slice(&bytes)?;
Ok(ListWorkspacesPage {
workspaces: wire.workspaces,
next_page_token: wire.next_page_token.filter(|t| !t.is_empty()),
})
}
}
fn first_nonempty_env(keys: &[&str]) -> Option<String> {
keys.iter()
.find_map(|k| std::env::var(k).ok().filter(|s| !s.is_empty()))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn workspace_parses_camelcase_wire() {
let json = serde_json::json!({
"workspaces": [{
"workspaceId": "ws_9f3a",
"name": "GPU vendor research",
"createTime": "2026-06-18T17:30:00Z",
"kind": "WORKSPACE_KIND_IMAGINE"
}],
"nextPageToken": "tok2"
});
let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap();
assert_eq!(wire.workspaces.len(), 1);
let w = &wire.workspaces[0];
assert_eq!(w.workspace_id, "ws_9f3a");
assert_eq!(w.name, "GPU vendor research");
assert_eq!(w.create_time.as_deref(), Some("2026-06-18T17:30:00Z"));
assert_eq!(w.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE"));
assert_eq!(wire.next_page_token.as_deref(), Some("tok2"));
}
#[test]
fn missing_fields_default_gracefully() {
let json = serde_json::json!({ "workspaces": [{ "workspaceId": "w1" }] });
let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap();
let w = &wire.workspaces[0];
assert_eq!(w.workspace_id, "w1");
assert!(w.name.is_empty());
assert!(w.create_time.is_none());
assert!(w.kind.is_none());
assert!(wire.next_page_token.is_none());
}
}
crates/codegen/kigi-shell/src/sampling/error.rs
+28 -27
View File
@@ -17,26 +17,36 @@ use agent_client_protocol as acp;
/// see this code and show a user-friendly upgrade message instead.
pub const RATE_LIMITED_ERROR_CODE: i32 = -32003;
/// OAuth / session rate-limit copy (personal plan upgrade path).
pub const RATE_LIMITED_USER_MESSAGE_OAUTH: &str =
"You\u{2019}ve hit the rate limit for your plan. Upgrade your account or try again later.";
/// Subscription (OAuth) rate-limit copy. PRD Q3: the official Kimi CLI and
/// Kigi draw on the SAME subscription quota, so the message says so — a user
/// who also runs `kimi` should understand why the limit arrived early.
/// Deliberately promises no reset duration; the quota window is server-side.
pub static RATE_LIMITED_USER_MESSAGE_OAUTH: std::sync::LazyLock<String> =
std::sync::LazyLock::new(|| {
format!(
"You\u{2019}ve hit the usage limit of your Kimi subscription. Note that Kigi and the \
official Kimi CLI share the same subscription quota. Upgrade your plan at {} or try \
again later.",
kigi_env::upgrade_page_url()
)
});
/// API key / team rate-limit copy. Personal grok.com upgrades do not raise API
/// team limits; admins purchase credits or a higher spend-based tier.
/// See https://docs.x.ai/developers/rate-limits#rate-limit-tiers
pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit your team\u{2019}s API rate limit. Ask a team admin to purchase more credits for higher limits, or try again later. See https://docs.x.ai/developers/rate-limits#rate-limit-tiers";
/// Moonshot API-key rate-limit copy. Platform keys are tier-limited (RPM/TPM);
/// raising the tier happens in the Moonshot Open Platform console, not via a
/// Kimi subscription.
pub const RATE_LIMITED_USER_MESSAGE_API_KEY: &str = "You\u{2019}ve hit the rate limit for your Moonshot API key. Check your tier\u{2019}s limits in the Moonshot Open Platform console (platform.moonshot.ai or platform.moonshot.cn), or try again later.";
/// Pick rate-limit copy from the *active* auth method.
///
/// Pass the real `is_api_key_auth` flag (pager `AppView`, `AuthMethodKind::is_api_key`
/// for the selected method). Do **not** decide from `has_xai_api_key_env()` alone:
/// for the selected method). Do **not** decide from the key env var alone:
/// when both an env key and a cached OAuth session exist, auth prefers the
/// cached session over the API key.
pub fn rate_limited_user_message(is_api_key_auth: bool) -> &'static str {
if is_api_key_auth {
RATE_LIMITED_USER_MESSAGE_API_KEY
} else {
RATE_LIMITED_USER_MESSAGE_OAUTH
RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()
}
}
@@ -318,20 +328,20 @@ mod tests {
fn rate_limited_user_message_oauth_vs_api_key() {
assert_eq!(
rate_limited_user_message(false),
RATE_LIMITED_USER_MESSAGE_OAUTH
RATE_LIMITED_USER_MESSAGE_OAUTH.as_str()
);
assert_eq!(
rate_limited_user_message(true),
RATE_LIMITED_USER_MESSAGE_API_KEY
);
assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("Upgrade your account"));
assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("team"));
assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("credits"));
assert!(
RATE_LIMITED_USER_MESSAGE_API_KEY
.contains("https://docs.x.ai/developers/rate-limits#rate-limit-tiers")
);
assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Upgrade your account"));
// PRD Q3: the subscription copy must state the shared quota with the
// official Kimi CLI and point at the upgrade page.
assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("official Kimi CLI"));
assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains("same subscription quota"));
assert!(RATE_LIMITED_USER_MESSAGE_OAUTH.contains(kigi_env::upgrade_page_url()));
// API-key copy points at the Moonshot platform, not the subscription.
assert!(RATE_LIMITED_USER_MESSAGE_API_KEY.contains("Moonshot"));
assert!(!RATE_LIMITED_USER_MESSAGE_API_KEY.contains("subscription quota"));
}
#[test]
@@ -341,7 +351,6 @@ mod tests {
message: "Rate limit exceeded".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
assert_eq!(acp_err.code, acp::ErrorCode::from(RATE_LIMITED_ERROR_CODE));
@@ -359,7 +368,6 @@ mod tests {
message: "Rate limit exceeded".into(),
model_metadata: None,
retry_after_secs: Some(60),
should_retry: None,
};
assert_eq!(err.retry_after(), Some(60));
let acp_err = map_sampling_err_to_acp(err);
@@ -374,14 +382,12 @@ mod tests {
message: "limited".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let server_err = SamplingError::Api {
status: StatusCode::INTERNAL_SERVER_ERROR,
message: "oops".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let rate_acp = map_sampling_err_to_acp(rate_err);
let server_acp = map_sampling_err_to_acp(server_err);
@@ -398,7 +404,6 @@ mod tests {
message: "bad token".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
assert_eq!(acp_err.code, acp::Error::auth_required().code);
@@ -420,7 +425,6 @@ mod tests {
.into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
assert_ne!(
@@ -476,7 +480,6 @@ mod tests {
message: "The model 'grok-build' requires a Grok subscription.".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
let data = acp_err.data.unwrap();
@@ -501,7 +504,6 @@ mod tests {
message: "The model 'grok-build' requires a Grok subscription.".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
let data = acp_err.data.unwrap();
@@ -522,7 +524,6 @@ mod tests {
message: "Content violates usage guidelines.".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
};
let acp_err = map_sampling_err_to_acp(err);
let data = acp_err.data.unwrap();
crates/codegen/kigi-shell/src/session/acp_session.rs
-2
View File
@@ -679,8 +679,6 @@ pub(crate) struct SessionActor {
pub(crate) origin_client: Option<crate::http::OriginClientInfo>,
/// Feedback manager for signal tracking and feedback request heuristics
pub(crate) feedback_manager: Arc<FeedbackManager>,
/// Cancellation token for the feedback sync loop (None if no feedback client)
pub(crate) sync_loop_cancel: Option<tokio_util::sync::CancellationToken>,
/// The fully-built Agent: owns the ToolBridge, system prompt, policies,
/// and the AgentDefinition. Replaces the old `tool_bridge` + `agent_definition` fields.
/// Wrapped in `RefCell` for mid-session mutation (skill refresh, prompt regen).
crates/codegen/kigi-shell/src/session/acp_session_impl/model_switch.rs
+1 -2
View File
@@ -1,5 +1,5 @@
use super::*;
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
use kigi_chat_state::conversation_util::replace_or_insert_system_head;
impl SessionActor {
pub(super) async fn handle_set_session_model(
@@ -72,7 +72,6 @@ impl SessionActor {
existing.auth_type,
),
alpha_test_key: existing.alpha_test_key,
client_version: sampling_config.client_version.clone(),
});
self.model_auth_facts.replace(None);
self.signals_handle()
crates/codegen/kigi-shell/src/session/acp_session_impl/prompt_build.rs
-1
View File
@@ -691,7 +691,6 @@ impl SessionActor {
crate::agent::config::finalize_image_describe_sampler_config(
resolved_describe,
&active_session_config,
self.client_identifier.clone(),
Some(self.max_retries),
);
let client = kigi_sampler::SamplingClient::new(sampler_config).map_err(|e| {
crates/codegen/kigi-shell/src/session/acp_session_impl/recap.rs
+1 -1
View File
@@ -3,7 +3,7 @@
use super::*;
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
impl SessionActor {
/// Handle a /btw side question — single-turn model call using the
crates/codegen/kigi-shell/src/session/acp_session_impl/reminders.rs
+1 -1
View File
@@ -134,7 +134,7 @@ pub(super) fn build_todo_gate_reminder(pending: &[&str], unbacked_in_progress: &
/// (which is disabled). Extracted from `spawn_session_actor` so the
/// precedence rules are unit-testable. Named `resolve_*` to match the
/// sibling precedence helpers in `crate::util::config`
/// (`resolve_zdr_access_enabled`, `resolve_restore_code`, …).
/// (`resolve_restore_code`, …).
pub(crate) fn resolve_reminder_policy(
remote: Option<&crate::util::config::RemoteSettings>,
todo_gate: bool,
crates/codegen/kigi-shell/src/session/acp_session_impl/run_loop.rs
+5 -9
View File
@@ -205,8 +205,7 @@ pub(super) async fn run_session(
.emit_buffered(notification). await; }
if let Some(tx) = respond_to { let _ =
tx.send(()); } } } } } maybe_completion = completion_rx.recv() => { let
Some((prompt_id, result)) = maybe_completion else { if let Some(cancel) = &
session.sync_loop_cancel { cancel.cancel(); } cleanup_session_scratch(&
Some((prompt_id, result)) = maybe_completion else { cleanup_session_scratch(&
session); return; }; if let Some(notification) = replay_buffer.flush() {
session.emit_buffered(notification). await; } let (turn_succeeded,
infra_pause_message) = SessionActor::post_turn_goal_degradation_plan(&
@@ -265,8 +264,7 @@ pub(super) async fn run_session(
{ let
model_id = session.current_model_id(). await; if let Some(signals) = session
.signals_handle().snapshot(). await {
} } if let
Some(cancel) = & session.sync_loop_cancel { cancel.cancel(); } session
} } session
.feedback_manager.shutdown(). await; if ! session
.startup_hints.is_subagent { session.persist_background_task_manifest().
await; } cleanup_session_scratch(& session); return; }; match cmd {
@@ -327,8 +325,7 @@ pub(super) async fn run_session(
::agent::config::try_resolve_model_credentials(model_name.as_str(), existing
.api_key.as_deref()) { session.chat_state_handle
.update_credentials(kigi_chat_state::Credentials { api_key : r.api_key,
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key,
client_version : existing.client_version, }); } session.model_auth_facts
auth_type : r.auth_type, alpha_test_key : existing.alpha_test_key, }); } session.model_auth_facts
.replace(None); } } SessionCommand::GetCurrentModel { responds_to } => { let
model = session.chat_state_handle.get_sampling_config(). await .map(| c | c
.model).unwrap_or_default(); let _ = responds_to.send(model); }
@@ -697,7 +694,7 @@ pub(super) async fn run_session(
await; session.send_hook_execution("session_start", None, None, & results).
await; } } SessionCommand::GetFeedbackContext { turn_number, responds_to } =>
{ let s = session.clone(); tokio::task::spawn_local(async move { use
prod_mc_cli_chat_proxy_types::feedback_types::FeedbackToolOutcome; let
crate::session::feedback_types::FeedbackToolOutcome; let
turn_idx = turn_number.and_then(| n | usize::try_from(n).ok()); let
(last_user_message, last_assistant_message) = match turn_idx { Some(n) => {
let conv = s.chat_state_handle.get_conversation(). await;
@@ -809,8 +806,7 @@ pub(super) async fn run_session(
"MEMORY_SUBAGENT_SKIP: skipping on_session_end for subagent session"); }
session.maybe_run_dream(). await; let telem = session.memory
.telemetry_snapshot(); session.emit_memory_session_summary(& telem,
total_chunks_at_end, session_end_result); if let Some(cancel) = & session
.sync_loop_cancel { cancel.cancel(); } session.feedback_manager
total_chunks_at_end, session_end_result); session.feedback_manager
.shutdown(). await; if ! session.startup_hints
.is_subagent { session.persist_background_task_manifest(). await; }
cleanup_session_scratch(& session); return; } } }
crates/codegen/kigi-shell/src/session/acp_session_impl/sampler_turn.rs
+1 -14
View File
@@ -49,7 +49,7 @@ impl SessionTokenAuthGate {
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_xai_url(base_url),
endpoint_is_first_party: crate::util::is_first_party_url(base_url),
}
}
fn active(self) -> bool {
@@ -314,22 +314,11 @@ impl SessionActor {
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: Some(self.max_retries),
stream_tool_calls: cfg.stream_tool_calls.unwrap_or(false),
idle_timeout_secs: None,
client_identifier: self.client_identifier.clone(),
deployment_id: crate::managed_config::resolve_deployment_id(
crate::managed_config::resolve_deployment_key().as_deref(),
),
user_id: self
.auth_manager
.as_ref()
.and_then(|am| am.current_or_expired())
.filter(|a| a.is_session_auth())
.map(|a| a.user_id),
origin_client: self.origin_client.clone(),
attribution_callback: self.attribution_callback.clone(),
bearer_resolver: if use_bearer_resolver {
@@ -482,7 +471,6 @@ impl SessionActor {
&endpoints,
session_key.as_deref(),
creds.alpha_test_key.clone(),
creds.client_version.clone(),
)
}
/// Resolve a dedicated sampler for the Auto-mode classifier model `slug`,
@@ -499,7 +487,6 @@ impl SessionActor {
crate::agent::config::stamp_session_local_sampler_fields(
&mut cfg,
&active_session_config,
self.client_identifier.clone(),
Some(self.max_retries),
);
let model = cfg.model.clone();
crates/codegen/kigi-shell/src/session/acp_session_impl/session_setup.rs
+7 -13
View File
@@ -323,10 +323,10 @@ impl SessionActor {
/// Check if the session has been idle and proactively refresh model metadata.
///
/// Called at the start of each turn. If idle > `IDLE_REFRESH_THRESHOLD_SECS`,
/// fetches `/models-v2` from cli-chat-proxy and updates the cached
/// fetches `/models` from cli-chat-proxy and updates the cached
/// context_window / max_completion_tokens if remote settings changed them.
///
/// Skipped for BYOK users (no remote settings, no `/models-v2`).
/// Skipped for BYOK users (no remote settings, no `/models`).
pub(super) async fn maybe_refresh_model_metadata_on_resume(&self) {
if !self.is_session_based_auth() {
return;
@@ -353,7 +353,7 @@ impl SessionActor {
tracing::info!(
idle_secs,
threshold_secs = Self::IDLE_REFRESH_THRESHOLD_SECS,
"Session resumed after idle — refreshing model metadata from cli-chat-proxy"
"Session resumed after idle — refreshing model metadata"
);
let creds = self.chat_state_handle.get_credentials().await;
let Some(ref am) = self.auth_manager else {
@@ -370,27 +370,21 @@ impl SessionActor {
);
let middleware_client =
crate::http::with_auth_retry(crate::http::shared_client(), provider);
let url = format!("{}/models-v2", base_url);
let url = format!("{}/models", base_url);
let parse_models_response =
|json: serde_json::Value| -> Option<(std::num::NonZeroU64, Option<u32>)> {
let data = json.get("data")?.as_array()?;
for entry in data {
let parsed = crate::remote::client::parse_remote_model_value(entry, base_url)?;
let parsed =
crate::agent::models_fetch::parse_remote_model_value(entry, base_url)?;
if parsed.model == *current_model {
return Some((parsed.context_window, parsed.max_completion_tokens));
}
}
None
};
#[allow(unused_mut)]
let mut request = middleware_client
let request = middleware_client
.get(&url)
.header("X-XAI-Token-Auth", "xai-grok-cli")
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
)
.timeout(std::time::Duration::from_secs(5));
let response = match request.send().await {
Ok(r) => r,
crates/codegen/kigi-shell/src/session/acp_session_impl/slash_exec.rs
+1 -1
View File
@@ -828,7 +828,7 @@ impl SessionActor {
);
let model_id = sampling_config.map(|c| c.model);
let resolved_model_id = model_metadata.resolved_model_id;
let client_version = credentials.client_version;
let client_version = Some(kigi_version::VERSION.to_string());
use crate::session::feedback_manager::{SessionFeedbackData, SubmitOutcome};
let outcome = self
crates/codegen/kigi-shell/src/session/acp_session_impl/spawn.rs
+18 -54
View File
@@ -3,7 +3,7 @@
//! the MCP auto-restart wiring (`SessionRestartActions`).
#![allow(clippy::items_after_test_module)]
use super::*;
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
/// Partition CLI `--allow` rules under the pin: blanket catch-all allows
/// (`Allow(Any)` `*` / `**`, plus bare/match-all Bash/MCP/WebFetch grants — see
/// `resolution::is_catchall_allow`) substitute for the blocked `--yolo`, so drop them when
@@ -123,10 +123,7 @@ pub(crate) async fn spawn_session_actor(
codebase_indexes: std::sync::Arc<parking_lot::Mutex<CodebaseIndexManager>>,
code_nav_enabled: bool,
fs_watch_caps: fs_watch::FsWatchCapabilities,
feedback_proxy_url: Option<String>,
feedback_user_token: Option<String>,
feedback_alpha_test_key: Option<String>,
deployment_key: Option<String>,
feedback_base_url: Option<String>,
client_terminal_capable: bool,
client_fs_capable: bool,
gateway_enabled: std::sync::Arc<std::sync::atomic::AtomicBool>,
@@ -141,7 +138,6 @@ pub(crate) async fn spawn_session_actor(
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>,
loc_tracking_enabled: bool,
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
@@ -879,36 +875,30 @@ pub(crate) async fn spawn_session_actor(
}
persist_chat_history_jsonl_sync(&session_info, &conversation);
chat_state_handle.replace_conversation(conversation);
let feedback_client = feedback_proxy_url.map(|base_url| {
let mut client =
crate::agent::feedback_client::FeedbackClient::new(base_url, feedback_user_token)
.with_alpha_test_key(feedback_alpha_test_key)
.with_deployment_key(deployment_key);
if let Some(am) = auth_manager.as_ref() {
client = client.with_auth_manager(am.clone());
}
client
});
let feedback_client = match (feedback_base_url, auth_manager.as_ref()) {
(Some(base_url), Some(am)) => Some(
crate::agent::feedback_client::FeedbackClient::new(base_url, am.clone())
.with_session_id(session_info.id.0.to_string()),
),
_ => None,
};
let has_feedback_client = feedback_client.is_some();
tracing::info!(
session_id = % session_info.id.0, has_feedback_client = has_feedback_client,
"Creating feedback manager"
);
let feedback_client_type = match client_type {
ClientType::GrokTUI => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui,
ClientType::GrokWeb => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Web,
ClientType::Nebula => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Nebula,
ClientType::Extension => {
prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Extension
}
ClientType::Generic => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Agent,
ClientType::Desktop => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop,
ClientType::GrokPager => prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui,
ClientType::GrokTUI => crate::session::feedback_types::ClientType::Tui,
ClientType::GrokWeb => crate::session::feedback_types::ClientType::Web,
ClientType::Nebula => crate::session::feedback_types::ClientType::Nebula,
ClientType::Extension => crate::session::feedback_types::ClientType::Extension,
ClientType::Generic => crate::session::feedback_types::ClientType::Agent,
ClientType::Desktop => crate::session::feedback_types::ClientType::Desktop,
ClientType::GrokPager => crate::session::feedback_types::ClientType::Tui,
};
let feedback_config = FeedbackManagerConfig {
feedback_enabled: feedback_flags.enabled,
client_type: feedback_client_type,
loc_tracking_enabled,
..Default::default()
};
let feedback_manager = Arc::new(FeedbackManager::new(
@@ -930,11 +920,6 @@ pub(crate) async fn spawn_session_actor(
}
signals_handle.set_primary_model(&primary_model_id);
signals_handle.set_tracing_config(inference_idle_timeout_secs);
let sync_loop_cancel = if has_feedback_client {
Some(tokio_util::sync::CancellationToken::new())
} else {
None
};
let force_compact = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false));
let resolved_workspace_root = kigi_workspace::session::git::find_git_root_from_path(
std::path::Path::new(&session_info.cwd),
@@ -1141,7 +1126,6 @@ pub(crate) async fn spawn_session_actor(
client_identifier: session_client_identifier.clone(),
origin_client: origin_client.clone(),
feedback_manager: feedback_manager.clone(),
sync_loop_cancel: sync_loop_cancel.clone(),
agent: std::cell::RefCell::new(agent),
last_reported_branch: Arc::new(Mutex::new(None)),
git_head_enabled: fs_watch_caps.git_head,
@@ -1375,18 +1359,6 @@ pub(crate) async fn spawn_session_actor(
}
});
}
if let Some(cancel) = sync_loop_cancel {
tracing::info!(session_id = % session_info.id.0, "Spawning feedback sync loop");
let fm = feedback_manager.clone();
tokio::spawn(async move {
fm.run_sync_loop(cancel).await;
});
} else {
tracing::debug!(
session_id = % session_info.id.0,
"No feedback client available, skipping sync loop"
);
}
{
use agent_client_protocol::Client as _;
use kigi_tools::implementations::grok_build::ask_user_question::{
@@ -1600,10 +1572,7 @@ pub(crate) async fn spawn_session_on_thread(
codebase_indexes: std::sync::Arc<parking_lot::Mutex<CodebaseIndexManager>>,
code_nav_enabled: bool,
fs_watch_caps: fs_watch::FsWatchCapabilities,
feedback_proxy_url: Option<String>,
feedback_user_token: Option<String>,
feedback_alpha_test_key: Option<String>,
deployment_key: Option<String>,
feedback_base_url: Option<String>,
client_terminal_capable: bool,
client_fs_capable: bool,
gateway_enabled: std::sync::Arc<std::sync::atomic::AtomicBool>,
@@ -1618,7 +1587,6 @@ pub(crate) async fn spawn_session_on_thread(
persisted_goal_mode: Option<crate::session::goal_tracker::GoalOrchestration>,
persisted_announcement_state: Option<crate::session::announcement_state::AnnouncementState>,
memory_config: Option<crate::config::MemoryConfig>,
loc_tracking_enabled: bool,
feedback_flags: crate::session::feedback_manager::FeedbackFlags,
managed_mcp_handle: crate::session::managed_mcp::ManagedMcpStateHandle,
managed_mcp_expires_at: Option<chrono::DateTime<chrono::Utc>>,
@@ -1751,10 +1719,7 @@ pub(crate) async fn spawn_session_on_thread(
codebase_indexes,
code_nav_enabled,
fs_watch_caps,
feedback_proxy_url,
feedback_user_token,
feedback_alpha_test_key,
deployment_key,
feedback_base_url,
client_terminal_capable,
client_fs_capable,
gateway_enabled,
@@ -1769,7 +1734,6 @@ pub(crate) async fn spawn_session_on_thread(
persisted_goal_mode,
persisted_announcement_state,
memory_config,
loc_tracking_enabled,
feedback_flags,
managed_mcp_handle,
managed_mcp_expires_at,
crates/codegen/kigi-shell/src/session/acp_session_tests/auth_error_no_retry_tests.rs
-4
View File
@@ -850,15 +850,11 @@ async fn set_session_model_invalidates_byok_memo_for_same_model_id() {
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
crates/codegen/kigi-shell/src/session/acp_session_tests/cancel_running_task_tests.rs
-20
View File
@@ -47,15 +47,11 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 100_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
@@ -193,7 +189,6 @@ async fn persist_ack_waits_for_disk_flush_before_success() {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -338,15 +333,11 @@ async fn first_turn_memory_injection_persists_to_chat_history() {
api_backend: Default::default(),
auth_scheme: Default::default(),
context_window: 100_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
@@ -470,15 +461,11 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
api_backend: Default::default(),
auth_scheme: Default::default(),
context_window: 100_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
@@ -641,7 +628,6 @@ async fn first_turn_memory_injection_disabled_does_not_persist_to_chat_history()
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -890,7 +876,6 @@ async fn cancel_running_task_teardown_clears_running_and_pending_work() {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(agent),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -1729,15 +1714,11 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 100_000,
client_version: None,
force_http1: false,
max_retries: Some(0),
stream_tool_calls: false,
idle_timeout_secs: Some(60),
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
@@ -1876,7 +1857,6 @@ async fn cancel_propagates_to_sampler_handle_so_no_further_emission() {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(agent),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
crates/codegen/kigi-shell/src/session/acp_session_tests/idle_resume_tests.rs
+4 -6
View File
@@ -43,7 +43,7 @@ async fn test_last_api_request_at_idle_detection() {
/// End-to-end test for `maybe_refresh_model_metadata_on_resume`.
///
/// Simulates a session idle for >10 minutes, then verifies the function
/// fetches `/models-v2`, parses the response, and updates `context_window`
/// fetches `/models`, parses the response, and updates `context_window`
/// and `max_completion_tokens` in the sampling config.
#[tokio::test(flavor = "current_thread")]
async fn test_e2e_idle_resume_refreshes_model_metadata() {
@@ -52,7 +52,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
local
.run_until(async {
let app = axum::Router::new().route(
"/v1/models-v2",
"/v1/models",
get(|| async {
axum::Json(serde_json::json!(
{ "data" : [{ "model" : "test-model", "name" : "Test Model",
@@ -117,7 +117,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
api_key: Some("test-key".to_string()),
auth_type: Default::default(),
alpha_test_key: None,
client_version: None,
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let actor = SessionActor {
@@ -219,7 +218,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -321,12 +319,12 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
assert_eq!(
cfg_after.context_window,
std::num::NonZeroU64::new(300_000).unwrap(),
"context_window should be updated to 300K from /models-v2"
"context_window should be updated to 300K from /models"
);
assert_eq!(
cfg_after.max_completion_tokens,
Some(16384),
"max_completion_tokens should be updated to 16384 from /models-v2"
"max_completion_tokens should be updated to 16384 from /models"
);
})
.await;
crates/codegen/kigi-shell/src/session/acp_session_tests/inline_auto_compact_flow_tests.rs
+4 -8
View File
@@ -152,7 +152,6 @@ async fn create_test_actor(
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -591,7 +590,6 @@ async fn create_test_actor_with_memory(
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-memory")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -1168,7 +1166,7 @@ async fn test_compact_on_error_no_trigger_when_tokens_within_new_window() {
/// End-to-end test for `maybe_refresh_model_metadata_on_resume`.
///
/// Simulates a session idle for >10 minutes, then verifies the function
/// fetches `/models-v2`, parses the response, and updates `context_window`
/// fetches `/models`, parses the response, and updates `context_window`
/// and `max_completion_tokens` in the sampling config.
#[tokio::test(flavor = "current_thread")]
async fn test_e2e_idle_resume_refreshes_model_metadata() {
@@ -1177,7 +1175,7 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
local
.run_until(async {
let app = axum::Router::new().route(
"/v1/models-v2",
"/v1/models",
get(|| async {
axum::Json(serde_json::json!(
{ "data" : [{ "model" : "test-model", "name" : "Test Model",
@@ -1241,7 +1239,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
api_key: Some("test-key".to_string()),
auth_type: Default::default(),
alpha_test_key: None,
client_version: None,
});
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let actor = SessionActor {
@@ -1346,7 +1343,6 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
@@ -1448,12 +1444,12 @@ async fn test_e2e_idle_resume_refreshes_model_metadata() {
assert_eq!(
cfg_after.context_window,
std::num::NonZeroU64::new(300_000).unwrap(),
"context_window should be updated to 300K from /models-v2"
"context_window should be updated to 300K from /models"
);
assert_eq!(
cfg_after.max_completion_tokens,
Some(16384),
"max_completion_tokens should be updated to 16384 from /models-v2"
"max_completion_tokens should be updated to 16384 from /models"
);
})
.await;
crates/codegen/kigi-shell/src/session/acp_session_tests/memory_config_tests.rs
-1
View File
@@ -211,7 +211,6 @@ async fn create_test_actor_with_memory(
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-memory")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
crates/codegen/kigi-shell/src/session/acp_session_tests/reactive_managed_reauth_e2e_tests.rs
+1 -1
View File
@@ -166,7 +166,7 @@ async fn actor_with_proxy(
let cfg = crate::agent::config::Config {
endpoints: crate::agent::config::EndpointsConfig {
cli_chat_proxy_base_url: Some(proxy_base.to_string()),
coding_api_base_url: Some(proxy_base.to_string()),
..Default::default()
},
..Default::default()
crates/codegen/kigi-shell/src/session/acp_session_tests/replay_buffer_send_update_tests.rs
-1
View File
@@ -157,7 +157,6 @@ pub(super) async fn make_replay_send_update_fixture() -> ReplaySendUpdateFixture
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
crates/codegen/kigi-shell/src/session/acp_session_tests/support.rs
-1
View File
@@ -271,7 +271,6 @@ pub(crate) async fn create_test_actor_ex(
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
crates/codegen/kigi-shell/src/session/acp_session_tests/web_search_e2e_tests.rs
-3
View File
@@ -64,9 +64,6 @@ async fn web_search_uses_model_override_from_config_end_to_end() {
entry,
crate::agent::config::resolve_credentials(entry, None),
None,
None,
None,
None,
);
let web_search_sampling = crate::tools::config::web_search_sampling_config(resolved);
crates/codegen/kigi-shell/src/session/acp_types.rs
+11 -11
View File
@@ -88,11 +88,11 @@ pub struct ClientFeedbackInput {
pub session_id: String,
/// Type of client submitting feedback
pub client_type: prod_mc_cli_chat_proxy_types::feedback_types::ClientType,
pub client_type: crate::session::feedback_types::ClientType,
/// Rating type (thumbs, stars, nps)
#[serde(default)]
pub rating_type: Option<prod_mc_cli_chat_proxy_types::feedback_types::RatingType>,
pub rating_type: Option<crate::session::feedback_types::RatingType>,
/// Rating value (interpretation depends on rating_type):
/// - thumbs: -1 (down), 0 (neutral), 1 (up)
@@ -113,7 +113,7 @@ pub struct ClientFeedbackInput {
/// Context type for the feedback
#[serde(default)]
pub context_type: Option<prod_mc_cli_chat_proxy_types::feedback_types::ContextType>,
pub context_type: Option<crate::session::feedback_types::ContextType>,
/// 0-based turn number this feedback is about.
#[serde(default, alias = "turnNumber")]
@@ -134,7 +134,7 @@ pub struct ClientFeedbackInput {
/// Terminal environment snapshot from the client.
#[serde(default)]
pub terminal_info: Option<prod_mc_cli_chat_proxy_types::feedback_types::FeedbackTerminalInfo>,
pub terminal_info: Option<crate::session::feedback_types::FeedbackTerminalInfo>,
}
impl ClientFeedbackInput {
@@ -144,10 +144,10 @@ impl ClientFeedbackInput {
/// - stars: 1 to 5
/// - nps: 0 to 10
fn clamp_rating_value(
rating_type: Option<prod_mc_cli_chat_proxy_types::feedback_types::RatingType>,
rating_type: Option<crate::session::feedback_types::RatingType>,
rating_value: Option<i32>,
) -> Option<i32> {
use prod_mc_cli_chat_proxy_types::feedback_types::RatingType;
use crate::session::feedback_types::RatingType;
match (rating_type, rating_value) {
(Some(RatingType::Thumbs), Some(v)) => Some(v.clamp(-1, 1)),
@@ -175,8 +175,8 @@ impl ClientFeedbackInput {
resolved_model_id: Option<String>,
model_fingerprint: Option<String>,
turn_number: Option<i64>,
) -> prod_mc_cli_chat_proxy_types::feedback_types::FeedbackSubmission {
use prod_mc_cli_chat_proxy_types::feedback_types::FeedbackContent;
) -> crate::session::feedback_types::FeedbackSubmission {
use crate::session::feedback_types::FeedbackContent;
let clamped_rating_value = Self::clamp_rating_value(self.rating_type, self.rating_value);
let content = match (
@@ -577,7 +577,7 @@ pub struct SessionInfoResponse {
pub struct FeedbackContext {
pub last_user_message: Option<String>,
pub last_assistant_message: Option<String>,
pub tool_outcomes: Vec<prod_mc_cli_chat_proxy_types::feedback_types::FeedbackToolOutcome>,
pub tool_outcomes: Vec<crate::session::feedback_types::FeedbackToolOutcome>,
pub compaction_count: i64,
pub context_window_usage: u8,
pub context_tokens_used: u64,
@@ -655,14 +655,14 @@ mod tests {
let input: ClientFeedbackInput = serde_json::from_str(json).unwrap();
assert_eq!(
input.client_type,
prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop
crate::session::feedback_types::ClientType::Desktop
);
assert_eq!(input.session_id, "sess-1");
let submission = input.to_submission(Some("grok-3".into()), None, None, Some(5));
assert_eq!(
submission.client_type,
prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Desktop
crate::session::feedback_types::ClientType::Desktop
);
assert_eq!(submission.client_type.to_string(), "desktop");
}
crates/codegen/kigi-shell/src/session/compaction.rs
+1 -2
View File
@@ -7,7 +7,7 @@
//! lives alongside the primary one in `acp_session.rs`.
use super::SessionActor;
use super::is_project_instructions;
use crate::remote::DEFAULT_CONTEXT_WINDOW;
use crate::agent::models_fetch::DEFAULT_CONTEXT_WINDOW;
use crate::session::compaction_config::{
AsyncCompactionCache, SUPPRESS_NONE, SUPPRESS_STICKY, SUPPRESS_TURN, SUPPRESS_UNTIL_SUCCESS,
};
@@ -2242,7 +2242,6 @@ mod inline_auto_compact_flow_tests {
client_identifier: None,
origin_client: None,
feedback_manager: Arc::new(FeedbackManager::local_only("test-session")),
sync_loop_cancel: None,
agent: std::cell::RefCell::new(test_agent_default().await),
last_reported_branch: std::sync::Arc::new(parking_lot::Mutex::new(None)),
git_head_enabled: false,
crates/codegen/kigi-shell/src/session/feedback.rs
+3 -5
View File
@@ -10,9 +10,7 @@ use super::signals::SessionSignals;
use crate::util::probabilistic_sample;
// Re-export shared feedback API wire types to avoid duplication
pub use prod_mc_cli_chat_proxy_types::feedback_types::{
FeedbackHeuristicsConfig, FeedbackMode, TierConfig,
};
pub use crate::session::feedback_types::{FeedbackHeuristicsConfig, FeedbackMode, TierConfig};
/// Feedback request tier with associated probability and criteria.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
@@ -288,7 +286,7 @@ impl FeedbackHeuristics {
/// Create a heuristics evaluator from a remote feedback-heuristics config.
pub fn from_config(config: &FeedbackHeuristicsConfig) -> Self {
use prod_mc_cli_chat_proxy_types::feedback_types::parse_feedback_mode_str;
use crate::session::feedback_types::parse_feedback_mode_str;
Self {
enabled: config.enabled,
@@ -343,7 +341,7 @@ impl FeedbackHeuristics {
/// Update the heuristics configuration from a loaded config.
/// Preserves the triggered_tiers state and request tracking.
pub fn update_config(&mut self, config: &FeedbackHeuristicsConfig) {
use prod_mc_cli_chat_proxy_types::feedback_types::parse_feedback_mode_str;
use crate::session::feedback_types::parse_feedback_mode_str;
self.enabled = config.enabled;
crates/codegen/kigi-shell/src/session/feedback_manager.rs
+54 -601
View File
@@ -3,57 +3,48 @@
//! This manager coordinates:
//! - Signal tracking via SessionSignalsHandle
//! - Heuristics evaluation to determine when to request feedback
//! - Periodic sync of signals to the feedback/analytics backend
//! - Background loading of feedback configuration from the backend
//! - Creating feedback request records when triggered
//! - Sending feedback request notifications to clients
//! - Local persistence of every feedback record
//! - Forwarding text feedback to the Kimi Code feedback endpoint for
//! subscription (OAuth) sessions
//!
//! ## Usage
//! ```ignore
//! // Create the manager when a session starts
//! let manager = FeedbackManager::new(session_id, feedback_api_url, user_token);
//! let manager = FeedbackManager::new(session_id, feedback_client, config);
//!
//! // Get the signals handle to pass around for event tracking
//! let signals = manager.signals_handle();
//!
//! // Spawn the background sync task (also loads config)
//! tokio::spawn(manager.run_sync_loop());
//!
//! // Track events
//! signals.increment_turn();
//! signals.record_tool_call("read_file");
//!
//! // Check for feedback after each turn
//! // This also records the request with the feedback API if triggered
//! if let Some(request) = manager.maybe_request_feedback(None).await {
//! // Send FeedbackRequest notification to client
//! }
//! ```
use std::ops::ControlFlow;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use tokio::sync::RwLock;
use crate::agent::feedback_client::{
FeedbackApiError, FeedbackClient, signals_to_update, snapshot_to_turn_delta,
};
use crate::agent::feedback_client::FeedbackClient;
use crate::session::feedback::{
FeedbackEvaluation, FeedbackHeuristics, FeedbackRequest, FeedbackTier, TriggerCondition,
};
use crate::session::signals::{SessionSignalsActor, SessionSignalsHandle, TurnDeltaSnapshot};
use crate::session::signals::{SessionSignalsActor, SessionSignalsHandle};
use prod_mc_cli_chat_proxy_types::feedback_types::{
ClientType, ContextType, CreateFeedbackRequestInput, FeedbackContent, FeedbackMode,
FeedbackSubmission, FeedbackToolOutcome,
use crate::session::feedback_types::{
ClientType, FeedbackContent, FeedbackMode, FeedbackSubmission, FeedbackToolOutcome,
};
use crate::session::persistence::{LocalFeedbackEntry, PersistenceMsg, UserFeedbackEntry};
pub(crate) enum SubmitOutcome {
Submitted,
/// No server configured for this session.
/// Persisted locally only: no subscription session, or a rating-only
/// record with no text content for the Kimi feedback endpoint.
LocalOnly,
/// Server request failed.
Failed(anyhow::Error),
@@ -70,8 +61,9 @@ pub(crate) fn new_submission(
s
}
/// Pipeline: persist → strip → submit. Callers merge `KIGI_USER_METADATA` and
/// set `submission.request_id`.
/// Pipeline: persist locally → forward text content to the Kimi feedback
/// endpoint (subscription sessions only). Callers merge `KIGI_USER_METADATA`
/// and set `submission.request_id`.
pub(crate) async fn submit_feedback_workflow(
submission: &mut FeedbackSubmission,
feedback_client: Option<&FeedbackClient>,
@@ -96,42 +88,33 @@ pub(crate) async fn submit_feedback_workflow(
}
}
let telemetry_model_id = submission.model_id.clone();
let telemetry_rating_value = submission.rating_value;
let telemetry_session_id = submission.session_id.clone();
let has_feedback_text = submission
.feedback_text
.as_ref()
.is_some_and(|t| !t.is_empty());
let request_id = submission.request_id.clone();
let appearance_id = request_id.clone();
let appearance_id = submission.request_id.clone();
// Keep client-enriched triage fields; do not strip_metadata (Slack shows Option fields when set).
let outcome = if let Some(client) = feedback_client {
let result = if let Some(req_id) = request_id {
with_one_shot_auth_retry(client, || async {
client
.complete_request(&req_id, submission)
.await
.map(|_| ())
})
.await
} else {
with_one_shot_auth_retry(client, || async {
client.submit_feedback(submission).await.map(|_| ())
})
.await
};
match result {
Ok(()) => SubmitOutcome::Submitted,
Err(e) => {
tracing::warn!(error = %e, "feedback submission failed");
SubmitOutcome::Failed(e)
// Only text-bearing feedback goes over the wire: the Kimi endpoint takes
// a `content` string (kimi-cli slash.py parity); ratings stay local.
let outcome = match (feedback_client, &submission.feedback_text) {
(Some(client), Some(text)) if !text.is_empty() => {
let model = submission
.model_id
.as_deref()
.or(submission.resolved_model_id.as_deref());
match client
.submit_feedback(&submission.session_id, text, model)
.await
{
Ok(()) => SubmitOutcome::Submitted,
Err(e) => {
tracing::warn!(error = %e, "feedback submission failed");
SubmitOutcome::Failed(e)
}
}
}
} else {
SubmitOutcome::LocalOnly
_ => SubmitOutcome::LocalOnly,
};
{
@@ -172,18 +155,13 @@ pub struct FeedbackFlags {
/// Configuration for the feedback manager.
#[derive(Debug, Clone)]
pub struct FeedbackManagerConfig {
/// Interval for syncing signals to the analytics backend (default: 30s)
/// Interval for the signals actor's periodic bookkeeping tick.
pub sync_interval: Duration,
/// Whether user-facing feedback features are enabled (popups, `/feedback`,
/// ratings). Gated by `KIGI_FEEDBACK_ENABLED`.
pub feedback_enabled: bool,
/// Client type (Agent, Tui, Web, Extension)
pub client_type: ClientType,
/// Whether LOC attribution tracking is enabled for this session.
/// Propagated into every `SessionTurnDelta` so the server can
/// distinguish "tracking off" (zeros are noise) from "tracking on,
/// no code changed" (zeros are real data).
pub loc_tracking_enabled: bool,
}
impl Default for FeedbackManagerConfig {
@@ -192,7 +170,6 @@ impl Default for FeedbackManagerConfig {
sync_interval: Duration::from_secs(60),
feedback_enabled: false,
client_type: ClientType::Agent,
loc_tracking_enabled: false,
}
}
}
@@ -205,19 +182,17 @@ pub struct FeedbackManager {
signals_handle: SessionSignalsHandle,
/// Feedback heuristics evaluator
heuristics: Arc<RwLock<FeedbackHeuristics>>,
/// REST client for the feedback/analytics backend
/// Client for the Kimi Code feedback endpoint (subscription sessions).
feedback_client: Option<FeedbackClient>,
/// Configuration
config: FeedbackManagerConfig,
/// Whether config has been loaded from server
config_loaded: Arc<AtomicBool>,
}
impl FeedbackManager {
/// Create a new feedback manager for a session.
///
/// If `feedback_client` is None, signal syncing is disabled but local
/// tracking and heuristics evaluation still work.
/// If `feedback_client` is None, submissions stay local but tracking and
/// heuristics evaluation still work.
pub fn new(
session_id: impl Into<String>,
feedback_client: Option<FeedbackClient>,
@@ -243,7 +218,6 @@ impl FeedbackManager {
heuristics: Arc::new(RwLock::new(FeedbackHeuristics::new())),
feedback_client,
config,
config_loaded: Arc::new(AtomicBool::new(false)),
}
}
@@ -267,13 +241,14 @@ impl FeedbackManager {
self.config.feedback_enabled
}
/// REST client for the feedback/analytics backend, if configured.
/// Client for the Kimi Code feedback endpoint, if this is a subscription
/// session.
pub fn feedback_client(&self) -> Option<&FeedbackClient> {
self.feedback_client.as_ref()
}
/// Client type for this session (Agent, Tui, Web, etc.).
pub fn client_type(&self) -> prod_mc_cli_chat_proxy_types::feedback_types::ClientType {
pub fn client_type(&self) -> ClientType {
self.config.client_type
}
@@ -330,47 +305,6 @@ impl FeedbackManager {
.await
}
/// Check if config has been loaded from the server.
pub fn is_config_loaded(&self) -> bool {
self.config_loaded.load(Ordering::Relaxed)
}
/// Load feedback heuristics config from the backend.
/// This is called automatically in run_sync_loop but can be called manually.
/// Does not block - errors are logged and defaults are used.
#[tracing::instrument(name = "feedback.load_config", skip_all, fields(
session_id = %self.session_id,
))]
pub async fn load_config(&self) {
let Some(client) = &self.feedback_client else {
return; // No client, use defaults
};
if self.config.feedback_enabled {
match client.get_feedback_config().await {
Ok(config) => {
let mut heuristics = self.heuristics.write().await;
heuristics.update_config(&config);
self.config_loaded.store(true, Ordering::Relaxed);
tracing::info!(
session_id = %self.session_id,
config_id = %config.config_id,
config_version = config.config_version,
enabled = config.enabled,
"Loaded feedback heuristics config from server"
);
}
Err(e) => {
tracing::warn!(
session_id = %self.session_id,
error = %e,
"Failed to load feedback heuristics config, using defaults"
);
}
}
}
}
/// Evaluate heuristics and return a FeedbackRequest if one should be sent.
///
/// Call this after each turn to check if feedback should be requested.
@@ -378,9 +312,6 @@ impl FeedbackManager {
/// - No tier criteria are met
/// - The tier was already triggered this session
/// - Probabilistic sampling says no
///
/// When a request is triggered, this method also creates a record via the
/// feedback API for tracking and analytics.
#[tracing::instrument(name = "feedback.maybe_request_feedback", skip_all, fields(
session_id = %self.session_id,
))]
@@ -395,7 +326,6 @@ impl FeedbackManager {
let signals = self.signals_handle.snapshot().await?;
let mut heuristics = self.heuristics.write().await;
// Check if heuristics are globally enabled (from server config)
if !heuristics.is_enabled() {
return None;
}
@@ -422,12 +352,10 @@ impl FeedbackManager {
tier = ?request.tier,
trigger_type = %request.trigger_type,
feedback_mode = ?request.feedback_mode,
prompt_id = ?prompt_id,
"Feedback request triggered"
);
self.record_feedback_request(&request, trigger_condition, feedback_mode, prompt_id)
.await;
return Some(request);
}
@@ -449,11 +377,6 @@ impl FeedbackManager {
/// `x.ai/debug/trigger_feedback` ACP extension method to exercise
/// the full feedback notification ↔ response flow without needing a
/// real session that meets tier criteria.
///
/// When a `feedback_client` is configured, the request is also recorded
/// via the feedback API — exactly like a real trigger — so that the
/// subsequent `complete_request` / `dismiss_request` round-trip from the
/// client works end-to-end.
#[tracing::instrument(name = "feedback.force_feedback_request", skip_all, fields(
session_id = %self.session_id,
))]
@@ -481,96 +404,7 @@ impl FeedbackManager {
// Manual/debug triggers are always dismissible regardless of tier config,
// since they exist for developer testing, not real user feedback collection.
let request = FeedbackRequest::with_mode(
self.session_id.clone(),
condition.clone(),
mode,
true,
None,
);
self.record_feedback_request(&request, &condition, mode, None)
.await;
request
}
/// Record a feedback request via the feedback API.
///
/// This is a best-effort operation — errors are logged but do not
/// prevent the request from being sent to the client.
#[tracing::instrument(name = "feedback.record_feedback_request", skip_all, fields(
session_id = %self.session_id,
))]
async fn record_feedback_request(
&self,
request: &FeedbackRequest,
trigger_condition: &TriggerCondition,
feedback_mode: FeedbackMode,
prompt_id: Option<String>,
) {
let Some(client) = &self.feedback_client else {
return;
};
let input = CreateFeedbackRequestInput {
request_id: request.request_id.clone(),
session_id: self.session_id.clone(),
client_type: self.config.client_type,
feedback_mode,
feedback_prompt: Some(request.prompt.clone()),
priority: tier_to_priority(trigger_condition.tier),
trigger_type: request.trigger_type.clone(),
trigger_reason: Some(trigger_condition.trigger_reason()),
context_type: Some(ContextType::Session),
context_message_ids: vec![],
expires_at: None,
experiment_id: None,
trigger_condition: serde_json::to_value(trigger_condition).ok(),
prompt_id,
};
match with_one_shot_auth_retry(client, || client.create_feedback_request(&input)).await {
Ok(response) => {
tracing::debug!(
request_id = %response.request_id,
"Feedback request recorded with feedback API"
);
}
Err(e) => {
tracing::warn!(
request_id = %request.request_id,
error = %e,
"Failed to record feedback request (continuing anyway)"
);
}
}
}
/// Capture a turn-end snapshot and send the delta to the analytics backend.
///
/// Call this once per user turn, after the agent has finished all tool-call
/// rounds and produced a final response (i.e. alongside `record_turn_complete`).
/// Intermediate tool-call steps within the same turn do NOT need their own
/// call — the signals actor accumulates tool calls, errors, and latency
/// continuously, so the single snapshot at turn end captures the full diff.
///
/// The caller provides a pre-captured `TurnDeltaSnapshot` (taken exactly
/// once inside the session actor). This avoids double-advancing the delta
/// baseline. If the snapshot is `None` (e.g. the signals actor was shut
/// down), this is a no-op.
///
/// The delta is converted and sent asynchronously to the backend. Errors
/// are logged but never block the turn flow.
///
/// Load feedback heuristics config on startup.
/// This should be spawned as a background task.
#[tracing::instrument(skip_all, fields(session_id = %self.session_id))]
pub async fn run_sync_loop(self: Arc<Self>, cancel: tokio_util::sync::CancellationToken) {
// Load config in background (non-blocking, errors logged)
self.load_config().await;
cancel.cancelled().await;
tracing::debug!("Feedback sync loop cancelled");
FeedbackRequest::with_mode(self.session_id.clone(), condition, mode, true, None)
}
/// Shutdown the manager: shuts down the signals actor.
@@ -579,173 +413,6 @@ impl FeedbackManager {
}
}
// Auth outcome handler used by run_sync_loop on 401.
/// Max consecutive failed sync ticks tolerated before stopping the loop.
/// ~10 minutes at the default 60s interval.
const MAX_CONSECUTIVE_AUTH_FAILURES: u8 = 10;
/// telemetry `reason` discriminators on the `signals sync loop stopped permanently`
/// event. Pinned because alerts filter on these strings.
const REASON_AUTH_PERMANENT_FAILURE: &str = "auth_permanent_failure";
const REASON_NO_CLIENT_OR_REFRESHER: &str = "no_client_or_refresher";
const LOG_TITLE_TRANSIENT: &str = "signals sync transient auth failure";
const LOG_TITLE_STOPPED_PERMANENTLY: &str = "signals sync loop stopped permanently";
/// Classification of one 401-recovery attempt.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum SyncAuthOutcome {
/// Refresh + retry succeeded.
Recovered,
/// Refresh or retry failed transiently (lock timeout, network, sibling
/// race, post-refresh 5xx). Increment the counter and retry next tick.
Transient,
/// IdP confirmed a terminal failure (`invalid_grant` / `invalid_client`).
/// Only re-login will recover.
Permanent,
/// No client or no refresher configured — nothing to retry.
Unrecoverable,
}
fn handle_auth_outcome(
outcome: SyncAuthOutcome,
consecutive_auth_failures: &mut u8,
session_id: &str,
) -> ControlFlow<()> {
match outcome {
SyncAuthOutcome::Recovered => {
*consecutive_auth_failures = 0;
tracing::info!(
session_id = %session_id,
"Signal sync recovered after token refresh"
);
ControlFlow::Continue(())
}
SyncAuthOutcome::Transient => {
*consecutive_auth_failures = consecutive_auth_failures.saturating_add(1);
tracing::warn!(
session_id = %session_id,
consecutive_failures = *consecutive_auth_failures,
max = MAX_CONSECUTIVE_AUTH_FAILURES,
"Signals sync transient auth failure"
);
kigi_log::unified_log::warn(
LOG_TITLE_TRANSIENT,
Some(session_id),
Some(serde_json::json!({
"consecutive_failures": *consecutive_auth_failures,
"max": MAX_CONSECUTIVE_AUTH_FAILURES,
})),
);
if *consecutive_auth_failures >= MAX_CONSECUTIVE_AUTH_FAILURES {
tracing::warn!(
session_id = %session_id,
consecutive_failures = *consecutive_auth_failures,
"Signals sync loop stopped: consecutive transient auth failures"
);
kigi_log::unified_log::warn(
"signals sync loop stopped: consecutive transient auth failures",
Some(session_id),
Some(serde_json::json!({
"consecutive_failures": *consecutive_auth_failures,
"max": MAX_CONSECUTIVE_AUTH_FAILURES,
})),
);
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
SyncAuthOutcome::Permanent => {
tracing::warn!(
session_id = %session_id,
reason = REASON_AUTH_PERMANENT_FAILURE,
"Signals sync loop stopped: IdP confirmed permanent auth failure"
);
kigi_log::unified_log::warn(
LOG_TITLE_STOPPED_PERMANENTLY,
Some(session_id),
Some(serde_json::json!({ "reason": REASON_AUTH_PERMANENT_FAILURE })),
);
ControlFlow::Break(())
}
SyncAuthOutcome::Unrecoverable => {
tracing::warn!(
session_id = %session_id,
reason = REASON_NO_CLIENT_OR_REFRESHER,
"Signals sync loop stopped: no client or no refresher configured"
);
kigi_log::unified_log::warn(
LOG_TITLE_STOPPED_PERMANENTLY,
Some(session_id),
Some(serde_json::json!({ "reason": REASON_NO_CLIENT_OR_REFRESHER })),
);
ControlFlow::Break(())
}
}
}
/// Check if an error is an HTTP 401 Unauthorized response.
///
/// Uses typed downcast on [`FeedbackApiError`] instead of string matching,
/// so it stays correct even if error messages change.
fn is_auth_error(error: &anyhow::Error) -> bool {
error
.downcast_ref::<FeedbackApiError>()
.is_some_and(|e| e.is_unauthorized())
}
/// Check if an error is an HTTP 403 Forbidden response.
///
/// 403 from the signals endpoint means the session does not belong to the
/// current user — a permanent condition that will never self-resolve.
fn is_forbidden_error(error: &anyhow::Error) -> bool {
error
.downcast_ref::<FeedbackApiError>()
.is_some_and(|e| e.is_forbidden())
}
/// Run `op` once; on 401, wait for an in-flight refresh to land, then
/// retry once. Prefers waiting for the proactive-refresh task or
/// main-request-path recovery over driving a `ServerRejected` refresh
/// itself, avoiding the 401-amplification pattern during token-expiry
/// windows.
async fn with_one_shot_auth_retry<T, F, Fut>(
client: &FeedbackClient,
mut op: F,
) -> anyhow::Result<T>
where
F: FnMut() -> Fut,
Fut: std::future::Future<Output = anyhow::Result<T>>,
{
match op().await {
Ok(v) => Ok(v),
Err(e) if is_auth_error(&e) => {
// 1. Wait briefly for the proactive refresh or main-path
// recovery to land a fresh token.
let refreshed = client.wait_for_token_refresh(Duration::from_secs(3)).await;
// 2. If nobody refreshed, drive our own recovery as fallback.
if refreshed || client.try_refresh_credentials().await {
op().await
} else {
Err(e)
}
}
Err(e) => Err(e),
}
}
/// Convert a FeedbackTier to a priority value (1-10, higher = more important).
fn tier_to_priority(tier: crate::session::feedback::FeedbackTier) -> i32 {
use crate::session::feedback::FeedbackTier;
match tier {
FeedbackTier::Tier1 => 5, // Standard engagement
FeedbackTier::Tier2 => 6, // Complex session with recovery
FeedbackTier::Tier3 => 7, // Recovery from friction
}
}
#[cfg(test)]
mod tests {
use super::*;
@@ -786,83 +453,6 @@ mod tests {
manager.shutdown().await;
}
#[test]
fn test_is_auth_error_detects_401() {
use crate::agent::feedback_client::FeedbackApiError;
let err: anyhow::Error = FeedbackApiError {
status: reqwest::StatusCode::UNAUTHORIZED,
context: "Signals update",
body: "Invalid or expired credentials".to_string(),
}
.into();
assert!(is_auth_error(&err));
}
#[test]
fn test_is_auth_error_ignores_other_statuses() {
use crate::agent::feedback_client::FeedbackApiError;
let err_500: anyhow::Error = FeedbackApiError {
status: reqwest::StatusCode::INTERNAL_SERVER_ERROR,
context: "Signals update",
body: "oops".to_string(),
}
.into();
assert!(!is_auth_error(&err_500));
let err_403: anyhow::Error = FeedbackApiError {
status: reqwest::StatusCode::FORBIDDEN,
context: "Signals update",
body: "ZDR team".to_string(),
}
.into();
assert!(!is_auth_error(&err_403));
}
#[test]
fn test_is_auth_error_ignores_non_api_errors() {
assert!(!is_auth_error(&anyhow::anyhow!("network timeout")));
assert!(!is_auth_error(&anyhow::anyhow!("connection refused")));
}
#[test]
fn test_is_forbidden_error_detects_403() {
use crate::agent::feedback_client::FeedbackApiError;
let err: anyhow::Error = FeedbackApiError {
status: reqwest::StatusCode::FORBIDDEN,
context: "Signals update",
body: "Access denied: session does not belong to this user".to_string(),
}
.into();
assert!(is_forbidden_error(&err));
}
#[test]
fn test_is_forbidden_error_ignores_other_statuses() {
use crate::agent::feedback_client::FeedbackApiError;
let err_401: anyhow::Error = FeedbackApiError {
status: reqwest::StatusCode::UNAUTHORIZED,
context: "Signals update",
body: "Invalid credentials".to_string(),
}
.into();
assert!(!is_forbidden_error(&err_401));
assert!(!is_forbidden_error(&anyhow::anyhow!("network error")));
}
#[test]
fn test_is_auth_error_works_through_anyhow_conversion() {
use crate::agent::feedback_client::FeedbackApiError;
// Verify the FeedbackApiError survives anyhow::Error round-trip
// (this is the actual path: send_json returns FeedbackApiError.into())
let api_err = FeedbackApiError {
status: reqwest::StatusCode::UNAUTHORIZED,
context: "Signals update",
body: "token expired".to_string(),
};
let anyhow_err: anyhow::Error = api_err.into();
assert!(is_auth_error(&anyhow_err));
}
#[tokio::test]
async fn test_feedback_manager_disabled() {
let config = FeedbackManagerConfig {
@@ -895,158 +485,21 @@ mod tests {
assert!(snapshot.is_none(), "Signals actor should be shut down");
}
// ── handle_auth_outcome tests ──────────────────────────────────────────
/// 9 transient failures must NOT break, and a subsequent `Recovered`
/// must reset the counter to 0.
#[test]
fn test_sync_loop_continues_through_transient_auth_failures() {
let mut counter: u8 = 0;
for _ in 0..(MAX_CONSECUTIVE_AUTH_FAILURES - 1) {
let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s");
assert_eq!(flow, ControlFlow::Continue(()));
}
assert_eq!(counter, MAX_CONSECUTIVE_AUTH_FAILURES - 1);
let flow = handle_auth_outcome(SyncAuthOutcome::Recovered, &mut counter, "s");
assert_eq!(flow, ControlFlow::Continue(()));
assert_eq!(counter, 0, "Recovered must reset the counter");
}
/// Exactly `MAX_CONSECUTIVE_AUTH_FAILURES` consecutive `Transient`
/// outcomes break the loop.
#[test]
fn test_sync_loop_breaks_after_max_transient_auth_failures() {
let mut counter: u8 = 0;
for i in 0..(MAX_CONSECUTIVE_AUTH_FAILURES - 1) {
let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s");
assert_eq!(
flow,
ControlFlow::Continue(()),
"iteration {i} should still continue"
);
}
// The 10th (== MAX_CONSECUTIVE_AUTH_FAILURES) transient breaks.
let flow = handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s");
assert_eq!(flow, ControlFlow::Break(()));
assert_eq!(counter, MAX_CONSECUTIVE_AUTH_FAILURES);
}
/// `Permanent` breaks immediately and does not bump the counter.
#[test]
fn test_sync_loop_breaks_immediately_on_permanent_failure() {
let mut counter: u8 = 0;
let flow = handle_auth_outcome(SyncAuthOutcome::Permanent, &mut counter, "s");
assert_eq!(flow, ControlFlow::Break(()));
assert_eq!(counter, 0);
}
/// 5 transient → 1 recovered → 5 transient must not break.
#[test]
fn test_sync_loop_counter_resets_on_successful_sync() {
let mut counter: u8 = 0;
for _ in 0..5 {
assert_eq!(
handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"),
ControlFlow::Continue(())
);
}
assert_eq!(counter, 5);
assert_eq!(
handle_auth_outcome(SyncAuthOutcome::Recovered, &mut counter, "s"),
ControlFlow::Continue(())
);
assert_eq!(counter, 0, "Recovered must reset the counter");
for _ in 0..5 {
assert_eq!(
handle_auth_outcome(SyncAuthOutcome::Transient, &mut counter, "s"),
ControlFlow::Continue(())
);
}
assert_eq!(counter, 5, "second burst should be re-counted from zero");
}
/// `Unrecoverable` breaks the loop and does not bump the counter.
#[test]
fn test_sync_loop_breaks_on_unrecoverable() {
let mut counter: u8 = 0;
let flow = handle_auth_outcome(SyncAuthOutcome::Unrecoverable, &mut counter, "s");
assert_eq!(flow, ControlFlow::Break(()));
assert_eq!(counter, 0);
}
/// `FeedbackClient::is_auth_permanently_failed` reflects the attached
/// `AuthManager`'s `permanent_failure()` cache (record → true,
/// age-out → false).
/// A rating-only submission (no text) must not hit the network even when
/// no client is configured — the workflow reports LocalOnly.
#[tokio::test]
async fn test_is_auth_permanently_failed_reads_auth_manager() {
use crate::agent::feedback_client::FeedbackClient;
use crate::auth::error::RefreshTokenFailedReason;
use crate::auth::{AuthManager, KimiAuth, KimiCodeConfig};
use std::sync::Arc;
async fn test_rating_only_submission_stays_local() {
use crate::session::feedback_types::RatingType;
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let client = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone());
assert!(!client.is_auth_permanently_failed());
// The tombstone is scoped to the live credential's refresh token.
am.hot_swap(KimiAuth {
key: "tok".into(),
refresh_token: Some("rt".into()),
expires_at: Some(chrono::Utc::now() - chrono::Duration::hours(1)),
..KimiAuth::test_default()
});
am.record_permanent_failure("rt".to_string(), RefreshTokenFailedReason::Other.into());
assert!(client.is_auth_permanently_failed());
am.force_permanent_failure_aged_out();
assert!(!client.is_auth_permanently_failed());
}
/// With no `AuthManager` attached, `is_auth_permanently_failed` is false.
#[test]
fn test_is_auth_permanently_failed_without_auth_manager() {
use crate::agent::feedback_client::FeedbackClient;
let client = FeedbackClient::new("http://example/v1", None);
assert!(!client.is_auth_permanently_failed());
}
/// `has_token_refresher` requires BOTH an `AuthManager` AND a refresher
/// wired in. Without this, a static-deployment-key session would be
/// mis-classified as recoverable.
#[tokio::test]
async fn test_has_token_refresher_requires_refresher_attached() {
use crate::agent::feedback_client::FeedbackClient;
use crate::auth::{AuthManager, KimiCodeConfig};
use std::sync::Arc;
struct NoOpRefresher;
#[async_trait::async_trait]
impl crate::auth::refresh::TokenRefresher for NoOpRefresher {
async fn refresh(
&self,
_reason: crate::auth::refresh::RefreshReason,
) -> crate::auth::refresh::RefreshOutcome {
crate::auth::refresh::RefreshOutcome::TransientFailure {
message: "noop".into(),
}
}
}
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), KimiCodeConfig::default()));
let bare = FeedbackClient::new("http://example/v1", None);
assert!(!bare.has_token_refresher());
let with_am = FeedbackClient::new("http://example/v1", None).with_auth_manager(am.clone());
assert!(
!with_am.has_token_refresher(),
"AuthManager without a refresher must NOT be reported as recoverable"
let mut submission = new_submission(
"sess-local".into(),
ClientType::Tui,
FeedbackContent::Rating {
rating_type: RatingType::Thumbs,
rating_value: 1,
},
);
am.set_refresher(std::sync::Arc::new(NoOpRefresher));
assert!(with_am.has_token_refresher());
let outcome = submit_feedback_workflow(&mut submission, None, None, false).await;
assert!(matches!(outcome, SubmitOutcome::LocalOnly));
}
}
crates/codegen/kigi-shell/src/session/feedback_types.rs
+763
View File
@@ -0,0 +1,763 @@
//! Local feedback data types.
//!
//! Formerly the wire contract with the deleted xAI cli-chat-proxy feedback
//! backend; now these types only back the LOCAL feedback records persisted in
//! the session store and the heuristics that decide when to solicit feedback.
//! The only remaining network surface is the Kimi Code `POST {base}/feedback`
//! call in [`crate::agent::feedback_client`], which sends a small flat JSON
//! body — none of these types go over the wire anymore, so the proxy-only
//! null-column fields (experiment/comparison/preference plumbing) are gone.
use serde::{Deserialize, Deserializer, Serialize};
pub use kigi_shared::session::FeedbackTerminalInfo;
/// Type of client submitting feedback.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ClientType {
/// Terminal/CLI agent
#[default]
Agent,
/// Terminal UI
Tui,
/// Web interface
Web,
/// IDE extension (VS Code, JetBrains, etc.)
Extension,
/// Remote workspace / hosted agent client (wire value `nebula`).
Nebula,
/// Desktop (Electron app)
Desktop,
}
impl std::fmt::Display for ClientType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ClientType::Agent => write!(f, "agent"),
ClientType::Tui => write!(f, "tui"),
ClientType::Web => write!(f, "web"),
ClientType::Extension => write!(f, "extension"),
ClientType::Nebula => write!(f, "nebula"),
ClientType::Desktop => write!(f, "desktop"),
}
}
}
/// Type of feedback being submitted.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FeedbackType {
/// Numeric rating only
#[default]
Rating,
/// Free-form text only
Text,
/// Both rating and text
RatingWithText,
/// Model preference comparison
ModelPreference,
/// Bug report
BugReport,
/// Feature request
FeatureRequest,
}
impl std::fmt::Display for FeedbackType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FeedbackType::Rating => write!(f, "rating"),
FeedbackType::Text => write!(f, "text"),
FeedbackType::RatingWithText => write!(f, "rating_with_text"),
FeedbackType::ModelPreference => write!(f, "model_preference"),
FeedbackType::BugReport => write!(f, "bug_report"),
FeedbackType::FeatureRequest => write!(f, "feature_request"),
}
}
}
/// Type of rating scale used.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RatingType {
/// Thumbs up/down (-1, 0, 1)
Thumbs,
/// Star rating (1-5)
Stars,
/// Net Promoter Score (0-10)
Nps,
}
impl std::fmt::Display for RatingType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RatingType::Thumbs => write!(f, "thumbs"),
RatingType::Stars => write!(f, "stars"),
RatingType::Nps => write!(f, "nps"),
}
}
}
/// Context type for what the feedback is about.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ContextType {
/// Feedback about a specific message
Message,
/// Feedback about the overall session/conversation
Session,
/// Feedback about a specific feature
Feature,
/// Feedback about tool usage
ToolUse,
/// General feedback
General,
}
impl std::fmt::Display for ContextType {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
ContextType::Message => write!(f, "message"),
ContextType::Session => write!(f, "session"),
ContextType::Feature => write!(f, "feature"),
ContextType::ToolUse => write!(f, "tool_use"),
ContextType::General => write!(f, "general"),
}
}
}
/// Type of feedback mode requested.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum FeedbackMode {
/// Thumbs up/down
Thumbs,
/// Star rating (1-5)
Stars,
/// Free-form text
Text,
/// Thumbs up/down with optional text comment
ThumbsText,
/// Star rating with optional text comment
StarsText,
/// Model comparison
Comparison,
/// Multi-question survey
Survey,
/// Net Promoter Score (0-10)
Nps,
/// NPS with optional text comment
NpsText,
}
impl std::fmt::Display for FeedbackMode {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
FeedbackMode::Thumbs => write!(f, "thumbs"),
FeedbackMode::Stars => write!(f, "stars"),
FeedbackMode::Text => write!(f, "text"),
FeedbackMode::ThumbsText => write!(f, "thumbs_text"),
FeedbackMode::StarsText => write!(f, "stars_text"),
FeedbackMode::Comparison => write!(f, "comparison"),
FeedbackMode::Survey => write!(f, "survey"),
FeedbackMode::Nps => write!(f, "nps"),
FeedbackMode::NpsText => write!(f, "nps_text"),
}
}
}
/// Parse a feedback mode string to FeedbackMode enum.
pub fn parse_feedback_mode_str(s: &str) -> FeedbackMode {
match s {
"thumbs" => FeedbackMode::Thumbs,
"stars" => FeedbackMode::Stars,
"text" => FeedbackMode::Text,
"thumbs_text" => FeedbackMode::ThumbsText,
"stars_text" => FeedbackMode::StarsText,
"comparison" => FeedbackMode::Comparison,
"survey" => FeedbackMode::Survey,
"nps" => FeedbackMode::Nps,
"nps_text" => FeedbackMode::NpsText,
_ => FeedbackMode::Thumbs,
}
}
/// Allowed `feedback_type` + value-field combinations. Construct submissions
/// via [`FeedbackSubmission::with_content`].
#[derive(Debug, Clone)]
#[non_exhaustive]
pub enum FeedbackContent {
Rating {
rating_type: RatingType,
rating_value: i32,
},
Text(String),
RatingWithText {
rating_type: RatingType,
rating_value: i32,
text: String,
},
}
impl FeedbackContent {
fn apply_to(self, s: &mut FeedbackSubmission) {
s.rating_type = None;
s.rating_value = None;
s.feedback_text = None;
match self {
Self::Rating {
rating_type,
rating_value,
} => {
s.feedback_type = FeedbackType::Rating;
s.rating_type = Some(rating_type);
s.rating_value = Some(rating_value);
}
Self::Text(text) => {
s.feedback_type = FeedbackType::Text;
s.feedback_text = Some(text);
}
Self::RatingWithText {
rating_type,
rating_value,
text,
} => {
s.feedback_type = FeedbackType::RatingWithText;
s.rating_type = Some(rating_type);
s.rating_value = Some(rating_value);
s.feedback_text = Some(text);
}
}
}
}
fn empty_string_as_none<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
D: Deserializer<'de>,
{
let opt = Option::<String>::deserialize(deserializer)?;
Ok(opt.filter(|s| !s.is_empty()))
}
/// A user feedback record. Persisted locally in the session store; the text
/// content is forwarded to the Kimi Code feedback endpoint for subscription
/// sessions. Construct via [`FeedbackSubmission::with_content`]; the `Default`
/// impl exists for builder-style construction and test fixtures and does not
/// produce a valid submission on its own (empty `session_id`).
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FeedbackSubmission {
/// Session ID this feedback is for
pub session_id: String,
/// Type of client submitting feedback
pub client_type: ClientType,
/// Type of feedback being submitted
pub feedback_type: FeedbackType,
/// Turn number within the session (optional)
#[serde(skip_serializing_if = "Option::is_none")]
pub turn_number: Option<i64>,
/// Rating type (if applicable)
#[serde(skip_serializing_if = "Option::is_none")]
pub rating_type: Option<RatingType>,
/// Rating value (interpretation depends on rating_type)
/// - thumbs: -1 (down), 0 (neutral), 1 (up)
/// - stars: 1-5
/// - nps: 0-10
#[serde(skip_serializing_if = "Option::is_none")]
pub rating_value: Option<i32>,
/// Free-form feedback text
#[serde(skip_serializing_if = "Option::is_none")]
pub feedback_text: Option<String>,
/// Feedback categories (e.g., ["accuracy", "speed", "helpfulness"])
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub feedback_categories: Vec<String>,
/// Model ID used for the response being rated
#[serde(skip_serializing_if = "Option::is_none")]
pub model_id: Option<String>,
/// Server-resolved model ID from the actual chat completion response.
#[serde(skip_serializing_if = "Option::is_none")]
pub resolved_model_id: Option<String>,
/// Checkpoint fingerprint from the inference provider (`system_fingerprint`).
#[serde(
default,
skip_serializing_if = "Option::is_none",
deserialize_with = "empty_string_as_none"
)]
pub model_fingerprint: Option<String>,
/// Context type for the feedback
#[serde(skip_serializing_if = "Option::is_none")]
pub context_type: Option<ContextType>,
/// Feedback request ID (set when responding to a solicited request)
#[serde(skip_serializing_if = "Option::is_none")]
pub request_id: Option<String>,
/// Client version
#[serde(skip_serializing_if = "Option::is_none")]
pub client_version: Option<String>,
/// Shell (kigi-shell) version
#[serde(skip_serializing_if = "Option::is_none")]
pub shell_version: Option<String>,
/// Additional metadata as JSON
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
/// Last user message at feedback time.
#[serde(
default,
skip_serializing_if = "Option::is_none",
alias = "last_user_turn"
)]
pub last_user_message: Option<String>,
/// Last assistant response at feedback time.
#[serde(
default,
skip_serializing_if = "Option::is_none",
alias = "last_assistant_turn"
)]
pub last_assistant_message: Option<String>,
/// Per-tool call counts for the rated turn.
#[serde(default, skip_serializing_if = "Vec::is_empty")]
pub tool_outcomes: Vec<FeedbackToolOutcome>,
/// Session working directory.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub session_cwd: Option<String>,
/// Number of compactions in the session.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub compaction_count: Option<i64>,
/// Context window usage percentage (0–100).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_window_usage: Option<u8>,
/// Raw context tokens used at feedback time.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_tokens_used: Option<u64>,
/// Raw model context window token limit at feedback time.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub context_window_tokens: Option<u64>,
/// Terminal environment snapshot at feedback time (brand, multiplexer, SSH, etc.).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub terminal_info: Option<FeedbackTerminalInfo>,
}
impl FeedbackSubmission {
/// Construct from typed content; set optional fields after.
pub fn with_content(
session_id: String,
client_type: ClientType,
content: FeedbackContent,
) -> Self {
let mut s = Self {
session_id,
client_type,
..Default::default()
};
content.apply_to(&mut s);
s
}
/// Merge a JSON object into `metadata`, inserting if absent.
pub fn merge_metadata(&mut self, extra: serde_json::Value) {
match &mut self.metadata {
Some(existing) if existing.is_object() => {
if let (Some(dst), Some(src)) = (existing.as_object_mut(), extra.as_object()) {
for (k, v) in src {
dst.insert(k.clone(), v.clone());
}
}
}
_ => {
self.metadata = Some(extra);
}
}
}
}
/// Per-tool call/failure counts for a single tool in a turn.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct FeedbackToolOutcome {
pub tool_name: String,
pub calls: u32,
pub failures: u32,
}
/// Configuration for a single feedback tier.
///
/// Each tier has specific thresholds and conditions that must be met
/// for feedback to be requested at that tier.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct TierConfig {
/// Whether this tier is enabled
pub enabled: bool,
/// Sample rate (0.0 to 1.0, e.g., 0.0005 = 0.05%)
pub sample_rate: f64,
/// Minimum turns required to trigger
pub min_turns: i64,
/// Minimum tool calls required (Tier 1 & 2)
#[serde(default)]
pub min_tool_calls: i64,
/// Minimum compactions required (Tier 1 & 2)
#[serde(default)]
pub min_compactions: i64,
/// Minimum errors required (Tier 2 only)
#[serde(default)]
pub min_errors: i64,
/// Whether cancellations disqualify this tier (Tier 1)
#[serde(default)]
pub no_cancellations: bool,
/// Whether cancellation is required (Tier 3)
#[serde(default)]
pub requires_cancellation: bool,
/// Whether revert is required (Tier 3)
#[serde(default)]
pub requires_revert: bool,
/// Whether at least one of cancellation/revert is required (Tier 3)
#[serde(default)]
pub requires_recovery: bool,
/// Feedback mode to use when this tier triggers
pub feedback_mode: FeedbackMode,
/// Whether feedback requests from this tier are dismissible (non-intrusive)
#[serde(default = "default_true")]
pub dismissible: bool,
/// Prompt text shown to users when this tier's feedback is requested
#[serde(default)]
pub prompt: String,
/// Max times this tier can trigger per session (0 = unlimited)
#[serde(default = "default_one")]
pub max_triggers: i32,
}
impl Default for TierConfig {
fn default() -> Self {
Self {
enabled: true,
sample_rate: 0.0005,
min_turns: 10,
min_tool_calls: 5,
min_compactions: 2,
min_errors: 0,
no_cancellations: false,
requires_cancellation: false,
requires_revert: false,
requires_recovery: false,
feedback_mode: FeedbackMode::Thumbs,
dismissible: true,
prompt: String::new(),
max_triggers: 1,
}
}
}
/// Configuration for feedback heuristics.
///
/// Formerly fetched from the proxy backend; now purely local — the built-in
/// [`Default`] is the only production source, kept as a struct so tests and
/// future config surfaces can tune the tiers.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct FeedbackHeuristicsConfig {
/// Unique configuration identifier
pub config_id: String,
/// Configuration version (monotonically increasing)
pub config_version: i64,
// === Global Settings ===
/// Master enable/disable switch for all feedback collection
pub enabled: bool,
/// Minimum seconds between feedback requests (cooldown period)
#[serde(default = "default_cooldown_seconds")]
pub cooldown_seconds: i64,
/// Maximum feedback requests per session
#[serde(default = "default_max_requests")]
pub max_requests_per_session: i64,
// === Tier 1: Standard Engagement ===
/// Whether Tier 1 is enabled
#[serde(default = "default_true")]
pub tier1_enabled: bool,
/// Sample rate for Tier 1 (0.0-1.0)
#[serde(default = "default_tier1_sample_rate")]
pub tier1_sample_rate: f64,
/// Minimum turns for Tier 1
#[serde(default = "default_tier1_min_turns")]
pub tier1_min_turns: i64,
/// Minimum tool calls for Tier 1
#[serde(default = "default_tier1_min_tool_calls")]
pub tier1_min_tool_calls: i64,
/// Minimum compactions for Tier 1
#[serde(default = "default_tier1_min_compactions")]
pub tier1_min_compactions: i64,
/// Whether Tier 1 requires no cancellations
#[serde(default = "default_true")]
pub tier1_no_cancellations: bool,
/// Feedback mode for Tier 1
#[serde(default = "default_feedback_mode_thumbs")]
pub tier1_feedback_mode: String,
/// Whether Tier 1 feedback requests are dismissible
#[serde(default = "default_true")]
pub tier1_dismissible: bool,
/// Prompt text shown to users when Tier 1 feedback is requested
#[serde(default = "default_tier1_prompt")]
pub tier1_prompt: String,
/// Max times Tier 1 can trigger per session (0 = unlimited)
#[serde(default = "default_one")]
pub tier1_max_triggers: i32,
// === Tier 2: Complex Session with Recovery ===
/// Whether Tier 2 is enabled
#[serde(default = "default_true")]
pub tier2_enabled: bool,
/// Sample rate for Tier 2 (0.0-1.0)
#[serde(default = "default_tier2_sample_rate")]
pub tier2_sample_rate: f64,
/// Minimum turns for Tier 2
#[serde(default = "default_tier2_min_turns")]
pub tier2_min_turns: i64,
/// Minimum tool calls for Tier 2
#[serde(default = "default_tier2_min_tool_calls")]
pub tier2_min_tool_calls: i64,
/// Minimum compactions for Tier 2
#[serde(default = "default_tier2_min_compactions")]
pub tier2_min_compactions: i64,
/// Minimum errors for Tier 2
#[serde(default = "default_tier2_min_errors")]
pub tier2_min_errors: i64,
/// Feedback mode for Tier 2
#[serde(default = "default_feedback_mode_thumbs_text")]
pub tier2_feedback_mode: String,
/// Whether Tier 2 feedback requests are dismissible
#[serde(default = "default_true")]
pub tier2_dismissible: bool,
/// Prompt text shown to users when Tier 2 feedback is requested
#[serde(default = "default_tier2_prompt")]
pub tier2_prompt: String,
/// Max times Tier 2 can trigger per session (0 = unlimited)
#[serde(default = "default_one")]
pub tier2_max_triggers: i32,
// === Tier 3: Recovery from Friction ===
/// Whether Tier 3 is enabled
#[serde(default = "default_true")]
pub tier3_enabled: bool,
/// Sample rate for Tier 3 (0.0-1.0)
#[serde(default = "default_tier3_sample_rate")]
pub tier3_sample_rate: f64,
/// Minimum turns for Tier 3
#[serde(default = "default_tier3_min_turns")]
pub tier3_min_turns: i64,
/// Whether Tier 3 requires at least one cancellation
#[serde(default)]
pub tier3_requires_cancellation: bool,
/// Whether Tier 3 requires at least one revert
#[serde(default)]
pub tier3_requires_revert: bool,
/// Whether Tier 3 requires recovery (cancellation OR revert)
#[serde(default = "default_true")]
pub tier3_requires_recovery: bool,
/// Feedback mode for Tier 3
#[serde(default = "default_feedback_mode_stars_text")]
pub tier3_feedback_mode: String,
/// Whether Tier 3 feedback requests are dismissible
#[serde(default = "default_true")]
pub tier3_dismissible: bool,
/// Prompt text shown to users when Tier 3 feedback is requested
#[serde(default = "default_tier3_prompt")]
pub tier3_prompt: String,
/// Max times Tier 3 can trigger per session (0 = unlimited)
#[serde(default = "default_one")]
pub tier3_max_triggers: i32,
}
impl Default for FeedbackHeuristicsConfig {
fn default() -> Self {
Self {
config_id: "default".to_string(),
config_version: 1,
enabled: true,
cooldown_seconds: 300,
max_requests_per_session: 3,
// Tier 1
tier1_enabled: true,
tier1_sample_rate: 0.0005,
tier1_min_turns: 10,
tier1_min_tool_calls: 5,
tier1_min_compactions: 2,
tier1_no_cancellations: true,
tier1_feedback_mode: "thumbs".to_string(),
tier1_dismissible: true,
tier1_prompt: default_tier1_prompt(),
tier1_max_triggers: 1,
// Tier 2
tier2_enabled: true,
tier2_sample_rate: 0.0002,
tier2_min_turns: 15,
tier2_min_tool_calls: 10,
tier2_min_compactions: 3,
tier2_min_errors: 1,
tier2_feedback_mode: "thumbs_text".to_string(),
tier2_dismissible: true,
tier2_prompt: default_tier2_prompt(),
tier2_max_triggers: 1,
// Tier 3
tier3_enabled: true,
tier3_sample_rate: 0.0001,
tier3_min_turns: 20,
tier3_requires_cancellation: false,
tier3_requires_revert: false,
tier3_requires_recovery: true,
tier3_feedback_mode: "stars_text".to_string(),
tier3_dismissible: true,
tier3_prompt: default_tier3_prompt(),
tier3_max_triggers: 1,
}
}
}
impl FeedbackHeuristicsConfig {
/// Get the Tier 1 configuration as a TierConfig.
pub fn tier1_config(&self) -> TierConfig {
TierConfig {
enabled: self.tier1_enabled,
sample_rate: self.tier1_sample_rate,
min_turns: self.tier1_min_turns,
min_tool_calls: self.tier1_min_tool_calls,
min_compactions: self.tier1_min_compactions,
min_errors: 0,
no_cancellations: self.tier1_no_cancellations,
requires_cancellation: false,
requires_revert: false,
requires_recovery: false,
feedback_mode: parse_feedback_mode_str(&self.tier1_feedback_mode),
dismissible: self.tier1_dismissible,
prompt: self.tier1_prompt.clone(),
max_triggers: self.tier1_max_triggers,
}
}
/// Get the Tier 2 configuration as a TierConfig.
pub fn tier2_config(&self) -> TierConfig {
TierConfig {
enabled: self.tier2_enabled,
sample_rate: self.tier2_sample_rate,
min_turns: self.tier2_min_turns,
min_tool_calls: self.tier2_min_tool_calls,
min_compactions: self.tier2_min_compactions,
min_errors: self.tier2_min_errors,
no_cancellations: false,
requires_cancellation: false,
requires_revert: false,
requires_recovery: false,
feedback_mode: parse_feedback_mode_str(&self.tier2_feedback_mode),
dismissible: self.tier2_dismissible,
prompt: self.tier2_prompt.clone(),
max_triggers: self.tier2_max_triggers,
}
}
/// Get the Tier 3 configuration as a TierConfig.
pub fn tier3_config(&self) -> TierConfig {
TierConfig {
enabled: self.tier3_enabled,
sample_rate: self.tier3_sample_rate,
min_turns: self.tier3_min_turns,
min_tool_calls: 0,
min_compactions: 0,
min_errors: 0,
no_cancellations: false,
requires_cancellation: self.tier3_requires_cancellation,
requires_revert: self.tier3_requires_revert,
requires_recovery: self.tier3_requires_recovery,
feedback_mode: parse_feedback_mode_str(&self.tier3_feedback_mode),
dismissible: self.tier3_dismissible,
prompt: self.tier3_prompt.clone(),
max_triggers: self.tier3_max_triggers,
}
}
}
// Default value functions for serde
fn default_true() -> bool {
true
}
fn default_cooldown_seconds() -> i64 {
300
}
fn default_max_requests() -> i64 {
3
}
fn default_tier1_sample_rate() -> f64 {
0.0005
}
fn default_tier1_min_turns() -> i64 {
10
}
fn default_tier1_min_tool_calls() -> i64 {
5
}
fn default_tier1_min_compactions() -> i64 {
2
}
fn default_tier2_sample_rate() -> f64 {
0.0002
}
fn default_tier2_min_turns() -> i64 {
15
}
fn default_tier2_min_tool_calls() -> i64 {
10
}
fn default_tier2_min_compactions() -> i64 {
3
}
fn default_tier2_min_errors() -> i64 {
1
}
fn default_tier3_sample_rate() -> f64 {
0.0001
}
fn default_tier3_min_turns() -> i64 {
20
}
fn default_feedback_mode_thumbs() -> String {
"thumbs".to_string()
}
fn default_feedback_mode_thumbs_text() -> String {
"thumbs_text".to_string()
}
fn default_feedback_mode_stars_text() -> String {
"stars_text".to_string()
}
fn default_tier1_prompt() -> String {
"You've been having a productive session! Would you mind sharing quick feedback?".to_string()
}
fn default_tier2_prompt() -> String {
"You've worked through a complex session. Your feedback would help us improve.".to_string()
}
fn default_tier3_prompt() -> String {
"Thanks for sticking with us through that session. Got a moment to share feedback?".to_string()
}
fn default_one() -> i32 {
1
}
crates/codegen/kigi-shell/src/session/fork.rs
+3 -72
View File
@@ -3,9 +3,7 @@
//! Forks a saved session to a new working directory with a new session ID.
//! This creates new session files but does not start the session.
use crate::remote::BackendClient;
const FORK_LOG: &str = "xai_fork";
use crate::session::export::ExportedMetadata;
const FORK_LOG: &str = "kigi_fork";
use crate::session::info::Info;
use crate::session::storage::{CopySessionOptions, JsonlStorageAdapter};
use crate::util::kigi_home::kigi_home;
@@ -62,11 +60,7 @@ fn generate_fork_session_id(_source_id: &str) -> String {
}
/// Fork a saved session to a new working directory.
pub async fn fork_session(
request: ForkSessionRequest,
agent_id: &str,
auth_manager: Option<std::sync::Arc<crate::auth::AuthManager>>,
) -> io::Result<ForkSessionResponse> {
pub async fn fork_session(request: ForkSessionRequest) -> io::Result<ForkSessionResponse> {
let t0 = std::time::Instant::now();
let root_dir = kigi_home();
@@ -114,32 +108,6 @@ pub async fn fork_session(
let copy_ms = t0.elapsed().as_millis() as u64;
// Writeback session to backend (fire-and-forget).
// This is telemetry-grade: the local fork works without it. All fork
// state lives locally (session files on disk), and the caller does not
// depend on synchronous backend registration. The backend eventually
// learns about the session when the background task completes.
// Spawning removes the network round-trip (~200-400ms) from the
// critical path.
if let Some(am) = auth_manager {
let sid = new_session_id.clone();
let cwd = request.new_cwd.clone();
let parent = request.source_session_id.clone();
let model = request.new_model_id.clone();
let aid = agent_id.to_string();
tokio::spawn(async move {
if let Err(e) =
sync_forked_session_to_backend(&sid, &cwd, parent, model, &aid, am).await
{
tracing::warn!(
session_id = %sid,
error = %e,
"Failed to register forked session with backend (background)"
);
}
});
}
let total_ms = t0.elapsed().as_millis() as u64;
tracing::info!(
target: FORK_LOG,
@@ -149,7 +117,7 @@ pub async fn fork_session(
total_ms,
chat_copied = result.chat_messages_copied,
updates_copied = result.updates_copied,
"FORK_COPY: session data copied (backend sync spawned in background)"
"FORK_COPY: session data copied"
);
Ok(ForkSessionResponse {
@@ -163,43 +131,6 @@ pub async fn fork_session(
})
}
/// Sync a forked session to the backend (for writeback mode).
async fn sync_forked_session_to_backend(
session_id: &str,
cwd: &str,
parent_session_id: String,
model_id: Option<String>,
agent_id: &str,
auth_manager: std::sync::Arc<crate::auth::AuthManager>,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let client = BackendClient::new().with_auth_manager(auth_manager);
let metadata = ExportedMetadata {
title: None, // Will be generated later when session runs
cwd: cwd.to_string(),
model_id,
created_at: Some(chrono::Utc::now().to_rfc3339()),
updated_at: Some(chrono::Utc::now().to_rfc3339()),
total_messages: Some(0),
parent_session_id: Some(parent_session_id),
session_kind: None,
subagent_type: None,
subagent_persona: None,
subagent_role: None,
fork_context_source: None,
subagent_depth: None,
};
client
.upsert_session(session_id, &metadata, agent_id)
.await?;
tracing::info!(
session_id = %session_id,
"Forked session registered with backend"
);
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
crates/codegen/kigi-shell/src/session/helpers/session_compact.rs
-8
View File
@@ -743,7 +743,6 @@ mod classify_tests {
message: "test".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
}))
};
assert!(det(StatusCode::BAD_REQUEST));
@@ -836,7 +835,6 @@ mod classify_tests {
.into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
})));
}
#[test]
@@ -866,7 +864,6 @@ mod classify_tests {
message: "bad payload".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
}) else {
panic!("expected Deterministic for 400");
};
@@ -878,7 +875,6 @@ mod classify_tests {
message: "upstream blip".into(),
model_metadata: None,
retry_after_secs: None,
should_retry: None,
}) else {
panic!("expected Transient for 500");
};
@@ -1595,15 +1591,11 @@ mod reasoning_compaction_regression_tests {
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
crates/codegen/kigi-shell/src/session/mod.rs
+2 -16
View File
@@ -12,6 +12,7 @@ pub mod two_pass;
pub use self::acp_session::*;
pub use self::acp_types::*;
pub use self::commands::*;
pub use self::feedback_types::{ClientType, FeedbackTerminalInfo, RatingType};
pub use self::fork::{ForkSessionRequest, ForkSessionResponse, fork_session};
pub use self::handle::*;
pub use self::persistence::{
@@ -19,13 +20,9 @@ pub use self::persistence::{
resolve_local_session_any_cwd, session_exists_by_id, session_exists_for_cwd,
};
pub use self::result::{Empty, ExtMethodResult};
pub use self::share::{ShareSessionRequest, ShareSessionResponse};
pub use kigi_fsnotify::{
FsConfig, FsEvent, FsEventKind, FsEventSource, FsNotifyError, GitMetaKind,
};
pub use prod_mc_cli_chat_proxy_types::feedback_types::{
ClientType, FeedbackTerminalInfo, RatingType,
};
/// `false` twin: this template is not compiled into this build, so no
/// template matches. Keeps ungated call sites compiling in both
/// configurations.
@@ -273,18 +270,6 @@ pub struct ClientFsConfig {
pub mode: ClientFsMode,
}
/// Share session request/response types
pub mod share {
/// Request to share a session via URL
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ShareSessionRequest {
pub session_id: String,
}
/// Response containing the shareable URL
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ShareSessionResponse {
pub share_url: String,
}
}
/// Proxy config for the session registry client.
/// Shared between `acp_session` (slash commands) and `persistence` (title generation).
#[derive(Clone)]
@@ -303,6 +288,7 @@ pub(crate) mod events;
pub mod export;
pub mod feedback;
pub mod feedback_manager;
pub mod feedback_types;
pub mod file_system;
pub mod fork;
pub(crate) mod fs_watch;
crates/codegen/kigi-shell/src/session/persistence.rs
+25 -326
View File
@@ -5,8 +5,6 @@ use std::sync::Arc;
use crate::config::StorageMode;
use crate::remote::RemoteSync;
use crate::sampling::Client as OaiCompatClient;
use crate::sampling::ConversationItem;
use crate::session::export::ExportedMetadata;
@@ -105,7 +103,7 @@ pub struct UserFeedbackEntry {
pub dismissed: bool,
/// The full submission payload (omitted when dismissed)
#[serde(skip_serializing_if = "Option::is_none")]
pub submission: Option<prod_mc_cli_chat_proxy_types::feedback_types::FeedbackSubmission>,
pub submission: Option<crate::session::feedback_types::FeedbackSubmission>,
}
/// Helper for `#[serde(skip_serializing_if)]` on bool fields.
@@ -116,14 +114,13 @@ pub(crate) fn is_false(v: &bool) -> bool {
#[cfg(test)]
mod feedback_tests {
use super::*;
use prod_mc_cli_chat_proxy_types::feedback_types::{
use crate::session::feedback_types::{
ClientType, FeedbackSubmission, FeedbackType, RatingType,
};
fn make_submission(thumbs_up: bool) -> FeedbackSubmission {
FeedbackSubmission {
session_id: "session-abc".into(),
user_id: None,
client_type: ClientType::Tui,
feedback_type: if thumbs_up {
FeedbackType::Rating
@@ -139,22 +136,13 @@ mod feedback_tests {
Some("could be better".into())
},
feedback_categories: vec![],
message_id: None,
model_id: Some("grok-3-fast".into()),
resolved_model_id: Some("grok-4.5".into()),
model_fingerprint: None,
context_type: None,
feature_name: None,
tool_name: None,
experiment_id: None,
comparison_id: None,
preferred_model_id: None,
preference_strength: None,
preference_reasons: vec![],
request_id: None,
client_version: None,
shell_version: None,
extension_host: None,
metadata: None,
last_user_message: None,
last_assistant_message: None,
@@ -165,7 +153,6 @@ mod feedback_tests {
context_tokens_used: None,
context_window_tokens: None,
terminal_info: None,
unified_log_url: None,
}
}
@@ -1392,7 +1379,6 @@ struct SessionPersistence {
/// Pending ACP notification for merging consecutive text chunks
pending_notification: Option<acp::SessionNotification>,
rx: mpsc::UnboundedReceiver<PersistenceMsg>,
remote_sync: Option<RemoteSync>,
/// Session title generation lifecycle.
summary: crate::session::summary::SummaryGenerator,
registry_title_sync: Option<RegistryGeneratedTitleSync>,
@@ -1497,20 +1483,12 @@ impl SessionPersistence {
}
}
/// Flush any pending merged ACP notification to disk and remote sync.
/// Flush any pending merged ACP notification to disk.
async fn flush_pending(&mut self) {
// Write any pending merged ACP notification
if let Some(notification) = self.pending_notification.take() {
self.write_update(&SessionUpdate::Acp(Box::new(notification.clone())))
self.write_update(&SessionUpdate::Acp(Box::new(notification)))
.await;
// HTTP-based remote sync (Writeback mode)
if let Some(sync) = &self.remote_sync {
sync.queue(notification);
}
}
// Flush HTTP sync
if let Some(sync) = &self.remote_sync {
sync.flush();
}
}
@@ -1549,12 +1527,8 @@ impl SessionPersistence {
SessionUpdate::Acp(notification) => {
// ACP notifications use merging to coalesce consecutive text chunks
if let Some(to_write) = self.maybe_merge_notification(&notification) {
self.write_update(&SessionUpdate::Acp(Box::new(to_write.clone())))
self.write_update(&SessionUpdate::Acp(Box::new(to_write)))
.await;
// HTTP-based remote sync (Writeback mode)
if let Some(sync) = &self.remote_sync {
sync.queue(to_write);
}
}
}
SessionUpdate::Xai(_) => {
@@ -1602,9 +1576,6 @@ impl SessionPersistence {
{
tracing::warn!(?e, "failed to update current model");
}
if let Some(sync) = &self.remote_sync {
sync.set_model_id(model_id.0.to_string());
}
}
PersistenceMsg::PlanState(state) => {
if let Err(e) = self.storage.write_plan_state(&self.info, &state).await {
@@ -1660,9 +1631,6 @@ impl SessionPersistence {
&self.info,
&title,
);
if let Some(sync) = &self.remote_sync {
sync.set_title(title.clone());
}
if let Some(reg) = self.registry_title_sync.as_ref()
&& !reg.suppress_for_zdr
{
@@ -1901,69 +1869,6 @@ fn collect_session_files_recursive(base: &Path, dir: &Path, files: &mut Vec<Copi
}
}
fn init_remote_sync(
summary: &Summary,
storage_mode: StorageMode,
auth_manager: Option<Arc<crate::auth::AuthManager>>,
) -> io::Result<Option<RemoteSync>> {
match storage_mode {
StorageMode::Local => Ok(None),
StorageMode::Writeback => {
let auth_manager = auth_manager.ok_or_else(|| {
io::Error::new(
io::ErrorKind::PermissionDenied,
"Writeback storage mode requires authentication. Run 'grok login' first.",
)
})?;
if auth_manager.current_or_expired().is_some() {
// ZDR was an xAI team concept; nothing gates remote sync here.
} else {
tracing::warn!(
"writeback: no auth loaded yet, ZDR check skipped (backend enforces server-side)"
);
}
tracing::info!("Writeback mode enabled, syncing to backend");
let client =
crate::remote::BackendClient::new().with_auth_manager(auth_manager.clone());
let metadata = ExportedMetadata::from_summary(summary);
Ok(Some(RemoteSync::new(
summary.info.id.to_string(),
metadata,
client,
)))
}
}
}
/// Pull a session from the backend if not found locally. Returns the pulled
/// session's [`Info`] (cwd may differ from caller's on different machines),
/// or `None` if not found or on error.
async fn try_pull_from_remote(info: &Info, client: &crate::remote::BackendClient) -> Option<Info> {
// BackendClient resolves auth internally via its auth_manager.
client.auth_manager.as_ref()?;
tracing::info!(session_id = %info.id, "Session not found locally, trying backend");
match crate::remote::pull_session_to_local(&info.id.0, client).await {
Ok(crate::remote::PullResult::Hydrated(pulled_info)) => {
tracing::info!(
session_id = %info.id,
pulled_cwd = %pulled_info.cwd,
"Pulled session from backend"
);
Some(pulled_info)
}
Ok(crate::remote::PullResult::NotFound) => {
tracing::debug!(session_id = %info.id, "Session not found on backend either");
None
}
Err(e) => {
tracing::warn!(session_id = %info.id, error = %e, "Backend pull failed");
None
}
}
}
/// Map a persistence `io::Error` into an `acp::Error` with a human-friendly
/// `message` and a stable `data.code` for log aggregation.
pub(crate) fn io_error_to_acp(e: &io::Error) -> acp::Error {
@@ -2066,7 +1971,6 @@ pub(crate) async fn new(
let info_clone = info.clone();
let storage: Arc<dyn StorageAdapter> = Arc::from(storage);
let remote_sync = init_remote_sync(&summary, storage_mode, auth_manager)?;
let handle = PersistenceHandle {
tx: tx.clone(),
noop: false,
@@ -2078,7 +1982,6 @@ pub(crate) async fn new(
storage: storage.clone(),
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
sampling_client,
@@ -2147,7 +2050,6 @@ pub async fn new_with_explicit_dir(
storage: storage.clone(),
pending_notification: None,
rx,
remote_sync: None,
summary: crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
sampling_client,
@@ -2195,101 +2097,12 @@ pub struct PersistedInfoLight {
pub goal_mode_state: Option<crate::session::goal_tracker::GoalOrchestration>,
}
/// On NotFound, try pulling from backend. Returns pulled info or the original error.
async fn pull_on_miss(
info: &Info,
client: &crate::remote::BackendClient,
err: io::Error,
) -> io::Result<Info> {
if err.kind() != io::ErrorKind::NotFound {
return Err(err);
}
try_pull_from_remote(info, client).await.ok_or(err)
}
#[expect(dead_code, reason = "wired when session restore flow calls load")]
pub(crate) async fn load(
info: &Info,
sampling_client: OaiCompatClient,
storage_mode: StorageMode,
auth_manager: Option<Arc<crate::auth::AuthManager>>,
backend: Option<&crate::remote::BackendClient>,
gateway: Option<GatewaySender>,
session_summary_model: String,
registry_title_sync: Option<RegistryGeneratedTitleSync>,
) -> io::Result<(PersistedInfo, PersistenceHandle)> {
let root_dir = kigi_home();
let storage: Box<dyn StorageAdapter> = Box::new(JsonlStorageAdapter::with_root(root_dir));
let (persisted, loaded_info) = match storage.load_session(info).await {
Ok(p) => (p, info.clone()),
Err(e) => match backend {
Some(client) => {
let pulled = pull_on_miss(info, client, e).await?;
let p = storage.load_session(&pulled).await?;
(p, pulled)
}
None => return Err(e),
},
};
// Touch on load too: resuming must reset the worktree's gc expiry clock.
touch_worktree_for_session(&loaded_info).await;
let persisted_info = PersistedInfo {
summary: persisted.summary,
chat_history: persisted.chat_history,
updates: persisted.updates,
plan_state: persisted.plan_state,
rewind_points: persisted.rewind_points,
signals: persisted.signals,
};
let (tx, rx) = mpsc::unbounded_channel::<PersistenceMsg>();
let storage: Arc<dyn StorageAdapter> = Arc::from(storage);
let remote_sync = init_remote_sync(&persisted_info.summary, storage_mode, auth_manager)?;
let has_title = !persisted_info.summary.display_title().is_empty();
let handle = PersistenceHandle {
tx: tx.clone(),
noop: false,
};
tokio::task::spawn(async move {
let mut summary_gen = crate::session::summary::SummaryGenerator::new(
crate::session::summary::SummaryConfig {
sampling_client,
model: session_summary_model,
persistence_tx: tx,
},
);
if has_title {
summary_gen.mark_done();
}
let persistence = SessionPersistence {
info: loaded_info,
storage: storage.clone(),
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
summary: summary_gen,
registry_title_sync,
gateway,
};
persistence.run().await;
});
Ok((persisted_info, handle))
}
/// Like `load`, but doesn't load updates into memory.
/// Loads a session for streaming updates without reading them into memory.
/// Instead, provides the path to the updates file for streaming reads.
/// Use this for memory-efficient session loading when replaying updates.
pub(crate) async fn load_light(
info: &Info,
sampling_client: OaiCompatClient,
storage_mode: StorageMode,
auth_manager: Option<Arc<crate::auth::AuthManager>>,
backend: Option<&crate::remote::BackendClient>,
gateway: Option<GatewaySender>,
session_summary_model: String,
registry_title_sync: Option<RegistryGeneratedTitleSync>,
@@ -2298,16 +2111,9 @@ pub(crate) async fn load_light(
let storage: Box<dyn StorageAdapter> =
Box::new(JsonlStorageAdapter::with_root(root_dir.clone()));
let (persisted, loaded_info) = match storage.load_session_without_updates(info).await {
Ok(p) => (p, info.clone()),
Err(e) => match backend {
Some(client) => {
let pulled = pull_on_miss(info, client, e).await?;
let p = storage.load_session_without_updates(&pulled).await?;
(p, pulled)
}
None => return Err(e),
},
let (persisted, loaded_info) = {
let p = storage.load_session_without_updates(info).await?;
(p, info.clone())
};
// Touch on load too: resuming must reset the worktree's gc expiry clock.
touch_worktree_for_session(&loaded_info).await;
@@ -2330,7 +2136,6 @@ pub(crate) async fn load_light(
let (tx, rx) = mpsc::unbounded_channel::<PersistenceMsg>();
let storage: Arc<dyn StorageAdapter> = Arc::from(storage);
let remote_sync = init_remote_sync(&persisted_info.summary, storage_mode, auth_manager)?;
let has_title = !persisted_info.summary.display_title().is_empty();
let handle = PersistenceHandle {
@@ -2353,7 +2158,6 @@ pub(crate) async fn load_light(
storage: storage.clone(),
pending_notification: None,
rx,
remote_sync: remote_sync.clone(),
summary: summary_gen,
registry_title_sync,
gateway,
@@ -2373,100 +2177,58 @@ pub async fn list_summaries(cwd: Option<&str>) -> io::Result<Vec<Summary>> {
}
/// Failure modes of [`delete_session_history`].
///
/// Kept distinct so callers can surface a precise message: a remote
/// failure is reported separately from a local-disk failure because the
/// remote delete runs first and aborts the whole operation (see the doc
/// on [`delete_session_history`]).
#[derive(Debug, thiserror::Error)]
pub enum DeleteSessionError {
/// Listing local summaries (to resolve the on-disk session dir) failed.
#[error("failed to list sessions: {0}")]
List(#[source] io::Error),
/// The remote (writeback) copy could not be deleted; local bits were
/// left untouched so the operation can be retried.
#[error("failed to delete remote session data: {0}")]
Remote(#[source] crate::remote::client::BackendError),
/// The local on-disk session directory could not be removed.
#[error("failed to delete session: {0}")]
Local(#[source] io::Error),
}
/// Where a session copy was actually removed by [`delete_session_history`].
/// Whether a session copy was removed by [`delete_session_history`].
///
/// Both fields are `false` when nothing existed to delete (still a
/// `local_removed` is `false` when nothing existed to delete (still a
/// success). Callers use [`Self::any_removed`] to decide between a
/// "deleted" and a "not found" message without conflating a remote-only
/// delete with a no-op.
/// "deleted" and a "not found" message.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct SessionDeletion {
/// A local on-disk session directory was found and removed.
pub local_removed: bool,
/// A remote (writeback) copy was found and removed. `false` when
/// `needs_remote` was not set, or the remote copy was already absent
/// (the backend returned `404`).
pub remote_removed: bool,
}
impl SessionDeletion {
/// `true` when a copy was removed from at least one location.
/// `true` when the local session directory was removed.
pub fn any_removed(self) -> bool {
self.local_removed || self.remote_removed
self.local_removed
}
}
/// Permanently delete a session's history: the remote (writeback) copy
/// when `needs_remote`, the local on-disk session directory, and the
/// FTS search-index entry.
/// Permanently delete a session's history: the local on-disk session
/// directory and the FTS search-index entry.
///
/// Idempotent: a session that is missing locally (e.g. remote-only)
/// still succeeds, and a remote `404` (copy already gone) is treated as
/// success rather than an error. When `needs_remote` is set the remote
/// delete runs *first* and is authoritative — only on its success (or a
/// `404`) are the local bits removed. This ordering prevents a partial
/// delete where the local copy is nuked but the remote copy lingers and
/// re-appears on the next session list.
/// Idempotent: a session that is missing locally still succeeds.
///
/// Returns a [`SessionDeletion`] recording which copies (local / remote)
/// were actually removed; both fields `false` means nothing existed
/// (still `Ok`).
/// Returns a [`SessionDeletion`] recording whether a local copy was
/// removed; `false` means nothing existed (still `Ok`).
pub async fn delete_session_history(
session_id: &str,
cwd: Option<&str>,
needs_remote: bool,
auth_manager: Arc<crate::auth::AuthManager>,
) -> Result<SessionDeletion, DeleteSessionError> {
let sid = acp::SessionId::new(Arc::from(session_id));
// Resolve the local session info, scoping to cwd if provided. A
// remote-only session won't be found here — that's fine, the remote
// delete (if applicable) still runs.
// Resolve the local session info, scoping to cwd if provided.
let summaries = list_summaries(cwd)
.await
.map_err(DeleteSessionError::List)?;
let local_info = summaries
let Some(info) = summaries
.iter()
.find(|s| s.info.id == sid)
.map(|s| s.info.clone());
// Remote delete first (authoritative for cloud history). A genuine
// failure aborts before any local mutation so the row does not
// reappear; a `404` means the copy is already gone, so deletion stays
// idempotent and falls through to local cleanup.
let remote_removed = if needs_remote {
let result = crate::remote::client::BackendClient::new()
.with_auth_manager(auth_manager)
.delete_session_data(session_id)
.await;
classify_remote_delete(result)?
} else {
false
};
let Some(info) = local_info else {
.map(|s| s.info.clone())
else {
return Ok(SessionDeletion {
local_removed: false,
remote_removed,
});
};
@@ -2481,85 +2243,22 @@ pub async fn delete_session_history(
Ok(SessionDeletion {
local_removed: true,
remote_removed,
})
}
/// Classify a remote `delete_session_data` result, reporting whether a
/// remote copy was actually removed: a `2xx` means a copy was deleted
/// (`Ok(true)`), a `404` means it was already gone so deletion stays
/// idempotent (`Ok(false)`), and any other backend error aborts the
/// delete (`Err`) so local bits are left untouched and it can be retried.
fn classify_remote_delete(
result: Result<(), crate::remote::client::BackendError>,
) -> Result<bool, DeleteSessionError> {
use crate::remote::client::BackendError;
match result {
Ok(()) => Ok(true),
Err(BackendError::RequestFailed { status: 404, .. }) => Ok(false),
Err(e) => Err(DeleteSessionError::Remote(e)),
}
}
#[cfg(test)]
mod delete_session_history_tests {
use super::{DeleteSessionError, SessionDeletion, classify_remote_delete};
use crate::remote::client::BackendError;
use super::SessionDeletion;
#[test]
fn remote_ok_reports_removed() {
assert!(
classify_remote_delete(Ok(())).unwrap(),
"a 2xx delete must report that a remote copy was removed"
);
}
#[test]
fn remote_404_is_treated_as_already_deleted() {
let removed = classify_remote_delete(Err(BackendError::RequestFailed {
status: 404,
body: "not found".into(),
}))
.expect("a 404 means the remote copy is gone — deletion must stay idempotent");
assert!(
!removed,
"a 404 must report that nothing was removed remotely"
);
}
#[test]
fn remote_non_404_request_failure_aborts() {
let res = classify_remote_delete(Err(BackendError::RequestFailed {
status: 500,
body: "boom".into(),
}));
assert!(matches!(res, Err(DeleteSessionError::Remote(_))));
}
#[test]
fn remote_auth_failure_aborts() {
let res = classify_remote_delete(Err(BackendError::Auth("denied".into())));
assert!(matches!(res, Err(DeleteSessionError::Remote(_))));
}
#[test]
fn any_removed_reflects_either_location() {
fn any_removed_reflects_local_removal() {
assert!(!SessionDeletion::default().any_removed());
assert!(
SessionDeletion {
local_removed: true,
remote_removed: false,
}
.any_removed()
);
assert!(
SessionDeletion {
local_removed: false,
remote_removed: true,
}
.any_removed(),
"a remote-only delete must count as removed"
);
}
}
crates/codegen/kigi-shell/src/session/storage/jsonl/tests.rs
+1 -12
View File
@@ -1085,7 +1085,7 @@ async fn test_load_prompts_only_large_session() {
#[tokio::test]
async fn test_append_feedback_creates_file_and_persists() {
use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry};
use prod_mc_cli_chat_proxy_types::feedback_types::{
use crate::session::feedback_types::{
ClientType, FeedbackSubmission, FeedbackType, RatingType,
};
let temp_dir = TempDir::new().unwrap();
@@ -1101,7 +1101,6 @@ async fn test_append_feedback_creates_file_and_persists() {
dismissed: false,
submission: Some(FeedbackSubmission {
session_id: "test-session-123".into(),
user_id: None,
client_type: ClientType::Tui,
feedback_type: FeedbackType::Rating,
turn_number: Some(3),
@@ -1109,22 +1108,13 @@ async fn test_append_feedback_creates_file_and_persists() {
rating_value: Some(1),
feedback_text: None,
feedback_categories: vec![],
message_id: None,
model_id: Some("grok-3-fast".into()),
resolved_model_id: Some("grok-4.5".into()),
model_fingerprint: None,
context_type: None,
feature_name: None,
tool_name: None,
experiment_id: None,
comparison_id: None,
preferred_model_id: None,
preference_strength: None,
preference_reasons: vec![],
request_id: None,
client_version: None,
shell_version: None,
extension_host: None,
metadata: None,
last_user_message: None,
last_assistant_message: None,
@@ -1135,7 +1125,6 @@ async fn test_append_feedback_creates_file_and_persists() {
context_tokens_used: None,
context_window_tokens: None,
terminal_info: None,
unified_log_url: None,
}),
});
adapter.append_feedback(&info, &user_entry).await.unwrap();
crates/codegen/kigi-shell/src/session/unified_list/cursor.rs
+72 -398
View File
@@ -3,7 +3,6 @@ use std::cmp::{Ordering, Reverse};
use base64::Engine as _;
use serde::{Deserialize, Serialize};
use super::PartialReason;
use super::envelope::SessionKind;
use super::row::UnifiedRow;
@@ -11,10 +10,6 @@ use super::row::UnifiedRow;
pub(super) struct CompositeCursor {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub boundary: Option<BoundaryKey>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub conv_page_token: Option<String>,
#[serde(default)]
pub conv_page_drained: bool,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
@@ -42,52 +37,21 @@ impl CompositeCursor {
}
}
pub(super) enum ConvLane {
Skipped,
Degraded(PartialReason),
Page {
rows: Vec<UnifiedRow>,
next_token: Option<String>,
frontier: Option<BoundaryKey>,
},
}
pub(super) fn conv_frontier(raw_rows: &[UnifiedRow], has_more: bool) -> Option<BoundaryKey> {
if !has_more {
return None;
}
raw_rows
.iter()
.max_by(|a, b| cmp_total_order(a, b))
.map(boundary_of)
}
pub(super) struct Paginated {
pub candidates: Vec<UnifiedRow>,
pub emit_count: usize,
pub next_cursor: Option<CompositeCursor>,
pub partial: Option<PartialReason>,
}
pub(super) fn merge_and_paginate(
/// Sort local rows newest-first, resume after the cursor boundary, and cut
/// one page. `next_cursor` is set only when rows remain past the page.
pub(super) fn paginate(
local: Vec<UnifiedRow>,
conv: ConvLane,
cursor: &CompositeCursor,
limit: usize,
) -> Paginated {
let (conv_rows, conv_next_token, conv_fetched, conv_frontier, partial) = match conv {
ConvLane::Skipped => (Vec::new(), None, false, None, None),
ConvLane::Degraded(reason) => (Vec::new(), None, false, None, Some(reason)),
ConvLane::Page {
rows,
next_token,
frontier,
} => (rows, next_token, true, frontier, None),
};
let mut keyed: Vec<(SortKey, UnifiedRow)> = local
.into_iter()
.chain(conv_rows)
.map(|row| (row_sort_key(&row), row))
.collect();
@@ -98,46 +62,12 @@ pub(super) fn merge_and_paginate(
keyed.sort_by(|(a, _), (b, _)| a.cmp(b));
let mut emit_count = keyed.len().min(limit);
if let Some(frontier) = &conv_frontier {
let fkey = boundary_sort_key(frontier);
let frontier_count = keyed
.iter()
.take_while(|(k, _)| k.cmp(&fkey) != Ordering::Greater)
.count();
emit_count = emit_count.min(frontier_count);
}
let emit_count = keyed.len().min(limit);
let new_boundary = (emit_count > 0).then(|| boundary_of(&keyed[emit_count - 1].1));
let has_more = keyed.len() > emit_count;
let tail = &keyed[emit_count..];
let local_has_more = tail.iter().any(|(_, r)| r.kind == SessionKind::Build);
let conv_in_tail = tail.iter().any(|(_, r)| r.kind == SessionKind::Chat);
let (next_conv_token, next_conv_drained, conv_has_more) = if conv_fetched {
if conv_in_tail {
(cursor.conv_page_token.clone(), false, true)
} else {
let has_more = conv_next_token.is_some();
(conv_next_token, true, has_more)
}
} else if partial.is_some() && cursor.conv_page_token.is_some() && new_boundary.is_some() {
(
cursor.conv_page_token.clone(),
cursor.conv_page_drained,
true,
)
} else {
(
cursor.conv_page_token.clone(),
cursor.conv_page_drained,
false,
)
};
let next_cursor = (local_has_more || conv_has_more).then(|| CompositeCursor {
let next_cursor = has_more.then(|| CompositeCursor {
boundary: new_boundary.or_else(|| cursor.boundary.clone()),
conv_page_token: next_conv_token,
conv_page_drained: next_conv_drained,
});
let candidates: Vec<UnifiedRow> = keyed.into_iter().map(|(_, row)| row).collect();
@@ -146,7 +76,6 @@ pub(super) fn merge_and_paginate(
candidates,
emit_count,
next_cursor,
partial,
}
}
@@ -205,12 +134,8 @@ pub(super) fn cmp_total_order(a: &UnifiedRow, b: &UnifiedRow) -> Ordering {
#[cfg(test)]
mod tests {
use super::*;
use crate::remote::Conversation;
use crate::session::merge::MergedSession;
use crate::session::unified_list::{
conversation_to_row, facet_registry, merged_session_to_row,
};
use std::collections::BTreeSet;
use crate::session::unified_list::{facet_registry, merged_session_to_row};
fn local(id: &str, ts: &str) -> UnifiedRow {
let m = MergedSession {
@@ -236,336 +161,85 @@ mod tests {
merged_session_to_row(m, facet_registry())
}
fn conv(id: &str, ts: &str) -> UnifiedRow {
let c = Conversation {
conversation_id: id.into(),
title: "t".into(),
modify_time: Some(ts.into()),
..Conversation::default()
};
conversation_to_row(c, facet_registry())
}
struct ConvSource {
rows: Vec<UnifiedRow>,
page_size: usize,
}
impl ConvSource {
fn new(mut rows: Vec<UnifiedRow>, page_size: usize) -> Self {
rows.sort_by(cmp_total_order);
Self { rows, page_size }
}
fn page(&self, token: Option<&str>) -> ConvLane {
if self.rows.is_empty() {
return ConvLane::Page {
rows: Vec::new(),
next_token: None,
frontier: None,
};
}
let idx = token
.and_then(|t| t.strip_prefix('p'))
.and_then(|n| n.parse::<usize>().ok())
.unwrap_or(0);
let start = idx * self.page_size;
let end = (start + self.page_size).min(self.rows.len());
let rows = self.rows.get(start..end).unwrap_or(&[]).to_vec();
let next_token = (end < self.rows.len()).then(|| format!("p{}", idx + 1));
let frontier = conv_frontier(&rows, next_token.is_some());
ConvLane::Page {
rows,
next_token,
frontier,
}
}
}
fn walk_all(local_window: &[UnifiedRow], conv: &ConvSource, limit: usize) -> Vec<String> {
let mut cursor = CompositeCursor::default();
let mut emitted: Vec<String> = Vec::new();
for _ in 0..1000 {
let lane = conv.page(cursor.conv_page_token.as_deref());
let result = merge_and_paginate(local_window.to_vec(), lane, &cursor, limit);
emitted.extend(
result.candidates[..result.emit_count]
.iter()
.map(|r| r.legacy.session_id.clone()),
);
match result.next_cursor {
Some(c) => cursor = c,
None => return emitted,
}
}
panic!("pagination did not terminate");
}
fn ids(rows: &[UnifiedRow]) -> Vec<String> {
rows.iter().map(|r| r.legacy.session_id.clone()).collect()
fn ids(p: &Paginated) -> Vec<String> {
p.candidates[..p.emit_count]
.iter()
.map(|r| r.legacy.session_id.clone())
.collect()
}
#[test]
fn cursor_round_trips() {
let cur = CompositeCursor {
fn cursor_roundtrip_boundary_only() {
let c = CompositeCursor {
boundary: Some(BoundaryKey {
updated_at: "2026-06-01T00:00:00Z".into(),
kind: SessionKind::Chat,
session_id: "conv_1".into(),
}),
conv_page_token: Some("p3".into()),
conv_page_drained: true,
};
let decoded = CompositeCursor::decode(Some(&cur.encode()));
assert_eq!(decoded.conv_page_token.as_deref(), Some("p3"));
assert!(decoded.conv_page_drained);
let b = decoded.boundary.unwrap();
assert_eq!(b.session_id, "conv_1");
assert_eq!(b.kind, SessionKind::Chat);
}
#[test]
fn malformed_cursor_decodes_to_fresh_first_page() {
for bad in [Some("not base64 !!!"), Some(""), None] {
let c = CompositeCursor::decode(bad);
assert!(c.boundary.is_none());
assert!(c.conv_page_token.is_none());
assert!(!c.conv_page_drained);
}
}
#[test]
fn multi_page_walk_equals_single_fetch_window() {
let local_window = vec![
local("l1", "2026-06-10T00:00:00Z"),
local("l2", "2026-06-08T00:00:00Z"),
local("l3", "2026-06-04T00:00:00Z"),
local("l4", "2026-05-30T00:00:00Z"),
];
let conv_rows = vec![
conv("c1", "2026-06-09T00:00:00Z"),
conv("c2", "2026-06-07T00:00:00Z"),
conv("c3", "2026-06-06T00:00:00Z"),
conv("c4", "2026-06-03T00:00:00Z"),
conv("c5", "2026-05-29T00:00:00Z"),
];
let mut expected_all = local_window.clone();
expected_all.extend(conv_rows.clone());
expected_all.sort_by(cmp_total_order);
let expected_ids = ids(&expected_all);
for &limit in &[1usize, 2, 3, 5, 7, 100] {
for &page_size in &[1usize, 2, 3] {
let source = ConvSource::new(conv_rows.clone(), page_size);
let got = walk_all(&local_window, &source, limit);
let unique: BTreeSet<&String> = got.iter().collect();
assert_eq!(
unique.len(),
got.len(),
"duplicate emitted (limit={limit}, page_size={page_size}): {got:?}"
);
assert_eq!(
got, expected_ids,
"walk != single fetch (limit={limit}, page_size={page_size})"
);
}
}
}
#[test]
fn equal_updated_at_tie_break_no_drop_or_dup() {
let ts = "2026-06-01T00:00:00Z";
let local_window = vec![local("l_same", ts), local("l_old", "2026-05-01T00:00:00Z")];
let conv_rows = vec![conv("c_same", ts), conv("c_old", "2026-05-15T00:00:00Z")];
let mut expected_all = local_window.clone();
expected_all.extend(conv_rows.clone());
expected_all.sort_by(cmp_total_order);
let expected_ids = ids(&expected_all);
assert_eq!(expected_ids[0], "l_same");
assert_eq!(expected_ids[1], "c_same");
for &limit in &[1usize, 2, 3] {
let source = ConvSource::new(conv_rows.clone(), 1);
let got = walk_all(&local_window, &source, limit);
let unique: BTreeSet<&String> = got.iter().collect();
assert_eq!(unique.len(), got.len(), "dup at limit={limit}: {got:?}");
assert_eq!(
got, expected_ids,
"tie-break walk mismatch at limit={limit}"
);
}
}
#[test]
fn partial_conv_page_is_not_advanced_until_drained() {
let local_window = vec![
local("l1", "2026-06-10T00:00:00Z"),
local("l2", "2026-06-08T00:00:00Z"),
];
let conv_rows = vec![
conv("c1", "2026-06-09T00:00:00Z"),
conv("c2", "2026-06-07T00:00:00Z"),
];
let source = ConvSource::new(conv_rows.clone(), 2);
let got = walk_all(&local_window, &source, 1);
let mut expected_all = local_window.clone();
expected_all.extend(conv_rows.clone());
expected_all.sort_by(cmp_total_order);
assert_eq!(got, ids(&expected_all));
}
#[test]
fn whole_page_filtered_out_does_not_drop_later_match() {
let local_window = vec![
local("l1", "2026-06-10T00:00:00Z"),
local("l2", "2026-06-01T00:00:00Z"),
];
let raw = vec![
conv("c1_drop", "2026-06-09T00:00:00Z"),
conv("c2_drop", "2026-06-05T00:00:00Z"),
conv("c3_ok", "2026-06-03T00:00:00Z"),
];
let source = ConvSource::new(raw, 1);
let mut cursor = CompositeCursor::default();
let mut emitted: Vec<String> = Vec::new();
for _ in 0..1000 {
let lane = match source.page(cursor.conv_page_token.as_deref()) {
ConvLane::Page {
rows,
next_token,
frontier,
} => ConvLane::Page {
rows: rows
.into_iter()
.filter(|r| r.legacy.session_id.contains("ok"))
.collect(),
next_token,
frontier,
},
other => other,
};
let result = merge_and_paginate(local_window.clone(), lane, &cursor, 2);
emitted.extend(
result.candidates[..result.emit_count]
.iter()
.map(|r| r.legacy.session_id.clone()),
);
match result.next_cursor {
Some(c) => cursor = c,
None => break,
}
}
let mut expected = local_window.clone();
expected.push(conv("c3_ok", "2026-06-03T00:00:00Z"));
expected.sort_by(cmp_total_order);
assert_eq!(emitted, ids(&expected));
assert!(
emitted.iter().any(|id| id == "c3_ok"),
"the later matching conversation must not be dropped"
);
}
#[test]
fn local_only_when_conversations_skipped() {
let local_window = vec![
local("l1", "2026-06-10T00:00:00Z"),
local("l2", "2026-06-08T00:00:00Z"),
];
let result = merge_and_paginate(
local_window.clone(),
ConvLane::Skipped,
&CompositeCursor::default(),
10,
);
assert_eq!(result.emit_count, 2);
assert!(result.partial.is_none());
assert!(result.next_cursor.is_none());
}
#[test]
fn degraded_lane_sets_partial_and_returns_local() {
let local_window = vec![local("l1", "2026-06-10T00:00:00Z")];
let result = merge_and_paginate(
local_window,
ConvLane::Degraded(PartialReason::Timeout),
&CompositeCursor::default(),
10,
);
assert_eq!(result.partial, Some(PartialReason::Timeout));
assert_eq!(result.emit_count, 1);
}
#[test]
fn degraded_mid_walk_with_progress_keeps_live_conv_token() {
let cursor = CompositeCursor {
boundary: Some(BoundaryKey {
updated_at: "2026-06-15T00:00:00Z".into(),
updated_at: "2026-02-01T00:00:00Z".into(),
kind: SessionKind::Build,
session_id: "z_newer".into(),
session_id: "a".into(),
}),
conv_page_token: Some("p2".into()),
conv_page_drained: true,
};
let result = merge_and_paginate(
vec![local("l1", "2026-06-10T00:00:00Z")],
ConvLane::Degraded(PartialReason::Timeout),
&cursor,
10,
);
assert_eq!(result.emit_count, 1, "the local row is emitted (progress)");
assert_eq!(result.partial, Some(PartialReason::Timeout));
let next = result
.next_cursor
.expect("progress + live conv token must keep the continuation");
assert_eq!(next.conv_page_token.as_deref(), Some("p2"));
assert_eq!(
next.boundary.as_ref().map(|b| b.session_id.as_str()),
Some("l1")
);
let decoded = CompositeCursor::decode(Some(&c.encode()));
let b = decoded.boundary.expect("boundary survives roundtrip");
assert_eq!(b.session_id, "a");
assert_eq!(b.updated_at, "2026-02-01T00:00:00Z");
}
#[test]
fn degraded_mid_walk_with_no_progress_terminates() {
let cursor = CompositeCursor {
boundary: Some(BoundaryKey {
updated_at: "2026-06-10T00:00:00Z".into(),
kind: SessionKind::Build,
session_id: "l1".into(),
}),
conv_page_token: Some("p2".into()),
conv_page_drained: true,
};
let result = merge_and_paginate(
vec![local("l1", "2026-06-10T00:00:00Z")],
ConvLane::Degraded(PartialReason::Timeout),
&cursor,
10,
);
assert_eq!(
result.emit_count, 0,
"local lane is exhausted (no progress)"
);
assert_eq!(result.partial, Some(PartialReason::Timeout));
fn decode_garbage_yields_default() {
assert!(
result.next_cursor.is_none(),
"a zero-progress degraded page must terminate, not re-emit an identical cursor"
CompositeCursor::decode(Some("!!!not-base64!!!"))
.boundary
.is_none()
);
assert!(CompositeCursor::decode(None).boundary.is_none());
assert!(CompositeCursor::decode(Some("")).boundary.is_none());
}
#[test]
fn degraded_first_page_with_no_token_does_not_fabricate_a_cursor() {
let result = merge_and_paginate(
Vec::new(),
ConvLane::Degraded(PartialReason::Error),
&CompositeCursor::default(),
10,
);
assert_eq!(result.partial, Some(PartialReason::Error));
assert!(result.next_cursor.is_none());
fn paginate_sorts_newest_first_and_cuts_page() {
let rows = vec![
local("old", "2026-01-01T00:00:00Z"),
local("new", "2026-03-01T00:00:00Z"),
local("mid", "2026-02-01T00:00:00Z"),
];
let page = paginate(rows, &CompositeCursor::default(), 2);
assert_eq!(ids(&page), vec!["new", "mid"]);
assert!(page.next_cursor.is_some(), "a third row remains");
}
#[test]
fn paginate_resumes_after_boundary_without_duplicates() {
let rows: Vec<UnifiedRow> = vec![
local("a", "2026-03-01T00:00:00Z"),
local("b", "2026-02-01T00:00:00Z"),
local("c", "2026-01-01T00:00:00Z"),
];
let first = paginate(rows.clone(), &CompositeCursor::default(), 2);
assert_eq!(ids(&first), vec!["a", "b"]);
let cursor = first.next_cursor.expect("more rows remain");
let second = paginate(rows, &cursor, 2);
assert_eq!(ids(&second), vec!["c"]);
assert!(second.next_cursor.is_none(), "list is exhausted");
}
#[test]
fn paginate_exact_page_has_no_next_cursor() {
let rows = vec![
local("a", "2026-03-01T00:00:00Z"),
local("b", "2026-02-01T00:00:00Z"),
];
let page = paginate(rows, &CompositeCursor::default(), 2);
assert_eq!(ids(&page).len(), 2);
assert!(page.next_cursor.is_none());
}
#[test]
fn paginate_ties_break_stably_by_session_id() {
let ts = "2026-02-01T00:00:00Z";
let rows = vec![local("b", ts), local("a", ts), local("c", ts)];
let first = paginate(rows.clone(), &CompositeCursor::default(), 2);
assert_eq!(ids(&first), vec!["a", "b"]);
let cursor = first.next_cursor.expect("one row remains");
let second = paginate(rows, &cursor, 2);
assert_eq!(ids(&second), vec!["c"]);
}
}
crates/codegen/kigi-shell/src/session/unified_list/facets.rs
+1 -247
View File
@@ -4,7 +4,6 @@ use serde::Serialize;
use super::envelope::{FacetMap, FacetValue, SessionKind};
use super::row::UnifiedRow;
use crate::remote::Conversation;
use crate::session::merge::MergedSession;
pub const KIND_FACET_KEY: &str = "kind";
@@ -44,25 +43,6 @@ impl NormalizedItem {
starred: false,
}
}
pub fn from_conversation(c: &Conversation) -> Self {
Self {
kind: SessionKind::Chat,
cwd: String::new(),
repo_name: None,
branch: None,
worktree_label: None,
git_root_dir: None,
source_workspace_dir: None,
workspace_ids: c
.workspaces
.iter()
.map(|w| w.workspace_id.clone())
.filter(|id| !id.is_empty())
.collect(),
starred: c.starred,
}
}
}
#[derive(Debug, Default)]
@@ -393,7 +373,7 @@ pub struct FacetSummaryValue {
#[cfg(test)]
mod tests {
use super::*;
use crate::session::unified_list::{conversation_to_row, merged_session_to_row};
use crate::session::unified_list::merged_session_to_row;
fn local_row(session_id: &str, repo: Option<&str>, branch: Option<&str>) -> UnifiedRow {
let m = MergedSession {
@@ -419,74 +399,6 @@ mod tests {
merged_session_to_row(m, &build_facet_registry())
}
fn conv_row(conversation_id: &str, workspaces: &[&str]) -> UnifiedRow {
let c = Conversation {
conversation_id: conversation_id.into(),
title: "t".into(),
modify_time: Some("2026-06-01T00:00:00Z".into()),
workspaces: workspaces
.iter()
.map(|w| crate::remote::conversations_client::Workspace {
workspace_id: (*w).into(),
})
.collect(),
..Conversation::default()
};
conversation_to_row(c, &build_facet_registry())
}
fn conv_row_starred(conversation_id: &str, starred: bool) -> UnifiedRow {
let c = Conversation {
conversation_id: conversation_id.into(),
title: "t".into(),
modify_time: Some("2026-06-01T00:00:00Z".into()),
starred,
..Conversation::default()
};
conversation_to_row(c, &build_facet_registry())
}
#[test]
fn project_facet_only_on_conversations() {
let reg = build_facet_registry();
let conv = NormalizedItem::from_conversation(&Conversation {
conversation_id: "c1".into(),
workspaces: vec![crate::remote::conversations_client::Workspace {
workspace_id: "ws_9f3a".into(),
}],
..Conversation::default()
});
let facets = reg.extract_all(&conv);
assert!(matches!(
facets.get(WORKSPACE_FACET_KEY),
Some(FacetValue::Many(v)) if v == &[serde_json::json!("ws_9f3a")]
));
let local = NormalizedItem::from_merged(&MergedSession {
session_id: "s".into(),
summary: String::new(),
first_prompt: None,
updated_at: String::new(),
created_at: String::new(),
cwd: "/x".into(),
hostname: None,
source: "local".into(),
model_id: None,
num_messages: 0,
last_active_at: None,
branch: Some("main".into()),
repo_name: Some("xai".into()),
worktree_label: None,
git_root_dir: None,
git_remotes: Vec::new(),
source_workspace_dir: None,
session_kind: None,
});
let lf = reg.extract_all(&local);
assert!(!lf.contains_key(WORKSPACE_FACET_KEY));
assert!(matches!(lf.get(REPO_FACET_KEY), Some(FacetValue::One(_))));
assert!(matches!(lf.get(BRANCH_FACET_KEY), Some(FacetValue::One(_))));
}
#[test]
fn project_pushdown_single_value_sets_workspace_id() {
let reg = build_facet_registry();
@@ -513,121 +425,6 @@ mod tests {
assert!(q.workspace_id.is_none());
}
#[test]
fn project_filter_is_partition_aware_keeps_local_rows() {
let reg = build_facet_registry();
let rows = vec![
local_row("local-1", Some("xai"), Some("main")),
conv_row("conv-match", &["ws_9f3a"]),
conv_row("conv-other", &["ws_zzz"]),
];
let mut filters = BTreeMap::new();
filters.insert(
WORKSPACE_FACET_KEY.to_owned(),
vec![serde_json::json!("ws_9f3a")],
);
let kept = reg.apply_in_memory_filters(&filters, rows);
let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect();
assert!(ids.contains(&"local-1"));
assert!(ids.contains(&"conv-match"));
assert!(!ids.contains(&"conv-other"));
}
#[test]
fn repo_filter_is_partition_aware_keeps_conversation_rows() {
let reg = build_facet_registry();
let rows = vec![
local_row("local-xai", Some("xai"), Some("main")),
local_row("local-other", Some("other"), Some("main")),
conv_row("conv-1", &["ws_9f3a"]),
];
let mut filters = BTreeMap::new();
filters.insert(REPO_FACET_KEY.to_owned(), vec![serde_json::json!("xai")]);
let kept = reg.apply_in_memory_filters(&filters, rows);
let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect();
assert!(ids.contains(&"local-xai"));
assert!(!ids.contains(&"local-other"));
assert!(ids.contains(&"conv-1"));
}
#[test]
fn pushdown_and_in_memory_project_filter_agree() {
let reg = build_facet_registry();
let convs = vec![conv_row("a", &["ws_1"]), conv_row("b", &["ws_2"])];
let mut filters = BTreeMap::new();
filters.insert(
WORKSPACE_FACET_KEY.to_owned(),
vec![serde_json::json!("ws_1")],
);
let in_memory = reg.apply_in_memory_filters(&filters, convs);
let ids: Vec<&str> = in_memory
.iter()
.map(|r| r.legacy.session_id.as_str())
.collect();
assert_eq!(ids, ["a"]);
let mut q = SourceQuery::default();
reg.apply_pushdown(&filters, &mut q);
assert_eq!(q.workspace_id.as_deref(), Some("ws_1"));
}
#[test]
fn starred_facet_present_only_for_starred_conversations() {
let reg = build_facet_registry();
let starred = NormalizedItem::from_conversation(&Conversation {
conversation_id: "c1".into(),
starred: true,
..Conversation::default()
});
assert!(matches!(
reg.extract_all(&starred).get(STARRED_FACET_KEY),
Some(FacetValue::One(serde_json::Value::Bool(true)))
));
let plain = NormalizedItem::from_conversation(&Conversation {
conversation_id: "c2".into(),
starred: false,
..Conversation::default()
});
assert!(!reg.extract_all(&plain).contains_key(STARRED_FACET_KEY));
let local = NormalizedItem::from_merged(&MergedSession {
session_id: "s".into(),
summary: String::new(),
first_prompt: None,
updated_at: String::new(),
created_at: String::new(),
cwd: "/x".into(),
hostname: None,
source: "local".into(),
model_id: None,
num_messages: 0,
last_active_at: None,
branch: None,
repo_name: None,
worktree_label: None,
git_root_dir: None,
git_remotes: Vec::new(),
source_workspace_dir: None,
session_kind: None,
});
assert!(!reg.extract_all(&local).contains_key(STARRED_FACET_KEY));
}
#[test]
fn starred_filter_is_partition_aware_keeps_local_rows() {
let reg = build_facet_registry();
let rows = vec![
local_row("local-1", Some("xai"), Some("main")),
conv_row_starred("conv-starred", true),
conv_row_starred("conv-plain", false),
];
let mut filters = BTreeMap::new();
filters.insert(STARRED_FACET_KEY.to_owned(), vec![serde_json::json!(true)]);
let kept = reg.apply_in_memory_filters(&filters, rows);
let ids: Vec<&str> = kept.iter().map(|r| r.legacy.session_id.as_str()).collect();
assert!(ids.contains(&"local-1"));
assert!(ids.contains(&"conv-starred"));
assert!(!ids.contains(&"conv-plain"));
}
fn local_row_with_git(
session_id: &str,
git_root: Option<&str>,
@@ -656,49 +453,6 @@ mod tests {
merged_session_to_row(m, &build_facet_registry())
}
#[test]
fn git_path_facets_present_only_for_local_rows() {
let reg = build_facet_registry();
let local = NormalizedItem::from_merged(&MergedSession {
session_id: "s".into(),
summary: String::new(),
first_prompt: None,
updated_at: String::new(),
created_at: String::new(),
cwd: "/x".into(),
hostname: None,
source: "local".into(),
model_id: None,
num_messages: 0,
last_active_at: None,
branch: None,
repo_name: None,
worktree_label: None,
git_root_dir: Some("/Users/me/xai".into()),
git_remotes: Vec::new(),
source_workspace_dir: Some("/Users/me/xai-main".into()),
session_kind: Some("worktree".into()),
});
let f = reg.extract_all(&local);
assert!(matches!(
f.get(GIT_ROOT_FACET_KEY),
Some(FacetValue::One(serde_json::Value::String(s))) if s == "/Users/me/xai"
));
assert!(matches!(
f.get(SOURCE_WORKSPACE_FACET_KEY),
Some(FacetValue::One(serde_json::Value::String(s))) if s == "/Users/me/xai-main"
));
// Conversations carry no local git enrichment.
let conv = NormalizedItem::from_conversation(&Conversation {
conversation_id: "c1".into(),
..Conversation::default()
});
let cf = reg.extract_all(&conv);
assert!(!cf.contains_key(GIT_ROOT_FACET_KEY));
assert!(!cf.contains_key(SOURCE_WORKSPACE_FACET_KEY));
}
#[test]
fn git_root_filter_keeps_matching_local_rows() {
let reg = build_facet_registry();
crates/codegen/kigi-shell/src/session/unified_list/mod.rs
+13 -430
View File
@@ -3,8 +3,7 @@ mod envelope;
mod facets;
mod row;
use crate::agent::session_registry_client::SessionRegistryClient;
use crate::remote::{ConvError, ConvQuery, ConversationsClient};
use cursor::{CompositeCursor, ConvLane, Paginated, merge_and_paginate};
use cursor::{CompositeCursor, Paginated, paginate};
pub use envelope::{FacetMap, FacetValue, SessionKind, SessionMetaEnvelope};
pub use facets::{
BRANCH_FACET_KEY, BranchFacet, CWD_FACET_KEY, CwdFacet, FacetProvider, FacetRegistry,
@@ -13,62 +12,18 @@ pub use facets::{
SOURCE_WORKSPACE_FACET_KEY, STARRED_FACET_KEY, SourceQuery, SourceWorkspaceFacet, StarredFacet,
WORKSPACE_FACET_KEY, WORKTREE_FACET_KEY, WorkspaceFacet, WorktreeFacet, build_facet_registry,
};
pub use row::{
ExtSupersetRow, RowMeta, SessionInfo, UnifiedRow, conversation_to_row, merged_session_to_row,
};
pub use row::{ExtSupersetRow, RowMeta, SessionInfo, UnifiedRow, merged_session_to_row};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::sync::LazyLock;
pub const DEFAULT_LIMIT: usize = 30;
const CONV_PAGE_HEADROOM: usize = 5;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PartialReason {
Timeout,
Error,
NoOauth,
}
impl PartialReason {
fn as_str(self) -> &'static str {
match self {
PartialReason::Timeout => "timeout",
PartialReason::Error => "error",
PartialReason::NoOauth => "no_oauth",
}
}
}
static FACET_REGISTRY: LazyLock<FacetRegistry> = LazyLock::new(build_facet_registry);
pub fn facet_registry() -> &'static FacetRegistry {
&FACET_REGISTRY
}
/// Hard-off in release builds so they can't enable the
/// conversations lane via env.
pub fn conversations_lane_enabled() -> bool {
if true {
return false;
}
std::env::var("KIGI_SESSION_LIST_CONVERSATIONS")
.ok()
.is_some_and(|v| {
!matches!(
v.trim().to_ascii_lowercase().as_str(),
"" | "0" | "false" | "off" | "no"
)
})
}
/// Env lane (desktop `KIGI_SESSION_LIST_CONVERSATIONS`) OR process-wide
/// `--chat` (`KIGI_CHAT_MODE`); hard-off in release builds.
/// The single predicate `MvpAgent::conversations_client()` keys on.
pub fn conversations_lane_active() -> bool {
conversations_lane_enabled() || crate::agent::chat_modes::process_chat_mode_enabled()
}
/// Parse `x.ai/session/list` params and, under process-wide chat mode, force
/// the conversations-only `kind` facet (see [`force_kind_chat`]).
pub fn parse_list_req(raw: &str) -> Result<ListReq, serde_json::Error> {
let mut req: ListReq = serde_json::from_str(raw)?;
if crate::agent::chat_modes::process_chat_mode_enabled() {
force_kind_chat(&mut req);
}
Ok(req)
serde_json::from_str(raw)
}
#[derive(Debug, Default, Deserialize)]
#[serde(rename_all = "camelCase")]
@@ -88,7 +43,6 @@ pub struct UnifiedListResult {
pub rows: Vec<UnifiedRow>,
pub next_cursor: Option<String>,
pub facets: FacetSummary,
pub conversations_partial: Option<PartialReason>,
}
#[derive(Debug, Default)]
struct ParsedMeta {
@@ -131,33 +85,8 @@ fn value_list(v: &serde_json::Value) -> Vec<serde_json::Value> {
other => vec![other.clone()],
}
}
/// Rewrite `req` so the `kind` facet filter is exactly `["chat"]`.
///
/// REPLACES any client-sent `kind` allow-list (a union with `"build"` would
/// re-enable the local lane); every other facet filter and `_meta` key is
/// left untouched.
pub fn force_kind_chat(req: &mut ListReq) {
let mut meta = match req.meta.take() {
Some(serde_json::Value::Object(map)) => map,
_ => serde_json::Map::new(),
};
let mut filters = match meta.remove("x.ai/facetFilters") {
Some(serde_json::Value::Object(map)) => map,
_ => serde_json::Map::new(),
};
filters.insert(
KIND_FACET_KEY.to_owned(),
serde_json::json!([SessionKind::Chat.as_str()]),
);
meta.insert(
"x.ai/facetFilters".to_owned(),
serde_json::Value::Object(filters),
);
req.meta = Some(serde_json::Value::Object(meta));
}
pub async fn build_unified_list(
registry_client: Option<&SessionRegistryClient>,
conversations_client: Option<&ConversationsClient>,
req: ListReq,
) -> UnifiedListResult {
let reg = facet_registry();
@@ -171,13 +100,11 @@ pub async fn build_unified_list(
let cursor = CompositeCursor::decode(req.cursor.as_deref());
let mut source_query = SourceQuery::default();
reg.apply_pushdown(&facet_filters, &mut source_query);
let exclude_conversations = excludes_conversations(&facet_filters);
let exclude_build = excludes_build(&facet_filters);
let over = (limit * 3).max(100);
let local_fut = async {
if exclude_build {
return Vec::new();
}
let local_rows = if exclude_build {
Vec::new()
} else {
crate::session::merge::fetch_merged(
registry_client,
req.cwd.as_deref(),
@@ -189,84 +116,17 @@ pub async fn build_unified_list(
.map(|m| merged_session_to_row(m, reg))
.collect::<Vec<UnifiedRow>>()
};
let conv_fut = async {
if exclude_conversations {
return ConvLane::Skipped;
}
let Some(client) = conversations_client else {
return ConvLane::Skipped;
};
let q = ConvQuery {
page_size: (limit + CONV_PAGE_HEADROOM) as i64,
page_token: cursor.conv_page_token.clone(),
search_query: query.clone(),
workspace_id: source_query.workspace_id.clone(),
};
match tokio::time::timeout(
crate::session::merge::REMOTE_TIMEOUT,
client.list_conversations(&q),
)
.await
{
Ok(Ok(page)) => {
let next_token = page.next_page_token;
let rows: Vec<UnifiedRow> = page
.conversations
.into_iter()
.map(|c| conversation_to_row(c, reg))
.collect();
let frontier = cursor::conv_frontier(&rows, next_token.is_some());
ConvLane::Page {
rows,
next_token,
frontier,
}
}
Ok(Err(ConvError::NoOauth)) => ConvLane::Degraded(PartialReason::NoOauth),
Ok(Err(e)) => {
tracing::warn!("conversation list failed: {e}");
ConvLane::Degraded(PartialReason::Error)
}
Err(_) => {
tracing::warn!("conversation list timed out");
ConvLane::Degraded(PartialReason::Timeout)
}
}
};
let (local_rows, conv_lane) = tokio::join!(local_fut, conv_fut);
{
let (conv_lane_status, conv_rows) = match &conv_lane {
ConvLane::Skipped => ("skipped", 0),
ConvLane::Degraded(reason) => (reason.as_str(), 0),
ConvLane::Page { rows, .. } => ("ok", rows.len()),
};
tracing::debug!(
local_lane_skipped = exclude_build,
local_rows = local_rows.len(),
conv_lane = conv_lane_status,
conv_rows,
"session list lanes"
);
}
tracing::debug!(
local_lane_skipped = exclude_build,
local_rows = local_rows.len(),
"session list"
);
let local_rows = reg.apply_in_memory_filters(&facet_filters, local_rows);
let conv_lane = match conv_lane {
ConvLane::Page {
rows,
next_token,
frontier,
} => ConvLane::Page {
rows: reg.apply_in_memory_filters(&facet_filters, rows),
next_token,
frontier,
},
other => other,
};
let Paginated {
candidates,
emit_count,
next_cursor,
partial,
} = merge_and_paginate(local_rows, conv_lane, &cursor, limit);
} = paginate(local_rows, &cursor, limit);
let mut rows = candidates;
rows.truncate(emit_count);
let facets = reg.summarize_window(&rows);
@@ -274,15 +134,6 @@ pub async fn build_unified_list(
rows,
next_cursor: next_cursor.map(|c| c.encode()),
facets,
conversations_partial: partial,
}
}
fn excludes_conversations(filters: &BTreeMap<String, Vec<serde_json::Value>>) -> bool {
match filters.get(KIND_FACET_KEY) {
Some(allowed) if !allowed.is_empty() => !allowed
.iter()
.any(|v| v.as_str() == Some(SessionKind::Chat.as_str())),
_ => false,
}
}
/// Mirror of [`excludes_conversations`]: `true` when a non-empty `kind`
@@ -307,21 +158,12 @@ pub struct ExtListResponse {
pub struct ExtListResponseMeta {
#[serde(rename = "x.ai/facets")]
pub facets: FacetSummary,
#[serde(rename = "x.ai/partial")]
pub partial: PartialInfo,
}
#[derive(Debug, Clone, Serialize)]
pub struct PartialInfo {
pub conversations: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub reason: Option<&'static str>,
}
pub fn ext_list_response(result: UnifiedListResult) -> ExtListResponse {
let UnifiedListResult {
rows,
next_cursor,
facets,
conversations_partial,
} = result;
ExtListResponse {
sessions: rows
@@ -329,13 +171,7 @@ pub fn ext_list_response(result: UnifiedListResult) -> ExtListResponse {
.map(UnifiedRow::into_ext_superset)
.collect(),
next_cursor,
meta: ExtListResponseMeta {
facets,
partial: PartialInfo {
conversations: conversations_partial.is_some(),
reason: conversations_partial.map(PartialReason::as_str),
},
},
meta: ExtListResponseMeta { facets },
}
}
#[cfg(test)]
@@ -494,74 +330,8 @@ mod tests {
);
filters
}
#[test]
fn excludes_build_mirrors_excludes_conversations() {
assert!(excludes_build(&kind_filter(&["chat"])));
assert!(!excludes_conversations(&kind_filter(&["chat"])));
assert!(!excludes_build(&kind_filter(&["build"])));
assert!(excludes_conversations(&kind_filter(&["build"])));
assert!(!excludes_build(&kind_filter(&["build", "chat"])));
assert!(!excludes_conversations(&kind_filter(&["build", "chat"])));
assert!(!excludes_build(&kind_filter(&[])));
assert!(!excludes_conversations(&kind_filter(&[])));
assert!(!excludes_build(&BTreeMap::new()));
assert!(!excludes_conversations(&BTreeMap::new()));
}
/// The forced `kind` REPLACES a client-sent `kind: ["build"]` (never
/// unions), so the local lane stays excluded.
#[test]
fn forced_kind_replaces_client_build_filter() {
let mut req = ListReq {
meta: Some(serde_json::json!({ "x.ai/facetFilters" : { "kind" : ["build"] }, })),
..ListReq::default()
};
force_kind_chat(&mut req);
let parsed = ParsedMeta::parse(req.meta.as_ref());
assert_eq!(
parsed.facet_filters.get(KIND_FACET_KEY),
Some(&vec![serde_json::json!("chat")]),
"forced kind must replace the client filter, not union with it"
);
assert!(excludes_build(&parsed.facet_filters));
assert!(!excludes_conversations(&parsed.facet_filters));
}
#[test]
fn forced_kind_preserves_other_facets() {
let mut req = ListReq {
meta: Some(serde_json::json!(
{ "x.ai/facetFilters" : { "kind" : ["build"], "starred" : [true],
"workspace" : ["w1"] }, "x.ai/query" : "antelope", "x.ai/limit" : 5,
}
)),
..ListReq::default()
};
force_kind_chat(&mut req);
let parsed = ParsedMeta::parse(req.meta.as_ref());
assert_eq!(
parsed.facet_filters.get(KIND_FACET_KEY),
Some(&vec![serde_json::json!("chat")])
);
assert_eq!(
parsed.facet_filters.get("starred"),
Some(&vec![serde_json::json!(true)])
);
assert_eq!(
parsed.facet_filters.get("workspace"),
Some(&vec![serde_json::json!("w1")])
);
assert_eq!(parsed.query.as_deref(), Some("antelope"));
assert_eq!(parsed.limit, Some(5));
}
#[test]
fn forced_kind_creates_facet_filters_when_meta_absent() {
let mut req = ListReq::default();
force_kind_chat(&mut req);
let parsed = ParsedMeta::parse(req.meta.as_ref());
assert_eq!(
parsed.facet_filters.get(KIND_FACET_KEY),
Some(&vec![serde_json::json!("chat")])
);
}
fn xai_auth_manager(dir: &std::path::Path) -> std::sync::Arc<crate::auth::AuthManager> {
let am = std::sync::Arc::new(crate::auth::AuthManager::new(
dir,
@@ -598,191 +368,4 @@ mod tests {
});
addr
}
/// A client-sent `kind: ["build"]` rewritten by [`force_kind_chat`]
/// yields conversations only.
#[tokio::test]
#[serial_test::serial]
async fn forced_kind_serves_conversations_only() {
let addr = spawn_conversations_stub(
serde_json::json!(
{ "conversations" : [{ "conversationId" : "c1", "title" : "Hello",
"modifyTime" : "2026-07-01T00:00:00Z" }, { "conversationId" : "c2",
"title" : "", "modifyTime" : "2026-07-02T00:00:00Z" },], }
)
.to_string(),
)
.await;
let _env = kigi_test_support::EnvGuard::set(
"KIGI_CONVERSATIONS_BASE_URL",
format!("http://{addr}"),
);
let home = tempfile::tempdir().expect("tempdir");
let client = ConversationsClient::new(xai_auth_manager(home.path()));
let mut req = ListReq {
meta: Some(serde_json::json!({ "x.ai/facetFilters" : { "kind" : ["build"] }, })),
..ListReq::default()
};
force_kind_chat(&mut req);
let result = build_unified_list(None, Some(&client), req).await;
let ids: Vec<&str> = result
.rows
.iter()
.map(|r| r.legacy.session_id.as_str())
.collect();
assert_eq!(ids, ["c2", "c1"], "conversations only, newest first");
assert!(
result
.rows
.iter()
.all(|r| r.legacy.source == "conversation"),
"no build row may survive the forced kind filter"
);
assert_eq!(result.conversations_partial, None);
}
/// A degraded conversations lane (no OAuth) surfaces through
/// `conversations_partial` instead of failing the list.
#[tokio::test]
#[serial_test::serial]
async fn degraded_conversations_lane_reports_no_oauth() {
let home = tempfile::tempdir().expect("tempdir");
let auth = std::sync::Arc::new(crate::auth::AuthManager::new(
home.path(),
crate::auth::KimiCodeConfig::default(),
));
let client = ConversationsClient::new(auth);
let mut req = ListReq::default();
force_kind_chat(&mut req);
let result = build_unified_list(None, Some(&client), req).await;
assert!(result.rows.is_empty());
assert_eq!(result.conversations_partial, Some(PartialReason::NoOauth));
}
/// Build-mode canary: with no conversations client the lane is skipped —
/// not degraded.
#[tokio::test]
async fn non_chat_list_without_client_skips_conversations_lane() {
let req = ListReq {
cwd: Some("/nonexistent/unified-list-canary".into()),
..ListReq::default()
};
let result = build_unified_list(None, None, req).await;
assert_eq!(
result.conversations_partial, None,
"no client ⇒ lane skipped, never reported as degraded"
);
assert!(result.rows.is_empty());
}
/// Desktop env lane stays env-gated; process chat mode is feature-gated.
#[test]
#[serial_test::serial]
fn conversations_lane_env_gating_matrix() {
{
let _off = kigi_test_support::EnvGuard::unset("KIGI_SESSION_LIST_CONVERSATIONS");
assert!(!conversations_lane_enabled());
}
{
let _on = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "1");
assert!(!conversations_lane_enabled());
}
{
let _off = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "0");
assert!(!conversations_lane_enabled());
}
}
/// Truth table for `conversations_lane_active`: desktop env lane OR
/// process chat mode, hard-off in release builds.
#[test]
#[serial_test::serial]
fn conversations_lane_active_truth_table() {
use crate::agent::chat_modes::KIGI_CHAT_MODE_ENV;
let _chat_off = kigi_test_support::EnvGuard::unset(KIGI_CHAT_MODE_ENV);
let _desktop_off = kigi_test_support::EnvGuard::unset("KIGI_SESSION_LIST_CONVERSATIONS");
assert!(
!conversations_lane_active(),
"no env ⇒ lane off (Build-mode default)"
);
{
let _desktop = kigi_test_support::EnvGuard::set("KIGI_SESSION_LIST_CONVERSATIONS", "1");
assert!(!conversations_lane_active());
}
{
let _chat = kigi_test_support::EnvGuard::set(KIGI_CHAT_MODE_ENV, "1");
assert!(
!conversations_lane_active(),
"process chat mode must enable the lane (chat feature only)"
);
}
}
/// `parse_list_req` forces the conversations-only `kind` exactly when
/// process chat mode is on; otherwise the client request is untouched.
#[test]
#[serial_test::serial]
fn parse_list_req_forces_kind_under_process_chat_mode_only() {
use crate::agent::chat_modes::KIGI_CHAT_MODE_ENV;
let raw = serde_json::json!(
{ "_meta" : { "x.ai/facetFilters" : { "kind" : ["build"], "starred" : [true]
} }, }
)
.to_string();
{
let _off = kigi_test_support::EnvGuard::unset(KIGI_CHAT_MODE_ENV);
let req = parse_list_req(&raw).expect("parse");
let parsed = ParsedMeta::parse(req.meta.as_ref());
assert_eq!(
parsed.facet_filters.get(KIND_FACET_KEY),
Some(&vec![serde_json::json!("build")]),
"non-chat: client kind filter untouched"
);
}
{
let _on = kigi_test_support::EnvGuard::set(KIGI_CHAT_MODE_ENV, "1");
let req = parse_list_req(&raw).expect("parse");
let parsed = ParsedMeta::parse(req.meta.as_ref());
let expected = if false { "chat" } else { "build" };
assert_eq!(
parsed.facet_filters.get(KIND_FACET_KEY),
Some(&vec![serde_json::json!(expected)])
);
assert_eq!(
parsed.facet_filters.get("starred"),
Some(&vec![serde_json::json!(true)]),
"other facets pass through"
);
}
}
/// Wire pin for the cross-crate `x.ai/partial` envelope the pager parses:
/// the serialized reason strings must not drift (the pager maps unknown
/// reasons to a generic retry notice, masking a rename).
#[test]
fn ext_list_response_serializes_partial_reasons() {
for (reason, wire) in [
(PartialReason::NoOauth, "no_oauth"),
(PartialReason::Timeout, "timeout"),
(PartialReason::Error, "error"),
] {
let value = serde_json::to_value(ext_list_response(UnifiedListResult {
rows: Vec::new(),
next_cursor: None,
facets: facet_registry().summarize_window(&[]),
conversations_partial: Some(reason),
}))
.expect("serialize");
assert_eq!(
value["_meta"]["x.ai/partial"],
serde_json::json!({ "conversations" :
true, "reason" : wire })
);
}
let healthy = serde_json::to_value(ext_list_response(UnifiedListResult {
rows: Vec::new(),
next_cursor: None,
facets: facet_registry().summarize_window(&[]),
conversations_partial: None,
}))
.expect("serialize");
assert_eq!(
healthy["_meta"]["x.ai/partial"],
serde_json::json!({ "conversations" :
false })
);
}
}
crates/codegen/kigi-shell/src/session/unified_list/row.rs
-82
View File
@@ -2,7 +2,6 @@ use serde::Serialize;
use super::envelope::{FacetMap, SessionKind, SessionMetaEnvelope};
use super::facets::{FacetRegistry, NormalizedItem};
use crate::remote::Conversation;
use crate::session::merge::MergedSession;
#[derive(Debug, Clone)]
@@ -73,44 +72,6 @@ pub fn merged_session_to_row(m: MergedSession, reg: &FacetRegistry) -> UnifiedRo
}
}
pub fn conversation_to_row(c: Conversation, reg: &FacetRegistry) -> UnifiedRow {
let facets = reg.extract_all(&NormalizedItem::from_conversation(&c));
let Conversation {
conversation_id,
title,
modify_time,
create_time,
..
} = c;
let legacy = MergedSession {
session_id: conversation_id,
summary: title.clone(),
first_prompt: None,
updated_at: modify_time.as_deref().unwrap_or_default().to_owned(),
created_at: create_time.unwrap_or_default(),
cwd: String::new(),
hostname: None,
source: "conversation".to_string(),
model_id: None,
num_messages: 0,
last_active_at: modify_time.clone(),
branch: None,
repo_name: None,
worktree_label: None,
git_root_dir: None,
git_remotes: Vec::new(),
source_workspace_dir: None,
session_kind: None,
};
UnifiedRow {
kind: SessionKind::Chat,
legacy,
title,
updated_at: modify_time,
facets,
}
}
fn effective_local_ts(m: &MergedSession) -> Option<String> {
m.last_active_at
.as_deref()
@@ -151,47 +112,4 @@ pub struct SessionInfo {
mod tests {
use super::*;
use crate::session::unified_list::facet_registry;
#[test]
fn conversation_row_uses_conversation_id_as_session_id() {
let c = Conversation {
conversation_id: "conv_abc123".into(),
title: "Compare GPU vendors".into(),
modify_time: Some("2026-06-18T18:02:00Z".into()),
create_time: Some("2026-06-18T17:30:00Z".into()),
..Conversation::default()
};
let row = conversation_to_row(c, facet_registry());
assert_eq!(row.legacy.session_id, "conv_abc123");
assert_eq!(row.kind, SessionKind::Chat);
assert_eq!(row.legacy.source, "conversation");
assert_eq!(row.legacy.cwd, "");
let ext = serde_json::to_value(row.clone().into_ext_superset()).unwrap();
assert_eq!(ext["sessionId"], "conv_abc123");
assert_eq!(ext["cwd"], "");
assert_eq!(ext["source"], "conversation");
assert_eq!(ext["_meta"]["x.ai/session"]["kind"], "chat");
// Chat rows have no local git enrichment (fields omitted).
assert!(ext.get("gitRootDir").is_none());
assert!(ext.get("gitRemotes").is_none());
assert!(ext.get("sourceWorkspaceDir").is_none());
assert!(ext.get("sessionKind").is_none());
let bare = serde_json::to_value(row.into_session_info()).unwrap();
assert_eq!(bare["sessionId"], "conv_abc123");
}
#[test]
fn conversation_missing_modify_time_still_resumable() {
let c = Conversation {
conversation_id: "conv_no_time".into(),
title: "Untitled".into(),
..Conversation::default()
};
let row = conversation_to_row(c, facet_registry());
assert_eq!(row.legacy.session_id, "conv_no_time");
assert!(row.updated_at.is_none());
assert_eq!(row.legacy.updated_at, "");
}
}
crates/codegen/kigi-shell/src/session/worktree.rs
+1 -1
View File
@@ -387,7 +387,7 @@ async fn resume_local_session_in_worktree(
source_workspace_dir: Some(resolved_source_cwd.to_owned()),
..Default::default()
};
let fork_resp = match fork_session(fork_req, agent_id, auth_manager).await {
let fork_resp = match fork_session(fork_req).await {
Ok(r) => r,
Err(e) => {
cleanup_worktree_on_failure(resolved_source_cwd, &wt_resp.worktree_path).await;
crates/codegen/kigi-shell/src/test_support/lsp_runtime.rs
-4
View File
@@ -47,15 +47,11 @@ pub(crate) fn ctx_with_toggle(toggle: HashMap<String, bool>) -> SubagentSpawnCon
auth_scheme: Default::default(),
extra_headers: Default::default(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
crates/codegen/kigi-shell/src/tools/config.rs
+1 -5
View File
@@ -205,7 +205,7 @@ impl ShellToolsetConfig {
pub fn new(base: Option<Self>, sampling_config: Option<SamplerConfig>) -> Self {
let default_base = SamplerConfig {
api_key: None,
base_url: "https://api.x.ai/v1".to_string(),
base_url: kigi_env::coding_api_base_url(),
model: String::new(),
max_completion_tokens: None,
temperature: None,
@@ -214,15 +214,11 @@ impl ShellToolsetConfig {
auth_scheme: Default::default(),
extra_headers: indexmap::IndexMap::new(),
context_window: 256_000,
client_version: None,
reasoning_effort: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
deployment_id: None,
user_id: None,
origin_client: None,
// Default base for the in-process web-search tool config.
// Real `SamplerConfig`s (e.g. from `sampling_config_for_model`)
crates/codegen/kigi-shell/src/util/config/resolve/auto_mode.rs
+1 -2
View File
@@ -67,8 +67,7 @@ fn resolve_auto_permission_mode_layers(
}
/// Resolve whether the **auto** permission mode feature (`PermissionMode::Auto`,
/// the LLM/heuristic classifier) is enabled. Full chain mirroring
/// [`resolve_zdr_access_enabled`](super::resolve_zdr_access_enabled):
/// the LLM/heuristic classifier) is enabled. Full precedence chain:
///
/// requirements > env (`KIGI_AUTO_PERMISSION_MODE`) > `[auto_mode] enabled` in
/// `config.toml` > managed > remote settings (`auto_mode.enabled`, coerced
crates/codegen/kigi-shell/src/util/config/resolve/features.rs
-22
View File
@@ -1,28 +1,6 @@
use crate::util::config::RemoteSettings;
use toml::Value as TomlValue;
/// Resolve whether ZDR users are allowed to use the product.
///
/// Precedence: requirements > env > config.toml > managed > remote settings > default (false).
pub fn resolve_zdr_access_enabled(
requirements: Option<&TomlValue>,
user: Option<&TomlValue>,
managed: Option<&TomlValue>,
remote: Option<&RemoteSettings>,
) -> bool {
use crate::agent::config::BoolFlag;
fn from_toml(v: Option<&TomlValue>) -> Option<bool> {
v?.get("features")?.get("zdr_access_enabled")?.as_bool()
}
BoolFlag::env("KIGI_ZDR_ACCESS_ENABLED")
.requirement(from_toml(requirements))
.config(from_toml(user))
.managed(from_toml(managed))
.feature_flag(remote.and_then(|r| r.zdr_access_enabled))
.resolve()
.value
}
/// Whether model-catalog (`/v1/models`) and remote-settings (`/v1/settings`)
/// fetches from xAI backends are allowed, including the deployment-config sync
/// bundled into the startup prefetch (the background managed-config sync has
crates/codegen/kigi-shell/tests/common/mod.rs
-4
View File
@@ -43,15 +43,11 @@ pub fn test_sampler_config(
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
context_window: 256_000,
client_version: None,
force_http1: false,
max_retries: None,
stream_tool_calls: false,
idle_timeout_secs: None,
client_identifier: None,
reasoning_effort: None,
deployment_id: None,
user_id: None,
origin_client: None,
attribution_callback: None,
bearer_resolver: None,
crates/codegen/kigi-shell/tests/git_contention_e2e.rs
+1 -1
View File
@@ -431,7 +431,7 @@ fn git_rebase_refresh_storm_e2e() {
// serve HTTP and never read the process environment.
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
crates/codegen/kigi-shell/tests/session_load_perf.rs
+1 -1
View File
@@ -625,7 +625,7 @@ async fn full_session_load_e2e() {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_INSTRUMENTATION", "log");
std::env::set_var("KIGI_INSTRUMENTATION_LOG", &instr_log);
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
crates/codegen/kigi-shell/tests/signed_managed_config/common.rs
+1 -1
View File
@@ -147,7 +147,7 @@ pub fn signed_dk_body(
requirements: Option<&str>,
) -> String {
let payload = SignedPayload {
version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION,
version: kigi_config::signed_policy::SIGNED_PAYLOAD_VERSION,
deployment_id: Some(deployment_id.to_owned()),
team_id: None,
managed_config: managed.map(str::to_owned),
crates/codegen/kigi-shell/tests/signed_managed_config_extended.rs
+1 -1
View File
@@ -38,7 +38,7 @@ async fn sync_fail_closed_policy(home: &std::path::Path, kp: &ring::signature::E
/// provisioned key with no config row.
fn signed_dk_empty_body(kp: &ring::signature::Ed25519KeyPair, deployment_id: &str) -> String {
let payload = SignedPayload {
version: prod_mc_cli_chat_proxy_types::SIGNED_PAYLOAD_VERSION,
version: kigi_config::signed_policy::SIGNED_PAYLOAD_VERSION,
deployment_id: Some(deployment_id.to_owned()),
team_id: None,
managed_config: None,
crates/codegen/kigi-shell/tests/test_built_binary_e2e.rs
+1 -1
View File
@@ -1267,7 +1267,7 @@ impl ConfigTestHarness {
home,
workdir: git_workdir(),
env: vec![
("KIGI_CLI_CHAT_PROXY_BASE_URL".into(), server.url()),
("KIGI_CODE_BASE_URL".into(), server.url()),
("KIGI_TELEMETRY_ENABLED".into(), "false".into()),
("KIGI_FEEDBACK_ENABLED".into(), "false".into()),
("KIGI_TRACE_UPLOAD".into(), "false".into()),
crates/codegen/kigi-shell/tests/test_leader_soak.rs
+1 -1
View File
@@ -136,7 +136,7 @@ async fn leader_soak_churning_clients_no_leaks_no_zombies() {
// code reads these process-globals (same pattern as session_load_perf).
unsafe {
std::env::set_var("KIGI_SHARE_DIR", kigi_home.path());
std::env::set_var("KIGI_CLI_CHAT_PROXY_BASE_URL", server.url());
std::env::set_var("KIGI_CODE_BASE_URL", server.url());
std::env::set_var("KIGI_XAI_API_BASE_URL", server.url());
std::env::set_var("XAI_API_KEY", "test-key-for-ci");
std::env::set_var("KIGI_TELEMETRY_ENABLED", "false");
crates/codegen/kigi-shell/tests/test_settings_refresh.rs
-258
View File
@@ -1,258 +0,0 @@
//! Integration test: MockInferenceServer `/v1/settings` endpoint and
//! remote settings settings refresh infrastructure.
//!
//! Tests the mock endpoint directly (no binary needed) and verifies
//! the `fetch_settings_blocking` client round-trips correctly with
//! runtime-mutated mock settings.
//!
//! Run locally:
//! ```bash
//! cargo test -p kigi-shell --test test_settings_refresh
//! ```
use std::future::Future;
use kigi_shell::util::config::RemoteSettings;
use kigi_test_support::*;
async fn with_local_set<F, Fut>(f: F)
where
F: FnOnce() -> Fut,
Fut: Future<Output = ()>,
{
tokio::task::LocalSet::new().run_until(f()).await;
}
/// Verify the mock `/v1/settings` endpoint returns 404 when no settings
/// are configured (the default). This preserves backward compatibility:
/// existing tests that never call `set_settings` see a 404, and
/// `fetch_settings_blocking` returns `None`.
#[tokio::test]
async fn test_settings_endpoint_returns_404_when_unconfigured() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
let resp = reqwest::get(format!("{}/settings", server.url()))
.await
.expect("request failed");
assert_eq!(resp.status(), 404);
})
.await;
}
/// Verify the mock `/v1/settings` endpoint returns configured settings
/// and that `set_settings` runtime mutation is reflected immediately.
#[tokio::test]
async fn test_settings_endpoint_returns_configured_settings() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Configure initial settings
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v1".into()]),
leader_mode: Some(false),
..Default::default()
});
// Fetch and verify
let resp = reqwest::get(format!("{}/settings", server.url()))
.await
.expect("request failed");
assert_eq!(resp.status(), 200);
let settings: RemoteSettings = resp.json().await.expect("parse failed");
assert_eq!(settings.tips, Some(vec!["tip_v1".into()]));
assert_eq!(settings.leader_mode, Some(false));
})
.await;
}
/// Verify that `set_settings` updates are visible to subsequent requests
/// (runtime mutation for multi-session test scenarios).
#[tokio::test]
async fn test_settings_endpoint_reflects_runtime_mutations() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Initial settings
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v1".into()]),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.tips, Some(vec!["tip_v1".into()]));
// Mutate settings (simulating a remote feature flag change)
server.set_settings(RemoteSettings {
tips: Some(vec!["tip_v2".into()]),
leader_mode: Some(true),
..Default::default()
});
// Subsequent request sees the updated values
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.tips, Some(vec!["tip_v2".into()]));
assert_eq!(settings.leader_mode, Some(true));
})
.await;
}
/// Verify `fetch_settings_blocking` round-trips through the mock server.
/// This is the actual client function used by `refresh_remote_settings`.
#[tokio::test]
async fn test_fetch_settings_blocking_round_trip() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Without settings configured: returns None (404 from mock)
let auth = kigi_shell::auth::KimiAuth {
key: "test-key".into(),
..Default::default()
};
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None)
})
.await
.unwrap();
assert!(
result.is_none(),
"Expected None when settings not configured"
);
// With settings configured: returns Some(settings)
server.set_settings(RemoteSettings {
tips: Some(vec!["fetched_tip".into()]),
..Default::default()
});
let result = tokio::task::spawn_blocking({
let url = server.url().to_string();
let auth = auth.clone();
move || kigi_shell::remote::fetch_settings_blocking(&url, &auth, None)
})
.await
.unwrap();
let settings = result.expect("Expected Some when settings are configured");
assert_eq!(settings.tips, Some(vec!["fetched_tip".into()]));
})
.await;
}
/// Verify the `doom_loop_recovery` settings object survives the
/// `/v1/settings` round-trip, that its absence deserializes to `None` (old
/// servers), and that a partial object keeps its unset fields `None`.
#[tokio::test]
async fn test_doom_loop_recovery_settings_round_trip() {
use kigi_shell::util::config::DoomLoopRecoverySettings;
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
// Absent from the payload ⇒ None on the client.
server.set_settings(RemoteSettings::default());
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
assert_eq!(settings.doom_loop_recovery, None);
server.set_settings(RemoteSettings {
doom_loop_recovery: Some(DoomLoopRecoverySettings {
enabled: Some(true),
max_threshold: Some(16),
max_retries: Some(1),
}),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
let recovery = settings.doom_loop_recovery.expect("object round-trips");
assert_eq!(recovery.enabled, Some(true));
assert_eq!(recovery.max_threshold, Some(16));
assert_eq!(recovery.max_retries, Some(1));
// Partial object: only the set field comes through; the rest stay
// None so the resolver falls through per-field.
server.set_settings(RemoteSettings {
doom_loop_recovery: Some(DoomLoopRecoverySettings {
max_threshold: Some(32),
..Default::default()
}),
..Default::default()
});
let settings: RemoteSettings = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap()
.json()
.await
.unwrap();
let recovery = settings.doom_loop_recovery.expect("object round-trips");
assert_eq!(recovery.enabled, None);
assert_eq!(recovery.max_threshold, Some(32));
assert_eq!(recovery.max_retries, None);
})
.await;
}
/// Verify that the mock server's request log correctly tracks
/// GET /v1/settings requests for assertion in multi-session tests.
#[tokio::test]
async fn test_settings_requests_appear_in_request_log() {
with_local_set(|| async {
let server = MockInferenceServer::start()
.await
.expect("start mock server");
server.set_settings(RemoteSettings::default());
assert_eq!(server.request_count(), 0);
// First request
let _ = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap();
let settings_reqs: Vec<_> = server
.requests()
.into_iter()
.filter(|r| r.method == "GET" && r.path.contains("/settings"))
.collect();
assert_eq!(settings_reqs.len(), 1, "Expected 1 settings request");
// Second request (simulating /new refresh)
let _ = reqwest::get(format!("{}/settings", server.url()))
.await
.unwrap();
let settings_reqs: Vec<_> = server
.requests()
.into_iter()
.filter(|r| r.method == "GET" && r.path.contains("/settings"))
.collect();
assert_eq!(settings_reqs.len(), 2, "Expected 2 settings requests");
})
.await;
}