M0: compilable skeleton — Kigi 0.1.0 fork surgery

Hard fork of xai-org/grok-build (Apache-2.0) re-targeted as Kigi, an
unofficial Kimi Code CLI community build.

Rename & identity
- 72 xai-*/xai-grok-* crates -> kigi-* (explicit: xai-grok-pager-bin ->
  kigi-bin [binary `kigi`], xai-grok-pager -> kigi-tui; rest mechanical);
  ptyctl, ptyctl-cli, third_party/ unchanged; proto package
  xai.grok.tools.v1 -> kigi.tools.v1
- Config home ~/.kigi (KIGI_SHARE_DIR override), env prefix GROK_* ->
  KIGI_*, `kigi --version` carries the unofficial-community-build notice
- clap identity, help text, startup banner, prompt templates rebranded
  (templates re-encrypted)

Deletions (PRD removal list #5/#6/#7/#9/#10)
- voice input (xai-grok-voice) and all TUI wiring
- telemetry: Mixpanel client, external OTel stream, Sentry, OTLP layers,
  trace/GCS/S3 upload queues (kigi-file-utils halved), workspace upload
  module & dc_log, heap-profile uploader, auth-diagnostics uploader,
  session-analytics halves of feedback; local zero-egress observability
  preserved in new kigi-log crate (unified log, --debug firehose,
  subsystem file logs, opt-in instrumentation)
- announcements (crate, remote-settings fields, TUI surfaces)
- plugin marketplace (crate, sources/browse/CTA/extensions-modal tab);
  direct plugin install/uninstall/update via kigi-agent git_install kept
- relay/gateway/assets endpoints and features (agent relay, headless
  relay transport, gateway bridge, LeaderEnvUrls); leader IPC socket now
  ~/.kigi/leader.sock + KIGI_LEADER_SOCKET, no ws-url derivation
- functional types rehomed instead of deleted: PermissionMode ->
  kigi-config-types, McpInitStrategy -> kigi-mcp, PrCreationSource ->
  session signals, TerminalDiagnostics -> kigi-pager-render, agent_id ->
  shell util

Endpoints
- kigi-env rewritten: single production KigiEndpoints {coding_api_base_url
  https://api.kimi.com/coding/v1 (KIGI_CODE_BASE_URL), oauth_host
  https://auth.kimi.com (KIGI_OAUTH_HOST), update_base_url (GitHub
  Releases API), upgrade_page_url}; GrokBuildEnvironment enum deleted

Toolchain & workspace hygiene
- Rust 1.97.0 pinned; edition 2024; full cargo update; git2 hoisted to
  workspace at 0.21 (Option->Result API migration), quick-xml 0.41
- Root Cargo.toml hand-maintained (PRD §8.1): version 0.1.0 inherited by
  all members, members sorted, unused deps pruned
- cargo-deny advisories gate (deny.toml with documented transitive
  exceptions); CI workflow (check/clippy/fmt/deny/test, macOS+Linux)
- cross-crate test seams re-gated behind `test-support` cargo feature;
  insta snapshot baselines renamed to the kigi_tui prefix
- clippy --workspace --all-targets: zero warnings; fmt clean

Fixes surfaced by the port
- updater probe/installer divergence (bin/kigi vs bin/grok symlink set)
- idle model-metadata refresh dead under KIGI_CODE_BASE_URL override
  (new is_effective_coding_endpoint_url, loopback+override aware)
- macOS symlinked-TMPDIR fixture canonicalization (foreign_sessions,
  fast-worktree); RSS measurement tests serialized via serial_test

Docs & legal (Apache §4)
- NOTICE added (upstream attribution + change statement); THIRD-PARTY
  notices sustained; kigi-tools ported-code notices extended; README,
  CONTRIBUTING, SECURITY, AGENTS.md rewritten

Out of scope for M0 (tracked): Kimi auth/inference (M1), search/fetch,
command parity, config import (M2), Computer Hub excision & final
brand-token sweep (M2), distribution & self-update rewrite (M3).
This commit is contained in:
2026-07-17 05:31:01 -04:00
commit d6c20fc13f
2612 changed files with 1353757 additions and 0 deletions
@@ -0,0 +1,411 @@
//! Send-safe view of the agent's in-flight work, shared with the leader's
//! auto-update checker and `RelaunchForUpdate` drain (`tokio::spawn` tasks
//! that cannot read the `!Send` `MvpAgent` state on the `LocalSet`).
//!
//! The leader's `agent_busy` flag only counts IPC (Unix-socket) requests;
//! relay (grok.com WebSocket) traffic is bridged straight into the agent's
//! ACP stdin and never sets it, so a relay-driven leader (devbox / remote)
//! always looked idle and got restarted mid-turn on every update —
//! surfacing as "Subagent result channel dropped".
//!
//! [`AgentActivity::is_busy`] derives busyness from agent state regardless
//! of transport, and [`AgentActivity::flush_all_sessions`] lets the shutdown
//! path end session actors gracefully instead of aborting them via
//! `LocalSet` drop.
//!
//! ## Lifecycle: entries expire with their actor, not with agent bookkeeping
//!
//! The agent only ever **registers** sessions (at handle creation). There is
//! deliberately no unregister: an entry is live exactly while its actor
//! holds the command receiver (`!cmd_tx.is_closed()`), and closed entries
//! are purged opportunistically. This sidesteps a whole class of races
//! between `MvpAgent`'s map bookkeeping and actor lifetime — an actor
//! removed from the agent's map but still winding down stays visible to
//! `is_busy`/`flush_all_sessions` until it actually exits, and a session id
//! rebuilt with a fresh actor is just a second (distinct) entry.
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use crate::session::pending_interaction::PendingInteractions;
use crate::session::{SessionCommand, SessionHandle};
/// How often [`AgentActivity::flush_all_sessions`] re-polls actors that have
/// not yet exited.
const FLUSH_POLL: Duration = Duration::from_millis(50);
/// Per-session slice of state shared with the session actor (the same `Arc`s
/// the actor mutates — see the matching `SessionHandle` fields).
struct SessionActivityEntry {
id: String,
cmd_tx: tokio::sync::mpsc::UnboundedSender<SessionCommand>,
/// `Some` while a turn is running (relay- or IPC-driven alike).
current_prompt_id: Arc<Mutex<Option<String>>>,
/// Non-empty while a blocking reverse-request (permission / question /
/// plan approval) is parked.
pending_interactions: PendingInteractions,
}
impl SessionActivityEntry {
/// The actor still holds the command receiver.
fn is_live(&self) -> bool {
!self.cmd_tx.is_closed()
}
/// A running turn or a parked blocking interaction.
fn is_busy(&self) -> bool {
self.current_prompt_id
.lock()
.unwrap_or_else(|p| p.into_inner())
.is_some()
|| !self
.pending_interactions
.lock()
.unwrap_or_else(|p| p.into_inner())
.is_empty()
}
}
#[derive(Default)]
struct ActivityInner {
/// Self-expiring: entries are dead once the actor drops its receiver
/// (see module docs), and are purged whenever the list is locked.
sessions: Mutex<Vec<SessionActivityEntry>>,
/// Subagents currently initializing or running; kept in sync by
/// `SubagentCoordinator::sync_running_gauge`.
subagents: Arc<AtomicUsize>,
}
/// Cheap-to-clone, `Send + Sync` handle. See module docs.
#[derive(Clone, Default)]
pub struct AgentActivity {
inner: Arc<ActivityInner>,
}
impl AgentActivity {
/// Register a session's shared state at handle-creation time. No
/// unregister exists — the entry expires when the actor exits.
pub(crate) fn register_session(&self, id: &str, handle: &SessionHandle) {
self.lock_live_sessions().push(SessionActivityEntry {
id: id.to_string(),
cmd_tx: handle.cmd_tx.clone(),
current_prompt_id: handle.current_prompt_id.clone(),
pending_interactions: handle.pending_interactions.clone(),
});
}
/// Shared gauge of initializing + running subagents; handed to the
/// `SubagentCoordinator`, which recomputes it on every state change.
pub(crate) fn subagent_gauge(&self) -> Arc<AtomicUsize> {
self.inner.subagents.clone()
}
/// Whether the agent has live work: a running turn, a parked blocking
/// interaction, or an initializing/running subagent.
///
/// Known sub-tick window: queued-but-not-started prompts
/// (`pending_inputs` in the actor) are not mirrored here, so a prompt
/// submitted exactly at a turn boundary can read as idle (the same
/// window `session_has_live_work` closes with an actor round-trip,
/// which a sync `Send` probe cannot do). The flush's quiesce loop
/// re-snapshots and still ends such an actor via its Shutdown arm.
pub fn is_busy(&self) -> bool {
self.inner.subagents.load(Ordering::Relaxed) > 0
|| self.lock_live_sessions().iter().any(|e| e.is_busy())
}
/// Number of live registered sessions (diagnostics/tests).
pub fn session_count(&self) -> usize {
self.lock_live_sessions().len()
}
/// Send [`SessionCommand::Shutdown`] to every live session actor
/// (replay-buffer flush → hooks → memory save → actor returns) and wait
/// up to `grace` for the actors to exit, observed via
/// `cmd_tx.is_closed()`.
///
/// This is a quiesce loop, not a one-shot broadcast: each poll
/// re-snapshots the registry and signals actors that appeared after the
/// flush started (deduped by channel identity, so a session id rebuilt
/// with a fresh actor gets its own signal), all against one deadline —
/// `grace` bounds the **total** shutdown delay.
///
/// Call **before** cancelling the leader's root token so session state
/// is durable before the `LocalSet` drop aborts remaining tasks. Actors
/// that miss the grace are logged and abandoned.
pub async fn flush_all_sessions(&self, grace: Duration) {
let deadline = tokio::time::Instant::now() + grace;
// Every distinct channel signaled so far (id kept for logging).
let mut signaled: Vec<(String, tokio::sync::mpsc::UnboundedSender<SessionCommand>)> =
Vec::new();
loop {
let snapshot: Vec<_> = self
.lock_live_sessions()
.iter()
.map(|e| (e.id.clone(), e.cmd_tx.clone()))
.collect();
for (id, tx) in snapshot {
if !signaled.iter().any(|(_, s)| s.same_channel(&tx)) {
tracing::info!(session_id = %id, "leader shutdown: flushing session");
let _ = tx.send(SessionCommand::Shutdown);
signaled.push((id, tx));
}
}
if signaled.iter().all(|(_, tx)| tx.is_closed()) {
return; // nothing to flush, or all actors exited
}
if tokio::time::Instant::now() >= deadline {
for (id, tx) in &signaled {
if !tx.is_closed() {
tracing::warn!(
session_id = %id,
"leader shutdown: session actor did not exit within grace; proceeding"
);
}
}
return;
}
tokio::time::sleep(FLUSH_POLL).await;
}
}
/// Lock the session list, dropping entries whose actor has exited.
///
/// Purging happens only here, so in modes with no periodic reader (no
/// auto-update checker) a dead entry lingers until the next register —
/// bounded and tiny (a sender handle + two `Arc`s per entry).
fn lock_live_sessions(&self) -> std::sync::MutexGuard<'_, Vec<SessionActivityEntry>> {
let mut guard = self
.inner
.sessions
.lock()
.unwrap_or_else(|p| p.into_inner());
guard.retain(SessionActivityEntry::is_live);
guard
}
/// Register a synthetic session from raw parts (no full `SessionHandle`).
/// Returns the command receiver (the "actor" side) plus the shared
/// running-turn and pending-interaction slots.
#[cfg(test)]
pub(crate) fn register_for_test(
&self,
id: &str,
) -> (
tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
Arc<Mutex<Option<String>>>,
PendingInteractions,
) {
let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
let current_prompt_id = Arc::new(Mutex::new(None));
let pending_interactions: PendingInteractions =
Arc::new(Mutex::new(std::collections::HashMap::new()));
self.lock_live_sessions().push(SessionActivityEntry {
id: id.to_string(),
cmd_tx,
current_prompt_id: current_prompt_id.clone(),
pending_interactions: pending_interactions.clone(),
});
(cmd_rx, current_prompt_id, pending_interactions)
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Build a registered entry from raw parts without a full SessionHandle.
fn register_raw(
activity: &AgentActivity,
id: &str,
) -> (
tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
Arc<Mutex<Option<String>>>,
PendingInteractions,
) {
activity.register_for_test(id)
}
/// Simulated session actor: exits (dropping its receiver) `delay` after
/// receiving `Shutdown`; resolves to whether Shutdown was received.
fn spawn_actor(
mut rx: tokio::sync::mpsc::UnboundedReceiver<SessionCommand>,
delay: Duration,
) -> tokio::task::JoinHandle<bool> {
tokio::spawn(async move {
while let Some(cmd) = rx.recv().await {
if matches!(cmd, SessionCommand::Shutdown) {
tokio::time::sleep(delay).await;
return true;
}
}
false
})
}
#[test]
fn idle_by_default() {
let activity = AgentActivity::default();
assert!(!activity.is_busy());
}
#[tokio::test]
async fn running_turn_marks_busy() {
let activity = AgentActivity::default();
let (_rx, prompt_id, _pending) = register_raw(&activity, "s1");
assert!(!activity.is_busy());
*prompt_id.lock().unwrap() = Some("prompt-1".to_string());
assert!(activity.is_busy());
*prompt_id.lock().unwrap() = None;
assert!(!activity.is_busy());
}
#[tokio::test]
async fn pending_interaction_marks_busy() {
let activity = AgentActivity::default();
let (_rx, _prompt_id, pending) = register_raw(&activity, "s1");
pending.lock().unwrap().insert(
"tc-1".to_string(),
crate::session::pending_interaction::PendingKind::Permission,
);
assert!(activity.is_busy());
pending.lock().unwrap().clear();
assert!(!activity.is_busy());
}
#[test]
fn subagent_gauge_marks_busy() {
let activity = AgentActivity::default();
let gauge = activity.subagent_gauge();
assert!(!activity.is_busy());
gauge.store(1, Ordering::Relaxed);
assert!(activity.is_busy());
gauge.store(0, Ordering::Relaxed);
assert!(!activity.is_busy());
}
/// An actor that is still running counts as busy even if the agent has
/// dropped its handle — liveness comes from the channel, not from agent
/// bookkeeping. Once the actor exits, the entry expires.
#[tokio::test]
async fn live_actor_counts_busy_until_it_exits() {
let activity = AgentActivity::default();
let (rx, prompt_id, _pending) = register_raw(&activity, "s1");
*prompt_id.lock().unwrap() = Some("prompt-1".to_string());
assert!(activity.is_busy());
assert_eq!(activity.session_count(), 1);
// Actor exits (receiver dropped) → entry expires, even though the
// shared prompt slot still says Some.
drop(rx);
assert!(!activity.is_busy());
assert_eq!(activity.session_count(), 0);
}
#[tokio::test(start_paused = true)]
async fn flush_sends_shutdown_and_waits_for_actor_exit() {
let activity = AgentActivity::default();
let (rx, _prompt_id, _pending) = register_raw(&activity, "s1");
// Simulated actor: exits (drops rx) when it receives Shutdown.
let actor = spawn_actor(rx, Duration::ZERO);
activity.flush_all_sessions(Duration::from_secs(5)).await;
assert!(actor.await.unwrap(), "actor should have received Shutdown");
}
#[tokio::test(start_paused = true)]
async fn flush_grace_bounds_total_delay_across_sessions() {
let activity = AgentActivity::default();
// One wedged actor (receiver kept open) and one healthy actor.
let (_wedged_rx, _p1, _i1) = register_raw(&activity, "wedged");
let (rx, _p2, _i2) = register_raw(&activity, "healthy");
let actor = spawn_actor(rx, Duration::ZERO);
// The wedged actor must not consume the healthy actor's budget, and
// the total wait must be ~one grace period, not one per session.
let start = tokio::time::Instant::now();
activity.flush_all_sessions(Duration::from_secs(2)).await;
let elapsed = start.elapsed();
assert!(elapsed >= Duration::from_secs(2));
assert!(
elapsed < Duration::from_secs(3),
"grace must be shared, not serial: {elapsed:?}"
);
assert!(actor.await.unwrap(), "healthy actor should get Shutdown");
}
/// A session id rebuilt with a fresh actor while the old actor is still
/// winding down: both channels must be signaled and awaited.
#[tokio::test(start_paused = true)]
async fn flush_awaits_both_channels_when_id_is_reused() {
let activity = AgentActivity::default();
let (old_rx, _p1, _i1) = register_raw(&activity, "s1");
let (new_rx, _p2, _i2) = register_raw(&activity, "s1");
let old_actor = spawn_actor(old_rx, Duration::from_millis(500));
let new_actor = spawn_actor(new_rx, Duration::ZERO);
activity.flush_all_sessions(Duration::from_secs(5)).await;
assert!(old_actor.is_finished(), "flush must wait for the old actor");
assert!(old_actor.await.unwrap());
assert!(new_actor.await.unwrap());
}
#[tokio::test(start_paused = true)]
async fn flush_signals_sessions_that_appear_mid_flush() {
let activity = AgentActivity::default();
// Actor 1: holds the flush open for a few polls, then exits.
let (rx1, _p1, _i1) = register_raw(&activity, "s1");
let actor1 = spawn_actor(rx1, Duration::from_millis(300));
// Actor 2 registers AFTER the flush has started (a relay-driven
// prompt racing the shutdown) — it must still receive Shutdown.
let activity_late = activity.clone();
let late = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(100)).await;
let (mut rx2, _p2, _i2) = activity_late.register_for_test("s2");
while let Some(cmd) = rx2.recv().await {
if matches!(cmd, SessionCommand::Shutdown) {
return true;
}
}
false
});
activity.flush_all_sessions(Duration::from_secs(5)).await;
assert!(actor1.await.unwrap());
assert!(
late.await.unwrap(),
"session registered mid-flush must receive Shutdown"
);
}
#[tokio::test(start_paused = true)]
async fn flush_gives_up_after_grace_when_actor_hangs() {
let activity = AgentActivity::default();
// Keep rx alive so the channel never closes (wedged actor).
let (_rx, _prompt_id, _pending) = register_raw(&activity, "s1");
let start = tokio::time::Instant::now();
activity.flush_all_sessions(Duration::from_secs(2)).await;
assert!(
start.elapsed() >= Duration::from_secs(2),
"flush should wait out the grace period"
);
// Returned rather than hanging forever — that's the assertion.
}
#[tokio::test]
async fn flush_with_no_sessions_is_noop() {
let activity = AgentActivity::default();
activity.flush_all_sessions(Duration::from_secs(1)).await;
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,334 @@
//! 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.
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.
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");
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,639 @@
//! Resilient parsing for `[model.<id>]` TOML overrides.
//!
//! A model entry must survive a bad field: warn and skip the field, never
//! drop the model (managed configs must not lose catalog entries).
//!
//! Every table is deserialized through `serde_ignored`, so unknown fields
//! warn on every path and [`ConfigModelOverride`] stays the single source of
//! truth for the field set. When the whole-table parse fails, fields that
//! fail to parse on their own are pruned (one warning each) and the table is
//! parsed again. Non-table values are dropped with a warning.
//!
//! Warnings are retained on `Config::model_override_warnings` and surfaced by
//! `grok inspect`.
use indexmap::IndexMap;
use serde::Serialize;
use super::config::ConfigModelOverride;
/// Category for a [`ModelOverrideWarning`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize)]
#[serde(rename_all = "kebab-case")]
pub enum ModelOverrideWarningKind {
/// Field name not recognized; field ignored.
UnknownField,
/// Value failed to parse; field skipped.
InvalidValue,
/// Legacy alias given alongside its canonical key; alias skipped.
DuplicateAlias,
/// Entry value is not a TOML table; entry dropped.
NotATable,
/// Entry failed to parse even after skipping invalid fields; the model
/// keeps an empty override.
UnparseableEntry,
}
/// One skipped field or dropped entry from `[model.*]` parsing.
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct ModelOverrideWarning {
/// `None` when the warning is about the `[model]` section itself.
#[serde(skip_serializing_if = "Option::is_none")]
pub model_key: Option<String>,
/// `None` for warnings about the entry as a whole.
#[serde(skip_serializing_if = "Option::is_none")]
pub field: Option<String>,
pub kind: ModelOverrideWarningKind,
pub reason: String,
}
/// Result of [`parse_model_overrides`].
pub(crate) struct ParsedModelOverrides {
pub models: IndexMap<String, ConfigModelOverride>,
pub warnings: Vec<ModelOverrideWarning>,
}
/// Parses every `[model.<id>]` entry in `raw_config`, returning the overrides
/// and a warning for each skipped field or dropped entry.
pub(crate) fn parse_model_overrides(raw_config: &toml::Value) -> ParsedModelOverrides {
let mut models = IndexMap::new();
let mut warnings = Vec::new();
let Some(section) = raw_config.get("model") else {
return ParsedModelOverrides { models, warnings };
};
let Some(table) = section.as_table() else {
warnings.push(ModelOverrideWarning {
model_key: None,
field: None,
kind: ModelOverrideWarningKind::NotATable,
reason: format!(
"`model` must be a table of [model.<id>] entries, got {}; all model overrides ignored",
section.type_str()
),
});
return ParsedModelOverrides { models, warnings };
};
for (model_key, value) in table {
let Some(entry_table) = value.as_table() else {
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.clone()),
field: None,
kind: ModelOverrideWarningKind::NotATable,
reason: format!(
"expected a table like [model.\"{model_key}\"], got {}; entry dropped",
value.type_str()
),
});
continue;
};
let (entry, entry_warnings) = parse_model_override_table(model_key, entry_table.clone());
warnings.extend(entry_warnings);
models.insert(model_key.clone(), entry);
}
ParsedModelOverrides { models, warnings }
}
/// Logs the warnings when they differ from the previous parse, so a
/// persistently broken config logs once per process instead of once per parse.
pub(crate) fn log_model_override_warnings(warnings: &[ModelOverrideWarning]) {
use std::hash::{Hash as _, Hasher as _};
use std::sync::atomic::{AtomicU64, Ordering};
static LAST_LOGGED: AtomicU64 = AtomicU64::new(0);
// 0 means "no warnings"; real hashes are clamped to nonzero.
let hash = if warnings.is_empty() {
0
} else {
let mut hasher = std::collections::hash_map::DefaultHasher::new();
warnings.hash(&mut hasher);
hasher.finish().max(1)
};
if LAST_LOGGED.swap(hash, Ordering::Relaxed) == hash {
return;
}
for warning in warnings {
tracing::warn!(
model = warning.model_key.as_deref().unwrap_or("(section)"),
field = warning.field.as_deref().unwrap_or("(entry)"),
kind = ?warning.kind,
reason = %warning.reason,
"model_override: skipped invalid config"
);
}
if !warnings.is_empty() {
tracing::warn!(
warnings = warnings.len(),
"model_override: parsed with warnings; run `grok inspect` for details"
);
}
}
fn parse_model_override_table(
model_key: &str,
mut table: toml::map::Map<String, toml::Value>,
) -> (ConfigModelOverride, Vec<ModelOverrideWarning>) {
let mut warnings = Vec::new();
dedupe_aliases(model_key, &mut table, &mut warnings);
// Unknown-field warnings come from whichever parse produces the returned
// entry, so both paths report them identically.
match deserialize_with_unknown_fields(table.clone()) {
Ok((entry, unknown)) => {
warnings.extend(unknown_field_warnings(model_key, unknown));
(entry, warnings)
}
Err(_) => {
prune_invalid_fields(model_key, &mut table, &mut warnings);
match deserialize_with_unknown_fields(table) {
Ok((entry, unknown)) => {
warnings.extend(unknown_field_warnings(model_key, unknown));
(entry, warnings)
}
Err(error) => {
// Reachable only when fields conflict jointly, e.g. an
// alias pair missing from `ALIASES`. Keep the model
// rather than dropping it.
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: None,
kind: ModelOverrideWarningKind::UnparseableEntry,
reason: format!(
"failed to parse after skipping invalid fields ({error}); using empty override"
),
});
(ConfigModelOverride::default(), warnings)
}
}
}
}
}
/// `(canonical, legacy)` key pairs that serde rejects as duplicate fields
/// when both appear in one table. Keep in sync with the `#[serde(alias)]`
/// attributes on [`ConfigModelOverride`].
const ALIASES: &[(&str, &str)] = &[("compactions_remaining", "send_compactions_remaining")];
/// Removes one key of each [`ALIASES`] pair that appears twice in `table`.
/// The canonical key wins; when its value doesn't parse, the legacy key is
/// kept instead.
fn dedupe_aliases(
model_key: &str,
table: &mut toml::map::Map<String, toml::Value>,
warnings: &mut Vec<ModelOverrideWarning>,
) {
for &(canonical, legacy) in ALIASES {
if !(table.contains_key(canonical) && table.contains_key(legacy)) {
continue;
}
match field_parse_error(canonical, &table[canonical]) {
None => {
table.remove(legacy);
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(legacy.to_owned()),
kind: ModelOverrideWarningKind::DuplicateAlias,
reason: format!("legacy alias of {canonical}; skipped in favor of {canonical}"),
});
}
Some(error) => {
table.remove(canonical);
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(canonical.to_owned()),
kind: ModelOverrideWarningKind::InvalidValue,
reason: format!("{error}; skipped in favor of {legacy}"),
});
}
}
}
}
/// Deserializes `table`, also returning the unknown field names that serde
/// would otherwise silently discard.
fn deserialize_with_unknown_fields(
table: toml::map::Map<String, toml::Value>,
) -> Result<(ConfigModelOverride, Vec<String>), toml::de::Error> {
let mut unknown = Vec::new();
let entry = serde_ignored::deserialize(toml::Value::Table(table), |path| {
unknown.push(path.to_string());
})?;
Ok((entry, unknown))
}
fn unknown_field_warnings(model_key: &str, unknown: Vec<String>) -> Vec<ModelOverrideWarning> {
unknown
.into_iter()
.map(|field| ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(field),
kind: ModelOverrideWarningKind::UnknownField,
reason: "unknown field".to_owned(),
})
.collect()
}
/// Removes each field that fails to parse on its own, one warning per field.
/// Unknown fields stay; the follow-up parse reports them.
fn prune_invalid_fields(
model_key: &str,
table: &mut toml::map::Map<String, toml::Value>,
warnings: &mut Vec<ModelOverrideWarning>,
) {
table.retain(|field, value| match field_parse_error(field, value) {
None => true,
Some(error) => {
warnings.push(ModelOverrideWarning {
model_key: Some(model_key.to_owned()),
field: Some(field.to_owned()),
kind: ModelOverrideWarningKind::InvalidValue,
reason: error.to_string(),
});
false
}
});
}
/// Parses `field` in isolation, returning the error if it fails.
fn field_parse_error(field: &str, value: &toml::Value) -> Option<toml::de::Error> {
let mut singleton = toml::map::Map::new();
singleton.insert(field.to_owned(), value.clone());
toml::Value::Table(singleton)
.try_into::<ConfigModelOverride>()
.err()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::sampling::ApiBackend;
use kigi_sampling_types::{
CompactionAtTokens, CompactionsRemaining, ReasoningEffort, ReasoningEffortOption,
};
fn parse_cfg(toml_str: &str) -> crate::agent::config::Config {
let raw: toml::Value = toml::from_str(toml_str).unwrap();
crate::agent::config::Config::new_from_toml_cfg(&raw).expect("config should parse")
}
fn parse_raw(
toml_str: &str,
) -> (
IndexMap<String, ConfigModelOverride>,
Vec<ModelOverrideWarning>,
) {
let raw: toml::Value = toml::from_str(toml_str).unwrap();
let ParsedModelOverrides { models, warnings } = parse_model_overrides(&raw);
(models, warnings)
}
#[test]
fn duplicate_compactions_keys_keeps_model() {
let cfg = parse_cfg(
r#"
[model."grok-4.5"]
model = "grok-4.5"
env_key = "ANTHROPIC_AUTH_TOKEN"
compactions_remaining = 1
send_compactions_remaining = true
"#,
);
let model = cfg
.config_models
.get("grok-4.5")
.expect("grok-4.5 must remain in catalog");
assert_eq!(
model.compactions_remaining,
Some(CompactionsRemaining::Fixed(1))
);
assert!(cfg.model_override_warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::DuplicateAlias
&& w.field.as_deref() == Some("send_compactions_remaining")
}));
let resolved = crate::agent::config::resolve_model_list(&cfg, None);
assert!(resolved.contains_key("grok-4.5"));
}
#[test]
fn legacy_alias_alone_parses_without_warning() {
let cfg = parse_cfg(
r#"
[model."grok-4.5"]
model = "grok-4.5"
send_compactions_remaining = 2
"#,
);
let model = cfg.config_models.get("grok-4.5").unwrap();
assert_eq!(
model.compactions_remaining,
Some(CompactionsRemaining::Fixed(2))
);
assert!(cfg.model_override_warnings.is_empty());
}
#[test]
fn invalid_reasoning_effort_skips_field_keeps_model() {
let cfg = parse_cfg(
r#"
[model."grok-4.5"]
model = "grok-4.5"
env_key = "ANTHROPIC_AUTH_TOKEN"
reasoning_effort = "not-a-level"
"#,
);
let model = cfg
.config_models
.get("grok-4.5")
.expect("grok-4.5 must remain in catalog");
assert_eq!(model.model.as_deref(), Some("grok-4.5"));
assert!(model.reasoning_effort.is_none());
assert!(cfg.model_override_warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::InvalidValue
&& w.field.as_deref() == Some("reasoning_effort")
}));
}
#[test]
fn unknown_field_warns_but_keeps_known_fields() {
let (models, warnings) = parse_raw(
r#"
[model."grok-4.5"]
model = "grok-4.5"
env_key = "TOKEN"
future_field = 1
"#,
);
let entry = models.get("grok-4.5").unwrap();
assert_eq!(entry.model.as_deref(), Some("grok-4.5"));
assert_eq!(
entry.env_key.as_ref().and_then(|k| k.primary()),
Some("TOKEN")
);
assert_eq!(
warnings,
vec![ModelOverrideWarning {
model_key: Some("grok-4.5".to_owned()),
field: Some("future_field".to_owned()),
kind: ModelOverrideWarningKind::UnknownField,
reason: "unknown field".to_owned(),
}]
);
}
/// An unknown field warns the same whether or not another field fails to
/// parse.
#[test]
fn unknown_field_warning_is_path_independent() {
let unknown_of = |toml_str: &str| {
let (_, warnings) = parse_raw(toml_str);
warnings
.into_iter()
.filter(|w| w.kind == ModelOverrideWarningKind::UnknownField)
.collect::<Vec<_>>()
};
let fast = unknown_of(
r#"
[model.m]
temprature = 0.5
"#,
);
let slow = unknown_of(
r#"
[model.m]
temprature = 0.5
reasoning_effort = "not-a-level"
"#,
);
assert_eq!(fast, slow);
assert_eq!(fast.len(), 1);
assert_eq!(fast[0].field.as_deref(), Some("temprature"));
}
#[test]
fn prune_skips_invalid_fields_and_keeps_the_rest() {
// A valid nested table survives an invalid sibling.
let (models, warnings) = parse_raw(
r#"
[model.m]
temperature = "hot"
[model.m.extra_headers]
x-team = "codegen"
"#,
);
let entry = models.get("m").unwrap();
assert_eq!(
entry.extra_headers.get("x-team").map(String::as_str),
Some("codegen")
);
assert!(entry.temperature.is_none());
assert!(warnings.iter().any(|w| {
w.kind == ModelOverrideWarningKind::InvalidValue
&& w.field.as_deref() == Some("temperature")
}));
// All fields invalid: the model stays, with an empty override.
let (models, warnings) = parse_raw(
r#"
[model.m]
temperature = "hot"
max_retries = "many"
"#,
);
let entry = models.get("m").expect("model must remain in catalog");
assert!(entry.temperature.is_none());
assert!(entry.max_retries.is_none());
assert_eq!(warnings.len(), 2);
assert!(
warnings
.iter()
.all(|w| w.kind == ModelOverrideWarningKind::InvalidValue)
);
}
#[test]
fn invalid_canonical_key_falls_back_to_legacy_alias() {
let (models, warnings) = parse_raw(
r#"
[model.m]
compactions_remaining = "bad"
send_compactions_remaining = 2
"#,
);
let entry = models.get("m").unwrap();
assert_eq!(
entry.compactions_remaining,
Some(CompactionsRemaining::Fixed(2))
);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::InvalidValue);
assert_eq!(warnings[0].field.as_deref(), Some("compactions_remaining"));
}
#[test]
fn non_table_model_section_warns_and_is_ignored() {
let (models, warnings) = parse_raw(r#"model = "grok-4""#);
assert!(models.is_empty());
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
assert_eq!(warnings[0].model_key, None);
assert_eq!(warnings[0].field, None);
}
#[test]
fn non_table_entry_is_dropped_with_warning() {
let (models, warnings) = parse_raw(
r#"
[model]
oops = 5
"#,
);
assert!(models.is_empty(), "a scalar cannot define a model");
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::NotATable);
assert_eq!(warnings[0].model_key.as_deref(), Some("oops"));
assert_eq!(warnings[0].field, None);
}
/// Exhaustive literal (no `..`): a new struct field is a compile error
/// here until the drift-guard tests cover it.
fn fully_populated_override() -> ConfigModelOverride {
ConfigModelOverride {
model: Some("m".into()),
base_url: Some("https://example.com".into()),
name: Some("Model M".into()),
description: Some("desc".into()),
api_key: Some("key".into()),
env_key: Some(crate::agent::config::EnvKeys::single("ENV_KEY")),
api_base_url: Some("https://api.example.com".into()),
max_completion_tokens: Some(1024),
temperature: Some(0.5),
top_p: Some(0.9),
api_backend: Some(ApiBackend::Messages),
extra_headers: [("x-team".to_owned(), "codegen".to_owned())]
.into_iter()
.collect(),
context_window: Some(200_000),
auto_compact_threshold_percent: Some(80),
system_prompt_label: Some("label".into()),
use_concise: Some(true),
agent_type: Some("agent".into()),
inference_idle_timeout_secs: Some(60),
max_retries: Some(3),
hidden: Some(false),
supported_in_api: Some(true),
reasoning_effort: Some(ReasoningEffort::High),
supports_reasoning_effort: Some(true),
reasoning_efforts: vec![ReasoningEffortOption {
id: "deep".to_string(),
value: ReasoningEffort::High,
label: "Deep".to_string(),
description: Some("Deep reasoning".to_string()),
default: true,
}],
supports_backend_search: Some(false),
compactions_remaining: Some(CompactionsRemaining::Fixed(1)),
compaction_at_tokens: Some(CompactionAtTokens::Fixed(100_000)),
show_model_fingerprint: Some(true),
stream_tool_calls: Some(false),
}
}
fn parse_single_entry(
entry: toml::map::Map<String, toml::Value>,
) -> (
IndexMap<String, ConfigModelOverride>,
Vec<ModelOverrideWarning>,
) {
let mut model_table = toml::map::Map::new();
model_table.insert("m".to_owned(), toml::Value::Table(entry));
let mut root = toml::map::Map::new();
root.insert("model".to_owned(), toml::Value::Table(model_table));
let ParsedModelOverrides { models, warnings } =
parse_model_overrides(&toml::Value::Table(root));
(models, warnings)
}
#[test]
fn fully_populated_override_round_trips_without_warnings() {
let serialized = toml::Value::try_from(fully_populated_override()).unwrap();
let (models, warnings) = parse_single_entry(serialized.as_table().unwrap().clone());
assert_eq!(warnings, Vec::new(), "no field may be skipped or unknown");
let reparsed = toml::Value::try_from(models.get("m").unwrap()).unwrap();
assert_eq!(reparsed, serialized, "round-trip must be lossless");
}
/// Drift guard: every `#[serde(alias)]` on [`ConfigModelOverride`] must
/// have a matching `ALIASES` pair, and vice versa. An unregistered alias
/// would send both-keys configs to the empty-override fallback.
#[test]
fn every_struct_alias_is_registered_in_aliases() {
let source = include_str!("config.rs");
let start = source
.find("pub struct ConfigModelOverride {")
.expect("ConfigModelOverride definition in config.rs");
let block = &source[start..];
let block = &block[..block.find("\n}").expect("struct end")];
let mut found = Vec::new();
let mut rest = block;
while let Some(pos) = rest.find("#[serde(alias = \"") {
let after = &rest[pos + "#[serde(alias = \"".len()..];
let legacy = &after[..after.find('"').expect("closing quote")];
let field = &after[after.find("pub ").expect("field after alias") + 4..];
let canonical = &field[..field.find(':').expect("field type colon")];
found.push((canonical.to_owned(), legacy.to_owned()));
rest = after;
}
assert_eq!(
block.matches("alias").count(),
found.len(),
"an alias on ConfigModelOverride was not recognized; write it as \
`#[serde(alias = \"...\")]` on its own line, or update this scan"
);
found.sort();
let mut registered: Vec<(String, String)> = ALIASES
.iter()
.map(|&(c, l)| (c.to_owned(), l.to_owned()))
.collect();
registered.sort();
assert_eq!(
found, registered,
"#[serde(alias)] attributes on ConfigModelOverride and ALIASES must match"
);
}
/// Drift guard for `ALIASES`, in both directions: every pair must be a
/// real serde alias (a both-keys table fails a plain parse), and the
/// parser must resolve it to the canonical key with a single warning.
#[test]
fn every_aliases_pair_is_a_real_serde_alias_and_dedupes() {
let reference = toml::Value::try_from(fully_populated_override()).unwrap();
for &(canonical, legacy) in ALIASES {
let value = reference
.get(canonical)
.unwrap_or_else(|| panic!("{canonical} missing from fully_populated_override"));
let mut entry = toml::map::Map::new();
entry.insert(canonical.to_owned(), value.clone());
entry.insert(legacy.to_owned(), value.clone());
assert!(
toml::Value::Table(entry.clone())
.try_into::<ConfigModelOverride>()
.is_err(),
"{canonical}/{legacy} is not a serde alias pair; remove it from ALIASES"
);
let (models, warnings) = parse_single_entry(entry);
let parsed = toml::Value::try_from(models.get("m").unwrap()).unwrap();
assert_eq!(
parsed.get(canonical),
Some(value),
"canonical value must be retained"
);
assert_eq!(warnings.len(), 1);
assert_eq!(warnings[0].kind, ModelOverrideWarningKind::DuplicateAlias);
assert_eq!(warnings[0].field.as_deref(), Some(legacy));
}
}
}
@@ -0,0 +1,259 @@
//! Wire-shape parsers for ext-notification params handled by `MvpAgent`.
//!
//! Pure parsing only (params JSON → `SessionCommand`); session lookup and
//! command dispatch stay in `mvp_agent::ext_notification`.
use crate::session::SessionCommand;
/// Parse a `x.ai/queue/{remove,reorder,clear,edit,interject}` ext-notification's
/// params into the corresponding [`SessionCommand`].
/// `owner` is the resolved attribution (params `owner`/`clientIdentifier`) used
/// to scope remove/clear to the requesting client's own items, and recorded as
/// `last_editor` for in-place text edits. Returns `None` for unrecognized
/// methods or for `edit` when `newText` is missing.
pub(super) fn parse_queue_edit_command(
method: &str,
params: &serde_json::Value,
owner: Option<String>,
) -> Option<SessionCommand> {
match method {
"x.ai/queue/remove" => {
let id = params.get("id").and_then(|v| v.as_str())?.to_string();
// The client supplies the version it last saw; the handler removes
// only on an exact match (stale = benign no-op + rebroadcast).
// Default 0 covers never-edited prompts (the common case).
let expected_version = params
.get("expectedVersion")
.and_then(|v| v.as_u64())
.unwrap_or(0);
Some(SessionCommand::RemoveQueuedPrompt {
id,
expected_version,
owner,
})
}
"x.ai/queue/reorder" => {
let ordered_ids = params
.get("orderedIds")
.and_then(|v| v.as_array())
.map(|arr| {
arr.iter()
.filter_map(|v| v.as_str().map(|s| s.to_string()))
.collect::<Vec<_>>()
})
.unwrap_or_default();
Some(SessionCommand::ReorderQueue { ordered_ids })
}
"x.ai/queue/clear" => Some(SessionCommand::ClearQueue { owner }),
"x.ai/queue/interject" => {
let id = params.get("id").and_then(|v| v.as_str())?.to_string();
// The client supplies the version it last saw; the handler acts
// only on an exact match (stale = benign no-op + rebroadcast).
let expected_version = params
.get("expectedVersion")
.and_then(|v| v.as_u64())
.unwrap_or(0);
// Optional client-edited replacement text (atomic edit+interject).
// Blank overrides are dropped (degrade to the stored queue text) —
// never interject an empty prompt on a malformed client param.
let new_text = params
.get("newText")
.and_then(|v| v.as_str())
.filter(|s| !s.trim().is_empty())
.map(str::to_string);
Some(SessionCommand::InterjectQueuedPrompt {
id,
expected_version,
owner,
new_text,
})
}
"x.ai/queue/edit" => {
let id = params.get("id").and_then(|v| v.as_str())?.to_string();
let new_text = params.get("newText").and_then(|v| v.as_str())?.to_string();
// `owner` is the resolved attribution; for edit it represents the
// most recent editor (recorded as `last_editor`), not the original
// enqueuer.
Some(SessionCommand::EditQueuedPrompt {
id,
new_text,
editor: owner,
})
}
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
/// Each `x.ai/queue/*` ext-notification maps to the
/// correct versioned/idempotent `SessionCommand`.
#[test]
fn parse_queue_edit_command_maps_each_method() {
// remove: id + expectedVersion + owner.
let p = serde_json::json!({
"sessionId": "s1", "id": "p7", "expectedVersion": 3
});
match parse_queue_edit_command("x.ai/queue/remove", &p, Some("grok-tui".into())) {
Some(SessionCommand::RemoveQueuedPrompt {
id,
expected_version,
owner,
}) => {
assert_eq!(id, "p7");
assert_eq!(expected_version, 3);
assert_eq!(owner.as_deref(), Some("grok-tui"));
}
_ => panic!("expected RemoveQueuedPrompt"),
}
// remove without expectedVersion defaults to 0.
let p = serde_json::json!({ "sessionId": "s1", "id": "p8" });
match parse_queue_edit_command("x.ai/queue/remove", &p, None) {
Some(SessionCommand::RemoveQueuedPrompt {
expected_version, ..
}) => assert_eq!(expected_version, 0),
_ => panic!("expected RemoveQueuedPrompt"),
}
// reorder: orderedIds array.
let p = serde_json::json!({ "sessionId": "s1", "orderedIds": ["a", "b", "c"] });
match parse_queue_edit_command("x.ai/queue/reorder", &p, None) {
Some(SessionCommand::ReorderQueue { ordered_ids }) => {
assert_eq!(ordered_ids, vec!["a", "b", "c"]);
}
_ => panic!("expected ReorderQueue"),
}
// clear: owner-scoped.
match parse_queue_edit_command(
"x.ai/queue/clear",
&serde_json::json!({ "sessionId": "s1" }),
Some("grok-tui".into()),
) {
Some(SessionCommand::ClearQueue { owner }) => {
assert_eq!(owner.as_deref(), Some("grok-tui"));
}
_ => panic!("expected ClearQueue"),
}
// edit: id + newText + editor (resolved via owner/clientIdentifier).
let p = serde_json::json!({
"sessionId": "s1", "id": "p9", "newText": "replacement text"
});
match parse_queue_edit_command("x.ai/queue/edit", &p, Some("grok-vscode".into())) {
Some(SessionCommand::EditQueuedPrompt {
id,
new_text,
editor,
}) => {
assert_eq!(id, "p9");
assert_eq!(new_text, "replacement text");
assert_eq!(editor.as_deref(), Some("grok-vscode"));
}
_ => panic!("expected EditQueuedPrompt"),
}
// edit without editor (no owner/clientIdentifier) → editor: None.
match parse_queue_edit_command(
"x.ai/queue/edit",
&serde_json::json!({ "sessionId": "s1", "id": "p9", "newText": "x" }),
None,
) {
Some(SessionCommand::EditQueuedPrompt { editor, .. }) => {
assert!(editor.is_none());
}
_ => panic!("expected EditQueuedPrompt"),
}
// edit without newText → None (can't replace text we don't have).
assert!(
parse_queue_edit_command(
"x.ai/queue/edit",
&serde_json::json!({ "sessionId": "s1", "id": "p9" }),
None,
)
.is_none()
);
// edit without id → None (can't target an entry).
assert!(
parse_queue_edit_command(
"x.ai/queue/edit",
&serde_json::json!({ "sessionId": "s1", "newText": "x" }),
None,
)
.is_none()
);
// interject: id + expectedVersion + owner (mirrors remove).
let p = serde_json::json!({
"sessionId": "s1", "id": "p10", "expectedVersion": 2
});
match parse_queue_edit_command("x.ai/queue/interject", &p, Some("grok-tui".into())) {
Some(SessionCommand::InterjectQueuedPrompt {
id,
expected_version,
owner,
new_text,
}) => {
assert_eq!(id, "p10");
assert_eq!(expected_version, 2);
assert_eq!(owner.as_deref(), Some("grok-tui"));
assert_eq!(new_text, None, "newText absent → None");
}
_ => panic!("expected InterjectQueuedPrompt"),
}
// interject with newText (client-edited row) carries the override.
let p = serde_json::json!({
"sessionId": "s1", "id": "p10", "expectedVersion": 2, "newText": "edited"
});
match parse_queue_edit_command("x.ai/queue/interject", &p, None) {
Some(SessionCommand::InterjectQueuedPrompt { new_text, .. }) => {
assert_eq!(new_text.as_deref(), Some("edited"));
}
_ => panic!("expected InterjectQueuedPrompt"),
}
// Blank newText is dropped → degrades to the stored queue text.
let p = serde_json::json!({
"sessionId": "s1", "id": "p10", "expectedVersion": 2, "newText": " "
});
match parse_queue_edit_command("x.ai/queue/interject", &p, None) {
Some(SessionCommand::InterjectQueuedPrompt { new_text, .. }) => {
assert_eq!(new_text, None, "blank override must be dropped");
}
_ => panic!("expected InterjectQueuedPrompt"),
}
// interject without expectedVersion defaults to 0.
match parse_queue_edit_command(
"x.ai/queue/interject",
&serde_json::json!({ "sessionId": "s1", "id": "p11" }),
None,
) {
Some(SessionCommand::InterjectQueuedPrompt {
expected_version, ..
}) => assert_eq!(expected_version, 0),
_ => panic!("expected InterjectQueuedPrompt"),
}
// interject without id → None (can't target an entry).
assert!(
parse_queue_edit_command("x.ai/queue/interject", &serde_json::json!({}), None)
.is_none()
);
// unknown method → None.
assert!(
parse_queue_edit_command("x.ai/queue/bogus", &serde_json::json!({}), None).is_none()
);
// remove without id → None (can't target an entry).
assert!(
parse_queue_edit_command("x.ai/queue/remove", &serde_json::json!({}), None).is_none()
);
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,3 @@
pub(crate) mod model_switch;
pub(crate) mod session;
pub(crate) mod workspaces;
@@ -0,0 +1,233 @@
//! Applies a model switch to a session — the ungated path. `set_session_model`
//! enforces the `allowed_models` gate before delegating here; internal callers
//! (`new_session`, `load_session`) call `apply` directly.
use crate::agent::config;
use crate::agent::mvp_agent::{
MvpAgent, agent_name_after_model_switch, harnesses_are_compatible, resolve_required_agent_type,
};
use crate::session::SessionCommand;
use agent_client_protocol::{self as acp};
use kigi_sampling_types::parse_reasoning_effort_meta;
use tokio::sync::oneshot;
/// Apply a model switch to a session (no gate — `set_session_model` gates first).
pub(crate) async fn apply(
agent: &MvpAgent,
args: acp::SetSessionModelRequest,
) -> Result<acp::SetSessionModelResponse, acp::Error> {
tracing::info!("Received set session model request {args:?}");
kigi_log::unified_log::info(
"model changed",
Some(args.session_id.0.as_ref()),
Some(serde_json::json!({ "model" : args.model_id.0.as_ref() })),
);
tracing::debug!("session_session_model::mvp_agent: {:?}", &args);
let effort_override = parse_reasoning_effort_meta(args.meta.as_ref());
let acp::SetSessionModelRequest {
session_id,
model_id,
..
} = args;
let handle = agent
.session_handle_waiting_for_load(&session_id)
.await
.ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
let model = agent.resolve_model_id(&model_id)?;
let use_concise = model.info().use_concise;
let session_default = handle
.session_default_agent_profile
.as_deref()
.unwrap_or(&handle.agent_name);
let required_agent_type =
resolve_required_agent_type(Some(model.info().agent_type.as_str()), session_default);
let previous_model_id = handle.model_id.0.clone();
let mut pending_rebuild_definition: Option<kigi_agent::AgentDefinition> = None;
{
let required = &required_agent_type;
let turn_count = handle
.signals_handle
.snapshot()
.await
.map(|s| s.turn_count)
.unwrap_or(0);
let (agent_tx, agent_rx) = oneshot::channel();
let _ = handle.cmd_tx.send(SessionCommand::GetActiveAgent {
responds_to: agent_tx,
});
let active_agent_type = agent_rx.await.ok().flatten();
let is_mismatch = active_agent_type
.as_ref()
.is_some_and(|active| !harnesses_are_compatible(active, required));
tracing::info!(
session_id = % session_id.0, model_id = % model_id.0, ? required_agent_type,
? active_agent_type, turn_count, is_mismatch,
"set_session_model: agent type compatibility check"
);
if is_mismatch && turn_count > 0 {
tracing::warn!(
session_id = % session_id.0, model_id = % model_id.0, active_agent = ?
active_agent_type, required_agent = % required, turn_count,
"set_session_model: agent type mismatch rejected"
);
let err_payload = config::ModelSwitchIncompatibleAgentError {
code: config::MODEL_SWITCH_INCOMPATIBLE_AGENT.to_string(),
active_agent_type: active_agent_type.unwrap_or_else(|| "unknown".to_owned()),
required_agent_type: required.clone(),
model_id: model_id.0.to_string(),
suggestion: "start_new_session".to_string(),
};
return Err(err_payload.into_acp_error());
}
if is_mismatch && turn_count == 0 {
let cwd = handle.tool_context.cwd.as_path();
let resolved = kigi_agent::discovery::by_name_in_cwd_with_plugins(
required,
cwd,
agent.plugin_registry_handle.snapshot().as_deref(),
);
match resolved {
Some(def) => {
tracing::info!(
session_id = % session_id.0, model_id = % model_id.0,
required_agent_type = % required, agent_def_name = % def.name,
"set_session_model: zero-turn harness switch — queued agent rebuild"
);
pending_rebuild_definition = Some(def);
}
None => {
tracing::warn!(
session_id = % session_id.0, model_id = % model_id.0,
required_agent_type = % required,
"set_session_model: zero-turn harness switch — could not resolve agent definition; proceeding with stale harness"
);
}
}
}
}
let mut model_sampling =
agent.prepare_sampling_config_for_model(&model, handle.origin_client.clone());
if let Some(eff) = effort_override {
if agent
.models_manager
.model_supports_reasoning_effort(model_id.0.as_ref())
{
tracing::info!(
session_id = % session_id.0, effort = % eff,
"set_session_model: applying reasoning_effort override from meta"
);
model_sampling.reasoning_effort = Some(eff);
} else {
tracing::warn!(
session_id = % session_id.0, model_id = % model_id.0, effort = % eff,
"set_session_model: ignoring reasoning_effort override — model does not support it"
);
}
}
let applied_effort = model_sampling.reasoning_effort;
let gate_closed = !handle
.gateway_enabled
.load(std::sync::atomic::Ordering::Relaxed);
let apply_prompt_override = !gate_closed;
if gate_closed {
tracing::info!(
session_id = % session_id.0, model_id = % model_id.0,
"set_session_model: gateway gate closed, prompt override suppressed"
);
pending_rebuild_definition = None;
}
let did_rebuild = if let Some(def) = pending_rebuild_definition {
let (rebuild_tx, rebuild_rx) = oneshot::channel();
let _ = handle
.cmd_tx
.send(SessionCommand::RebuildAgentForDefinition {
definition: def,
responds_to: rebuild_tx,
});
let rebuild_result = rebuild_rx
.await
.map_err(|_| acp::Error::internal_error().data("rebuild_agent: actor closed"))?;
match rebuild_result {
Ok(()) => true,
Err(e) => {
tracing::error!(
session_id = % session_id.0, model_id = % model_id.0, error = ? e,
"set_session_model: zero-turn harness rebuild failed; aborting model switch"
);
return Err(e);
}
}
} else {
false
};
let model_unchanged = previous_model_id == model_id.0;
let new_threshold = {
let cfg = agent.cfg.borrow();
let models = agent.models_manager.models();
let model = config::find_model_by_id(&models, model_sampling.model.as_str());
crate::util::config::resolve_auto_compact_threshold_percent(
&cfg,
model_sampling.model.as_str(),
model.map(|e| &e.info),
)
};
let (tx, rx) = oneshot::channel();
let _ = handle.cmd_tx.send(SessionCommand::SetSessionModel {
sampling_config: model_sampling,
use_concise,
apply_prompt_override,
skip_prompt_rewrite: did_rebuild || model_unchanged,
auto_compact_threshold_percent: new_threshold,
responds_to: tx,
});
let updated_model = rx
.await
.map_err(|_| acp::Error::internal_error().data("failed to set session model"))?;
if let Some(handle) = agent.sessions.borrow_mut().get_mut(&session_id) {
handle.model_id = model_id.clone();
handle.reasoning_effort = applied_effort;
handle.agent_name =
agent_name_after_model_switch(did_rebuild, &required_agent_type, &handle.agent_name);
}
broadcast_model_changed(
agent,
&session_id,
model_id.0.as_ref(),
applied_effort.map(|eff| eff.to_string()),
);
if agent.cfg.borrow().mode != config::AgentMode::Leader {
agent.models_manager.set_current_model_id(model_id);
agent
.models_manager
.set_current_reasoning_effort(applied_effort);
}
Ok(acp::SetSessionModelResponse::new().meta(
serde_json::json!({ "model" : updated_model, })
.as_object()
.cloned(),
))
}
/// Broadcast a `ModelChanged` to every client subscribed to this session so
/// followers mirror the new model. The originating client ignores its own echo
/// (gated by `model_switch_pending`). Broadcast-only — no eventId, not persisted.
fn broadcast_model_changed(
agent: &MvpAgent,
session_id: &acp::SessionId,
model_id: &str,
reasoning_effort: Option<String>,
) {
let notification = crate::extensions::notification::SessionNotification {
session_id: session_id.clone(),
update: crate::extensions::notification::SessionUpdate::ModelChanged {
model_id: model_id.to_owned(),
reasoning_effort,
},
meta: None,
};
if let Ok(params) = serde_json::value::to_raw_value(&notification) {
agent
.gateway
.forward_fire_and_forget(acp::ExtNotification::new(
"x.ai/session_notification",
params.into(),
));
}
}
@@ -0,0 +1,291 @@
//! Session meta-information handlers.
//!
//! Router pattern: single `handle()` dispatches by method name.
//! Business logic delegates to pure functions or MvpAgent methods.
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Arc;
use agent_client_protocol::{self as acp};
use serde::Deserialize;
use super::super::mvp_agent::MvpAgent;
use crate::session::persistence::{Summary, list_recent_summaries, list_summaries};
use crate::session::{
AllSessionOverviewRequest, AllSessionOverviewResponse, ContextInfo, ExtMethodResult,
SessionCommand, SessionInfoData, SessionInfoResponse, SessionListRequest, SessionListResponse,
};
/// Mirrors the display title (`generated_title`, else `session_summary`) into
/// `session_summary` so clients that only read that field show the same title
/// as `display_title()` — including after a `/rename` that updated only
/// `generated_title`. Mutates the response copy only; never persisted.
fn backfill_session_summary(summary: &mut Summary) {
let display = summary.display_title().to_owned();
if !display.is_empty() && display != summary.session_summary {
summary.session_summary = display;
}
}
/// Router for x.ai/session/* and x.ai/session_summaries/* methods.
pub async fn handle(
agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
match args.method.as_ref() {
"x.ai/session/info" => handle_session_info(agent, args).await,
"x.ai/session/close" => handle_session_close(agent, args).await,
"x.ai/session/list" => handle_session_list(agent, args).await,
"x.ai/sessions/list" => handle_roster_list(agent, args).await,
m if m.starts_with("x.ai/session_summaries/") => {
handle_session_summaries(agent, args).await
}
_ => Err(acp::Error::method_not_found()),
}
}
/// `x.ai/sessions/list` — the FleetView roster. Returns every
/// resident session plus recently-touched on-disk `Dormant` sessions. Clients
/// poll this while the dashboard is open and reconcile against the
/// `x.ai/sessions/changed` broadcast.
async fn handle_roster_list(
agent: &MvpAgent,
_args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
let sessions = agent.build_roster().await;
ExtMethodResult::success(crate::agent::roster::RosterListResponse { sessions })
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct SessionInfoRequest {
session_id: Option<String>,
}
#[derive(Deserialize)]
struct RecentSessionsRequest {
limit: usize,
}
async fn handle_session_info(
agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
let req: SessionInfoRequest = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
let session_id = req.session_id.or_else(|| {
agent
.sessions
.borrow()
.keys()
.next()
.map(|id| id.0.to_string())
});
let Some(session_id) = session_id else {
return ExtMethodResult::success(serde_json::json!({}))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()));
};
let sid = acp::SessionId::new(session_id.clone());
let Some(session) = agent.sessions.borrow().get(&sid).cloned() else {
return ExtMethodResult::success(serde_json::json!({}))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()));
};
let (tx, rx) = tokio::sync::oneshot::channel();
let _ = session
.cmd_tx
.send(SessionCommand::GetSessionInfo { responds_to: tx });
let info = rx.await.ok();
// Construct display data for `/session-info`.
let mut data = info.unwrap_or_else(|| SessionInfoData {
agent_name: None,
model: None,
model_display_name: None,
resolved_model_id: None,
model_fingerprint: None,
show_model_fingerprint: false,
api_backend: None,
conversation_id: None,
turns: 0,
turn_index: 0,
context: ContextInfo {
auto_compact_threshold_percent:
crate::util::config::DEFAULT_AUTO_COMPACT_THRESHOLD_PERCENT,
..ContextInfo::default()
},
});
// Calculate the model's display name.
data.model_display_name = agent
.models_manager
.models()
.get(session.model_id.0.as_ref())
.and_then(|entry| entry.info.name.clone());
// Construct `SessionInfoResponse`.
let response = SessionInfoResponse {
session_id,
cwd: session.info.cwd.clone(),
data,
};
// Wrap `SessionInfoResponse` in `ExtMethodResult` and return it.
ExtMethodResult::success(serde_json::to_value(&response).unwrap_or_default())
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
async fn handle_session_close(
agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct CloseRequest {
session_id: String,
}
let req: CloseRequest = serde_json::from_str(args.params.get())
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
let sid = acp::SessionId::new(req.session_id.clone());
let existed = agent.sessions.borrow().contains_key(&sid);
if existed {
// Explicit terminal close: shut the actor down and finalize the cloud
// replica (genuine session end). Distinct from a mere client disconnect,
// which detaches but keeps the session resumable and never finalizes
// (see `MvpAgent::handle_evict_sessions` / `close_session_explicit`).
agent.request_session_shutdown(&sid);
agent.close_session_explicit(&sid);
tracing::info!(session_id = %req.session_id, "session closed via x.ai/session/close");
} else {
tracing::debug!(session_id = %req.session_id, "session/close: session not found (already closed)");
}
ExtMethodResult::success(serde_json::json!({ "success": true }))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
async fn handle_session_summaries(
_agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
match args.method.as_ref() {
"x.ai/session_summaries/session_list" => {
let req = serde_json::from_str::<SessionListRequest>(args.params.get())?;
let cwd = req.workspace_directory.to_string_lossy().to_string();
let _timer = crate::instrumentation_timer!("session.list_sessions_for_workspace");
let mut summaries = list_summaries(Some(&cwd)).await.map_err(|e| {
acp::Error::internal_error().data(format!("failed to list sessions: {e}"))
})?;
for s in &mut summaries {
backfill_session_summary(s);
}
let value = serde_json::to_value(SessionListResponse {
session_summaries: summaries,
})
.map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
.expect("to work")
.expect("to work");
Ok(acp::ExtResponse::new(value))
}
"x.ai/session_summaries/workspace_list" => {
tracing::debug!("xai/session_summaries/workspace_list is working");
let _req = serde_json::from_str::<AllSessionOverviewRequest>(args.params.get())?;
let _timer = crate::instrumentation_timer!("session.list_sessions_for_load");
let summaries = list_summaries(None).await.map_err(|e| {
acp::Error::internal_error().data(format!("failed to list workspaces: {e}"))
})?;
summaries_to_overview_response(summaries)
}
"x.ai/session_summaries/workspace_list_recent" => {
let req = serde_json::from_str::<RecentSessionsRequest>(args.params.get())?;
let _timer = crate::instrumentation_timer!("session.list_sessions_recent");
let limit = req.limit.min(10_000);
let mut summaries = list_recent_summaries(limit).await.map_err(|e| {
acp::Error::internal_error().data(format!("failed to list workspaces: {e}"))
})?;
for s in &mut summaries {
backfill_session_summary(s);
}
let value = serde_json::to_value(&summaries)
.map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
.expect("to work")
.expect("to work");
Ok(acp::ExtResponse::new(value))
}
_ => Err(acp::Error::method_not_found()),
}
}
/// Group summaries by cwd and serialize into an [`AllSessionOverviewResponse`].
fn summaries_to_overview_response(summaries: Vec<Summary>) -> Result<acp::ExtResponse, acp::Error> {
let mut by_cwd: BTreeMap<String, Vec<Summary>> = Default::default();
for mut s in summaries {
backfill_session_summary(&mut s);
by_cwd.entry(s.info.cwd.clone()).or_default().push(s);
}
let value = serde_json::to_value(AllSessionOverviewResponse {
all_sessions: by_cwd
.into_iter()
.map(|(k, v)| (PathBuf::from(k), v))
.collect(),
})
.map(|v| serde_json::value::to_raw_value(&v).map(Arc::from))
.expect("to work")
.expect("to work");
Ok(acp::ExtResponse::new(value))
}
// ── Merged session list (local + remote) ─────────────────────────────
async fn handle_session_list(
agent: &MvpAgent,
args: &acp::ExtRequest,
) -> Result<acp::ExtResponse, acp::Error> {
use crate::session::unified_list;
// Under chat mode `parse_list_req` REPLACES any client-sent `kind` facet
// (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"
);
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;
ExtMethodResult::success(unified_list::ext_list_response(result))
.to_ext_response()
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
}
@@ -0,0 +1,174 @@
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");
}
}
+147
View File
@@ -0,0 +1,147 @@
//! Agent bootstrap and lifecycle hooks.
//!
//! [`bootstrap`] runs the full init sequence (config resolution, process
//! singletons, model catalog) and returns a resolved config + `ModelsManager`.
use std::sync::Arc;
use indexmap::IndexMap;
use crate::agent::config::{self, Config as AgentConfig, ModelEntry};
use crate::agent::models::ModelsManager;
use crate::auth::AuthManager;
use crate::config::StorageMode;
/// Resolve config, init process singletons, build the model catalog.
///
/// The `ModelsManager` is `Clone + Send`, so callers that need a handle
/// for the config watcher can clone it before passing it to
/// `MvpAgent::with_models`.
pub fn bootstrap(
cfg: &AgentConfig,
auth_manager: &Arc<AuthManager>,
prefetched: Option<IndexMap<String, ModelEntry>>,
) -> Result<(AgentConfig, ModelsManager), String> {
// Fail closed before any policy is read: a tampered managed policy must not run unmanaged.
crate::managed_config::managed_policy_gate()?;
let cfg = resolve_config(cfg, auth_manager);
cfg.validate_model_filters()?;
init_process(&cfg, auth_manager);
let models_manager = ModelsManager::from_config(&cfg, prefetched, auth_manager.clone())?;
// Refresh on every auth refresh — the FSEvents watcher can silently die after
// macOS sleep, stranding the catalog on bundled defaults.
models_manager.start_auth_refresh_watcher(auth_manager.refresh_notifier());
Ok((cfg, models_manager))
}
/// Print a `bootstrap`/`MvpAgent::new` config error and exit (process boundary).
///
/// Restores native stderr first: a managed-policy refusal on the ACP/server path reaches here
/// while fd 2 may still point at the `/dev/null` the TUI's `redirect_native_stderr()` set, which
/// would swallow the message. No-op when stderr was never redirected (headless).
pub(crate) fn exit_on_config_error<T>(e: String) -> T {
kigi_tty_utils::restore_native_stderr();
eprintln!("\nConfiguration error:\n\n {e}\n");
std::process::exit(1);
}
/// Config transform: apply managed settings, fetch remote settings,
/// resolve storage mode.
fn resolve_config(cfg: &AgentConfig, auth_manager: &AuthManager) -> AgentConfig {
let mut cfg = cfg.clone();
if let Ok(layers) = crate::config::ConfigLayers::load()
&& layers.has_managed()
{
let origins = crate::config::config_origins(&layers);
let managed_keys: Vec<&str> = origins
.iter()
.filter(|(_, s)| matches!(s, config::ConfigSource::ManagedConfig))
.map(|(k, _)| k.as_str())
.collect();
if !managed_keys.is_empty() {
tracing::info!(keys = ?managed_keys, "managed_config.toml fields");
}
}
let managed_enforced = crate::config::apply_managed_settings_features(&mut cfg);
let requirements_enforced = crate::config::apply_requirements(&mut cfg);
for e in managed_enforced.iter().chain(&requirements_enforced) {
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.grok_com_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());
// env var > remote settings > Local. Skip remote settings for Generic (grok -p, subagents).
if cfg.storage_mode == StorageMode::Local
&& cfg.mode != crate::agent::config::AgentMode::Generic
{
cfg.storage_mode = StorageMode::resolve(None, cfg.remote_settings.as_ref());
}
// Writeback talks to the code backend; requires grok.com auth.
if cfg.storage_mode == StorageMode::Writeback
&& !auth_manager.current().is_some_and(|a| a.is_xai_auth())
{
tracing::info!("Writeback is disabled: requires auth with grok.com");
cfg.storage_mode = StorageMode::Local;
}
if let Some(rs) = cfg.remote_settings.as_ref()
&& let Some(v) = rs.path_not_found_hints
{
cfg.path_not_found_hints = v;
}
cfg
}
/// Initialize process-level singletons (deployment sync, bundled files).
/// `Once`-guarded: only the first call takes effect.
fn init_process(cfg: &AgentConfig, auth_manager: &AuthManager) {
use std::sync::Once;
static INIT: Once = Once::new();
INIT.call_once(|| {
if !cfg!(test) {
// Clear a logged-out team's files before the background sync runs.
crate::managed_config::clear_orphan();
crate::managed_config::spawn_sync(tokio_util::sync::CancellationToken::new());
}
let kigi_home = crate::util::kigi_home::kigi_home();
crate::builtin::extract_bundled_files(&kigi_home);
let feedback = cfg.resolve_feedback();
let feedback_url = cfg.endpoints.resolve_feedback_base_url();
tracing::info!(
feedback = %feedback,
feedback_url = %feedback_url,
feedback_url_custom = cfg.endpoints.feedback_base_url.is_some(),
"data capture config resolved",
);
});
}
@@ -0,0 +1,25 @@
pub mod activity;
pub mod app;
pub mod auth_method;
pub mod chat_modes;
pub mod config;
pub mod config_model_override_parse;
mod ext_parsers;
pub mod feedback_client;
pub mod folder_trust;
pub(crate) mod handlers;
pub mod init;
pub mod models;
pub mod mvp_agent;
pub(crate) mod proxy;
pub(crate) mod restore_code;
pub mod roster;
pub mod server;
pub mod session_config;
pub mod session_registry_client;
pub(crate) mod subagent;
pub(crate) mod subscription_check;
pub(crate) mod update_chunk_merge;
pub use mvp_agent::MvpAgent;
pub use server::{ServerConfig, run_agent_server};
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,261 @@
//! Code-navigation eligibility gating and codebase-index management for [`MvpAgent`].
//! Co-located child of `mvp_agent` (`use super::*`).
use super::*;
impl MvpAgent {
/// Parse the `x.ai/codeNavigation.enabled` capability from an initialize
/// request. Returns `false` if the field is absent or not `true`.
pub(crate) fn parse_code_nav_capability(init: &acp::InitializeRequest) -> bool {
init.client_capabilities
.meta
.as_ref()
.and_then(|m| m.get("x.ai/codeNavigation"))
.and_then(|v| v.get("enabled"))
.and_then(|v| v.as_bool())
.unwrap_or(false)
}
/// Start (or reuse) the codebase index for an eligible code-nav request.
///
/// Returns `Some((handle, was_newly_started))` on success or `None` when
/// config/git-root checks prevent starting. The bool is the authoritative
/// "first spawn vs reuse" signal threaded up from `CodebaseIndexManager`.
///
/// This is the narrow `pub(crate)` entry point for lazy index startup
/// from `extensions/code_nav.rs`. Callers must verify eligibility with
/// [`code_nav_eligibility_for_request`] before calling this.
pub(crate) fn start_codebase_index_for_code_nav(
&self,
session_id: Option<&acp::SessionId>,
cwd: &std::path::Path,
) -> Option<(std::sync::Arc<kigi_codebase_graph::IndexManagerHandle>, bool)> {
let (handle, was_newly_started) = self.resolve_codebase_index(cwd)?;
// Pin the index to the requesting session so the Weak in
// CodebaseIndexManager doesn't orphan it immediately.
if let Some(sid) = session_id {
self.session_index_claims
.borrow_mut()
.insert(sid.clone(), std::sync::Arc::clone(&handle));
}
Some((handle, was_newly_started))
}
/// Core eligibility check — pure function that accepts explicit client
/// context rather than reading global agent state.
///
/// This is the single place that applies all four gates. Call it via
/// [`code_nav_eligibility_for_request`] (leader-mode safe) or
/// [`code_nav_eligibility`] (global state, non-leader use only).
pub(super) fn code_nav_eligibility_inner(
&self,
cwd: &std::path::Path,
client_type: ClientType,
code_nav_enabled: bool,
) -> Result<(), CodeNavEligibility> {
use crate::agent::config::CodebaseIndexingSetting;
// Gate 1: client type
if !matches!(client_type, ClientType::GrokWeb) {
tracing::info!(
client_type = ?client_type,
gate = "client_type",
skip_reason = "client_not_web",
"code-nav eligibility check: skipping (client type not eligible)"
);
return Err(CodeNavEligibility::ClientNotWeb);
}
// Gate 2: capability advertised
if !code_nav_enabled {
tracing::info!(
gate = "capability",
skip_reason = "capability_not_advertised",
"code-nav eligibility check: skipping (x.ai/codeNavigation.enabled not advertised)"
);
return Err(CodeNavEligibility::CapabilityNotAdvertised);
}
// Gate 3: config
let setting = self.cfg.borrow().features.codebase_indexing.clone();
if let CodebaseIndexingSetting::Enabled(false) = &setting {
tracing::info!(
gate = "config",
skip_reason = "disabled_by_config",
"code-nav eligibility check: skipping (codebase_indexing disabled in config)"
);
return Err(CodeNavEligibility::DisabledByConfig);
}
// Gate 4: git root / config globs
let git_root = kigi_workspace::session::git::find_git_root_from_path(cwd).ok();
match &setting {
CodebaseIndexingSetting::Enabled(true) => {
if git_root.is_none() {
tracing::info!(
cwd = %cwd.display(),
gate = "git_root",
skip_reason = "not_git_repo",
"code-nav eligibility check: skipping (not inside a git repo)"
);
return Err(CodeNavEligibility::NotGitRepo);
}
}
CodebaseIndexingSetting::Patterns(_) => {
let check_path = git_root.as_deref().unwrap_or(cwd);
if !setting.should_index(check_path) {
tracing::info!(
cwd = %cwd.display(),
gate = "config_globs",
skip_reason = "disabled_by_config",
"code-nav eligibility check: skipping (not matched by config globs)"
);
return Err(CodeNavEligibility::DisabledByConfig);
}
}
CodebaseIndexingSetting::Enabled(false) => {} // handled above
}
Ok(())
}
/// Check eligibility using per-session context (leader-mode safe).
///
/// When `session_id` is provided, reads the session's own client type
/// and code-nav capability — the values that were in effect when that
/// specific client created the session. This is correct in leader mode
/// where multiple clients share one agent process and `initialize()` is
/// called once per connection; the global fields on `MvpAgent` reflect
/// only the **last** client to call `initialize()`.
///
/// Falls back to global agent state when no session_id is given.
pub fn code_nav_eligibility_for_request(
&self,
session_id: Option<&acp::SessionId>,
cwd: &std::path::Path,
) -> Result<(), CodeNavEligibility> {
let session_id = match session_id {
Some(sid) => sid,
// No session_id: per-client capability cannot be determined without a
// session. Reject with SessionRequired rather than fall back to shared
// global state. Callers must provide sessionId for x.ai/code/* requests.
None => return Err(CodeNavEligibility::SessionRequired),
};
let sessions = self.sessions.borrow();
let (client_type, code_nav_enabled) = if let Some(handle) = sessions.get(session_id) {
let ct = crate::http::client_type_from_origin(handle.origin_client.as_ref());
(ct, handle.code_nav_enabled)
} else {
// Session not found (evicted/unknown): reject rather than silently
// falling back to shared global state — that would reintroduce the
// last-client-wins bug for stale session IDs in leader mode.
return Err(CodeNavEligibility::SessionRequired);
};
drop(sessions);
self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
}
/// Check eligibility using the stored initialize_request context.
///
/// **Not safe in leader mode** — reads the last `initialize()` call's
/// client_type and capability. Prefer [`code_nav_eligibility_for_request`]
/// when a session_id is available.
pub fn code_nav_eligibility(&self, cwd: &std::path::Path) -> Result<(), CodeNavEligibility> {
let client_type = *self.client_type.borrow();
let code_nav_enabled = self.code_nav_enabled.get();
self.code_nav_eligibility_inner(cwd, client_type, code_nav_enabled)
}
/// Resolve and get-or-create the codebase index for `cwd`, applying config
/// and git-root eligibility checks.
///
/// Returns `Some((handle, was_newly_started))` when an index is available,
/// `None` when config or git-root checks rule it out. The bool is the
/// authoritative "was this a first spawn?" signal from the manager.
pub(super) fn resolve_codebase_index(
&self,
cwd: &std::path::Path,
) -> Option<(std::sync::Arc<kigi_codebase_graph::IndexManagerHandle>, bool)> {
use crate::agent::config::CodebaseIndexingSetting;
let setting = self.cfg.borrow().features.codebase_indexing.clone();
let git_root = kigi_workspace::session::git::find_git_root_from_path(cwd).ok();
match (&setting, &git_root) {
(CodebaseIndexingSetting::Enabled(false), _) => {
tracing::info!(
cwd = %cwd.display(),
skip_reason = "disabled_by_config",
"code-nav: skipping index creation (disabled in config)"
);
return None;
}
(CodebaseIndexingSetting::Enabled(true), None) => {
tracing::info!(
cwd = %cwd.display(),
skip_reason = "not_git_repo",
"code-nav: skipping index creation (not inside a git repo)"
);
return None;
}
(CodebaseIndexingSetting::Patterns(_), _) => {
let check_path = git_root.as_deref().unwrap_or(cwd);
if !setting.should_index(check_path) {
tracing::info!(
cwd = %cwd.display(),
skip_reason = "disabled_by_config",
"code-nav: skipping index creation (not matched by config globs)"
);
return None;
}
}
(CodebaseIndexingSetting::Enabled(true), Some(_)) => {}
}
let target = git_root.unwrap_or_else(|| cwd.to_path_buf());
// get_or_create returns the authoritative (handle, was_newly_started) pair.
// Log only on actual first spawn so reuse requests are not misleadingly
// labelled as "starting".
let (handle, was_newly_started) = self.get_or_create_codebase_index(target.clone());
if was_newly_started {
tracing::info!(
cwd = %cwd.display(),
index_target = %target.display(),
event = "index_first_spawn",
"code-nav: first lazy spawn of codebase index"
);
}
Some((handle, was_newly_started))
}
pub(super) fn indexed_roots_for(&self, cwd: &std::path::Path) -> Vec<String> {
if self.get_codebase_index(cwd).is_some() {
return vec![cwd.to_string_lossy().into_owned()];
}
if let Ok(git_root) = kigi_workspace::session::git::find_git_root_from_path(cwd)
&& self.get_codebase_index(&git_root).is_some()
{
return vec![git_root.to_string_lossy().into_owned()];
}
Vec::new()
}
/// Returns `(handle, was_newly_started)` — the bool is the authoritative
/// "did this call spawn a new actor?" bit from `CodebaseIndexManager::get_or_create`.
pub(super) fn get_or_create_codebase_index(
&self,
cwd: PathBuf,
) -> (std::sync::Arc<kigi_codebase_graph::IndexManagerHandle>, bool) {
self.codebase_indexes.lock().get_or_create(cwd)
}
/// Get an existing codebase index for the given cwd.
/// Returns None if no index exists for this cwd.
pub fn get_codebase_index(
&self,
cwd: &std::path::Path,
) -> Option<std::sync::Arc<kigi_codebase_graph::IndexManagerHandle>> {
self.codebase_indexes.lock().get(cwd)
}
}
@@ -0,0 +1,453 @@
//! Interactive folder-trust prompt: a dormant agent→GUI-client ACP round-trip
//! (`x.ai/folder_trust/request`) that asks a GUI client (grok-desktop) to decide
//! trust for an untrusted-with-configs workspace, then grants + reloads the
//! now-trusted project servers without a restart.
//!
//! DORMANT in production: it only fires when the connected client advertised
//! `x.ai/folderTrust.interactive` AND the folder-trust feature flag is on AND the
//! verdict is [`kigi_workspace::folder_trust::TrustOutcome::Prompt`]. No
//! client advertises the capability until the desktop UI ships — so this is
//! inert by default even with the feature flag on. The TUI/headless clients never
//! advertise it (they self-gate trust client-side), so they are never
//! double-prompted. Co-located child of `mvp_agent` (`use super::*`).
//!
//! Post-grant reload scope: MCP, plugins, and each session's own project hooks
//! are hot-reloaded in place — for EVERY session sharing the granted workspace
//! (same `workspace_key`), each reloaded against its OWN cwd. Project LSP is NOT
//! hot-reloaded — the LSP backend is baked into the agent's tool bridge at build
//! time (one-shot startup coordinator, no in-place reconfigure API), so repo-local
//! `.kigi/lsp.json` servers start on the NEXT session open (the durable grant
//! makes the re-spawn trusted). `lsp` is still REPORTED in the prompt's
//! `configKinds` (it is a real reason the folder is gated) — only the post-grant
//! hot-reload skips it.
use super::*;
/// Max wait for a GUI client's trust decision before giving up (fail-closed).
/// Generous because it is a human decision, but bounds the detached task so a
/// connected-but-silent client (modal left open / client bug) can't leak it for
/// the whole connection lifetime.
const TRUST_PROMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30 * 60);
/// ACP `x.ai/folder_trust/request` payload (agent → GUI client). Serialized as
/// `camelCase` for the ACP JSON-RPC wire format.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct FolderTrustRequest {
/// The session this prompt belongs to. REQUIRED for leader Tier-2 routing:
/// non-interaction reverse-requests are delivered to the driver keyed on
/// `params.sessionId`; omitting it makes the leader silently drop the message,
/// so the prompt would never reach the client.
pub session_id: String,
/// The session cwd whose workspace is being gated.
pub cwd: String,
/// Display path of the canonical workspace key (the trust grant's scope).
pub workspace: String,
/// Detected repo-local config kinds (e.g. `mcp`, `hooks`, `lsp`) — the
/// reasons the folder is gated — for the prompt UI. Display-only, NOT the
/// trust gate; derived from the same scan as the gate. `lsp` may appear: it
/// is a real reason to prompt, but project LSP applies on the NEXT session
/// open rather than hot-reloading on grant (see module docs).
pub config_kinds: Vec<String>,
}
/// Outcome of the trust prompt (GUI client → agent). Fail-closed: any value
/// other than `"trust"` (including unknown strings, via `#[serde(other)]`)
/// decodes to [`FolderTrustOutcome::Reject`], so only an explicit grant unblocks.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub(crate) enum FolderTrustOutcome {
Trust,
#[serde(other)]
Reject,
}
/// ACP `x.ai/folder_trust/request` response (GUI client → agent).
#[derive(Debug, Clone, serde::Deserialize)]
pub(crate) struct FolderTrustResponse {
pub outcome: FolderTrustOutcome,
}
impl MvpAgent {
/// Parse the `x.ai/folderTrust.interactive` capability from an initialize
/// request. Returns `false` if absent or not `true`. Mirrors
/// [`Self::parse_code_nav_capability`].
pub(crate) fn parse_interactive_trust_capability(init: &acp::InitializeRequest) -> bool {
init.client_capabilities
.meta
.as_ref()
.and_then(|m| m.get("x.ai/folderTrust"))
.and_then(|v| v.get("interactive"))
.and_then(|v| v.as_bool())
.unwrap_or(false)
}
/// Ask a GUI client to decide trust for `session_id`'s workspace, then grant
/// + reload on accept. DORMANT no-op unless the client advertised
/// `x.ai/folderTrust.interactive` AND [`folder_trust::prompt_warranted`]
/// (feature on + untrusted + repo configs present).
///
/// Non-blocking: the session was already created with project servers GATED
/// (the untrusted resolve in `new_session`/`load_session`), so nothing
/// repo-local spawns while the prompt is open. The round-trip + reload run in
/// a detached `spawn_local` task, so the `new_session` response is not
/// delayed by the (potentially long) user decision. At most one outstanding
/// request per workspace per process (dedup), and the await is bounded by
/// [`TRUST_PROMPT_TIMEOUT`].
pub(crate) fn maybe_spawn_interactive_trust_prompt(
&self,
session_id: &acp::SessionId,
cwd: &std::path::Path,
remote: Option<&crate::util::config::RemoteSettings>,
) {
if !self.interactive_trust_client.get() {
return;
}
if !folder_trust::prompt_warranted(cwd, remote) {
return;
}
let key = kigi_workspace::trust::workspace_key(cwd);
// Dedup: skip if this workspace was already prompted/decided (reconnect)
// or has a prompt in flight (concurrent same-workspace session). `insert`
// returns false when already present. Agent-owned set (no process
// global), captured into the task for release on failure/timeout.
let prompted = self.interactive_trust_prompted.clone();
if !prompted.borrow_mut().insert(key.clone()) {
return;
}
// Capture EVERY session sharing the GRANTED WORKSPACE (same
// `workspace_key` — the grant's actual scope, aligned with the dedup key),
// each with its OWN cwd, so a grant reloads every sibling against its own
// project config — exactly like the per-cwd `handle_reload_project_mcp_servers`
// / `broadcast_plugin_registry_to_sessions`. `&self` can't be borrowed
// across the `spawn_local` boundary, so capture owned clones now.
//
// INTENTIONAL fail-safe limitation: this is a one-time snapshot taken at
// prompt-spawn. A same-workspace session created WHILE the modal is open is
// deduped (no second prompt) and is not in this set, so the grant won't
// reload it — it stays GATED until its own next session (secure, never
// over-exposed). Re-querying at grant time would need the `sessions` map
// (a non-`Rc` `RefCell` field) shared into the detached task, which isn't
// available here; the fail-safe stale-session window is accepted instead.
let targets: Vec<ReloadTarget> = self
.sessions
.borrow()
.values()
.filter(|h| {
kigi_workspace::trust::workspace_key(std::path::Path::new(&h.info.cwd)) == key
})
.map(|h| ReloadTarget {
cmd_tx: h.cmd_tx.clone(),
initial_client_mcp_servers: h.initial_client_mcp_servers.clone(),
cwd: PathBuf::from(&h.info.cwd),
})
.collect();
if targets.is_empty() {
prompted.borrow_mut().remove(&key);
return;
}
let gateway = self.gateway.clone();
let plugin_handle = self.plugin_registry_handle.clone();
let managed_mcp_cache = self.managed_mcp_cache.clone();
let auth_manager = self.auth_manager.clone();
let can_fetch_managed = self.can_fetch_managed_mcps();
let proxy_url = self.cfg.borrow().endpoints.proxy_url();
let compat = self.cfg.borrow().compat_resolved;
let remote = remote.cloned();
let cwd = cwd.to_path_buf();
let workspace = key.display().to_string();
let config_kinds = folder_trust::detected_config_kinds(&cwd);
let session_id = session_id.0.to_string();
// Regression guard: every reverse-request must carry a
// non-empty sessionId or leader Tier-2 routing silently drops it.
debug_assert!(
!session_id.is_empty(),
"folder_trust reverse-request must carry a non-empty sessionId (design §5.4)"
);
tokio::task::spawn_local(async move {
let request = FolderTrustRequest {
session_id,
cwd: cwd.to_string_lossy().into_owned(),
workspace,
config_kinds,
};
// Non-panicking: a struct of String/Vec<String> can't fail to
// serialize, but avoid `expect` in prod — bail (and release the dedup
// key) on the impossible error rather than aborting the task thread.
let raw_params = match serde_json::value::to_raw_value(&request) {
Ok(p) => p,
Err(e) => {
tracing::error!(error = %e, "folder trust: request serialization failed");
prompted.borrow_mut().remove(&key);
return;
}
};
let ext_request = acp::ExtRequest::new("x.ai/folder_trust/request", raw_params.into());
use agent_client_protocol::Client as _;
let outcome = match tokio::time::timeout(
TRUST_PROMPT_TIMEOUT,
gateway.ext_method(ext_request),
)
.await
{
// A decodable response carries the user's decision. An
// undecodable success payload is a client/protocol error, not a
// decision: stay gated (fail-closed) but release the dedup key so
// a later session can re-prompt — same as transport/timeout below.
Ok(Ok(raw)) => match serde_json::from_str::<FolderTrustResponse>(raw.0.get()) {
Ok(r) => r.outcome,
Err(e) => {
tracing::debug!(error = %e, "folder trust: undecodable trust response; staying gated, releasing dedup key");
prompted.borrow_mut().remove(&key);
return;
}
},
Ok(Err(e)) => {
// Client disconnected / transport error: not a decision —
// release the key so a later session can re-prompt.
tracing::debug!(error = %e, "folder trust: client trust request failed");
prompted.borrow_mut().remove(&key);
return;
}
Err(_elapsed) => {
// Connected but silent past the deadline: stay gated, release
// the key so a future session may re-prompt.
tracing::info!(
cwd = %cwd.display(),
"folder trust: no client decision before timeout; staying gated"
);
prompted.borrow_mut().remove(&key);
return;
}
};
if outcome != FolderTrustOutcome::Trust {
// Decided "reject": keep the dedup key so the user is not
// re-prompted for this workspace on every reconnect.
tracing::info!(
cwd = %cwd.display(),
"folder trust: GUI client declined; workspace stays gated"
);
return;
}
// Re-check the dedup key before granting. `HooksAction::Untrust`
// removes this workspace's key (and revokes asynchronously) when the
// user untrusts. If that fired while the modal was open, the key is
// gone — honor the untrust and drop this now-stale "trust" rather
// than re-persisting a grant the user just revoked. The single-
// threaded LocalSet makes this check + grant atomic w.r.t. the
// untrust task (no await in between).
if !prompted.borrow().contains(&key) {
tracing::info!(
cwd = %cwd.display(),
"folder trust: workspace untrusted while prompt was open; ignoring stale grant"
);
return;
}
// Persist the grant, then flip the cached untrusted verdict to trusted
// (the `Some(false)` arm of `resolve_and_record` re-reads the store).
folder_trust::grant_folder_trust(&cwd);
folder_trust::resolve_and_record(&cwd, remote.as_ref(), false);
reload_project_servers_after_grant(ReloadAfterGrant {
gateway: &gateway,
targets,
plugin_handle: &plugin_handle,
managed_mcp_cache: &managed_mcp_cache,
auth_manager: &auth_manager,
can_fetch_managed,
proxy_url: &proxy_url,
compat: &compat,
prompt_cwd: &cwd,
})
.await;
tracing::info!(
cwd = %cwd.display(),
"folder trust: granted via GUI client; reloaded project servers"
);
});
}
}
/// One session to reload after a grant, with ITS OWN cwd (so the MCP merge +
/// plugin build use the session's own project config — matching the per-cwd
/// canonical reloaders, not the prompt's cwd).
struct ReloadTarget {
cmd_tx: tokio::sync::mpsc::UnboundedSender<crate::session::SessionCommand>,
initial_client_mcp_servers: Vec<acp::McpServer>,
cwd: PathBuf,
}
/// Inputs for [`reload_project_servers_after_grant`], bundled to keep the
/// orchestrator free of a long positional arg list.
struct ReloadAfterGrant<'a> {
gateway: &'a GatewaySender,
/// Every session sharing the granted workspace, each with its own cwd.
targets: Vec<ReloadTarget>,
plugin_handle: &'a kigi_agent::plugins::SharedPluginRegistryHandle,
managed_mcp_cache: &'a crate::session::managed_mcp::ManagedMcpStateHandle,
auth_manager: &'a std::sync::Arc<AuthManager>,
can_fetch_managed: bool,
proxy_url: &'a str,
compat: &'a kigi_tools::types::CompatConfig,
/// The prompting session's cwd — used only for the client catalog push.
prompt_cwd: &'a std::path::Path,
}
/// Reload each granted-workspace session's now-trusted project servers in place
/// (no restart), driving the canonical primitives the normal spawn/reload paths
/// use — PER SESSION CWD, like `handle_reload_project_mcp_servers` /
/// `broadcast_plugin_registry_to_sessions`: `fetch_managed_mcp_configs` +
/// `merge_managed_mcp_servers` (`SessionCommand::UpdateMcpServers`), `build_for_cwd`
/// (`SessionCommand::ReloadPlugins`), and `reload_hooks_impl`
/// (`SessionCommand::ReloadHooks`), then push the refreshed MCP catalog. LSP is
/// spawn-baked and applies on the next session open (see module docs). Caller
/// must have granted + recorded trust first.
async fn reload_project_servers_after_grant(ctx: ReloadAfterGrant<'_>) {
// Managed (gateway/Toolbox) servers must survive the re-merge; fetch them once
// (cwd-independent) via the shared helper (single-sources the auth-key dance
// with `MvpAgent::get_managed_mcp_configs`). The plugin MCP snapshot is also
// global, so it is fine to reuse across cwds for the merge.
let managed = if ctx.can_fetch_managed {
crate::session::managed_mcp::fetch_managed_mcp_configs(
ctx.managed_mcp_cache,
ctx.proxy_url,
ctx.auth_manager,
)
.await
} else {
vec![]
};
let plugin_snapshot = ctx.plugin_handle.snapshot();
for target in ctx.targets {
// Per-session cwd: a sibling session in a subdir of the granted workspace
// must get ITS OWN project config, not the prompt's.
let session_cwd = target.cwd.as_path();
// MCP: `merge_managed_mcp_servers` re-reads disk + runs
// `filter_untrusted_project_mcp`, which now KEEPS project servers because
// the cached verdict was flipped to trusted (same workspace key).
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
target.initial_client_mcp_servers,
session_cwd,
&managed,
plugin_snapshot.as_deref(),
ctx.compat,
);
let (tx, _rx) = tokio::sync::oneshot::channel();
let _ = target
.cmd_tx
.send(crate::session::SessionCommand::UpdateMcpServers {
mcp_servers: merged,
respond_to: tx,
});
// Plugins (+ plugin-contributed hooks) built for this session's own cwd
// on the folder-trust verdict (mirrors `broadcast_plugin_registry_to_sessions`);
// the grant + resolve_and_record above flipped the cached verdict to trusted.
let disk_cfg =
crate::config::resolve_effective_plugins_config(session_cwd).to_discovery_config();
let project_trusted = folder_trust::project_scope_allowed(session_cwd);
// Session `_meta.pluginDirs` are re-merged by the receiving actor
// (`preserve_session_plugin_dirs` on `ReloadPlugins`).
let registry =
ctx.plugin_handle
.build_for_cwd(session_cwd, &disk_cfg, &[], project_trusted);
let _ = target
.cmd_tx
.send(crate::session::SessionCommand::ReloadPlugins { registry });
// The session's OWN project hooks (`.kigi/hooks`, `.cursor/hooks.json`),
// which `ReloadPlugins` does NOT touch — re-discovered against the actor's
// own `session_info.cwd` on the now-trusted verdict by `reload_hooks_impl`.
let _ = target
.cmd_tx
.send(crate::session::SessionCommand::ReloadHooks);
}
// Push the refreshed MCP catalog (for the prompting session's cwd) so the
// client UI reflects the now-trusted repo-local servers.
let local = folder_trust::filter_untrusted_project_mcp(
ctx.prompt_cwd,
crate::util::config::load_mcp_servers(ctx.prompt_cwd, ctx.compat),
);
crate::extensions::mcp::notify_servers_updated(ctx.gateway, &managed, &local).await;
}
#[cfg(test)]
mod tests {
use super::*;
fn init_with_meta(meta: Option<serde_json::Value>) -> acp::InitializeRequest {
// Production reads `client_capabilities.meta`, not top-level request meta.
let mut caps = acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false);
if let Some(m) = meta
&& let Some(map) = m.as_object().cloned()
{
caps = caps.meta(map);
}
acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(caps)
}
#[test]
fn parse_interactive_trust_capability_present_and_true() {
let mut meta = serde_json::Map::new();
meta.insert(
"x.ai/folderTrust".to_string(),
serde_json::json!({ "interactive": true }),
);
let init = init_with_meta(Some(serde_json::Value::Object(meta)));
assert!(MvpAgent::parse_interactive_trust_capability(&init));
}
#[test]
fn parse_interactive_trust_capability_absent_returns_false() {
let init = init_with_meta(None);
assert!(!MvpAgent::parse_interactive_trust_capability(&init));
}
#[test]
fn parse_interactive_trust_capability_false_returns_false() {
let mut meta = serde_json::Map::new();
meta.insert(
"x.ai/folderTrust".to_string(),
serde_json::json!({ "interactive": false }),
);
let init = init_with_meta(Some(serde_json::Value::Object(meta)));
assert!(!MvpAgent::parse_interactive_trust_capability(&init));
}
#[test]
fn request_serializes_camel_case_with_session_id() {
let req = FolderTrustRequest {
session_id: "sess-1".into(),
cwd: "/repo".into(),
workspace: "/repo".into(),
config_kinds: vec!["mcp".into()],
};
let json = serde_json::to_value(&req).unwrap();
assert!(json.get("configKinds").is_some());
assert!(json.get("config_kinds").is_none());
// Leader Tier-2 routing reads `params.sessionId`; it must be present and
// non-empty (regression guard for the silently-dropped-in-leader bug).
assert_eq!(json["sessionId"], "sess-1");
assert!(!json["sessionId"].as_str().unwrap().is_empty());
}
#[test]
fn response_decodes_trust_reject_and_unknown_fail_closed() {
let trust: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"trust"}"#).unwrap();
assert_eq!(trust.outcome, FolderTrustOutcome::Trust);
let reject: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"reject"}"#).unwrap();
assert_eq!(reject.outcome, FolderTrustOutcome::Reject);
// Unknown outcome must fail closed to Reject (never silently "trust").
let unknown: FolderTrustResponse = serde_json::from_str(r#"{"outcome":"banana"}"#).unwrap();
assert_eq!(unknown.outcome, FolderTrustOutcome::Reject);
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,147 @@
use super::{PromptResponseMetaArgs, build_prompt_response_meta};
use kigi_sampling_types::TokenUsage;
/// Baseline args with no usage, cancellation, or structured output.
fn args<'a>(
session_id: &'a str,
prompt_id: &'a str,
total_tokens: u64,
model_id: &'a str,
) -> PromptResponseMetaArgs<'a> {
PromptResponseMetaArgs {
session_id,
prompt_id,
total_tokens,
model_id,
last_turn_usage: None,
prompt_usage: None,
cancellation_category: None,
cancel_trigger: None,
structured_output: None,
}
}
#[test]
fn includes_baseline_keys_without_usage() {
let meta = build_prompt_response_meta(args("sess-1", "prompt-1", 42_000, "grok-4.5"));
assert_eq!(meta["sessionId"], "sess-1");
assert_eq!(meta["requestId"], "prompt-1");
assert_eq!(meta["promptId"], "prompt-1");
assert_eq!(meta["totalTokens"], 42_000);
assert_eq!(meta["modelId"], "grok-4.5");
// No per-turn keys when usage is absent.
assert!(meta.get("inputTokens").is_none());
assert!(meta.get("outputTokens").is_none());
assert!(meta.get("cachedReadTokens").is_none());
}
#[test]
fn enriches_meta_with_camelcase_token_keys() {
let usage = TokenUsage {
prompt_tokens: 1500,
completion_tokens: 200,
total_tokens: 1700,
reasoning_tokens: 75,
cached_prompt_tokens: 1000,
};
let meta = build_prompt_response_meta(PromptResponseMetaArgs {
last_turn_usage: Some(&usage),
..args("sess-1", "prompt-1", 1_700, "grok-4.5")
});
// Bot's _META_TOKEN_KEY_MAP expects exactly these camelCase keys.
assert_eq!(meta["inputTokens"], 1500);
assert_eq!(meta["outputTokens"], 200);
assert_eq!(meta["cachedReadTokens"], 1000);
// Reasoning tokens carried through for diagnostic visibility.
assert_eq!(meta["reasoningTokens"], 75);
}
#[test]
fn preserves_zero_token_values() {
// Responses API hits with no cache return cached_prompt_tokens=0.
// The key is still emitted as 0 so the bot can distinguish "no cache
// hit" from "no usage data". (The bot's _merge_meta_usage requires
// the key to be present and integer-typed.)
let usage = TokenUsage {
prompt_tokens: 100,
completion_tokens: 10,
total_tokens: 110,
reasoning_tokens: 0,
cached_prompt_tokens: 0,
};
let meta = build_prompt_response_meta(PromptResponseMetaArgs {
last_turn_usage: Some(&usage),
..args("s", "p", 110, "m")
});
assert_eq!(meta["cachedReadTokens"], 0);
assert_eq!(meta["reasoningTokens"], 0);
}
#[test]
fn usage_object_lands_on_meta() {
let mut ledger = kigi_chat_state::UsageLedger::default();
ledger.record_main_loop_call(
"m",
&TokenUsage {
prompt_tokens: 100,
completion_tokens: 10,
total_tokens: 999_999,
reasoning_tokens: 0,
cached_prompt_tokens: 0,
},
None,
None,
);
let meta = build_prompt_response_meta(PromptResponseMetaArgs {
prompt_usage: Some(crate::extensions::notification::PromptUsage::from(&ledger)),
..args("s", "p", 110, "m")
});
assert_eq!(meta["usage"]["totalTokens"], 110);
assert_eq!(meta["usage"]["modelUsage"]["m"]["inputTokens"], 100);
assert!(
build_prompt_response_meta(args("s", "p", 0, "m"))
.get("usage")
.is_none()
);
}
#[test]
fn cancel_trigger_lands_as_camelcase_meta_key() {
// A send-now cancelled turn's PromptResponse `_meta` carries `cancelTrigger: "send_now"`.
let meta = build_prompt_response_meta(PromptResponseMetaArgs {
cancel_trigger: Some("send_now".to_string()),
..args("s", "p", 0, "m")
});
assert_eq!(meta["cancelTrigger"], "send_now");
// Absent for non-cancel completions — the key must not appear.
let none = build_prompt_response_meta(args("s", "p", 0, "m"));
assert!(none.get("cancelTrigger").is_none());
}
#[test]
fn structured_output_maps_to_camelcase_meta_keys() {
// Success carries the validated value under `structuredOutput`; no error key.
let ok = build_prompt_response_meta(PromptResponseMetaArgs {
structured_output: Some(Ok(serde_json::json!({"name": "ada"}))),
..args("s", "p", 0, "m")
});
assert_eq!(ok["structuredOutput"]["name"], "ada");
assert!(ok.get("structuredOutputError").is_none());
// Failure carries the message under `structuredOutputError`; no value key.
let err = build_prompt_response_meta(PromptResponseMetaArgs {
structured_output: Some(Err("output does not match the required schema".to_string())),
..args("s", "p", 0, "m")
});
assert_eq!(
err["structuredOutputError"],
"output does not match the required schema"
);
assert!(err.get("structuredOutput").is_none());
// No schema requested → neither key present.
let none = build_prompt_response_meta(args("s", "p", 0, "m"));
assert!(none.get("structuredOutput").is_none());
assert!(none.get("structuredOutputError").is_none());
}
@@ -0,0 +1,406 @@
//! Session lifecycle, roster deltas, and the idle-session supervisor for [`MvpAgent`].
//! Co-located `#[path]`-style child of `mvp_agent` (`use super::*`) so the `impl`
//! block keeps access to `MvpAgent`'s private fields.
use super::*;
impl MvpAgent {
/// Ask a live session actor to shut down.
pub(crate) fn request_session_shutdown(&self, id: &acp::SessionId) {
if let Some(handle) = self.sessions.borrow().get(id) {
let _ = handle.cmd_tx.send(SessionCommand::Shutdown);
}
}
/// Finalize the cloud session replica (fire-and-forget, "Hook 4").
///
/// Marks the session **done** upstream, so this MUST only run on a genuine
/// session end — a terminal/explicit close (`x.ai/session/close`). It must
/// NOT run on a mere client disconnect or a dead-actor reap: those leave the
/// conversation resumable on disk, and finalizing would wrongly mark a still
/// running/resumable session "done".
pub(super) fn finalize_session_replica(&self, id: &acp::SessionId) {
#[cfg(test)]
self.finalize_spy.borrow_mut().push(id.0.to_string());
if let Some(client) = self.session_registry_client() {
let sid = id.0.to_string();
tokio::spawn(async move {
if let Err(e) = client.finalize(&sid).await {
tracing::warn!(
error = % e, "session registry finalize failed (non-fatal)"
);
}
});
}
}
/// Remove a session and its thread handle **without** finalizing the cloud
/// replica.
///
/// Used for dead-actor reaping and idle-unload: the conversation stays
/// resumable on disk, so it must NOT be marked "done" upstream. Genuine
/// terminal closes go through [`MvpAgent::close_session_explicit`]. Also
/// drops the `session_live_state` entry so that map stays bounded.
pub(crate) fn remove_session(&self, id: &acp::SessionId) {
self.sessions.borrow_mut().remove(id);
self.prompt_intake_locks.borrow_mut().remove(id);
self.session_threads.borrow_mut().remove(id);
self.session_index_claims.borrow_mut().remove(id);
self.require_gateway_sessions.borrow_mut().remove(id);
self.session_live_state.borrow_mut().remove(id);
}
/// Get-or-create the per-session prompt-intake lock (see
/// [`Self::prompt_intake_locks`]). Cheap clone of the shared `Rc`.
pub(super) fn prompt_intake_lock(
&self,
id: &acp::SessionId,
) -> std::rc::Rc<tokio::sync::Mutex<()>> {
self.prompt_intake_locks
.borrow_mut()
.entry(id.clone())
.or_default()
.clone()
}
/// Close a session in response to an **explicit** terminal close
/// (`x.ai/session/close`). Finalizes the cloud replica (genuine session
/// end), then removes the session terminally as `Completed`.
pub(crate) fn close_session_explicit(&self, id: &acp::SessionId) {
self.finalize_session_replica(id);
self.remove_session_terminal(id, SessionLiveState::Completed);
}
/// Record the coarse lifecycle state for a session.
pub(super) fn set_session_live_state(&self, id: &acp::SessionId, state: SessionLiveState) {
self.session_live_state
.borrow_mut()
.insert(id.clone(), state);
}
/// Read the recorded lifecycle state for a session (test observability).
#[cfg(test)]
pub(super) fn session_live_state_for(&self, id: &acp::SessionId) -> Option<SessionLiveState> {
self.session_live_state.borrow().get(id).copied()
}
/// Roster-delta hook for a terminally removed session. Broadcasts an
/// `x.ai/sessions/changed` notification with the session in `removed` so
/// every attached dashboard drops the row promptly. Also
/// records the call site (and the terminal state) for test observability,
/// since the `session_live_state` entry is dropped on removal.
pub(super) fn record_roster_delta(&self, id: &acp::SessionId, final_state: SessionLiveState) {
#[cfg(test)]
self.roster_delta_spy
.borrow_mut()
.push((id.0.to_string(), final_state));
tracing::debug!(
session_id = % id.0, ? final_state, "roster delta: session removed"
);
self.emit_roster_changed(Vec::new(), vec![id.0.to_string()]);
}
/// Roster-delta hook for a newly-resident / changed session. Broadcasts an
/// `x.ai/sessions/changed` notification with the current entry in
/// `upserted` so dashboards add/refresh the row.
pub(crate) fn push_roster_delta_upserted(&self, id: &acp::SessionId) {
if let Some(entry) = self.resident_roster_entry(id) {
self.emit_roster_changed(vec![entry], Vec::new());
}
}
/// Emit an `x.ai/sessions/changed` upsert for a resident session with an
/// explicit `activity`, so every attached dashboard reflects a
/// turn-boundary transition (Working / Idle / NeedsInput) *immediately*
/// rather than waiting for the ≤1s roster poll (deltas are emitted
/// at turn-start/turn-end). Without this, a viewer client that holds no
/// local `AgentView` for the session only learns its activity from the
/// poll, so a turn driven by another client shows as `Idle` for up to a
/// poll interval — and not at all while that viewer's poll is dormant.
///
/// The `activity` is supplied by the caller rather than read from
/// `resident_activity` because at turn-start the actor may not have
/// published `current_prompt_id` yet (it is set asynchronously once the
/// actor dequeues the `SessionCommand::Prompt`), so a natural read would
/// still observe `Idle`. The authoritative entry (cwd / worktree / model /
/// yolo) is built by `resident_roster_entry`, so it never diverges from
/// the polled entry; only the `activity` field is overridden.
pub(super) fn push_roster_activity_delta(
&self,
id: &acp::SessionId,
activity: crate::agent::roster::RosterActivity,
) {
if let Some(mut entry) = self.resident_roster_entry(id) {
entry.activity = activity;
self.emit_roster_changed(vec![entry], Vec::new());
}
}
/// Fan an `x.ai/sessions/changed` delta out to every attached client.
///
/// This is a roster-wide notification (no `sessionId`), so the leader IPC
/// server broadcasts it to all clients rather than routing by session (see
/// the `x.ai/sessions/changed` special-case in `leader/server.rs`).
pub(super) fn emit_roster_changed(
&self,
upserted: Vec<crate::agent::roster::RosterEntry>,
removed: Vec<String>,
) {
if upserted.is_empty() && removed.is_empty() {
return;
}
let payload = crate::agent::roster::RosterChanged { upserted, removed };
if let Ok(params) = serde_json::value::to_raw_value(&payload) {
self.gateway
.forward_fire_and_forget(acp::ExtNotification::new(
crate::agent::roster::SESSIONS_CHANGED_METHOD,
params.into(),
));
}
}
/// Coarse activity of a resident session for the dashboard status column.
///
/// Precedence: a non-empty pending-interaction map →
/// `NeedsInput` (wins even over a running turn — a session awaiting a
/// permission *mid-turn* is "needs input"); else a running turn →
/// `Working`; else map the coarse `SessionLiveState`.
pub(super) fn resident_activity(
&self,
id: &acp::SessionId,
) -> crate::agent::roster::RosterActivity {
use crate::agent::roster::RosterActivity;
let (needs_input, turn_running) = self
.sessions
.borrow()
.get(id)
.map(|h| {
let needs_input = h
.pending_interactions
.lock()
.map(|g| !g.is_empty())
.unwrap_or(false);
let turn_running = h
.current_prompt_id
.lock()
.map(|g| g.is_some())
.unwrap_or(false);
(needs_input, turn_running)
})
.unwrap_or((false, false));
if needs_input {
return RosterActivity::NeedsInput;
}
if turn_running {
return RosterActivity::Working;
}
match self.session_live_state.borrow().get(id).copied() {
Some(SessionLiveState::Completed) => RosterActivity::Completed,
Some(SessionLiveState::DeadFailed) => RosterActivity::Dead,
Some(SessionLiveState::Dormant) => RosterActivity::Dormant,
_ => RosterActivity::Idle,
}
}
/// Build a single roster entry for a resident session, or `None` if it is
/// not currently resident.
pub(super) fn resident_roster_entry(
&self,
id: &acp::SessionId,
) -> Option<crate::agent::roster::RosterEntry> {
let session_id = id.0.to_string();
let (cwd, is_worktree, model_id, reasoning_effort, yolo) = {
let sessions = self.sessions.borrow();
let h = sessions.get(id)?;
(
h.display_cwd.clone().unwrap_or_else(|| h.info.cwd.clone()),
h.display_cwd.is_some(),
Some(h.model_id.0.to_string()),
h.reasoning_effort,
h.yolo_mode,
)
};
Some(crate::agent::roster::RosterEntry {
title: self
.resident_roster_titles
.borrow()
.get(&session_id)
.cloned(),
session_id,
cwd,
is_worktree,
model_id,
reasoning_effort,
yolo,
activity: self.resident_activity(id),
resident: true,
last_change_unix_ms: chrono::Utc::now().timestamp_millis(),
origin: crate::agent::roster::RosterOrigin::Local,
})
}
/// Snapshot all resident sessions as roster entries (synchronous; no disk).
pub(super) fn resident_roster_entries(&self) -> Vec<crate::agent::roster::RosterEntry> {
let ids: Vec<acp::SessionId> = self.sessions.borrow().keys().cloned().collect();
ids.iter()
.filter_map(|id| self.resident_roster_entry(id))
.collect()
}
/// Build the full roster: resident actors plus recently-touched on-disk
/// (`Dormant`) sessions. Resident wins on an id collision; hidden sessions
/// are excluded.
pub(crate) async fn build_roster(&self) -> Vec<crate::agent::roster::RosterEntry> {
let resident = self.resident_roster_entries();
let summaries = crate::session::persistence::list_recent_summaries(200)
.await
.unwrap_or_default();
let entries = crate::agent::roster::merge_roster(resident, summaries);
self.cache_resident_titles(&entries);
entries
}
/// Refresh `resident_roster_titles` from the freshly-built roster.
pub(super) fn cache_resident_titles(&self, entries: &[crate::agent::roster::RosterEntry]) {
*self.resident_roster_titles.borrow_mut() = entries
.iter()
.filter(|e| e.resident)
.filter_map(|e| Some((e.session_id.clone(), e.title.clone()?)))
.collect();
}
/// Terminally remove a session: emit the roster delta with its final state,
/// then drop it from all maps (no finalize — callers that need finalize do
/// it first, see `close_session_explicit`).
pub(super) fn remove_session_terminal(
&self,
id: &acp::SessionId,
final_state: SessionLiveState,
) {
self.record_roster_delta(id, final_state);
self.remove_session(id);
}
/// Reap a session whose **resident** actor thread exited unexpectedly
/// (panic / load failure). Demotes it to `DeadFailed`, emits the roster
/// delta, and removes it WITHOUT finalize — the conversation persists on
/// disk and stays resumable (reaping a dead actor is harmless;
/// it demotes to Dormant).
pub(super) fn reap_dead_session(&self, id: &acp::SessionId) {
self.remove_session_terminal(id, SessionLiveState::DeadFailed);
}
/// Sweep `session_threads` for finished threads and clean them up.
///
/// A finished thread has two distinct meanings, and conflating them
/// corrupts the `SessionLiveState` roster source:
///
/// - **Still resident in `sessions`** → the actor exited unexpectedly while
/// the session was hosted (panic / load failure). Reap as `DeadFailed`.
/// - **Not resident** (already idle-unloaded → `Dormant`, or explicitly
/// closed) → this is the *expected* clean exit. The `SessionThread` was
/// kept only so `drain_old_session_thread` could wait on it; now that it
/// has finished there is nothing left to drain, so just drop the leftover
/// `SessionThread`/state entries. Do **not** demote to `DeadFailed` and do
/// **not** emit a second roster delta.
///
/// `JoinHandle::is_finished()` is non-blocking and cannot distinguish a
/// clean exit from a panic on its own, which is exactly why the residency
/// check is required. Runs both opportunistically and from the join-handle
/// supervisor (`ensure_session_supervisor`).
pub(super) fn sweep_dead_sessions(&self) {
let dead: Vec<acp::SessionId> = self
.session_threads
.borrow()
.iter()
.filter(|(_, t)| t.is_finished())
.map(|(id, _)| id.clone())
.collect();
for id in dead {
if self.sessions.borrow().contains_key(&id) {
tracing::warn!(
session_id = % id.0,
"Resident session actor exited unexpectedly; reaping as DeadFailed"
);
self.reap_dead_session(&id);
} else {
self.session_threads.borrow_mut().remove(&id);
self.session_live_state.borrow_mut().remove(&id);
tracing::debug!(
session_id = % id.0,
"Reaped finished thread for non-resident session (clean exit)"
);
}
}
}
/// Start the join-handle supervisor. **Idempotent.**
///
/// A single `spawn_local` task periodically reaps actor threads that have
/// exited (panicked or finished) so a dead actor never lingers as a roster
/// zombie. `std::thread::JoinHandle` is not awaitable, so we poll
/// `is_finished()` on a tick — the same mechanism `drain_old_session_thread`
/// and `sweep_dead_sessions` already use. A panicked actor is therefore
/// reaped within one [`SESSION_SUPERVISOR_TICK`].
///
/// The sweep body is wrapped in `catch_unwind` so a single panicking sweep
/// can never terminate the loop (which would silently disable reaping for
/// the rest of the process). The task holds a `LocalRef` (raw pointer) to
/// `self` for the lifetime of the `LocalSet`; this is sound because the
/// agent owns the `LocalSet` and outlives it (same contract as
/// `start_subagent_coordinator`), and `LocalRef` is `!Send`.
pub(super) fn ensure_session_supervisor(&self) {
if self.supervisor_started.replace(true) {
return;
}
#[cfg(test)]
self.supervisor_spawn_count
.set(self.supervisor_spawn_count.get() + 1);
let agent_ref = LocalRef::new(self);
tokio::task::spawn_local(async move {
loop {
tokio::time::sleep(SESSION_SUPERVISOR_TICK).await;
let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
agent_ref.get().sweep_dead_sessions();
}));
if result.is_err() {
tracing::error!("session supervisor sweep panicked; continuing supervision");
}
}
});
}
/// Coarse "any work pending" check for the idle-unload stub.
/// Returns `true` while the session has work in flight.
///
/// Three layers:
/// 1. **Fast path (sync):** the shared `current_prompt_id` slot, which the
/// actor sets while a turn is running (`maybe_start_running_task`) and
/// clears via its RAII guard. A poisoned lock is treated as busy → never
/// unload.
/// 1b. **Parked plan-approval (sync):** the shared `pending_interactions`
/// slot. The parked plan-approval resume re-park is the one outstanding work with no
/// running turn, so it needs its own sync check (the same shared-`Arc`
/// idiom as `current_prompt_id`) rather than the async round-trip below.
/// 2. **Queue check (async):** when no turn is running, the actor is between
/// turns and responsive, so we ask it whether `pending_inputs` is
/// non-empty (a prompt queued at the turn boundary). This closes the
/// sub-tick window where `current_prompt_id` is momentarily `None` but a
/// queued input is about to be drained. On timeout we keep the session
/// resident (conservative).
///
/// TODO(PR-4): once the aggregate `SessionActivity` signal exists, also
/// consult the autonomous background sources so a detached session is never
/// idle-unloaded (→ `Shutdown` → `KillOnDrop`) while they are live:
/// `monitor_event_buffer`, pending scheduler fires,
/// `ToolContext.background_tasks`, and background subagent sessions. Until
/// then those background-only sessions rely on the keep-resident default and
/// the `current_prompt_id` auto-wake turn being active.
///
/// TODO(PR-4): this is also inherently a *check-then-act* across the
/// actor-thread boundary — work can arrive (a new `Prompt`/auto-wake) in the
/// gap between this `IsBusy` answer and the caller's subsequent `Shutdown`,
/// so an idle-unload can still race a just-arrived turn. The actor processes
/// its mailbox in order, so the lost work is bounded and recoverable on
/// reload; PR-4 closes the gap properly by gating the unload inside the
/// actor (a single `Unload`-if-idle command) rather than check-then-send.
pub(super) async fn session_has_live_work(&self, id: &acp::SessionId) -> bool {
let Some(handle) = self.sessions.borrow().get(id).cloned() else {
return false;
};
let turn_running = handle
.current_prompt_id
.lock()
.map(|g| g.is_some())
.unwrap_or(true);
if turn_running {
return true;
}
if crate::session::pending_interaction::has_parked_plan_approval(
&handle.pending_interactions,
) {
return true;
}
tokio::time::timeout(IDLE_QUERY_TIMEOUT, handle.is_busy())
.await
.unwrap_or(true)
}
}
@@ -0,0 +1,552 @@
//! Subagent coordinator drain task and spawn-context construction for [`MvpAgent`].
//! Co-located child of `mvp_agent` (`use super::*`); tested by `tests/subagent_spawn_context_tests.rs`.
use super::*;
impl MvpAgent {
/// Start the subagent coordinator drain task.
///
/// Takes the `subagent_event_rx` receiver (once) and spawns a `spawn_local` task
/// that receives `SubagentRequest`s and delegates each to
/// `handle_subagent_request()` on its own `spawn_local` task.
///
/// Uses `LocalRef` to reference `self` from
/// `spawn_local` closures. Idempotent: subsequent calls are no-ops.
pub(super) fn start_subagent_coordinator(&self) {
let Some(mut rx) = self.subagent_event_rx.borrow_mut().take() else {
return;
};
let agent_ref = LocalRef::new(self);
use crate::agent::subagent::{BlockWaitSlot, is_running, resolve_snapshot};
use kigi_tools::implementations::grok_build::task::types::{
SubagentCancelOutcome, SubagentCancelTarget, SubagentEvent,
};
tokio::task::spawn_local({
let agent_ref = agent_ref.clone();
async move {
while let Some(event) = rx.recv().await {
match event {
SubagentEvent::Spawn(boxed) => {
let request = *boxed;
let agent_ref = agent_ref.clone();
tokio::task::spawn_local(async move {
let this = agent_ref.get();
let parent_sid = request.parent_session_id.clone();
let mut ctx = this.build_subagent_spawn_context(&parent_sid);
let parent_handle = {
let parent_sid_acp = acp::SessionId::new(parent_sid.clone());
this.sessions.borrow().get(&parent_sid_acp).cloned()
};
if let Some(handle) = parent_handle {
ctx.parent_mcp_pool = handle.snapshot_mcp_pool().await;
ctx.client_hooks = handle.snapshot_client_hooks().await;
let parent_tools = handle.snapshot_tool_definitions().await;
ctx.parent_tool_snapshot =
(!parent_tools.is_empty()).then_some(parent_tools);
}
crate::agent::subagent::handle_subagent_request(
request,
ctx,
&this.subagent_coordinator,
&this.gateway,
)
.await;
});
}
SubagentEvent::Query(query) => {
let agent_ref = agent_ref.clone();
tokio::task::spawn_local(async move {
let subagent_id = query.subagent_id;
let block = query.block;
let timeout_ms = query.timeout_ms;
let slot: BlockWaitSlot = std::rc::Rc::new(
std::cell::RefCell::new(Some(query.respond_to)),
);
let send_via_slot =
|slot: &BlockWaitSlot, snap| match slot.borrow_mut().take() {
Some(tx) => tx.send(snap).is_ok(),
None => false,
};
let lookup = {
let this = agent_ref.get();
let result =
this.subagent_coordinator.borrow().lookup(&subagent_id);
if block && result.is_some() {
this.subagent_coordinator
.borrow_mut()
.register_block_wait(&subagent_id, slot.clone());
}
this.subagent_coordinator
.borrow_mut()
.evict_stale_completed();
result
};
let snapshot = resolve_snapshot(lookup).await;
let should_block =
block && snapshot.as_ref().is_some_and(is_running);
if should_block {
let timeout_ms = timeout_ms.unwrap_or(30_000);
let deadline = tokio::time::Instant::now()
+ tokio::time::Duration::from_millis(timeout_ms);
loop {
tokio::time::sleep(tokio::time::Duration::from_millis(200))
.await;
let receiver_gone =
slot.borrow().as_ref().is_none_or(|tx| tx.is_closed());
if receiver_gone {
let this = agent_ref.get();
let mut coord = this.subagent_coordinator.borrow_mut();
coord.clear_block_waited(&subagent_id);
coord.unregister_block_wait(&subagent_id, &slot);
return;
}
let lookup = {
let this = agent_ref.get();
this.subagent_coordinator.borrow().lookup(&subagent_id)
};
let snap = resolve_snapshot(lookup).await;
let still_running = snap.as_ref().is_some_and(is_running);
if !still_running || tokio::time::Instant::now() >= deadline
{
{
let this = agent_ref.get();
let mut coord =
this.subagent_coordinator.borrow_mut();
if still_running {
coord.clear_block_waited(&subagent_id);
}
coord.unregister_block_wait(&subagent_id, &slot);
}
if !send_via_slot(&slot, snap) && !still_running {
let this = agent_ref.get();
this.subagent_coordinator
.borrow_mut()
.clear_block_waited(&subagent_id);
}
return;
}
}
} else {
let delivered = send_via_slot(&slot, snapshot);
if block {
let this = agent_ref.get();
let mut coord = this.subagent_coordinator.borrow_mut();
coord.unregister_block_wait(&subagent_id, &slot);
if !delivered {
coord.clear_block_waited(&subagent_id);
}
}
}
});
}
SubagentEvent::Cancel(request) => {
let this = agent_ref.get();
let outcome = {
let mut coord = this.subagent_coordinator.borrow_mut();
match request.target {
SubagentCancelTarget::SubagentId(ref subagent_id) => {
coord.mark_explicitly_killed(subagent_id);
coord.cancel_with_outcome(subagent_id)
}
SubagentCancelTarget::ParentPromptId(ref parent_prompt_id) => {
coord.cancel_by_parent_prompt_id(parent_prompt_id);
SubagentCancelOutcome::Cancelled
}
}
};
let _ = request.respond_to.send(outcome);
}
SubagentEvent::ListActive(request) => {
let this = agent_ref.get();
let summaries = this
.subagent_coordinator
.borrow()
.active_summaries_for(&request.parent_session_id);
let _ = request.respond_to.send(summaries);
}
SubagentEvent::Completions(request) => {
let this = agent_ref.get();
let mut completions = this
.subagent_coordinator
.borrow_mut()
.drain_pending_completions();
completions.retain(|c| !request.suppress_ids.contains(&c.subagent_id));
let _ = request.respond_to.send(completions);
}
SubagentEvent::Outstanding(request) => {
let this = agent_ref.get();
let reply = this
.subagent_coordinator
.borrow()
.outstanding_reply_for_prompt(&request.prompt_id);
let _ = request.respond_to.send(reply);
}
SubagentEvent::ClearUsageNotApplied(request) => {
let this = agent_ref.get();
this.subagent_coordinator
.borrow_mut()
.clear_subagent_usage_not_applied(&request.prompt_id);
}
SubagentEvent::MarkUsageNotApplied(request) => {
let this = agent_ref.get();
this.subagent_coordinator
.borrow_mut()
.mark_subagent_usage_not_applied(&request.prompt_id);
let _ = request.respond_to.send(());
}
SubagentEvent::ValidateType(request) => {
let agent_ref = agent_ref.clone();
tokio::task::spawn_local(async move {
let this = agent_ref.get();
let ctx = this
.build_subagent_validation_context(&request.parent_session_id);
let outcome = crate::agent::subagent::validate_subagent_type(
&request.subagent_type,
&ctx,
);
let _ = request.respond_to.send(outcome);
});
}
SubagentEvent::DescribeType(request) => {
let agent_ref = agent_ref.clone();
tokio::task::spawn_local(async move {
use kigi_tools::implementations::grok_build::task::types::SubagentDescribeOutcome;
let this = agent_ref.get();
let outcome = match this
.try_build_subagent_spawn_context(&request.parent_session_id)
{
Some(ctx) => crate::agent::subagent::describe_subagent_type(
&request.subagent_type,
request.harness_agent_type.as_deref(),
&ctx,
),
None => {
tracing::warn!(
parent_session_id = % request.parent_session_id,
subagent_type = % request.subagent_type,
"DescribeType for unknown/evicted parent session, replying Unavailable",
);
SubagentDescribeOutcome::Unavailable
}
};
let _ = request.respond_to.send(outcome);
});
}
}
}
}
});
}
/// Lightweight context for the `SubagentEvent::ValidateType` drain arm;
/// tolerates evicted parent sessions (returns built-in defaults + warns).
pub(super) fn build_subagent_validation_context(
&self,
parent_session_id: &str,
) -> crate::agent::subagent::SubagentValidationContext {
let parent_sid = acp::SessionId::new(parent_session_id);
let (parent_cwd, allowed_subagent_types) = {
let sessions = self.sessions.borrow();
let ps = sessions.get(&parent_sid);
warn_on_missing_parent_session_for_validate_type(parent_session_id, ps.is_some());
(
ps.map(|h| std::path::PathBuf::from(&h.info.cwd))
.unwrap_or_default(),
ps.and_then(|h| h.allowed_subagent_types.clone()),
)
};
let cli_agent_names: Vec<String> = {
let cfg = self.cfg.borrow();
cfg.cli_agents.iter().map(|d| d.name.clone()).collect()
};
crate::agent::subagent::SubagentValidationContext {
parent_cwd,
plugin_registry: self.plugin_registry_handle.snapshot(),
subagent_toggle: self.subagent_toggle.clone(),
allowed_subagent_types,
cli_agent_names,
}
}
/// Build a `SubagentSpawnContext` from the current agent state and the
/// parent session's shared resources.
///
/// This is the ONLY subagent-related method on MvpAgent besides the
/// coordinator startup.
/// Build a spawn context for a real subagent spawn. The parent session is
/// guaranteed present here because the parent just issued the spawn request,
/// so a missing parent is a real invariant violation and panics. Read-only
/// callers that can race a parent teardown (e.g. `DescribeType`) must use
/// [`Self::try_build_subagent_spawn_context`] instead.
pub(super) fn build_subagent_spawn_context(
&self,
parent_session_id: &str,
) -> crate::agent::subagent::SubagentSpawnContext {
self.try_build_subagent_spawn_context(parent_session_id)
.expect("parent session must exist when spawning subagents")
}
/// Fallible variant of [`Self::build_subagent_spawn_context`]: returns
/// `None` when the parent `SessionHandle` is absent (evicted / torn down)
/// instead of panicking, so read-only paths that can race a teardown can
/// fail open.
pub(super) fn try_build_subagent_spawn_context(
&self,
parent_session_id: &str,
) -> Option<crate::agent::subagent::SubagentSpawnContext> {
let parent_sid = acp::SessionId::new(parent_session_id);
let (
parent_model_id,
parent_chat_state,
parent_cmd_tx,
parent_cwd,
yolo_mode,
parent_depth,
hunk_tracker_handle,
hunk_tracking_enabled,
fs,
terminal,
session_env,
parent_attribution_callback,
parent_agent_name,
parent_managed_mcp_proxy_base_url,
) = {
let sessions = self.sessions.borrow();
let ps = sessions.get(&parent_sid);
(
ps.map(|h| h.model_id.clone())
.unwrap_or_else(|| self.models_manager.current_model_id()),
ps.map(|h| h.chat_state_handle.clone()),
ps.map(|h| h.cmd_tx.clone()),
ps.map(|h| std::path::PathBuf::from(&h.info.cwd))
.unwrap_or_default(),
ps.map(|h| h.yolo_mode).unwrap_or(self.default_yolo_mode),
ps.map(|h| h.tool_context.subagent_depth).unwrap_or(0),
ps.map(|h| h.tool_context.hunk_tracker_handle.clone())
.unwrap_or_else(kigi_hunk_tracker::HunkTrackerHandle::noop),
ps.map(|h| h.tool_context.hunk_tracking_enabled)
.unwrap_or(false),
ps.map(|h| h.tool_context.fs.inner().clone())
.unwrap_or_else(|| {
let cwd = ps
.map(|h| std::path::PathBuf::from(&h.info.cwd))
.unwrap_or_default();
std::sync::Arc::new(kigi_workspace::file_system::LocalFs::new(cwd))
}),
ps.map(|h| h.tool_context.terminal.clone())
.unwrap_or_else(|| {
std::sync::Arc::new(crate::terminal::TerminalRunner::new(
std::sync::Arc::new(self.gateway.clone()),
parent_sid.clone(),
))
}),
ps.map(|h| h.tool_context.session_env.clone())
.unwrap_or_else(|| std::sync::Arc::new(std::collections::HashMap::new())),
ps.and_then(|h| h.attribution_callback.clone()),
ps.map(|h| h.agent_name.clone()),
ps.map(|h| h.managed_mcp_proxy_base_url.clone()),
)
};
let (
parent_workspace_ops,
parent_terminal_backend,
parent_notification_handle,
parent_scheduler_handle,
) = {
let sessions = self.sessions.borrow();
sessions.get(&parent_sid).map(|ps| {
(
ps.workspace_ops.clone(),
ps.terminal_backend.clone(),
ps.tools_notification_handle.clone(),
ps.scheduler_handle.clone(),
)
})
}?;
let available_models = self.models_manager.models();
let parent_lsp = {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.and_then(|h| h.tool_context.lsp.clone())
};
let am = self.auth_manager.clone();
let inference_idle_timeout_secs = {
let per_model = config::find_model_by_id(&available_models, parent_model_id.0.as_ref())
.and_then(|e| e.info.inference_idle_timeout_secs);
let cfg = self.cfg.borrow();
let remote = cfg
.remote_settings
.as_ref()
.and_then(|s| s.inference_idle_timeout_secs);
per_model.or(remote).unwrap_or(600).max(10)
};
let parent_hook_registry = {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.and_then(|h| h.hook_registry.clone())
};
let parent_max_turns = {
let sessions = self.sessions.borrow();
sessions.get(&parent_sid).and_then(|h| h.max_turns)
};
let parent_model_agent_type =
config::find_model_by_id(&available_models, parent_model_id.0.as_ref())
.map(|e| e.info.agent_type.clone());
let ask_user_question_enabled = {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.ask_user_question_enabled)
.unwrap_or_else(|| self.cfg.borrow().resolve_ask_user_question().value)
};
Some(crate::agent::subagent::SubagentSpawnContext {
lsp: parent_lsp,
gateway: self.gateway.clone(),
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()),
alpha_test_key: self.alpha_test_key(),
auth_method_id: self
.auth_method_id
.load()
.as_deref()
.cloned()
.unwrap_or_else(|| acp::AuthMethodId::new("default")),
model_id: parent_model_id,
storage_mode: self.storage_mode,
auth: self.current_or_buffered_auth(),
parent_cwd: parent_cwd.clone(),
parent_session_id: parent_session_id.to_string(),
yolo_mode,
subagent_event_tx: self.subagent_event_tx.clone(),
parent_depth,
inference_idle_timeout_secs,
auto_compact_threshold_tiers:
crate::agent::subagent::AutoCompactThresholdTiers::capture(&self.cfg.borrow()),
hunk_tracker_handle,
hunk_tracking_enabled,
fs,
terminal,
session_env,
memory_config: self.memory_config.clone(),
web_search_sampling_config: self.prepare_web_search_sampling_config(),
web_fetch_config: self.prepare_web_fetch_config(),
image_gen_config: self.prepare_image_gen_config(),
video_gen_config: self.prepare_video_gen_config(),
app_builder_deployer_config: self.prepare_app_builder_deployer_config(),
write_file_enabled: self.cfg.borrow().resolve_write_file().value,
goal_enabled: self.cfg.borrow().resolve_goal().value,
ask_user_question_enabled,
parent_cmd_tx: parent_cmd_tx.clone(),
parent_session_info: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| crate::session::info::Info {
id: parent_sid.clone(),
cwd: h.info.cwd.clone(),
})
},
parent_chat_state,
parent_max_turns,
available_models,
subagent_model_overrides: self.subagent_model_overrides.clone(),
subagent_toggle: self.subagent_toggle.clone(),
subagent_roles: self.subagent_roles.clone(),
subagent_personas: self.subagent_personas.clone(),
persona_io_summaries: self.persona_io_summaries.clone(),
disable_web_search: self.cfg.borrow().disable_web_search,
todo_gate: self.cfg.borrow().todo_gate,
remote_settings: self.cfg.borrow().remote_settings.clone(),
laziness_debug_log: self.cfg.borrow().laziness_debug_log.clone(),
backend_tools_enabled: self.cfg.borrow().resolve_backend_tools().value,
respect_gitignore: self.cfg.borrow().respect_gitignore,
path_not_found_hints: self.cfg.borrow().path_not_found_hints,
plugin_registry: self.plugin_registry_handle.snapshot(),
models_manager: self.models_manager.clone(),
file_tool_overrides: {
let cfg = self.cfg.borrow();
let effective = cfg
.toolset
.resolve_file_toolset(cfg.remote_settings.as_ref());
if effective != crate::tools::FileToolset::Standard {
effective.tool_configs(&cfg.toolset.hashline).ok()
} else {
None
}
},
agent_config: Some(self.cfg.borrow().clone()),
hook_registry: parent_hook_registry,
hook_workspace_root: String::new(),
permission_handle: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.permission_handle.clone())
},
worktree_type: self.worktree_type,
api_key_provider: Some(Arc::new(crate::auth::manager::SharedAuthKeyProvider(
am.clone(),
))),
image_description_model: self.resolve_image_description_model(),
workspace_ops: parent_workspace_ops.clone(),
auth_manager: am.clone(),
attribution_callback: parent_attribution_callback,
parent_agent_name,
parent_model_agent_type,
allowed_subagent_types: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.and_then(|h| h.allowed_subagent_types.clone())
},
parent_mcp_configs: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.mcp_servers.clone())
.unwrap_or_default()
},
managed_mcp_state: self.managed_mcp_cache.clone(),
parent_mcp_pool: None,
parent_tool_snapshot: None,
parent_skills: None,
parent_skills_config: self.cfg.borrow().skills.clone(),
parent_compat: self.cfg.borrow().compat_resolved,
auto_wake_delivered: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.and_then(|h| h.tool_context.auto_wake_delivered.clone())
},
task_output_tool_name: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.tool_context.task_output_tool_name.clone())
.unwrap_or_else(|| {
kigi_tools::reminders::task_completion::DEFAULT_TASK_OUTPUT_TOOL
.to_string()
})
},
auto_wake_enabled: self.cfg.borrow().auto_wake_enabled,
goal_loop_active: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.tool_context.goal_loop_active_gate.clone())
.unwrap_or_else(|| {
std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false))
})
},
parent_blocking_wait_depth: {
let sessions = self.sessions.borrow();
sessions
.get(&parent_sid)
.map(|h| h.tool_context.blocking_wait_depth.clone())
.unwrap_or_else(|| std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)))
},
parent_terminal_backend: parent_terminal_backend.clone(),
parent_notification_handle: parent_notification_handle.clone(),
parent_scheduler_handle: parent_scheduler_handle.clone(),
})
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,135 @@
//! Subagent spawn-context inheritance: a child session must inherit the parent's
//! permission handle and goal-loop gate so policy and run-state can't be bypassed
//! by delegating to a subagent.
use super::{build_minimal_agent_for_tests, make_test_handle};
use agent_client_protocol as acp;
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
/// Subagents inherit the parent permission handle, so a managed `Read(**/.env)`
/// deny still blocks the child — direct read and the `cat .env` shell equivalent.
#[tokio::test]
async fn subagent_spawn_context_inherits_parent_permission_handle() {
use kigi_workspace::permission::types::{
PatternMode, PermissionConfig, PermissionRule, RuleAction, ToolFilter,
};
let local = tokio::task::LocalSet::new();
local
.run_until(async {
let agent = build_minimal_agent_for_tests();
let sid = acp::SessionId::new("parent-permission");
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
let gateway = GatewaySender::new(tx);
let cwd = kigi_paths::AbsPathBuf::new(std::path::PathBuf::from("/tmp"))
.expect("absolute cwd");
let (permission_handle, _events_rx) =
kigi_workspace::permission::spawn_permission_manager(
sid.clone(),
gateway,
cwd,
kigi_workspace::permission::types::ClientType::Generic,
Some(PermissionConfig::new(vec![PermissionRule {
action: RuleAction::Deny,
tool: ToolFilter::Read,
pattern: Some("**/.env".to_owned()),
pattern_mode: PatternMode::Glob,
}])),
Vec::new(), // deny_read_globs
Vec::new(),
false,
None,
);
let mut handle = make_test_handle("test-model", false, None);
handle.permission_handle = permission_handle;
agent.sessions.borrow_mut().insert(sid.clone(), handle);
let ctx = agent.build_subagent_spawn_context(sid.0.as_ref());
let inherited = ctx
.permission_handle
.expect("subagent context must inherit parent permission handle");
// Direct file read and the shell equivalent both hit the parent deny.
for access in [
kigi_workspace::permission::AccessKind::Read(Some(".env".into())),
kigi_workspace::permission::AccessKind::Bash("cat .env".into()),
] {
let decision = inherited
.request(
access.clone(),
acp::ToolCallUpdate::new(acp::ToolCallId::new("tc"), Default::default()),
Some("child-session".to_owned()),
Some("general-purpose".to_owned()),
Some("permission inheritance regression".to_owned()),
)
.await;
assert!(
matches!(
decision,
kigi_workspace::permission::Decision::PolicyDeny(_)
),
"subagent-inherited handle must enforce parent deny for {access:?}, got {decision:?}"
);
}
})
.await;
}
/// A subagent shares the parent's `goal_loop_active_gate` Arc, so flipping the
/// parent gate is observed through the child context (same allocation).
#[tokio::test]
async fn subagent_spawn_context_shares_parent_goal_loop_gate() {
use std::sync::atomic::Ordering::Relaxed;
let agent = build_minimal_agent_for_tests();
let sid = acp::SessionId::new("parent-goal");
let handle = make_test_handle("test-model", false, None);
// Clone the parent's live gate before the handle moves into `sessions`.
let parent_gate = handle.tool_context.goal_loop_active_gate.clone();
agent.sessions.borrow_mut().insert(sid.clone(), handle);
let ctx = agent.build_subagent_spawn_context(sid.0.as_ref());
// Flipping the parent gate must surface through the child flag (shared Arc).
assert!(!ctx.goal_loop_active.load(Relaxed));
parent_gate.store(true, Relaxed);
assert!(
ctx.goal_loop_active.load(Relaxed),
"subagent context must observe the parent's goal-loop gate (same Arc)"
);
}
/// A subagent inherits the parent session's `ask_user_question` gate, so
/// `--no-ask-user` strips the tool from subagents too, while the default keeps it.
#[tokio::test]
async fn subagent_spawn_context_inherits_parent_ask_user_question_gate() {
let agent = build_minimal_agent_for_tests();
// Parent with the tool disabled (the `--no-ask-user` case) → child off.
let sid_off = acp::SessionId::new("parent-no-ask");
let mut handle_off = make_test_handle("test-model", false, None);
handle_off.ask_user_question_enabled = false;
agent
.sessions
.borrow_mut()
.insert(sid_off.clone(), handle_off);
let ctx_off = agent.build_subagent_spawn_context(sid_off.0.as_ref());
assert!(
!ctx_off.ask_user_question_enabled,
"subagent must inherit the parent's disabled ask_user_question gate (--no-ask-user)"
);
// Parent with the tool enabled (the default) → child on.
let sid_on = acp::SessionId::new("parent-ask");
let handle_on = make_test_handle("test-model", false, None);
agent
.sessions
.borrow_mut()
.insert(sid_on.clone(), handle_on);
let ctx_on = agent.build_subagent_spawn_context(sid_on.0.as_ref());
assert!(
ctx_on.ask_user_question_enabled,
"subagent must inherit the parent's enabled ask_user_question gate"
);
}
@@ -0,0 +1,619 @@
//! HTTP CONNECT proxy support for WebSocket connections.
//!
//! When running behind a corporate egress proxy,
//! `tokio-tungstenite`'s `connect_async` cannot reach external
//! hosts directly because it does not read the standard `HTTPS_PROXY` /
//! `HTTP_PROXY` environment variables.
//!
//! This module provides:
//! - [`resolve_proxy_for_host`]: reads proxy env vars and `NO_PROXY`, returning
//! the proxy URL to use for a given target host (or `None` for direct).
//! - [`connect_via_proxy`]: opens a TCP connection to the proxy, sends an HTTP
//! CONNECT request to create a tunnel, wraps the result in TLS, and returns a
//! stream suitable for `tokio_tungstenite::client_async`.
use std::sync::{Arc, OnceLock};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use tokio::net::TcpStream;
use tokio_tungstenite::MaybeTlsStream;
use tracing::debug;
// ---------------------------------------------------------------------------
// Environment-variable resolution
// ---------------------------------------------------------------------------
/// Read proxy configuration from the environment and decide whether `target_host`
/// should be connected through a proxy.
///
/// Resolution order (matches `curl` / `reqwest` behaviour):
/// 1. If `NO_PROXY` contains `target_host` (or a matching domain suffix / CIDR),
/// return `None`.
/// 2. If `HTTPS_PROXY` (or `https_proxy`) is set, return its value.
/// 3. If `HTTP_PROXY` (or `http_proxy`) is set, return its value.
/// 4. Otherwise return `None`.
pub fn resolve_proxy_for_host(target_host: &str) -> Option<String> {
resolve_proxy_for_host_with(target_host, |key| std::env::var(key))
}
/// Testable inner implementation that accepts a custom env-var reader.
fn resolve_proxy_for_host_with<F>(target_host: &str, env: F) -> Option<String>
where
F: for<'a> Fn(&'a str) -> Result<String, std::env::VarError>,
{
// Check NO_PROXY / no_proxy.
let no_proxy = env("NO_PROXY")
.or_else(|_| env("no_proxy"))
.unwrap_or_default();
if is_host_bypassed(target_host, &no_proxy) {
return None;
}
// HTTPS_PROXY takes precedence (our target is always wss://).
if let Ok(url) = env("HTTPS_PROXY").or_else(|_| env("https_proxy")) {
let url = url.trim().to_string();
if !url.is_empty() {
return Some(url);
}
}
// Fall back to HTTP_PROXY.
if let Ok(url) = env("HTTP_PROXY").or_else(|_| env("http_proxy")) {
let url = url.trim().to_string();
if !url.is_empty() {
return Some(url);
}
}
None
}
/// Check whether `host` is in the `no_proxy` list.
///
/// The `no_proxy` value is a comma-separated list of hostnames, domain
/// suffixes (with or without a leading dot), IP addresses, or CIDR ranges.
/// The special value `*` matches everything.
fn is_host_bypassed(host: &str, no_proxy: &str) -> bool {
let host_lower = host.to_ascii_lowercase();
for entry in no_proxy.split(',') {
let entry = entry.trim().to_ascii_lowercase();
if entry.is_empty() {
continue;
}
// Wildcard — bypass all hosts.
if entry == "*" {
return true;
}
// Exact match.
if host_lower == entry {
return true;
}
// Domain suffix match: ".example.com" matches "foo.example.com".
// Also handle the common convention of omitting the leading dot:
// "example.com" in NO_PROXY should match "sub.example.com".
let matches_suffix = if entry.starts_with('.') {
host_lower.ends_with(entry.as_str())
} else {
host_lower.len() > entry.len()
&& host_lower.ends_with(entry.as_str())
&& host_lower.as_bytes()[host_lower.len() - entry.len() - 1] == b'.'
};
if matches_suffix {
return true;
}
// CIDR / IP matching is intentionally omitted here — our target host
// is always a DNS name, not an IP literal. Keeping this simple avoids
// pulling in a CIDR parsing dependency.
}
false
}
// ---------------------------------------------------------------------------
// HTTP CONNECT tunnel
// ---------------------------------------------------------------------------
/// Establish a TLS-wrapped TCP stream through an HTTP CONNECT proxy.
///
/// Steps:
/// 1. Parse the proxy URL to get host + port.
/// 2. Open a TCP connection to the proxy and perform the CONNECT handshake.
/// 3. Wrap the tunnel in TLS (using rustls with native root certificates).
/// 4. Return the stream as `MaybeTlsStream<TcpStream>` so it is compatible
/// with `tokio_tungstenite::client_async`.
pub async fn connect_via_proxy(
proxy_url: &str,
target_host: &str,
target_port: u16,
) -> anyhow::Result<MaybeTlsStream<TcpStream>> {
let stream = open_connect_tunnel(proxy_url, target_host, target_port).await?;
let tls_stream = tls_wrap(stream, target_host).await?;
Ok(MaybeTlsStream::Rustls(tls_stream))
}
/// Open a raw TCP tunnel through an HTTP CONNECT proxy (no TLS).
///
/// 1. Parse the proxy URL to get host + port.
/// 2. Open a plain TCP connection to the proxy.
/// 3. Send `CONNECT target_host:target_port HTTP/1.1\r\n\r\n`.
/// 4. Read the proxy's response; expect `HTTP/1.x 200 …`.
/// 5. Return the raw `TcpStream` positioned after the CONNECT response.
async fn open_connect_tunnel(
proxy_url: &str,
target_host: &str,
target_port: u16,
) -> anyhow::Result<TcpStream> {
// 1. Parse proxy URL.
let (proxy_host, proxy_port) = parse_proxy_url(proxy_url)?;
// 2. TCP connect to proxy.
let proxy_addr = format!("{proxy_host}:{proxy_port}");
debug!(proxy_addr = %proxy_addr, "Opening TCP to proxy");
let stream = TcpStream::connect(&proxy_addr)
.await
.map_err(|e| anyhow::anyhow!("Failed to connect to proxy at {proxy_addr}: {e}"))?;
// 3. Send HTTP CONNECT.
let connect_req = format!(
"CONNECT {target_host}:{target_port} HTTP/1.1\r\n\
Host: {target_host}:{target_port}\r\n\
\r\n"
);
let (reader_half, mut writer_half) = stream.into_split();
writer_half.write_all(connect_req.as_bytes()).await?;
writer_half.flush().await?;
// 4. Read the status line from the proxy.
let mut reader = BufReader::new(reader_half);
let mut status_line = String::new();
reader.read_line(&mut status_line).await?;
debug!(status_line = %status_line.trim(), "Proxy CONNECT response");
if !status_line.starts_with("HTTP/1.1 200") && !status_line.starts_with("HTTP/1.0 200") {
anyhow::bail!("Proxy CONNECT failed: {}", status_line.trim());
}
// Consume remaining response headers (until empty line).
loop {
let mut line = String::new();
reader.read_line(&mut line).await?;
if line.trim().is_empty() {
break;
}
}
// 5. Assert the BufReader's internal buffer is empty before reuniting.
// BufReader::read_line may have read ahead into its buffer. If extra
// bytes were consumed beyond the HTTP headers (e.g., from a proxy that
// eagerly forwards data or coalesced TCP segments), dropping them would
// corrupt the subsequent TLS handshake.
let remaining = reader.buffer();
if !remaining.is_empty() {
anyhow::bail!(
"Proxy sent {} unexpected byte(s) after CONNECT response headers",
remaining.len()
);
}
// 6. Reunite the split halves back into a TcpStream.
let stream = reader.into_inner().reunite(writer_half)?;
Ok(stream)
}
/// Lazily-initialized TLS client configuration.
///
/// Loading native root certificates involves syscalls (reading `/etc/ssl/certs/`
/// or the macOS Keychain) and the cert store never changes at runtime. We build
/// the `ClientConfig` once and reuse it across all proxy connections / reconnects.
///
/// Stores `Ok(config)` on success or `Err(message)` if cert loading fails.
static TLS_CONFIG: OnceLock<Result<Arc<rustls::ClientConfig>, String>> = OnceLock::new();
/// Build (or return the cached) TLS client configuration.
fn get_tls_config() -> anyhow::Result<Arc<rustls::ClientConfig>> {
let result = TLS_CONFIG.get_or_init(|| {
let mut root_store = rustls::RootCertStore::empty();
let cert_result = rustls_native_certs::load_native_certs();
if cert_result.certs.is_empty() {
let errors: Vec<_> = cert_result.errors.iter().map(|e| e.to_string()).collect();
return Err(format!(
"No native root certificates found. Errors: {}",
if errors.is_empty() {
"(none)".to_string()
} else {
errors.join("; ")
}
));
}
for cert in cert_result.certs {
if let Err(e) = root_store.add(cert) {
tracing::warn!(error = %e, "Skipping unparseable native root certificate");
}
}
let config = rustls::ClientConfig::builder()
.with_root_certificates(root_store)
.with_no_client_auth();
Ok(Arc::new(config))
});
match result {
Ok(config) => Ok(config.clone()),
Err(msg) => anyhow::bail!("{msg}"),
}
}
/// Perform a TLS handshake over an existing TCP stream using rustls with
/// native root certificates (cached via [`TLS_CONFIG`]).
async fn tls_wrap(
stream: TcpStream,
server_name: &str,
) -> anyhow::Result<tokio_rustls::client::TlsStream<TcpStream>> {
let tls_config = get_tls_config()?;
let connector = tokio_rustls::TlsConnector::from(tls_config);
let dns_name = rustls::pki_types::ServerName::try_from(server_name.to_string())
.map_err(|e| anyhow::anyhow!("Invalid TLS server name '{server_name}': {e}"))?;
let tls_stream = connector
.connect(dns_name, stream)
.await
.map_err(|e| anyhow::anyhow!("TLS handshake through proxy failed: {e}"))?;
Ok(tls_stream)
}
/// Parse a proxy URL into (host, port).
///
/// Accepted formats:
/// - `http://host:port`
/// - `http://host` (defaults to port 80)
/// - `host:port`
fn parse_proxy_url(url: &str) -> anyhow::Result<(String, u16)> {
// Strip scheme if present.
let without_scheme = url
.strip_prefix("http://")
.or_else(|| url.strip_prefix("https://"))
.unwrap_or(url);
// Strip trailing path/slash.
let authority = without_scheme.split('/').next().unwrap_or(without_scheme);
if let Some((host, port_str)) = authority.rsplit_once(':') {
let port: u16 = port_str
.parse()
.map_err(|_| anyhow::anyhow!("Invalid proxy port in '{url}'"))?;
Ok((host.to_string(), port))
} else {
// No port — default to 80 for HTTP proxies.
Ok((authority.to_string(), 80))
}
}
// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------
#[cfg(test)]
mod tests {
use super::*;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
// ===== parse_proxy_url =====
#[test]
fn test_parse_proxy_url_with_scheme_and_port() {
let (host, port) = parse_proxy_url("http://proxy.example.com:3140").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 3140);
}
#[test]
fn test_parse_proxy_url_without_scheme() {
let (host, port) = parse_proxy_url("proxy.example.com:8080").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 8080);
}
#[test]
fn test_parse_proxy_url_without_port() {
let (host, port) = parse_proxy_url("http://proxy.example.com").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 80);
}
#[test]
fn test_parse_proxy_url_with_trailing_slash() {
let (host, port) = parse_proxy_url("http://proxy.example.com:3140/").unwrap();
assert_eq!(host, "proxy.example.com");
assert_eq!(port, 3140);
}
#[test]
fn test_parse_proxy_url_https_scheme() {
let (host, port) = parse_proxy_url("https://secure-proxy:443").unwrap();
assert_eq!(host, "secure-proxy");
assert_eq!(port, 443);
}
#[test]
fn test_parse_proxy_url_multi_label_host() {
let (host, port) =
parse_proxy_url("http://http-proxy.services.internal.example:3128").unwrap();
assert_eq!(host, "http-proxy.services.internal.example");
assert_eq!(port, 3128);
}
#[test]
fn test_parse_proxy_url_invalid_port() {
assert!(parse_proxy_url("http://proxy:notaport").is_err());
}
// ===== is_host_bypassed =====
#[test]
fn test_bypass_exact_match() {
assert!(is_host_bypassed("localhost", "localhost,127.0.0.1"));
}
#[test]
fn test_bypass_domain_suffix_with_dot() {
assert!(is_host_bypassed(
"api.corp.example",
"localhost,.corp.example"
));
}
#[test]
fn test_bypass_domain_suffix_without_dot() {
// Common convention: "example.com" in NO_PROXY matches "api.example.com".
assert!(is_host_bypassed("api.example.com", "localhost,example.com"));
}
#[test]
fn test_bypass_wildcard() {
assert!(is_host_bypassed("anything.example.com", "*"));
}
#[test]
fn test_no_bypass_when_not_listed() {
assert!(!is_host_bypassed(
"api.external.example",
"localhost,127.0.0.1,.corp.example,.internal.example"
));
}
#[test]
fn test_bypass_case_insensitive() {
assert!(is_host_bypassed("API.Corp.EXAMPLE", ".corp.example"));
}
#[test]
fn test_bypass_empty_no_proxy() {
assert!(!is_host_bypassed("api.external.example", ""));
}
#[test]
fn test_bypass_spaces_in_entries() {
assert!(is_host_bypassed(
"foo.example.com",
" localhost , .example.com , .other.com "
));
}
#[test]
fn test_bypass_cidr_not_matched_for_dns_names() {
// CIDR entries like 10.0.0.0/8 should not match DNS names.
assert!(!is_host_bypassed("api.external.example", "10.0.0.0/8"));
}
#[test]
fn test_bypass_combined_no_proxy_list() {
// A typical corporate NO_PROXY mixes loopback, private CIDRs, and domain suffixes.
let no_proxy = "localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example";
assert!(!is_host_bypassed("api.external.example", no_proxy));
assert!(is_host_bypassed("db.internal.example", no_proxy));
assert!(is_host_bypassed("git.corp.example", no_proxy));
assert!(is_host_bypassed("localhost", no_proxy));
}
// ===== resolve_proxy_for_host_with =====
#[test]
fn test_resolve_no_proxy_vars_set() {
let result = resolve_proxy_for_host_with("api.external.example", |_| {
Err(std::env::VarError::NotPresent)
});
assert_eq!(result, None);
}
#[test]
fn test_resolve_https_proxy_used() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
#[test]
fn test_resolve_http_proxy_fallback() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
}
#[test]
fn test_resolve_no_proxy_bypasses() {
let result = resolve_proxy_for_host_with("api.corp.example", |key| match key {
"HTTPS_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => Ok("localhost,.corp.example".to_string()),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, None);
}
#[test]
fn test_resolve_https_proxy_takes_precedence() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok("http://https-proxy.example.com:443".to_string()),
"HTTP_PROXY" => Ok("http://http-proxy.example.com:80".to_string()),
"NO_PROXY" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(
result,
Some("http://https-proxy.example.com:443".to_string())
);
}
#[test]
fn test_resolve_lowercase_env_vars() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"https_proxy" => Ok("http://proxy.example.com:3128".to_string()),
"no_proxy" => Err(std::env::VarError::NotPresent),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
#[test]
fn test_resolve_empty_proxy_ignored() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" => Ok(" ".to_string()),
"HTTP_PROXY" => Ok("http://proxy.example.com:8080".to_string()),
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:8080".to_string()));
}
#[test]
fn test_resolve_respects_no_proxy_when_proxy_set() {
let result = resolve_proxy_for_host_with("api.external.example", |key| match key {
"HTTPS_PROXY" | "HTTP_PROXY" => Ok("http://proxy.example.com:3128".to_string()),
"NO_PROXY" => {
Ok("localhost,127.0.0.1,10.0.0.0/8,.internal.example,.corp.example".to_string())
}
_ => Err(std::env::VarError::NotPresent),
});
assert_eq!(result, Some("http://proxy.example.com:3128".to_string()));
}
// ===== HTTP CONNECT tunnel (integration-style) =====
/// Helper: spawn a mock HTTP CONNECT proxy that accepts one connection.
///
/// On receiving a CONNECT request, it validates the request format,
/// replies with `status_line`, and then echoes data (simulating a tunnel).
/// Returns the proxy's listen address.
async fn spawn_mock_proxy(status_line: &'static str) -> std::net::SocketAddr {
let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr = listener.local_addr().unwrap();
tokio::spawn(async move {
let (mut stream, _) = listener.accept().await.unwrap();
// Read CONNECT request (read until \r\n\r\n).
let mut buf = vec![0u8; 4096];
let mut total = 0;
loop {
let n = stream.read(&mut buf[total..]).await.unwrap();
if n == 0 {
return;
}
total += n;
let so_far = std::str::from_utf8(&buf[..total]).unwrap_or("");
if so_far.contains("\r\n\r\n") {
break;
}
}
let request = std::str::from_utf8(&buf[..total]).unwrap().to_string();
assert!(
request.contains("CONNECT ") && request.contains(" HTTP/1.1"),
"Expected CONNECT request, got: {request}"
);
// Reply with the provided status line.
stream.write_all(status_line.as_bytes()).await.unwrap();
// Echo loop (simulates the transparent tunnel).
let mut echo_buf = [0u8; 1024];
loop {
let n = match stream.read(&mut echo_buf).await {
Ok(0) | Err(_) => break,
Ok(n) => n,
};
if stream.write_all(&echo_buf[..n]).await.is_err() {
break;
}
}
});
addr
}
/// Tests that `open_connect_tunnel` sends a correct CONNECT request,
/// parses the proxy's 200 response, and returns a usable tunnel stream.
#[tokio::test]
async fn test_open_connect_tunnel_success() {
let addr =
spawn_mock_proxy("HTTP/1.1 200 Connection Established\r\nServer: mock\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
// Call the real function under test.
let mut stream = open_connect_tunnel(&proxy_url, "example.com", 443)
.await
.expect("tunnel should succeed");
// Verify the tunnel works by echoing data through it.
stream.write_all(b"hello tunnel").await.unwrap();
stream.flush().await.unwrap();
let mut response = vec![0u8; 12];
stream.read_exact(&mut response).await.unwrap();
assert_eq!(&response, b"hello tunnel");
}
/// Tests that `open_connect_tunnel` with a non-default port sends the
/// correct CONNECT target.
#[tokio::test]
async fn test_open_connect_tunnel_custom_port() {
let addr = spawn_mock_proxy("HTTP/1.1 200 OK\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
let stream = open_connect_tunnel(&proxy_url, "internal.example.com", 8443).await;
assert!(stream.is_ok(), "tunnel should succeed for custom port");
}
/// Tests that `open_connect_tunnel` returns an error when the proxy
/// rejects the CONNECT request with a non-200 status.
#[tokio::test]
async fn test_open_connect_tunnel_proxy_rejects() {
let addr = spawn_mock_proxy("HTTP/1.1 403 Forbidden\r\n\r\n").await;
let proxy_url = format!("http://{addr}");
let result = open_connect_tunnel(&proxy_url, "blocked.example.com", 443).await;
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("403"),
"Error should mention 403: {err_msg}"
);
}
/// Tests that `open_connect_tunnel` returns an error when connecting
/// to a proxy that isn't listening.
#[tokio::test]
async fn test_open_connect_tunnel_proxy_unreachable() {
let result = open_connect_tunnel("http://127.0.0.1:1", "example.com", 443).await;
assert!(result.is_err());
let err_msg = result.unwrap_err().to_string();
assert!(
err_msg.contains("Failed to connect to proxy"),
"Error should mention proxy connection failure: {err_msg}"
);
}
}
@@ -0,0 +1,51 @@
//! Thin wire-format adapter that wraps the shared
//! [`kigi_workspace::session::git::build_restore_decision`] helper
//! into the JSON shape emitted by `LoadSession` on `_meta.codeRestore`.
use kigi_workspace::session::git::{CheckoutSessionOutcome, RestoreKind, build_restore_decision};
use serde_json::Value;
/// Build the `codeRestore` JSON meta, or `None` when no restore should
/// be reported (no checkout AND no archive applied). The shared
/// [`build_restore_decision`] is the source of truth; this function
/// only adapts the result into the wire JSON shape used by the
/// non-worktree path.
pub(crate) fn build_code_restore_meta(
target_sha: &str,
outcome: &CheckoutSessionOutcome,
kind: RestoreKind,
) -> Option<Value> {
let decision = build_restore_decision(Some(target_sha), outcome, kind);
let summary = decision.summary?;
Some(serde_json::json!(
{ "restored" : decision.restored, "summary" : summary, "degree" : decision
.degree, }
))
}
#[cfg(test)]
mod tests {
use super::*;
fn outcome(
checked_out: bool,
stash_ref: Option<&str>,
skipped: Option<&str>,
) -> CheckoutSessionOutcome {
CheckoutSessionOutcome {
checked_out,
stash_ref: stash_ref.map(str::to_owned),
stash_skipped_reason: skipped.map(str::to_owned),
}
}
#[test]
fn checkout_failed_emits_restored_false_meta() {
let meta = build_code_restore_meta(
"0123456789abcdef",
&outcome(false, None, Some("MERGE_HEAD present")),
RestoreKind::RegistryOff,
)
.unwrap();
assert_eq!(meta["restored"], false);
assert!(meta["degree"].is_null());
let s = meta["summary"].as_str().unwrap();
assert!(s.contains("restore aborted"));
assert!(s.contains("MERGE_HEAD present"));
}
}
@@ -0,0 +1,311 @@
//! Roster types for the multi-client FleetView dashboard.
//!
//! The roster is a list
//! of dashboard-sized summaries of every session the leader hosts (resident
//! actors) plus recently-touched on-disk (`Dormant`) sessions. Clients read it
//! two ways:
//!
//! - request/response `x.ai/sessions/list` → `{ "sessions": [RosterEntry, …] }`
//! - broadcast notification `x.ai/sessions/changed` →
//! `{ "upserted": [RosterEntry, …], "removed": ["sess-abc", …] }`
//!
//! The wire shape is intentionally small and current-state only — no event
//! fold or materialized snapshot is required (the snapshot is deferred).
use kigi_sampling_types::ReasoningEffort;
use serde::{Deserialize, Serialize};
use crate::session::persistence::Summary;
/// Coarse activity of a session as rendered in the dashboard's status column.
///
/// Mirrors the design's `SessionActivity` at dashboard granularity. A full
/// background-work breakdown (bg tasks / monitors / scheduler / subagents)
/// lands with a richer `SessionActivity`; the dashboard only needs this
/// coarse signal to pick a status glyph.
#[derive(Serialize, Deserialize, Clone, Copy, Debug, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum RosterActivity {
/// A turn (user-originated or autonomous) is running.
Working,
/// Resident, no turn in flight.
Idle,
/// A permission / question / plan-approval is pending.
NeedsInput,
/// On disk, not resident.
Dormant,
/// Finished and resumable.
Completed,
/// Actor panicked / load failed.
Dead,
}
/// Where the session lives. Only `Local` is produced today; `Remote` is
/// reserved for cross-machine roster aggregation.
#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)]
#[serde(rename_all = "snake_case", tag = "kind")]
pub enum RosterOrigin {
Local,
Remote { host: String },
}
/// One dashboard row.
#[derive(Serialize, Deserialize, Clone, Debug)]
#[serde(rename_all = "camelCase")]
pub struct RosterEntry {
pub session_id: String,
/// Generated/display title, if known. Clients fall back to cwd / id.
#[serde(default)]
pub title: Option<String>,
pub cwd: String,
pub is_worktree: bool,
#[serde(default)]
pub model_id: Option<String>,
/// Per-session reasoning effort for `model_id`. Carried alongside the model
/// so clients can render the session's effort in the roster without a
/// separate `model_state` fetch. `None` means "use the model/global
/// default" (or the session predates per-session effort persistence).
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_effort: Option<ReasoningEffort>,
pub yolo: bool,
pub activity: RosterActivity,
/// `true` while a resident actor hosts the session (vs. read from disk).
pub resident: bool,
/// Best-effort last-change timestamp (unix millis). Used for sort order.
pub last_change_unix_ms: i64,
pub origin: RosterOrigin,
}
/// Response payload for `x.ai/sessions/list`.
#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct RosterListResponse {
pub sessions: Vec<RosterEntry>,
}
/// Params payload for the `x.ai/sessions/changed` broadcast notification.
#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct RosterChanged {
#[serde(default)]
pub upserted: Vec<RosterEntry>,
#[serde(default)]
pub removed: Vec<String>,
}
/// JSON-RPC method names for the roster API.
pub const SESSIONS_LIST_METHOD: &str = "x.ai/sessions/list";
pub const SESSIONS_CHANGED_METHOD: &str = "x.ai/sessions/changed";
/// Merge live `resident` rows with on-disk `summaries` into the sorted roster.
/// Pure, so it is unit-testable without disk or a live actor.
///
/// Resident rows own the live state but carry no title or last-active time, so
/// each adopts those from its summary — except a `Working` row keeps its "now"
/// timestamp. Summaries with no resident row become `Dormant`; keying by id
/// dedups them. Hidden summaries are excluded.
pub(crate) fn merge_roster(
mut entries: Vec<RosterEntry>,
summaries: Vec<Summary>,
) -> Vec<RosterEntry> {
let mut by_id: std::collections::HashMap<String, Summary> = summaries
.into_iter()
.filter(|s| !s.is_hidden())
.map(|s| (s.info.id.0.to_string(), s))
.collect();
// Backfill resident rows; remove the summary so it isn't re-emitted below.
for entry in &mut entries {
let Some(summary) = by_id.remove(&entry.session_id) else {
continue;
};
if let Some(title) = summary.display_title_opt() {
entry.title = Some(title);
}
if entry.activity != RosterActivity::Working {
entry.last_change_unix_ms = summary.last_change_unix_ms();
}
}
// Remaining summaries have no resident row: emit them as dormant.
entries.extend(by_id.into_values().map(|summary| RosterEntry {
session_id: summary.info.id.0.to_string(),
title: summary.display_title_opt(),
cwd: summary.info.cwd.clone(),
is_worktree: summary.session_kind.as_deref() == Some("worktree")
|| summary.source_workspace_dir.is_some(),
model_id: Some(summary.current_model_id.0.to_string()),
reasoning_effort: summary.reasoning_effort,
yolo: false,
activity: RosterActivity::Dormant,
resident: false,
last_change_unix_ms: summary.last_change_unix_ms(),
origin: RosterOrigin::Local,
}));
// Most-recently-changed first.
entries.sort_by_key(|entry| std::cmp::Reverse(entry.last_change_unix_ms));
entries
}
#[cfg(test)]
mod merge_roster_tests {
use super::*;
use crate::session::info::Info;
use crate::session::persistence::default_model_id;
use agent_client_protocol as acp;
fn summary(id: &str, title: Option<&str>, last_active_ms: i64) -> Summary {
let mut s = Summary::new(
&Info {
id: acp::SessionId::new(id),
cwd: format!("/repo/{id}"),
},
default_model_id(),
)
.expect("summary");
s.generated_title = title.map(String::from);
s.last_active_at = chrono::DateTime::from_timestamp_millis(last_active_ms);
s
}
fn resident(id: &str, activity: RosterActivity, last_change_unix_ms: i64) -> RosterEntry {
RosterEntry {
session_id: id.to_string(),
title: None,
cwd: format!("/live/{id}"),
is_worktree: false,
model_id: Some("grok-4".into()),
reasoning_effort: None,
yolo: false,
activity,
resident: true,
last_change_unix_ms,
origin: RosterOrigin::Local,
}
}
#[test]
fn idle_resident_adopts_persisted_title_and_last_active() {
let now = 9_000;
let out = merge_roster(
vec![resident("a", RosterActivity::Idle, now)],
vec![summary("a", Some("Fix the roster"), 1_234)],
);
assert_eq!(out.len(), 1, "resident must not be duplicated as dormant");
assert_eq!(out[0].title.as_deref(), Some("Fix the roster"));
assert_eq!(out[0].last_change_unix_ms, 1_234, "idle adopts last-active");
assert!(out[0].resident);
assert_eq!(out[0].cwd, "/live/a", "live cwd is preserved");
assert_eq!(out[0].activity, RosterActivity::Idle);
}
#[test]
fn working_resident_keeps_now_but_adopts_title() {
let now = 9_000;
let out = merge_roster(
vec![resident("a", RosterActivity::Working, now)],
vec![summary("a", Some("Busy turn"), 1_234)],
);
assert_eq!(out[0].title.as_deref(), Some("Busy turn"));
assert_eq!(out[0].last_change_unix_ms, now, "Working stays 'now'");
}
#[test]
fn new_resident_without_summary_stays_titleless_now() {
let now = 9_000;
let out = merge_roster(vec![resident("a", RosterActivity::Idle, now)], vec![]);
assert_eq!(out[0].title, None);
assert_eq!(out[0].last_change_unix_ms, now);
}
#[test]
fn blank_persisted_title_leaves_row_untitled() {
let out = merge_roster(
vec![resident("a", RosterActivity::Idle, 9_000)],
vec![summary("a", Some(" "), 1_234)],
);
assert_eq!(out[0].title, None, "blank title normalizes to None");
assert_eq!(out[0].last_change_unix_ms, 1_234);
}
#[test]
fn dormant_sessions_are_emitted_and_sorted_after_residents() {
let out = merge_roster(
vec![resident("live", RosterActivity::Idle, 5_000)],
vec![
summary("live", Some("Live one"), 4_000),
summary("old", Some("Dormant one"), 1_000),
summary("new", Some("Newer dormant"), 8_000),
],
);
let ids: Vec<&str> = out.iter().map(|e| e.session_id.as_str()).collect();
assert_eq!(
ids,
vec!["new", "live", "old"],
"sorted by last-change desc"
);
let dormant = out.iter().find(|e| e.session_id == "new").unwrap();
assert_eq!(dormant.activity, RosterActivity::Dormant);
assert!(!dormant.resident);
}
#[test]
fn duplicate_summaries_are_deduped() {
let out = merge_roster(
vec![],
vec![
summary("dup", Some("First"), 1_000),
summary("dup", Some("Second"), 2_000),
],
);
assert_eq!(out.len(), 1, "duplicate ids collapse to one row");
}
#[test]
fn hidden_summaries_are_excluded() {
let mut hidden = summary("sub", Some("Subagent"), 5_000);
hidden.session_kind = Some("subagent".into());
let out = merge_roster(vec![], vec![hidden]);
assert!(out.is_empty(), "hidden/subagent summaries are dropped");
}
#[test]
fn dormant_row_carries_persisted_reasoning_effort() {
let mut s = summary("dorm", Some("Dormant"), 1_000);
s.reasoning_effort = Some(ReasoningEffort::Xhigh);
let out = merge_roster(vec![], vec![s]);
assert_eq!(out[0].reasoning_effort, Some(ReasoningEffort::Xhigh));
}
#[test]
fn resident_effort_is_taken_from_the_live_row_not_the_summary() {
// The live handle is authoritative for a resident session, so the
// resident row's effort must survive the summary backfill.
let mut live = resident("a", RosterActivity::Idle, 9_000);
live.reasoning_effort = Some(ReasoningEffort::High);
let mut s = summary("a", Some("Title"), 1_234);
s.reasoning_effort = Some(ReasoningEffort::Low);
let out = merge_roster(vec![live], vec![s]);
assert_eq!(out.len(), 1);
assert_eq!(out[0].reasoning_effort, Some(ReasoningEffort::High));
}
#[test]
fn reasoning_effort_serializes_as_camel_case_and_skips_when_none() {
let with_effort = RosterEntry {
reasoning_effort: Some(ReasoningEffort::Xhigh),
..resident("a", RosterActivity::Idle, 1)
};
let json = serde_json::to_string(&with_effort).unwrap();
assert!(
json.contains("\"reasoningEffort\":\"xhigh\""),
"effort must be camelCase and snake_case-valued: {json}"
);
let without = resident("b", RosterActivity::Idle, 1);
let json = serde_json::to_string(&without).unwrap();
assert!(
!json.contains("reasoningEffort"),
"a None effort must not be serialized: {json}"
);
}
}
@@ -0,0 +1,487 @@
//! WebSocket server for remote agent connections.
//!
//! This module provides a WebSocket server that allows remote TUI clients to
//! connect to a grok agent running on a different machine.
//!
//! The agent persists across WebSocket reconnections: a single MvpAgent instance
//! is created on first connection and reused for all subsequent connections. This
//! ensures that session actors (and any in-flight prompts) survive client
//! disconnects — when a client reconnects and loads an existing session, ongoing
//! work continues to stream to the new connection.
use std::cell::RefCell;
use std::net::SocketAddr;
use std::rc::Rc;
use std::sync::Arc;
use std::thread;
use axum::{
Router,
extract::{
ConnectInfo, Query, State,
ws::{Message, WebSocket, WebSocketUpgrade},
},
http::{HeaderMap, StatusCode},
response::{IntoResponse, Response},
routing::get,
};
use futures::{SinkExt, StreamExt};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, simplex};
use tokio::net::TcpListener;
use tokio::sync::mpsc;
use tokio::time::Duration;
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
use tracing::{info, warn};
use agent_client_protocol as acp;
use kigi_acp_lib::{
AcpAgentGatewayReceiver as GatewayReceiver, AcpAgentGatewaySender as GatewaySender,
AcpClientMessage, LineBufferedRead,
};
use crate::agent::config::{Config as AgentConfig, ModelEntry};
use crate::agent::models::{ModelFetchAuth, prefetch_models_blocking};
use crate::agent::mvp_agent::MvpAgent;
use indexmap::IndexMap;
/// Swappable destination for the relay task.
///
/// Points at the current ACP connection's gateway sender. When no client is
/// connected, the value is `None` and outbound messages are silently dropped
/// (matching the old behaviour where the gateway channel's receiver was simply
/// gone).
type RelayDest = Rc<RefCell<Option<mpsc::UnboundedSender<AcpClientMessage>>>>;
const MAX_BUFFER_SIZE: usize = 8 * 1024 * 1024;
const KEEPALIVE_INTERVAL_SECS: u64 = 15;
/// Configuration for the agent WebSocket server.
#[derive(Debug, Clone)]
pub struct ServerConfig {
/// Address to bind the server to
pub bind_addr: SocketAddr,
/// Secret token for client authentication (required)
pub secret: String,
}
/// Shared state for the WebSocket server.
struct ServerState {
agent_config: AgentConfig,
secret: String,
/// Channel to send new WebSocket connections to the persistent agent thread.
/// Lazily initialised on first connection; protected by a tokio Mutex so the
/// axum handler (which is `Send`) can acquire it.
agent_conn_tx: tokio::sync::Mutex<Option<mpsc::UnboundedSender<NewConnectionChannels>>>,
}
/// Channels bridging a single WebSocket connection to the agent thread.
struct NewConnectionChannels {
from_ws_rx: mpsc::UnboundedReceiver<String>,
to_ws_tx: mpsc::UnboundedSender<String>,
}
/// Query parameters for WebSocket connection.
#[derive(Debug, serde::Deserialize, Default)]
pub struct WsQueryParams {
#[serde(rename = "server-key")]
pub server_key: Option<String>,
}
/// Validate the bearer token from request headers or query parameters.
fn validate_auth(headers: &HeaderMap, query: &WsQueryParams, expected_secret: &str) -> bool {
// Try Authorization header
if let Some(token) = headers
.get("authorization")
.and_then(|v| v.to_str().ok())
.and_then(|v| v.strip_prefix("Bearer "))
{
return token == expected_secret;
}
// Fall back to query parameter for browser connections
if let Some(ref key) = query.server_key {
return key == expected_secret;
}
false
}
/// WebSocket upgrade handler with authentication.
async fn ws_handler(
ws: WebSocketUpgrade,
State(state): State<Arc<ServerState>>,
ConnectInfo(addr): ConnectInfo<SocketAddr>,
headers: HeaderMap,
Query(query): Query<WsQueryParams>,
) -> Response {
// Validate secret token from header or query param
if !validate_auth(&headers, &query, &state.secret) {
warn!("Unauthorized connection attempt from {}", addr);
return (
StatusCode::UNAUTHORIZED,
"Invalid or missing authorization token",
)
.into_response();
}
info!("Authenticated WebSocket connection from {}", addr);
ws.on_upgrade(move |socket| handle_connection(socket, state, addr))
}
/// Handle an authenticated WebSocket connection.
///
/// On first connection, spawns a persistent agent thread that owns the MvpAgent.
/// On subsequent connections (reconnects), sends new WS channels to the existing
/// agent thread so that session actors can continue streaming to the new client.
async fn handle_connection(ws: WebSocket, state: Arc<ServerState>, peer_addr: SocketAddr) {
info!("New WebSocket connection from {}", peer_addr);
let (mut ws_write, mut ws_read) = ws.split();
// Channels for bridging WS <-> Agent thread
let (to_agent_tx, to_agent_rx) = mpsc::unbounded_channel::<String>();
let (from_agent_tx, mut from_agent_rx) = mpsc::unbounded_channel::<String>();
// Ensure the persistent agent thread is running (lazy init on first connection).
// If the previous agent thread died (panic, etc.), clear the stale sender so we
// respawn a fresh one.
{
let mut agent_tx_guard = state.agent_conn_tx.lock().await;
// Check if existing sender is still alive (receiver not dropped)
if let Some(ref tx) = *agent_tx_guard
&& tx.is_closed()
{
warn!("Persistent agent thread died — will respawn");
*agent_tx_guard = None;
}
if agent_tx_guard.is_none() {
let (conn_tx, conn_rx) = mpsc::unbounded_channel();
let agent_config = state.agent_config.clone();
let _agent_thread = thread::Builder::new()
.name("agent-persistent".to_string())
.spawn(move || {
// Prefetch models before creating the runtime (blocking is OK here)
let auth = agent_config.create_auth_manager().current();
let fetch_auth =
ModelFetchAuth::resolve(&agent_config.endpoints, auth.is_some());
let prefetched_models = if auth.is_some()
|| agent_config.endpoints.has_custom_endpoint()
|| fetch_auth != ModelFetchAuth::Session
{
prefetch_models_blocking(&agent_config.endpoints, auth.as_ref(), fetch_auth)
} else {
None
};
info!("Prefetched models: {:?}", prefetched_models);
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("Failed to create runtime for agent");
let local_set = tokio::task::LocalSet::new();
local_set.block_on(&rt, async move {
run_persistent_agent(agent_config, conn_rx, prefetched_models).await
});
warn!("Persistent agent thread exiting");
});
*agent_tx_guard = Some(conn_tx);
info!("Persistent agent thread spawned");
}
// Send new WS channels to the agent thread
if let Some(ref tx) = *agent_tx_guard
&& tx
.send(NewConnectionChannels {
from_ws_rx: to_agent_rx,
to_ws_tx: from_agent_tx,
})
.is_err()
{
warn!("Failed to send connection channels to agent thread");
}
}
// Task: Read from WS, send to agent thread
let read_task = tokio::spawn(async move {
while let Some(msg) = ws_read.next().await {
match msg {
Ok(Message::Text(text)) => {
let text_str: &str = text.as_ref();
let trimmed = text_str.trim_end_matches(['\r', '\n']);
// Skip browser keepalive pings (non-JSON text)
if trimmed == "ping" || trimmed.is_empty() {
continue;
}
if to_agent_tx.send(trimmed.to_string()).is_err() {
break;
}
}
Ok(Message::Binary(bin)) => {
if let Ok(s) = std::str::from_utf8(&bin) {
let trimmed = s.trim_end_matches(['\r', '\n']);
if trimmed == "ping" || trimmed.is_empty() {
continue;
}
if to_agent_tx.send(trimmed.to_string()).is_err() {
break;
}
}
}
Ok(Message::Close(frame)) => {
if let Some(f) = frame {
info!(
"WebSocket close from {}: {} {}",
peer_addr, f.code, f.reason
);
}
break;
}
Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {}
Err(e) => {
warn!("WebSocket read error from {}: {:?}", peer_addr, e);
break;
}
}
}
});
// Task: Read from agent thread, send to WS (with keepalive)
let write_task = tokio::spawn(async move {
let mut keepalive = tokio::time::interval(Duration::from_secs(KEEPALIVE_INTERVAL_SECS));
loop {
tokio::select! {
Some(msg) = from_agent_rx.recv() => {
if ws_write.send(Message::Text(msg.into())).await.is_err() {
break;
}
}
_ = keepalive.tick() => {
if ws_write.send(Message::Ping(vec![].into())).await.is_err() {
break;
}
}
else => break,
}
}
});
// Wait for either task to complete
tokio::select! {
_ = read_task => {}
_ = write_task => {}
}
info!("WebSocket connection ended for {}", peer_addr);
}
/// Run the persistent agent on a dedicated thread with LocalSet.
///
/// The MvpAgent is created **once** and reused across WebSocket reconnections.
/// A persistent gateway channel ensures that session actors (which hold cloned
/// `GatewaySender` handles) can always send notifications. A relay task forwards
/// messages from the persistent channel to the *current* ACP connection's channel,
/// so notifications reach whichever client is currently connected.
async fn run_persistent_agent(
agent_config: AgentConfig,
mut connection_rx: mpsc::UnboundedReceiver<NewConnectionChannels>,
prefetched_models: Option<IndexMap<String, ModelEntry>>,
) {
// Persistent gateway channel — the MvpAgent and all session actors hold
// clones of `gw_tx`. This channel survives across reconnections.
let (gw_tx, mut gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
let gateway = GatewaySender::new(gw_tx);
// Create MvpAgent ONCE -- it persists for the lifetime of the server.
let auth_manager = Arc::new(agent_config.create_auth_manager());
// Proactive token refresh; runs until process exit.
auth_manager.start_proactive_refresh(tokio_util::sync::CancellationToken::new());
// Restore managed policy right before bootstrap reads it — the agent is created lazily here,
// so an earlier restore could go stale before the gate.
crate::managed_config::ensure_managed_policy_present(&auth_manager).await;
let agent = Rc::new(
MvpAgent::new(gateway, &agent_config, auth_manager, prefetched_models)
.unwrap_or_else(crate::agent::init::exit_on_config_error),
);
let relay_dest: RelayDest = Rc::new(RefCell::new(None));
// Relay task: reads from the persistent gateway channel and forwards to
// whichever ACP connection is currently active.
let relay_dest_for_task = relay_dest.clone();
tokio::task::spawn_local(async move {
while let Some(msg) = gw_rx.recv().await {
let maybe_tx = relay_dest_for_task.borrow().clone();
if let Some(tx) = maybe_tx
&& tx.send(msg).is_err()
{
// Connection's gateway receiver was dropped — clear it.
*relay_dest_for_task.borrow_mut() = None;
}
// If no connection, the message (and its response_tx) is dropped.
// The caller (session actor) gets a send error which is already
// handled with `let _ = ...`.
}
});
// Accept new connections in a loop
while let Some(channels) = connection_rx.recv().await {
info!("Agent thread: setting up new ACP connection (reconnect)");
setup_acp_connection(agent.clone(), channels, relay_dest.clone());
}
info!("Agent thread: connection channel closed, exiting");
}
/// Set up a new ACP connection for a WebSocket connection, reusing the existing
/// MvpAgent. The relay destination is updated so that session actor notifications
/// flow to the new client.
fn setup_acp_connection(
agent: Rc<MvpAgent>,
channels: NewConnectionChannels,
relay_dest: RelayDest,
) {
let NewConnectionChannels {
mut from_ws_rx,
to_ws_tx,
} = channels;
// Create new simplex IO streams for this ACP connection
let (agent_read_rx, mut agent_read_tx) = simplex(MAX_BUFFER_SIZE);
let (agent_write_rx, agent_write_tx) = simplex(MAX_BUFFER_SIZE);
let incoming = agent_read_rx.compat();
let outgoing = agent_write_tx.compat_write();
// Create a per-connection gateway channel for the GatewayReceiver.
// The relay task will forward persistent-channel messages here.
let (conn_gw_tx, conn_gw_rx) = tokio::sync::mpsc::unbounded_channel::<AcpClientMessage>();
// Point the relay at this new connection's channel
*relay_dest.borrow_mut() = Some(conn_gw_tx);
// Create new ACP connection reusing the same MvpAgent (via Rc clone).
// `Agent` is implemented for `Rc<T: Agent>` so this works.
let incoming = LineBufferedRead::spawn_local(incoming);
let (conn, handle_io) = acp::AgentSideConnection::new(agent, outgoing, incoming, |fut| {
tokio::task::spawn_local(fut);
});
tokio::task::spawn_local(
GatewayReceiver::new(conn_gw_rx, conn)
.with_on_meta(kigi_file_utils::trace_context::span_from_meta_traceparent)
.run(),
);
// Task: Forward WS messages → agent (incoming ACP bytes)
tokio::task::spawn_local(async move {
while let Some(msg) = from_ws_rx.recv().await {
// Log messages that lack both `id` and `method` — the ACP layer
// only prints "received message with neither id nor method" without
// the payload, making debugging impossible.
if let Ok(v) = serde_json::from_str::<serde_json::Value>(&msg)
&& v.get("id").is_none()
&& v.get("method").is_none()
{
warn!(
len = msg.len(),
"incoming WS message has neither id nor method"
);
}
if agent_read_tx.write_all(msg.as_bytes()).await.is_err() {
break;
}
if agent_read_tx.write_all(b"\n").await.is_err() {
break;
}
}
// WS disconnected — the simplex writer is dropped, causing `handle_io`
// to complete. The GatewayReceiver for this connection will also stop.
// But the MvpAgent and session actors stay alive, ready for the next
// connection.
});
// Task: Forward agent messages → WS (outgoing ACP bytes)
tokio::task::spawn_local(async move {
let mut reader = BufReader::new(agent_write_rx);
let mut line = String::new();
loop {
line.clear();
match reader.read_line(&mut line).await {
Ok(0) => break,
Ok(_) => {
let msg = line.trim_end_matches(['\r', '\n']);
if !msg.is_empty() && to_ws_tx.send(msg.to_string()).is_err() {
break;
}
}
Err(_) => break,
}
}
});
// Run the ACP IO handler — fire-and-forget since we don't block the
// connection loop. It completes when the WS disconnects.
tokio::task::spawn_local(async move {
let _ = handle_io.await;
info!("ACP connection IO handler completed");
});
}
/// Run the agent WebSocket server.
///
/// This starts a WebSocket server that accepts authenticated connections from
/// remote TUI clients. A single agent instance is shared across all connections
/// (persisted across reconnections) so that in-flight session work survives
/// client disconnects.
///
/// # Arguments
/// * `config` - Server configuration (bind address and secret)
/// * `agent_config` - Agent configuration to use for each connection
///
/// # Example
/// ```ignore
/// let server_config = ServerConfig {
/// bind_addr: "0.0.0.0:9000".parse().unwrap(),
/// secret: "my-secret-token".to_string(),
/// };
/// run_agent_server(server_config, agent_config).await?;
/// ```
pub async fn run_agent_server(
config: ServerConfig,
agent_config: AgentConfig,
) -> anyhow::Result<()> {
let state = Arc::new(ServerState {
agent_config,
secret: config.secret,
agent_conn_tx: tokio::sync::Mutex::new(None),
});
let app = Router::new()
.route("/ws", get(ws_handler))
.with_state(state);
let listener = TcpListener::bind(config.bind_addr).await?;
info!("Agent server listening on ws://{}/ws", config.bind_addr);
info!(
"Clients should connect with: --remote ws://{}:{}/ws --secret <token>",
config.bind_addr.ip(),
config.bind_addr.port()
);
axum::serve(
listener,
app.into_make_service_with_connect_info::<SocketAddr>(),
)
.await?;
Ok(())
}
@@ -0,0 +1,218 @@
use agent_client_protocol as acp;
use kigi_sampling_types::{ReasoningEffort, ReasoningEffortOption};
use serde::Serialize;
use crate::session::unified_list::SessionKind;
pub(crate) const SELECTABLE_REASONING_EFFORTS: [ReasoningEffort; 5] = [
ReasoningEffort::Minimal,
ReasoningEffort::Low,
ReasoningEffort::Medium,
ReasoningEffort::High,
ReasoningEffort::Xhigh,
];
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionConfigOption {
pub id: String,
pub category: String,
pub label: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub description: Option<String>,
pub selected: bool,
}
#[derive(Debug, Clone, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GrokSessionDetail {
pub session_id: String,
pub kind: String,
pub cwd: String,
pub current_model_id: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub title: Option<String>,
}
impl GrokSessionDetail {
pub fn build(
session_id: String,
cwd: String,
current_model_id: String,
title: Option<String>,
) -> Self {
Self {
session_id,
kind: SessionKind::Build.as_str().to_string(),
cwd,
current_model_id,
title,
}
}
}
fn effort_label(effort: ReasoningEffort) -> String {
match effort {
ReasoningEffort::None => "None",
ReasoningEffort::Minimal => "Minimal",
ReasoningEffort::Low => "Low",
ReasoningEffort::Medium => "Medium",
ReasoningEffort::High => "High",
ReasoningEffort::Xhigh => "X-High",
}
.to_string()
}
/// The built-in session-picker modes used when the model has no server list.
/// Reproduces the historical five rows and their labels.
pub(crate) fn legacy_session_effort_options() -> Vec<ReasoningEffortOption> {
SELECTABLE_REASONING_EFFORTS
.iter()
.map(|&effort| ReasoningEffortOption {
id: effort.as_str().to_string(),
value: effort,
label: effort_label(effort),
description: None,
default: false,
})
.collect()
}
pub(crate) fn build_session_config_options(
available_models: &[acp::ModelInfo],
current_model_id: &acp::ModelId,
effort_options: &[ReasoningEffortOption],
current_effort: Option<ReasoningEffort>,
) -> Vec<SessionConfigOption> {
let mut options = Vec::with_capacity(available_models.len() + effort_options.len());
for model in available_models {
let label = if model.name.is_empty() {
model.model_id.0.to_string()
} else {
model.name.clone()
};
options.push(SessionConfigOption {
id: model.model_id.0.to_string(),
category: "model".to_string(),
label,
description: None,
selected: model.model_id == *current_model_id,
});
}
for effort in effort_options {
options.push(SessionConfigOption {
id: effort.id.clone(),
category: "mode".to_string(),
label: effort.label.clone(),
description: effort.description.clone(),
selected: Some(effort.value) == current_effort,
});
}
options
}
#[cfg(test)]
mod tests {
use super::*;
fn model(id: &'static str, name: &str) -> acp::ModelInfo {
acp::ModelInfo::new(acp::ModelId::new(id), name.to_string())
}
#[test]
fn options_have_one_selected_model_and_a_mode_per_effort() {
let models = [
model("grok-build", "Grok Build"),
model("grok-4.5", "Grok 4.5"),
];
let current = acp::ModelId::from("grok-build");
let opts = build_session_config_options(
&models,
&current,
&legacy_session_effort_options(),
Some(ReasoningEffort::High),
);
let model_opts: Vec<_> = opts.iter().filter(|o| o.category == "model").collect();
assert_eq!(model_opts.len(), 2);
let selected_models: Vec<_> = model_opts.iter().filter(|o| o.selected).collect();
assert_eq!(selected_models.len(), 1);
assert_eq!(selected_models[0].id, "grok-build");
let mode_opts: Vec<_> = opts.iter().filter(|o| o.category == "mode").collect();
assert_eq!(mode_opts.len(), SELECTABLE_REASONING_EFFORTS.len());
let selected_modes: Vec<_> = mode_opts.iter().filter(|o| o.selected).collect();
assert_eq!(selected_modes.len(), 1);
assert_eq!(selected_modes[0].id, "high");
assert_eq!(selected_modes[0].label, "High");
}
#[test]
fn none_effort_is_not_a_user_selectable_mode() {
assert!(!SELECTABLE_REASONING_EFFORTS.contains(&ReasoningEffort::None));
let models = [model("grok-build", "Grok Build")];
let current = acp::ModelId::from("grok-build");
let opts = build_session_config_options(
&models,
&current,
&legacy_session_effort_options(),
Some(ReasoningEffort::None),
);
let modes: Vec<_> = opts.iter().filter(|o| o.category == "mode").collect();
assert!(modes.iter().all(|o| o.id != "none"));
assert!(modes.iter().all(|o| !o.selected));
}
#[test]
fn no_mode_options_when_model_lacks_effort_support() {
let models = [model("grok-build", "Grok Build")];
let current = acp::ModelId::from("grok-build");
let opts = build_session_config_options(&models, &current, &[], None);
assert_eq!(opts.len(), 1);
assert!(opts.iter().all(|o| o.category == "model"));
}
#[test]
fn model_label_falls_back_to_id_when_name_empty() {
let models = [model("grok-build", "")];
let current = acp::ModelId::from("grok-build");
let opts = build_session_config_options(&models, &current, &[], None);
assert_eq!(opts[0].label, "grok-build");
}
#[test]
fn session_config_option_serializes_camel_case() {
let opt = SessionConfigOption {
id: "grok-build".to_string(),
category: "model".to_string(),
label: "Grok Build".to_string(),
description: None,
selected: true,
};
let v = serde_json::to_value(&opt).expect("serialize");
assert_eq!(v["id"], "grok-build");
assert_eq!(v["category"], "model");
assert_eq!(v["label"], "Grok Build");
assert_eq!(v["selected"], true);
assert!(v.get("description").is_none());
}
#[test]
fn grok_session_detail_serializes_camel_case() {
let detail = GrokSessionDetail::build(
"sess-1".to_string(),
"/Users/me/xai".to_string(),
"grok-build".to_string(),
None,
);
let v = serde_json::to_value(&detail).expect("serialize");
assert_eq!(v["sessionId"], "sess-1");
assert_eq!(v["kind"], "build");
assert_eq!(v["cwd"], "/Users/me/xai");
assert_eq!(v["currentModelId"], "grok-build");
assert!(v.get("title").is_none());
}
}
@@ -0,0 +1,641 @@
//! REST client for the session replicas registry (cli-chat-proxy).
//!
//! Handles registering, updating, finalizing, searching, and downloading
//! session replicas for cross-host session replication. Write methods
//! (register/update/finalize) are fire-and-forget safe. Read methods
//! (search/get/download_file) return typed results.
use anyhow::{Context, Result};
use reqwest::RequestBuilder;
use serde::{Deserialize, Serialize};
// ============================================================================
// Request / response types (local — not in cli-chat-proxy since these
// are only used by the agent, not consumed by other crates)
// ============================================================================
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct RegisterRequest {
pub session_id: String,
pub cwd: String,
pub gcs_trace_prefix: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub model_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo_remote_url: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo_branch: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo_head_at_start: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub hostname: Option<String>,
/// Opaque per-machine device id (telemetry `agent_id()`) for machine disambiguation.
#[serde(skip_serializing_if = "Option::is_none")]
pub device_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub parent_session_id: Option<String>,
// --- Subagent-specific fields (optional, backward-compatible) ---
#[serde(skip_serializing_if = "Option::is_none")]
pub session_kind: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub subagent_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub subagent_persona: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub subagent_role: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub fork_context_source: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub subagent_depth: Option<u32>,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct UpdateRequest {
#[serde(skip_serializing_if = "Option::is_none")]
pub summary: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub first_prompt: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub last_turn_number: Option<i32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub repo_head_at_end: Option<String>,
/// Latest turn whose restore artifacts are confirmed durable.
/// Omitted from the wire when `None` — old servers ignore unknown fields.
#[serde(skip_serializing_if = "Option::is_none")]
pub restorable_turn_number: Option<i32>,
}
// ============================================================================
// Response types
// ============================================================================
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SessionRecord {
pub session_id: String,
pub summary: String,
pub first_prompt: Option<String>,
pub model_id: Option<String>,
pub created_at: String,
pub updated_at: String,
pub last_turn_number: i32,
/// Present on servers that have applied the restorable-turn migration.
/// `None` when talking to an older server — callers should fall back to
/// `last_turn_number` in that case.
#[serde(default)]
pub restorable_turn_number: Option<i32>,
pub cwd: String,
pub repo_remote_url: Option<String>,
pub hostname: Option<String>,
pub status: String,
pub gcs_trace_prefix: String,
pub gcs_bucket: String,
#[serde(default)]
pub last_active_at: Option<String>,
}
impl From<crate::session::persistence::Summary> for SessionRecord {
fn from(s: crate::session::persistence::Summary) -> Self {
Self {
session_id: s.info.id.to_string(),
summary: s.session_summary,
first_prompt: None,
model_id: Some(s.current_model_id.to_string()),
created_at: s.created_at.to_rfc3339(),
updated_at: s.updated_at.to_rfc3339(),
last_turn_number: s.num_messages as i32,
restorable_turn_number: None,
cwd: s.info.cwd,
repo_remote_url: None,
hostname: None,
status: "local".to_string(),
gcs_trace_prefix: String::new(),
gcs_bucket: String::new(),
last_active_at: s.last_active_at.map(|t| t.to_rfc3339()),
}
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct SearchResponse {
pub sessions: Vec<SessionRecord>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct DownloadResponse {
pub download_url: String,
pub file: String,
pub turn: i32,
}
// ============================================================================
// Client
// ============================================================================
#[derive(Clone)]
pub struct SessionRegistryClient {
raw_client: reqwest::Client,
client: reqwest_middleware::ClientWithMiddleware,
base_url: String,
credentials: crate::util::grok_auth_credentials::GrokAuthCredentials,
session_id: Option<String>,
}
impl SessionRegistryClient {
pub fn new(base_url: impl Into<String>, user_token: impl Into<String>) -> Self {
let http_client = crate::http::shared_client();
Self {
raw_client: http_client.clone(),
client: reqwest_middleware::ClientBuilder::new(http_client).build(),
base_url: base_url.into(),
credentials: crate::util::grok_auth_credentials::GrokAuthCredentials::new(Some(
user_token.into(),
)),
session_id: None,
}
}
pub fn with_deployment_key(mut self, key: Option<String>) -> Self {
self.credentials.deployment_key = key;
self
}
pub fn with_alpha_test_key(mut self, key: Option<String>) -> Self {
self.credentials.alpha_test_key = key;
self
}
pub fn with_session_id(mut self, session_id: impl Into<String>) -> Self {
self.session_id = Some(session_id.into());
self
}
/// Attach an `AuthManager` so the request signing and 401
/// recovery go through the consolidated auth path.
pub fn with_auth(mut self, auth_manager: std::sync::Arc<crate::auth::AuthManager>) -> Self {
let provider: std::sync::Arc<dyn kigi_auth::AuthCredentialProvider> = std::sync::Arc::new(
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
auth_manager.clone(),
self.credentials.deployment_key.clone(),
self.credentials.alpha_test_key.clone(),
),
);
self.credentials = self.credentials.with_auth_manager(auth_manager);
self.client = crate::http::with_auth_retry(self.raw_client.clone(), provider);
self
}
async fn send_authed(
&self,
builder: RequestBuilder,
op: &'static str,
) -> Result<reqwest::Response> {
let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder);
let request = builder.build().context(op)?;
self.client.execute(request).await.map_err(|e| match e {
reqwest_middleware::Error::Middleware(e) => e.context(op),
reqwest_middleware::Error::Reqwest(e) => anyhow::Error::from(e).context(op),
})
}
/// Non-auth headers only -- the `Authorization` header lives in
/// `send_authed` so it picks up freshly-refreshed tokens.
fn add_common_headers(&self, builder: RequestBuilder) -> RequestBuilder {
builder
}
fn check_response(&self, response: reqwest::Response, op: &str) -> anyhow::Error {
if response.status() == reqwest::StatusCode::UNAUTHORIZED {
self.record_401_attribution(op);
anyhow::anyhow!("{op}: {}", self.credentials.auth_error_hint())
} else {
anyhow::anyhow!("{op} failed: {}", response.status())
}
}
/// Emit a single `auth 401 attribution` log entry tagged with
/// `consumer = "SessionRegistryClient.<op>"`. The op string is the
/// operation name passed to `check_response` (e.g.,
/// `"session register"`).
fn record_401_attribution(&self, op: &str) {
if let Some(manager) = self.credentials.auth_manager() {
let resolved = self.credentials.resolve();
let sent = resolved
.deployment_key
.clone()
.or(resolved.user_token.clone());
crate::auth::attribution::record_consumer_401(
manager.as_ref(),
self.session_id.as_deref(),
crate::auth::attribution::ConsumerKind::SessionRegistryClient,
op,
sent.as_deref(),
);
}
}
fn post(&self, url: &str) -> RequestBuilder {
self.add_common_headers(self.raw_client.post(url))
}
fn get(&self, url: &str) -> RequestBuilder {
self.add_common_headers(self.raw_client.get(url))
}
/// POST /v1/sessions/register (idempotent via ON CONFLICT)
pub async fn register(&self, req: &RegisterRequest) -> Result<()> {
let url = format!("{}/sessions/register", self.base_url);
let response = self
.send_authed(self.post(&url).json(req), "session register")
.await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session register"));
}
Ok(())
}
/// POST /v1/sessions/{id}/replicas/update
pub async fn update(&self, session_id: &str, req: &UpdateRequest) -> Result<()> {
let url = format!("{}/sessions/{}/replicas/update", self.base_url, session_id);
let response = self
.send_authed(self.post(&url).json(req), "session update")
.await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session update"));
}
Ok(())
}
/// POST /v1/sessions/{id}/replicas/finalize
pub async fn finalize(&self, session_id: &str) -> Result<()> {
let url = format!(
"{}/sessions/{}/replicas/finalize",
self.base_url, session_id
);
let response = self
.send_authed(self.post(&url), "session finalize")
.await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session finalize"));
}
Ok(())
}
/// GET /v1/sessions/search
pub async fn search(&self, query: Option<&str>, limit: i64) -> Result<Vec<SessionRecord>> {
let url = format!("{}/sessions/search", self.base_url);
let mut builder = self.get(&url).query(&[("limit", limit.to_string())]);
if let Some(q) = query {
builder = builder.query(&[("query", q)]);
}
let response = self.send_authed(builder, "session search").await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session search"));
}
let resp: SearchResponse = response.json().await.context("parse search response")?;
Ok(resp.sessions)
}
/// GET /v1/sessions/{id}/replicas
pub async fn get_session(&self, session_id: &str) -> Result<SessionRecord> {
let url = format!("{}/sessions/{}/replicas", self.base_url, session_id);
let response = self.send_authed(self.get(&url), "session get").await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session get"));
}
response.json().await.context("parse session response")
}
/// GET /v1/sessions/{id}/download — returns a signed GCS URL without downloading.
pub async fn get_download_url(
&self,
session_id: &str,
file: &str,
turn: i32,
) -> Result<String> {
let url = format!("{}/sessions/{}/download", self.base_url, session_id);
let builder = self
.get(&url)
.query(&[("file", file), ("turn", &turn.to_string())]);
let response = self.send_authed(builder, "session download url").await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session download url"));
}
let resp: DownloadResponse = response.json().await.context("parse download response")?;
Ok(resp.download_url)
}
/// GET /v1/sessions/{id}/download — returns a signed URL, then streams to dest file.
pub async fn download_file(
&self,
session_id: &str,
file: &str,
turn: i32,
dest: &std::path::Path,
) -> Result<()> {
let url = format!("{}/sessions/{}/download", self.base_url, session_id);
let builder = self
.get(&url)
.query(&[("file", file), ("turn", &turn.to_string())]);
let response = self.send_authed(builder, "session download").await?;
if !response.status().is_success() {
return Err(self.check_response(response, "session download"));
}
let resp: DownloadResponse = response.json().await.context("parse download response")?;
// Stream from the signed GCS URL directly to disk (archives can be hundreds of MB)
let mut gcs_response = self
.raw_client
.get(&resp.download_url)
.send()
.await
.context("download from GCS")?;
if !gcs_response.status().is_success() {
anyhow::bail!("GCS download failed: {}", gcs_response.status());
}
if let Some(parent) = dest.parent() {
tokio::fs::create_dir_all(parent).await?;
}
let mut out = tokio::fs::File::create(dest)
.await
.context("create dest file")?;
let chunk_timeout = std::time::Duration::from_secs(60);
loop {
match tokio::time::timeout(chunk_timeout, gcs_response.chunk()).await {
Ok(Ok(Some(chunk))) => {
tokio::io::AsyncWriteExt::write_all(&mut out, &chunk)
.await
.context("write chunk to disk")?;
}
Ok(Ok(None)) => break,
Ok(Err(e)) => return Err(e).context("read GCS chunk"),
Err(_) => anyhow::bail!(
"GCS download stalled: no data received for {chunk_timeout:?} \
while downloading {file}"
),
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
// ── UpdateRequest wire shapes ────────────────────────────────────────────
//
// The writer split relies on two distinct update payloads being sent at
// different times:
//
// 1. Immediate post-turn: `last_turn_number` + `repo_head_at_end`
// 2. Artifact-ready: `restorable_turn_number` only
//
// These tests verify that `skip_serializing_if = "Option::is_none"` does the
// right thing for each shape, so old servers silently ignore the new field and
// clients don't accidentally overwrite unrelated fields with nulls.
#[test]
fn immediate_turn_update_omits_restorable_field() {
let req = UpdateRequest {
summary: None,
first_prompt: None,
last_turn_number: Some(5),
repo_head_at_end: Some("abc123".into()),
restorable_turn_number: None,
};
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["lastTurnNumber"], 5);
assert_eq!(json["repoHeadAtEnd"], "abc123");
assert!(json.get("restorableTurnNumber").is_none());
assert!(json.get("summary").is_none());
assert!(json.get("firstPrompt").is_none());
}
#[test]
fn restorable_turn_update_omits_last_turn_and_head_fields() {
let req = UpdateRequest {
summary: None,
first_prompt: None,
last_turn_number: None,
repo_head_at_end: None,
restorable_turn_number: Some(5),
};
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["restorableTurnNumber"], 5);
assert!(json.get("lastTurnNumber").is_none());
assert!(json.get("repoHeadAtEnd").is_none());
assert!(json.get("summary").is_none());
assert!(json.get("firstPrompt").is_none());
}
#[test]
fn summary_update_omits_all_turn_fields() {
let req = UpdateRequest {
summary: Some("My session summary".into()),
first_prompt: None,
last_turn_number: None,
repo_head_at_end: None,
restorable_turn_number: None,
};
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["summary"], "My session summary");
assert!(json.get("lastTurnNumber").is_none());
assert!(json.get("restorableTurnNumber").is_none());
assert!(json.get("repoHeadAtEnd").is_none());
}
// Wire-contract tests: server reads the camelCase `deviceId` key.
fn minimal_register_request(device_id: Option<String>) -> RegisterRequest {
RegisterRequest {
session_id: "s1".into(),
cwd: "/x".into(),
gcs_trace_prefix: "t".into(),
model_id: None,
repo_remote_url: None,
repo_branch: None,
repo_head_at_start: None,
hostname: None,
device_id,
parent_session_id: None,
session_kind: None,
subagent_type: None,
subagent_persona: None,
subagent_role: None,
fork_context_source: None,
subagent_depth: None,
}
}
#[test]
fn register_request_serializes_device_id_as_camel_case() {
let req = minimal_register_request(Some("machine-uuid-123".into()));
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["deviceId"], "machine-uuid-123");
assert!(json.get("device_id").is_none());
}
#[test]
fn register_request_serializes_empty_device_id_as_present() {
let req = minimal_register_request(Some(String::new()));
let json = serde_json::to_value(&req).unwrap();
assert_eq!(json["deviceId"], "");
}
#[test]
fn register_request_omits_device_id_when_none() {
let req = minimal_register_request(None);
let json = serde_json::to_value(&req).unwrap();
assert!(json.get("deviceId").is_none());
assert!(json.get("device_id").is_none());
}
// ── SessionRecord backward compatibility ─────────────────────────────────
//
// Older servers do not include `restorable_turn_number` in their response.
// The field is `#[serde(default)]` so it must deserialize as `None` when
// absent, keeping new clients compatible with old servers.
#[test]
fn session_record_without_restorable_turn_deserializes_as_none() {
let json = serde_json::json!({
"sessionId": "sess-abc",
"summary": "hello",
"firstPrompt": null,
"modelId": null,
"createdAt": "2026-01-01T00:00:00Z",
"updatedAt": "2026-01-01T00:00:00Z",
"lastTurnNumber": 3,
"cwd": "/home/user/repo",
"repoRemoteUrl": null,
"hostname": null,
"status": "active",
"gcsTracePrefix": "sessions/sess-abc",
"gcsBucket": "my-bucket"
});
let record: SessionRecord = serde_json::from_value(json).unwrap();
assert_eq!(record.last_turn_number, 3);
assert_eq!(record.restorable_turn_number, None);
}
#[test]
fn session_record_with_restorable_turn_deserializes_correctly() {
let json = serde_json::json!({
"sessionId": "sess-xyz",
"summary": "hello",
"firstPrompt": null,
"modelId": null,
"createdAt": "2026-01-01T00:00:00Z",
"updatedAt": "2026-01-01T00:00:00Z",
"lastTurnNumber": 7,
"restorableTurnNumber": 6,
"cwd": "/home/user/repo",
"repoRemoteUrl": null,
"hostname": null,
"status": "active",
"gcsTracePrefix": "sessions/sess-xyz",
"gcsBucket": "my-bucket"
});
let record: SessionRecord = serde_json::from_value(json).unwrap();
assert_eq!(record.last_turn_number, 7);
assert_eq!(record.restorable_turn_number, Some(6));
}
/// Verify per-request auth resolve picks up rotated tokens.
#[tokio::test]
async fn session_registry_client_uses_active_auth_for_each_request() {
use crate::auth::{AuthManager, AuthMode, GrokAuth, GrokComConfig};
use axum::{Router, response::IntoResponse, routing::post};
use chrono::{Duration, Utc};
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
let captured = Arc::new(parking_lot::Mutex::new(None::<String>));
let captured_for_handler = captured.clone();
let router = Router::new().route(
"/sessions/register",
post(move |headers: axum::http::HeaderMap, _body: String| {
let captured = captured_for_handler.clone();
async move {
if let Some(auth) = headers.get(axum::http::header::AUTHORIZATION) {
*captured.lock() = Some(auth.to_str().unwrap_or("").to_owned());
}
(axum::http::StatusCode::OK, "").into_response()
}
}),
);
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
let addr: SocketAddr = listener.local_addr().unwrap();
tokio::spawn(async move { axum::serve(listener, router).await.unwrap() });
let dir = tempfile::tempdir().unwrap();
let am = Arc::new(AuthManager::new(dir.path(), GrokComConfig::default()));
am.hot_swap(GrokAuth {
key: "fresh-from-auth-manager".into(),
auth_mode: AuthMode::ApiKey,
create_time: Utc::now(),
user_id: "user-42".into(),
expires_at: Some(Utc::now() + Duration::hours(1)),
..GrokAuth::test_default()
});
let client = SessionRegistryClient::new(format!("http://{addr}"), "STALE-build-time-token")
.with_auth(am);
let req = RegisterRequest {
session_id: "s1".into(),
cwd: "/x".into(),
gcs_trace_prefix: "t".into(),
model_id: None,
repo_remote_url: None,
repo_branch: None,
repo_head_at_start: None,
hostname: None,
device_id: None,
parent_session_id: None,
session_kind: None,
subagent_type: None,
subagent_persona: None,
subagent_role: None,
fork_context_source: None,
subagent_depth: None,
};
client.register(&req).await.unwrap();
let sent = captured.lock().clone().expect("server saw the request");
assert_eq!(
sent, "Bearer fresh-from-auth-manager",
"outgoing bearer must come from AuthManager (not the build-time token)"
);
}
// Verify the split-pointer invariant: last_turn_number can be ahead of
// restorable_turn_number (codebase best-effort means a turn may be "done"
// but not yet restorable if session-state upload is still in flight).
#[test]
fn session_record_allows_last_turn_ahead_of_restorable() {
let json = serde_json::json!({
"sessionId": "sess-lag",
"summary": "",
"firstPrompt": null,
"modelId": null,
"createdAt": "2026-01-01T00:00:00Z",
"updatedAt": "2026-01-01T00:00:00Z",
"lastTurnNumber": 10,
"restorableTurnNumber": 8,
"cwd": "/repo",
"repoRemoteUrl": null,
"hostname": null,
"status": "active",
"gcsTracePrefix": "sessions/sess-lag",
"gcsBucket": "bucket"
});
let record: SessionRecord = serde_json::from_value(json).unwrap();
assert!(record.last_turn_number > record.restorable_turn_number.unwrap_or(0));
}
}
@@ -0,0 +1,465 @@
#![cfg_attr(rustfmt, rustfmt::skip)]
#![allow(unused_imports)]
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use agent_client_protocol as acp;
use tokio::sync::{Notify, mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use crate::extensions::notification::{SessionNotification, SessionUpdate};
use crate::session::{
self, SessionCommand, SessionHandle, SessionThread,
commands::{PromptCompletionKind, PromptTurnResult as SubagentPromptTurnResult},
fs_watch::FsWatchCapabilities, info::Info as SessionInfo,
};
use crate::terminal::AsyncTerminalRunner;
use crate::tools::ToolContext;
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
use kigi_tools::implementations::grok_build::task::types::*;
use kigi_workspace::file_system::AsyncFileSystem;
use kigi_hunk_tracker::HunkTrackerHandle;
use super::*;
impl SubagentCoordinator {
pub fn new() -> Self {
Self {
pending: HashMap::new(),
active: HashMap::new(),
completed: HashMap::new(),
completion_notify: Arc::new(Notify::new()),
pending_completions: Vec::new(),
is_turn_active: Arc::new(std::sync::atomic::AtomicBool::new(false)),
running_gauge: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
block_wait_slots: HashMap::new(),
subagent_usage_not_applied_prompts: std::collections::HashSet::new(),
}
}
pub fn mark_subagent_usage_not_applied(&mut self, prompt_id: &str) {
self.subagent_usage_not_applied_prompts.insert(prompt_id.to_string());
}
pub fn subagent_usage_not_applied(&self, prompt_id: &str) -> bool {
self.subagent_usage_not_applied_prompts.contains(prompt_id)
}
pub fn clear_subagent_usage_not_applied(&mut self, prompt_id: &str) {
self.subagent_usage_not_applied_prompts.remove(prompt_id);
}
pub fn parent_prompt_id_for(&self, subagent_id: &str) -> Option<String> {
self.active
.get(subagent_id)
.and_then(|t| t.parent_prompt_id.clone())
.or_else(|| {
self.pending.get(subagent_id).and_then(|p| p.parent_prompt_id.clone())
})
}
/// Rebind the running-subagent gauge, copying the current count so a
/// late rebind cannot under-report.
pub fn set_running_gauge(&mut self, gauge: Arc<std::sync::atomic::AtomicUsize>) {
gauge
.store(
self.pending.len() + self.active.len(),
std::sync::atomic::Ordering::Relaxed,
);
self.running_gauge = gauge;
}
/// Recompute the gauge from `pending` + `active` after every mutation of
/// either map — recomputing (rather than incrementing) prevents drift.
fn sync_running_gauge(&self) {
self.running_gauge
.store(
self.pending.len() + self.active.len(),
std::sync::atomic::Ordering::Relaxed,
);
}
/// Returns a handle to the completion [`Notify`].
#[cfg_attr(
not(test),
expect(
dead_code,
reason = "used from tests only; remove expect when wired in production"
)
)]
pub fn completion_notify(&self) -> Arc<Notify> {
Arc::clone(&self.completion_notify)
}
/// Returns a shared handle to the turn-active flag.
pub fn turn_active_flag(&self) -> Arc<std::sync::atomic::AtomicBool> {
Arc::clone(&self.is_turn_active)
}
/// Whether the model's turn is currently active.
#[cfg_attr(
not(test),
expect(
dead_code,
reason = "used from tests only; remove expect when wired in production"
)
)]
pub fn is_turn_active(&self) -> bool {
self.is_turn_active.load(std::sync::atomic::Ordering::Relaxed)
}
/// Pending + active turn-blocking subagent IDs for `prompt_id`.
/// Background children are excluded: they outlive the turn by design, so
/// the freeze drain must not wait on them (their spend reaches the session
/// ledger when they finish; the prompt report flags them via
/// `background_live`).
pub fn outstanding_for_prompt(&self, prompt_id: &str) -> Vec<String> {
let mut ids: Vec<String> = self
.pending
.values()
.filter(|p| {
p.parent_prompt_id.as_deref() == Some(prompt_id) && !p.run_in_background
})
.map(|p| p.subagent_id.clone())
.chain(
self
.active
.values()
.filter(|t| {
t.parent_prompt_id.as_deref() == Some(prompt_id)
&& !t.run_in_background
})
.map(|t| t.subagent_id.clone()),
)
.collect();
ids.sort();
ids
}
/// True while any background child of `prompt_id` is pending or active.
/// Their spend is missing from the prompt report (it lands on the session
/// ledger at completion), so the report is incomplete — without waiting.
pub fn background_live_for_prompt(&self, prompt_id: &str) -> bool {
self
.pending
.values()
.any(|p| {
p.parent_prompt_id.as_deref() == Some(prompt_id) && p.run_in_background
})
|| self
.active
.values()
.any(|t| {
t.parent_prompt_id.as_deref() == Some(prompt_id)
&& t.run_in_background
})
}
/// Record that a foreground child was auto-backgrounded (await budget
/// expired): it no longer blocks the turn, so the freeze drain must stop
/// waiting on it.
pub fn mark_backgrounded(&mut self, subagent_id: &str) {
if let Some(t) = self.active.values_mut().find(|t| t.subagent_id == subagent_id)
{
t.run_in_background = true;
}
if let Some(p) = self.pending.values_mut().find(|p| p.subagent_id == subagent_id)
{
p.run_in_background = true;
}
}
pub fn outstanding_reply_for_prompt(
&self,
prompt_id: &str,
) -> kigi_tools::implementations::grok_build::task::types::SubagentOutstandingReply {
kigi_tools::implementations::grok_build::task::types::SubagentOutstandingReply {
live_ids: self.outstanding_for_prompt(prompt_id),
background_live: self.background_live_for_prompt(prompt_id),
subagent_usage_not_applied: self.subagent_usage_not_applied(prompt_id),
}
}
/// Drain all buffered completion summaries, returning them and clearing the buffer.
pub fn drain_pending_completions(&mut self) -> Vec<SubagentCompletionSummary> {
std::mem::take(&mut self.pending_completions)
}
/// Register a subagent as pending (initializing). Call this early,
/// before any blocking work like worktree creation, so that
/// `get_task_output` can report the subagent as initializing instead
/// of "not found".
pub fn insert_pending(&mut self, entry: PendingSubagent) {
self.pending.insert(entry.subagent_id.clone(), entry);
self.sync_running_gauge();
}
/// Remove a pending subagent without recording a failure.
/// Used by cancel flows where the subagent was intentionally stopped.
#[cfg(test)]
pub fn remove_pending(&mut self, id: &str) {
self.pending.remove(id);
self.sync_running_gauge();
}
/// Move a pending subagent directly to `completed` so it stays queryable via
/// `get_task_output`. `cancelled` stamps `"cancelled"` vs `"failed"`.
fn move_pending_to_terminal(&mut self, id: &str, error: &str, cancelled: bool) {
let Some(pending) = self.pending.remove(id) else {
return;
};
self.record_failure_completion(FailureCompletion {
subagent_id: pending.subagent_id,
subagent_type: pending.subagent_type,
description: pending.description,
parent_prompt_id: pending.parent_prompt_id,
parent_session_id: pending.parent_session_id,
persona: pending.persona,
started_at: pending.started_at,
error,
surface_completion: pending.surface_completion,
cancelled,
});
}
/// Move a pending subagent to `completed` as a failure so it stays queryable
/// via `get_task_output`.
pub fn move_pending_to_failed(&mut self, id: &str, error: &str) {
self.move_pending_to_terminal(id, error, false);
}
/// Like [`Self::move_pending_to_failed`] but stamps `"cancelled"` — a pending
/// subagent killed while initializing.
pub fn move_pending_to_cancelled(&mut self, id: &str, error: &str) {
self.move_pending_to_terminal(id, error, true);
}
/// Record a synthetic failure for a subagent that never reached `pending`.
pub fn record_pre_spawn_failure(
&mut self,
subagent_id: String,
subagent_type: String,
description: String,
parent_prompt_id: Option<String>,
parent_session_id: String,
error: &str,
surface_completion: bool,
) {
self.record_failure_completion(FailureCompletion {
subagent_id,
subagent_type,
description,
parent_prompt_id,
parent_session_id,
persona: None,
started_at: std::time::Instant::now(),
error,
surface_completion,
cancelled: false,
});
}
/// Insert a synthetic failed entry, push a completion summary, notify waiters.
/// Clears any stale pending entry for the same id.
fn record_failure_completion(&mut self, c: FailureCompletion<'_>) {
self.pending.remove(&c.subagent_id);
self.sync_running_gauge();
let FailureCompletion {
subagent_id,
subagent_type,
description,
parent_prompt_id,
parent_session_id,
persona,
started_at,
error,
surface_completion,
cancelled,
} = c;
let result = SubagentResult {
success: false,
cancelled,
error: Some(error.to_string()),
subagent_id: subagent_id.clone(),
..Default::default()
};
let summary_output = result.output.clone();
self.completed
.insert(
subagent_id.clone(),
CompletedSubagent {
subagent_id: subagent_id.clone(),
parent_session_id,
parent_prompt_id,
child_session_id: String::new(),
description: description.clone(),
subagent_type: subagent_type.clone(),
persona,
started_at,
completed_at: std::time::Instant::now(),
result,
resumed_from: None,
child_cwd: String::new(),
worktree_path: None,
snapshot_ref: None,
effective_model_id: String::new(),
block_waited: false,
explicitly_killed: false,
},
);
if surface_completion {
self.pending_completions
.push(SubagentCompletionSummary {
subagent_id,
subagent_type,
description,
success: false,
duration_ms: 0,
tool_calls: 0,
turns: 0,
output: summary_output,
});
}
self.completion_notify.notify_waiters();
}
pub fn insert(&mut self, tracker: SubagentTracker) {
self.pending.remove(&tracker.subagent_id);
self.active.insert(tracker.subagent_id.clone(), tracker);
self.sync_running_gauge();
}
/// Move a finished subagent from `active` to `completed`.
/// Returns the tracker if it was active.
pub fn move_to_completed(
&mut self,
id: &str,
description: String,
subagent_type: String,
result: SubagentResult,
) -> Option<SubagentTracker> {
let tracker = self.active.remove(id);
self.sync_running_gauge();
let started_at = tracker
.as_ref()
.map(|t| t.started_at)
.unwrap_or_else(std::time::Instant::now);
let parent_session_id = tracker
.as_ref()
.map(|t| t.parent_session_id.clone())
.unwrap_or_default();
let child_session_id = tracker
.as_ref()
.map(|t| t.child_session_id.0.to_string())
.unwrap_or_default();
let parent_prompt_id = tracker.as_ref().and_then(|t| t.parent_prompt_id.clone());
let persona = tracker.as_ref().and_then(|t| t.persona.clone());
let child_cwd = tracker
.as_ref()
.map(|t| t.child_cwd.clone())
.unwrap_or_default();
let worktree_path = tracker.as_ref().and_then(|t| t.worktree_path.clone());
let resumed_from = tracker.as_ref().and_then(|t| t.resumed_from.clone());
let effective_model_id = tracker
.as_ref()
.map(|t| t.effective_model_id.clone())
.unwrap_or_default();
let block_waited = tracker.as_ref().is_some_and(|t| t.block_waited);
let explicitly_killed = tracker.as_ref().is_some_and(|t| t.explicitly_killed);
let surface_completion = tracker.as_ref().is_none_or(|t| t.surface_completion);
self.completed
.insert(
id.to_string(),
CompletedSubagent {
subagent_id: id.to_string(),
parent_session_id,
parent_prompt_id,
child_session_id,
description,
subagent_type,
persona,
started_at,
completed_at: std::time::Instant::now(),
result,
resumed_from,
child_cwd,
worktree_path,
snapshot_ref: None,
effective_model_id,
block_waited,
explicitly_killed,
},
);
let completed = self.completed.get(id).expect("just inserted");
let success = completed.result.success && !completed.result.cancelled;
{
let preview = crate::util::truncate(&completed.result.output, 200);
let level_fn = if success {
kigi_log::unified_log::info
} else {
kigi_log::unified_log::error
};
level_fn(
if success { "subagent completed" } else { "subagent failed" },
None,
Some(
serde_json::json!(
{ "subagent_id" : & completed.subagent_id, "subagent_type" : &
completed.subagent_type, "effective_model" : & completed
.effective_model_id, "success" : success, "cancelled" : completed
.result.cancelled, "duration_ms" : completed.result.duration_ms,
"turns" : completed.result.turns, "tool_calls" : completed.result
.tool_calls, "output_preview" : preview, "error" : & completed
.result.error, }
),
),
);
}
if surface_completion {
self.pending_completions
.push(SubagentCompletionSummary {
subagent_id: id.to_string(),
subagent_type: completed.subagent_type.clone(),
description: completed.description.clone(),
success,
duration_ms: completed.result.duration_ms,
tool_calls: completed.result.tool_calls,
turns: completed.result.turns,
output: completed.result.output.clone(),
});
}
self.completion_notify.notify_waiters();
tracker
}
/// Record the durable worktree snapshot ref on a completed subagent so
/// in-memory `resume_from` resolution can rehydrate the disposed worktree.
/// No-op if the entry was already evicted (the on-disk meta.json still has it).
pub fn set_completed_snapshot_ref(&mut self, id: &str, snapshot_ref: String) {
if let Some(completed) = self.completed.get_mut(id) {
completed.snapshot_ref = Some(snapshot_ref);
}
}
/// Cancel all active subagents that were launched by a specific parent turn,
/// including `run_in_background: true` subagents.
pub fn cancel_by_parent_prompt_id(&mut self, parent_prompt_id: &str) {
for tracker in self.active.values() {
if tracker.parent_prompt_id.as_deref() == Some(parent_prompt_id) {
Self::cancel_tracker(tracker);
}
}
for pending in self.pending.values() {
if pending.parent_prompt_id.as_deref() == Some(parent_prompt_id) {
pending.cancel_token.cancel();
}
}
}
/// Attempt to cancel a subagent. Returns a typed outcome covering all cases:
/// - Active → cancel it, return Cancelled
/// - Pending (initializing) → fire its spawn token, return Cancelled
/// - Already finished → return AlreadyFinished with terminal status
/// - Unknown ID → return NotFound
pub fn cancel_with_outcome(&mut self, subagent_id: &str) -> SubagentCancelOutcome {
if let Some(tracker) = self.active.get(subagent_id) {
Self::cancel_tracker(tracker);
return SubagentCancelOutcome::Cancelled;
}
if let Some(pending) = self.pending.get(subagent_id) {
pending.cancel_token.cancel();
return SubagentCancelOutcome::Cancelled;
}
if let Some(entry) = self.completed.get(subagent_id) {
return SubagentCancelOutcome::AlreadyFinished {
status: entry.result.status().to_string(),
};
}
SubagentCancelOutcome::NotFound
}
/// Internal: send Cancel + Shutdown to a tracked subagent.
fn cancel_tracker(tracker: &SubagentTracker) {
tracker.cancel_token.cancel();
let _ = tracker
.child_handle
.cmd_tx
.send(SessionCommand::Cancel {
cancel_subagents: true,
kill_background_tasks: true,
rewind_if_pristine: false,
trigger: None,
});
let _ = tracker.child_handle.cmd_tx.send(SessionCommand::Shutdown);
}
}
@@ -0,0 +1,348 @@
#![cfg_attr(rustfmt, rustfmt::skip)]
#![allow(unused_imports)]
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use agent_client_protocol as acp;
use tokio::sync::{Notify, mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use crate::extensions::notification::{SessionNotification, SessionUpdate};
use crate::session::{
self, SessionCommand, SessionHandle, SessionThread,
commands::{PromptCompletionKind, PromptTurnResult as SubagentPromptTurnResult},
fs_watch::FsWatchCapabilities, info::Info as SessionInfo,
};
use crate::terminal::AsyncTerminalRunner;
use crate::tools::ToolContext;
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
use kigi_tools::implementations::grok_build::task::types::*;
use kigi_workspace::file_system::AsyncFileSystem;
use kigi_hunk_tracker::HunkTrackerHandle;
use super::*;
impl SubagentCoordinator {
/// Synchronous lookup of a subagent by ID.
///
/// Returns a three-way result so the caller can drop the `RefCell` borrow
/// before awaiting the signals handle for running subagents.
///
/// - `Ready` — completed/failed/cancelled snapshot, no async work needed.
/// - `NeedsSignals` — subagent is running; caller must await
/// `resolve_snapshot()` after dropping the coordinator borrow.
/// - `None` — ID not found in active, completed, or pending maps.
pub(crate) fn lookup(&self, id: &str) -> Option<SnapshotLookup> {
if let Some(tracker) = self.active.get(id) {
return Some(
SnapshotLookup::NeedsSignals(RunningSnapshotSeed {
subagent_id: tracker.subagent_id.clone(),
description: tracker.description.clone(),
subagent_type: tracker.subagent_type.clone(),
started_at_epoch_ms: instant_to_epoch_ms(tracker.started_at),
duration_ms: tracker.started_at.elapsed().as_millis() as u64,
persona: tracker.persona.clone(),
signals_handle: tracker.child_handle.signals_handle.clone(),
}),
);
}
if let Some(completed) = self.completed.get(id) {
let status = if completed.result.cancelled {
SubagentSnapshotStatus::Cancelled {
reason: completed.result.error.clone(),
}
} else if completed.result.success {
SubagentSnapshotStatus::Completed {
output: completed.result.output.to_string(),
tool_calls: completed.result.tool_calls,
turns: completed.result.turns,
worktree_path: completed.result.worktree_path.clone(),
}
} else {
SubagentSnapshotStatus::Failed {
error: completed
.result
.error
.clone()
.unwrap_or_else(|| "Unknown error".to_string()),
}
};
return Some(
SnapshotLookup::Ready(SubagentSnapshot {
subagent_id: completed.subagent_id.clone(),
description: completed.description.clone(),
subagent_type: completed.subagent_type.clone(),
status,
started_at_epoch_ms: instant_to_epoch_ms(completed.started_at),
duration_ms: completed.result.duration_ms,
persona: completed.persona.clone(),
}),
);
}
if let Some(pending) = self.pending.get(id) {
return Some(
SnapshotLookup::Ready(SubagentSnapshot {
subagent_id: pending.subagent_id.clone(),
description: pending.description.clone(),
subagent_type: pending.subagent_type.clone(),
status: SubagentSnapshotStatus::Initializing,
started_at_epoch_ms: instant_to_epoch_ms(pending.started_at),
duration_ms: pending.started_at.elapsed().as_millis() as u64,
persona: pending.persona.clone(),
}),
);
}
None
}
/// Return `(parent_session_id, child_session_id)` for a given subagent.
///
/// Checks active first, then completed. Returns `None` if not found.
pub(crate) fn session_ids_for(&self, id: &str) -> Option<(String, String)> {
if let Some(t) = self.active.get(id) {
return Some((t.parent_session_id.clone(), t.child_session_id.0.to_string()));
}
if let Some(c) = self.completed.get(id) {
return Some((c.parent_session_id.clone(), c.child_session_id.clone()));
}
None
}
/// Mark a subagent as block-waited so auto-wake is suppressed on completion.
pub(crate) fn mark_block_waited(&mut self, id: &str) {
if let Some(t) = self.active.get_mut(id) {
t.block_waited = true;
} else if let Some(c) = self.completed.get_mut(id) {
c.block_waited = true;
}
}
/// Clear the block-waited flag after a block timed out without receiving
/// the completion, so auto-wake can still fire when the subagent finishes.
pub(crate) fn clear_block_waited(&mut self, id: &str) {
if let Some(t) = self.active.get_mut(id) {
t.block_waited = false;
} else if let Some(c) = self.completed.get_mut(id) {
c.block_waited = false;
}
}
/// Whether a block-waiter already consumed this subagent's result.
pub(crate) fn is_block_waited(&self, id: &str) -> bool {
self.active.get(id).is_some_and(|t| t.block_waited)
|| self.completed.get(id).is_some_and(|c| c.block_waited)
}
/// Register a live blocking-query reply slot and mark `block_waited`.
///
/// The slot lets `block_wait_delivered_or_live` verify at completion
/// time that the waiter can still receive the result — the flag alone
/// can be stale when the waiting turn was cancelled moments before the
/// subagent finished.
pub(crate) fn register_block_wait(&mut self, id: &str, slot: BlockWaitSlot) {
self.mark_block_waited(id);
self.block_wait_slots.entry(id.to_string()).or_default().push(slot);
}
/// Drop a previously registered reply slot (query poll loop exited).
pub(crate) fn unregister_block_wait(&mut self, id: &str, slot: &BlockWaitSlot) {
if let Some(slots) = self.block_wait_slots.get_mut(id) {
slots.retain(|s| !std::rc::Rc::ptr_eq(s, slot));
if slots.is_empty() {
self.block_wait_slots.remove(id);
}
}
}
/// Decision-time gate for the completion auto-wake: returns true when
/// the result was already delivered to a blocking waiter, or a live
/// waiter is still parked and will receive it. When every registered
/// waiter is gone (receivers dropped by a cancelled turn), clears
/// `block_waited` and returns false so the auto-wake fires.
///
/// This closes the race where the query poll loop clears the flag up to
/// one poll interval *after* the caller cancelled — the completion
/// handler could read the stale flag in that window and skip the wake.
/// Consumes the id's slot registrations (completion is terminal).
pub(crate) fn block_wait_delivered_or_live(&mut self, id: &str) -> bool {
let slots = self.block_wait_slots.remove(id).unwrap_or_default();
if !self.is_block_waited(id) {
return false;
}
let delivered_or_live = slots.is_empty()
|| slots
.iter()
.any(|s| s.borrow().as_ref().is_none_or(|tx| !tx.is_closed()));
if !delivered_or_live {
self.clear_block_waited(id);
}
delivered_or_live
}
/// Mark a subagent as explicitly killed so auto-wake is suppressed on completion.
pub(crate) fn mark_explicitly_killed(&mut self, id: &str) {
if let Some(t) = self.active.get_mut(id) {
t.explicitly_killed = true;
} else if let Some(c) = self.completed.get_mut(id) {
c.explicitly_killed = true;
}
}
/// Whether the model explicitly killed this subagent via the kill tool.
pub(crate) fn is_explicitly_killed(&self, id: &str) -> bool {
self.active.get(id).is_some_and(|t| t.explicitly_killed)
|| self.completed.get(id).is_some_and(|c| c.explicitly_killed)
}
/// Return fork provenance for a given subagent.
pub(crate) fn provenance_for(&self, id: &str) -> SubagentProvenance {
if let Some(t) = self.active.get(id) {
return SubagentProvenance {
fork_parent_prompt_id: t.parent_prompt_id.clone(),
resumed_from: t.resumed_from.clone(),
};
}
if let Some(c) = self.completed.get(id) {
return SubagentProvenance {
fork_parent_prompt_id: c.parent_prompt_id.clone(),
resumed_from: c.resumed_from.clone(),
};
}
SubagentProvenance::default()
}
/// Resolve a completed subagent scoped to the requesting parent session.
///
/// Returns `None` if the subagent is not found, still active, or belongs
/// to a different parent session (prevents cross-session context bleed).
///
/// Fast path: checks the in-memory `completed` map first. When that
/// misses (e.g. after TTL eviction), falls back to on-disk metadata
/// in `{parent_session_dir}/subagents/{id}/meta.json`.
pub(crate) fn resumable_source_for(
&self,
id: &str,
parent_session_id: &str,
parent_cwd: &Path,
) -> Option<ResumeSourceData> {
if let Some(completed) = self.completed.get(id) {
if completed.parent_session_id != parent_session_id {
return None;
}
return Some(ResumeSourceData {
subagent_id: completed.subagent_id.clone(),
child_session_id: completed.child_session_id.clone(),
child_cwd: completed.child_cwd.clone(),
worktree_path: completed.worktree_path.clone(),
snapshot_ref: completed.snapshot_ref.clone(),
subagent_type: completed.subagent_type.clone(),
persona: completed.persona.clone(),
model_id: Some(completed.effective_model_id.clone()),
});
}
let parent_info = SessionInfo {
id: acp::SessionId::new(parent_session_id),
cwd: parent_cwd.to_string_lossy().to_string(),
};
let meta_path = session::persistence::session_dir(&parent_info)
.join("subagents")
.join(id)
.join("meta.json");
let data = std::fs::read_to_string(&meta_path).ok()?;
let meta: SubagentMeta = serde_json::from_str(&data).ok()?;
if meta.parent_session_id != parent_session_id {
return None;
}
match meta.status.as_str() {
"completed" | "failed" | "cancelled" => {}
_ => return None,
}
Some(ResumeSourceData {
subagent_id: meta.subagent_id,
child_session_id: meta.child_session_id,
child_cwd: meta.child_cwd.unwrap_or_default(),
worktree_path: meta.worktree_path.map(PathBuf::from),
snapshot_ref: meta.snapshot_ref,
subagent_type: meta.subagent_type,
persona: meta.persona,
model_id: meta.effective_model_id,
})
}
/// Check whether an ID refers to a currently-active (running) subagent.
pub(crate) fn is_active(&self, id: &str) -> bool {
self.active.contains_key(id)
}
/// Whether the coordinator still has this id in flight (spawning or running).
/// Orphan reconcile skips these — there is nothing stuck to heal.
pub(crate) fn is_active_or_pending(&self, id: &str) -> bool {
self.active.contains_key(id) || self.pending.contains_key(id)
}
/// The terminal `SubagentFinished` for an id the coordinator already holds in
/// `completed`, else `None`. Lets orphan reconcile re-emit a subagent's real
/// outcome when only its terminal meta write was lost (reconnect race: entry
/// in `completed` but the on-disk meta is still `running`) instead of
/// force-cancelling it and discarding the result.
pub(crate) fn completed_finish(&self, id: &str) -> Option<SessionUpdate> {
let c = self.completed.get(id)?;
let duration_ms = c
.completed_at
.saturating_duration_since(c.started_at)
.as_millis() as u64;
Some(SessionUpdate::SubagentFinished {
subagent_id: c.subagent_id.clone(),
child_session_id: c.child_session_id.clone(),
status: c.result.status().to_string(),
error: c.result.error.clone(),
tool_calls: c.result.tool_calls,
turns: c.result.turns,
duration_ms,
tokens_used: 0,
output: None,
will_wake: false,
})
}
/// TTL cleanup: remove completed entries older than 30 minutes.
pub fn evict_stale_completed(&mut self) {
let cutoff = std::time::Duration::from_secs(30 * 60);
self.completed.retain(|_, entry| entry.completed_at.elapsed() < cutoff);
}
/// Snapshot all currently-running subagents for compaction state context.
///
/// Returns one `ActiveSubagentSummary` per entry in the `active` map.
/// Completed/failed/cancelled subagents are NOT included — they live in
/// the `completed` map and are irrelevant for post-compaction reminders
/// (the model already saw their tool results before compaction).
///
/// The `elapsed_ms` field is computed from `started_at.elapsed()` at call
/// time, so the values are a snapshot of "right now" — appropriate for
/// compaction since it happens once and the reminder is static.
#[cfg(test)]
pub fn active_summaries(&self) -> Vec<ActiveSubagentSummary> {
self.active.values().map(tracker_to_summary).collect()
}
pub fn active_summaries_for(
&self,
parent_session_id: &str,
) -> Vec<ActiveSubagentSummary> {
self.active
.values()
.filter(|t| t.parent_session_id == parent_session_id)
.map(tracker_to_summary)
.collect()
}
/// Return seeds for all running subagents belonging to `parent_session_id`.
///
/// Each seed carries copied identity metadata plus a cloned
/// `SessionSignalsHandle` so the caller can resolve live progress
/// asynchronously after dropping the coordinator borrow.
///
/// Returns an empty `Vec` if no active subagents match the given
/// parent session ID. Callers (e.g. the `x.ai/subagent/list_running`
/// ACP handler) should treat an empty result as a normal "no running
/// subagents" response, not an error.
pub(crate) fn list_running_for_parent(
&self,
parent_session_id: &str,
) -> Vec<RunningSubagentListSeed> {
self.active
.values()
.filter(|t| t.parent_session_id == parent_session_id)
.map(|t| RunningSubagentListSeed {
subagent_id: t.subagent_id.clone(),
parent_session_id: t.parent_session_id.clone(),
child_session_id: t.child_session_id.0.to_string(),
subagent_type: t.subagent_type.clone(),
description: t.description.clone(),
started_at_epoch_ms: instant_to_epoch_ms(t.started_at),
duration_ms: t.started_at.elapsed().as_millis() as u64,
signals_handle: t.child_handle.signals_handle.clone(),
})
.collect()
}
}
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,191 @@
//! Subscription check for paywall gate lift.
//!
//! Provides `single_check()` which queries `GET /user?include=subscription`
//! for the live subscription tier from the backend, independent of the JWT.
//! If a qualifying tier is detected, does a best-effort JWT refresh and
//! settings re-fetch, then returns an `UnblockResult` so the agent can
//! lift the gate.
//!
//! The pager drives the polling via `x.ai/auth/check_subscription`: the 5s
//! paywall chain, the free-tier watch, the refocus check, and
//! verify-before-paywall gate deferral (see the pager's `app::subscription`
//! module).
use crate::auth::AuthManager;
use crate::auth::UserInfo;
use crate::auth::manager::RefreshReason;
use crate::auth::token_type::TokenType;
use std::sync::Arc;
use std::time::Duration;
/// Subscription tiers that qualify for Grok Build access.
/// Any active subscription qualifies -- the access gate in remote settings
/// controls which tiers are actually allowed.
const QUALIFYING_TIERS: &[&str] = &[
"SuperGrokPro",
"GrokPro",
"SuperGrokLite",
"XPremiumPlus",
"XPremium",
"XBasic",
];
/// Successful subscription check result: confirmed qualifying tier +
/// optionally refreshed settings.
pub(crate) struct UnblockResult {
pub(crate) new_tier: String,
pub(crate) settings: Option<crate::util::config::RemoteSettings>,
}
/// Fetch `/user?include=subscription` and return the parsed `UserInfo`.
async fn fetch_user_info(
http_client: &reqwest::Client,
url: &str,
auth: &crate::auth::GrokAuth,
auth_manager: &AuthManager,
alpha_test_key: Option<&str>,
) -> Result<UserInfo, &'static str> {
let request = http_client
.get(url)
.timeout(Duration::from_secs(10))
.header("Authorization", format!("Bearer {}", auth.key))
.header(
"X-XAI-Token-Auth",
auth_manager.grok_com_config().token_header.as_str(),
)
.header("x-grok-client-version", kigi_version::VERSION)
.header(
crate::http::CLIENT_MODE_HEADER,
crate::http::process_client_mode(),
);
let _ = alpha_test_key;
match request.send().await {
Ok(resp) if resp.status().is_success() => {
resp.json::<UserInfo>().await.map_err(|_| "parse")
}
Ok(_resp) => Err("http_status"),
Err(e) if e.is_timeout() => Err("timeout"),
Err(_) => Err("transport"),
}
}
/// Single-shot subscription check. Called by the pager every 5s while
/// the paywall is shown (`x.ai/auth/check_subscription`).
///
/// Queries `/user?include=subscription` for the live tier. If a qualifying
/// tier is found, does a best-effort JWT refresh + settings re-fetch and
/// returns `Some(UnblockResult)`. Returns `None` if no qualifying
/// subscription exists or the request fails.
#[tracing::instrument(name = "paywall_check", skip_all, fields(user_id = %user_id))]
pub(crate) async fn single_check(
auth_manager: Arc<AuthManager>,
proxy_base_url: &str,
alpha_test_key: Option<&str>,
user_id: &str,
) -> Option<UnblockResult> {
let user_url = format!("{}/user?include=subscription", proxy_base_url);
let http_client = crate::http::shared_client();
let auth = auth_manager.current()?;
let user_info = match fetch_user_info(
&http_client,
&user_url,
&auth,
&auth_manager,
alpha_test_key,
)
.await
{
Ok(ui) => ui,
Err(kind) => {
kigi_log::unified_log::warn(
"paywall_check_error",
None,
Some(serde_json::json!({ "user_id" : user_id, "kind" : kind })),
);
return None;
}
};
kigi_log::unified_log::info(
"paywall_check_result",
None,
Some(serde_json::json!(
{ "user_id" : user_id, "subscription_tier" : user_info.subscription_tier,
}
)),
);
let new_tier = match &user_info.subscription_tier {
Some(tier) if !tier.is_empty() => tier.clone(),
_ => return None,
};
if !QUALIFYING_TIERS.contains(&new_tier.as_str()) {
return None;
}
kigi_log::unified_log::info(
"paywall_check_subscription_detected",
None,
Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier, })),
);
if let Err(e) = auth_manager
.refresh_chain(TokenType::OidcSession, RefreshReason::ServerRejected)
.await
{
kigi_log::unified_log::warn(
"paywall_check_error",
None,
Some(serde_json::json!(
{ "user_id" : user_id, "kind" : "refresh_failed", "detail" : e
.to_string(), }
)),
);
}
let settings = if crate::util::config::resolve_remote_fetch_enabled() {
let base_url = proxy_base_url.to_string();
let auth_for_settings = auth_manager.current().unwrap_or(auth);
let atk = alpha_test_key.map(str::to_string);
tokio::task::spawn_blocking(move || {
crate::remote::fetch_settings_blocking(&base_url, &auth_for_settings, atk.as_deref())
})
.await
.ok()
.flatten()
} else {
None
};
kigi_log::unified_log::info(
"paywall_check_unblocked",
None,
Some(serde_json::json!({ "user_id" : user_id, "new_tier" : new_tier })),
);
Some(UnblockResult { new_tier, settings })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn qualifying_tiers_includes_all_paid_tiers() {
for tier in &[
"SuperGrokPro",
"GrokPro",
"SuperGrokLite",
"XPremiumPlus",
"XPremium",
"XBasic",
] {
assert!(
QUALIFYING_TIERS.contains(tier),
"{tier} must be in QUALIFYING_TIERS"
);
}
}
#[test]
fn free_tier_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&"Free"));
}
#[test]
fn empty_tier_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&""));
}
/// The subscription check only returns `Some` when `/user` reports a
/// qualifying tier. Verify the tier matching is exact (no prefix match).
#[test]
fn partial_tier_name_is_not_qualifying() {
assert!(!QUALIFYING_TIERS.contains(&"Super"));
assert!(!QUALIFYING_TIERS.contains(&"Grok"));
assert!(!QUALIFYING_TIERS.contains(&"XPremium+"));
}
}
File diff suppressed because it is too large Load Diff