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,479 @@
//! Bridge a leader IPC connection into an `AcpClientChannel`.
//!
//! Adapts the leader's raw JSON string channels into the typed ACP channel
//! interface, reusing `ClientSideConnection` from `agent_client_protocol`
//! for JSON-RPC ser/deser.
use std::sync::Arc;
use std::thread;
use anyhow::Result;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, simplex};
use tokio::sync::{Mutex as TokioMutex, mpsc};
use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
use tokio_util::sync::CancellationToken;
use agent_client_protocol as acp;
use kigi_acp_lib::{
AcpClientChannel, AcpGatewayReceiver, AcpGatewaySender, LineBufferedRead, acp_channels,
};
pub use kigi_shell::leader::ConnectionStatus;
use kigi_shell::leader::{LeaderConnection, LeaderReconnector, ReconnectPolicy};
const MAX_BUF: usize = 8 * 1024 * 1024;
pub struct LeaderBridge {
pub channel: AcpClientChannel,
pub cancel: CancellationToken,
pub thread_handle: thread::JoinHandle<Result<()>>,
}
/// How [`forward_outbound_line`] resolved one outbound line.
#[derive(Debug, PartialEq, Eq)]
enum ForwardOutcome {
Sent,
/// The connection the line was composed for died and a new one replaced
/// it; the line was dropped.
DroppedStale,
Cancelled,
}
/// Send one outbound line to the (swappable) leader tx.
///
/// A failed send means the connection is dead; the line is held — blocking
/// the lines queued behind it — until the reader task installs a fresh tx,
/// then dropped. Replaying it would be worse: a stale `session/load`
/// re-delivered on the new connection triggers a second full replay into the
/// same reload window (duplicated transcript); the reconnect re-init
/// re-establishes state explicitly instead.
///
/// Scoping is by FIRST OBSERVED send failure, a best-effort heuristic: a
/// pre-disconnect line whose first send happens after the swap never fails
/// and goes out on the new connection, and lines queued behind a held one
/// are forwarded post-swap regardless of when they were composed.
async fn forward_outbound_line(
leader_tx: &TokioMutex<mpsc::UnboundedSender<String>>,
cancel: &CancellationToken,
mut pending: String,
) -> ForwardOutcome {
let mut failed_on: Option<mpsc::UnboundedSender<String>> = None;
loop {
{
let tx = leader_tx.lock().await;
if let Some(ref dead) = failed_on
&& !tx.same_channel(dead)
{
return ForwardOutcome::DroppedStale;
}
pending = match tx.send(pending) {
Ok(()) => return ForwardOutcome::Sent,
Err(mpsc::error::SendError(returned)) => returned,
};
if failed_on.is_none() {
failed_on = Some(tx.clone());
}
}
tracing::debug!("Writer send failed; holding line until the reconnect swap");
tokio::select! {
biased;
_ = cancel.cancelled() => return ForwardOutcome::Cancelled,
_ = tokio::time::sleep(std::time::Duration::from_millis(100)) => {}
}
}
}
/// Bridge a `LeaderConnection` into an `AcpClientChannel`.
///
/// When `reconnector` is `Some`, the bridge automatically attempts to reconnect
/// on leader disconnect using the given `policy`. On reconnection failure (or if
/// `reconnector` is `None`), the cancel token fires so the caller can exit.
pub fn bridge_leader_connection(
conn: LeaderConnection,
cancel: CancellationToken,
reconnector: Option<LeaderReconnector>,
policy: ReconnectPolicy,
) -> Result<LeaderBridge> {
let (leader_tx, leader_rx) = conn.into_channels();
bridge_channels(leader_tx, leader_rx, cancel, reconnector, policy)
}
/// Bridge raw IPC channels into an `AcpClientChannel`.
///
/// Spawns a dedicated thread with a `LocalSet` because `ClientSideConnection`
/// uses `spawn_local` internally. On leader disconnect, reconnects via
/// `reconnector` (if provided) or fires the cancel token.
pub(crate) fn bridge_channels(
leader_tx: mpsc::UnboundedSender<String>,
leader_rx: mpsc::UnboundedReceiver<String>,
cancel: CancellationToken,
reconnector: Option<LeaderReconnector>,
policy: ReconnectPolicy,
) -> Result<LeaderBridge> {
let (client_channel, agent_channel) = acp_channels();
let (incoming_read, incoming_write) = simplex(MAX_BUF);
let (outgoing_read, outgoing_write) = simplex(MAX_BUF);
let bridge_cancel = cancel.clone();
let thread_handle = thread::Builder::new()
.name("pager-leader-bridge".into())
.spawn(move || -> Result<()> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let local = tokio::task::LocalSet::new();
local.block_on(&rt, async move {
let leader_tx_shared = Arc::new(TokioMutex::new(leader_tx));
// Reader: leader IPC -> incoming simplex pipe -> ClientSideConnection
let cancel_r = bridge_cancel.clone();
let leader_tx_for_reader = leader_tx_shared.clone();
let reader_task = tokio::task::spawn_local(async move {
let mut incoming_write = incoming_write;
let mut leader_rx = leader_rx;
loop {
tokio::select! {
biased;
_ = cancel_r.cancelled() => break,
msg = leader_rx.recv() => {
match msg {
Some(json_line) => {
if incoming_write.write_all(json_line.as_bytes()).await.is_err()
|| incoming_write.write_all(b"\n").await.is_err()
{
break;
}
}
None => {
tracing::warn!("Leader connection closed");
if let Some(ref reconnector) = reconnector {
tracing::info!("Attempting to reconnect to leader...");
match reconnector.reconnect(policy, &cancel_r).await {
Ok((new_tx, new_rx, _disconnect_rx)) => {
tracing::info!("Reconnected to leader IPC");
leader_rx = new_rx;
*leader_tx_for_reader.lock().await = new_tx;
// Swap first, notify second — see
// `LeaderReconnector::notify_connected`.
reconnector.notify_connected();
continue;
}
Err(e) => {
tracing::error!(error = %e, "Failed to reconnect to leader");
cancel_r.cancel();
break;
}
}
} else {
cancel_r.cancel();
break;
}
}
}
}
}
}
});
// Writer: ClientSideConnection -> outgoing simplex pipe -> leader IPC
let cancel_w = bridge_cancel.clone();
let leader_tx_for_writer = leader_tx_shared;
let writer_task = tokio::task::spawn_local(async move {
let mut reader = BufReader::new(outgoing_read);
let mut line = String::new();
loop {
line.clear();
tokio::select! {
biased;
_ = cancel_w.cancelled() => break,
result = reader.read_line(&mut line) => {
match result {
Ok(0) => break,
Ok(_) => {
let pending = line.trim_end();
if pending.is_empty() {
continue;
}
match forward_outbound_line(
&leader_tx_for_writer,
&cancel_w,
pending.to_string(),
)
.await
{
ForwardOutcome::Sent => {}
ForwardOutcome::DroppedStale => {
// Unified-log marker: this drop is deliberate
// (replaying a stale `session/load` would
// double-replay the transcript), but it can eat
// one-shot notifications like `session/cancel`, a
// known stuck-cancel failure mode. Record WHAT was
// dropped so the next
// investigation sees it in the unified log.
let method = serde_json::from_str::<serde_json::Value>(pending)
.ok()
.and_then(|j| {
j.get("method").and_then(|m| m.as_str()).map(str::to_owned)
});
crate::unified_log::warn(
"leader.ipc.outbound_dropped_stale",
None,
Some(serde_json::json!({
"method": method,
"len": pending.len(),
})),
);
tracing::debug!(
"Dropped outbound line composed for a replaced leader connection"
);
}
ForwardOutcome::Cancelled => break,
}
}
Err(_) => break,
}
}
}
}
});
// Wire ClientSideConnection for JSON-RPC ser/deser.
let gw_tx = AcpGatewaySender::new(agent_channel.tx).with_tracing(true);
let incoming = LineBufferedRead::spawn_local(incoming_read.compat());
let (conn, handle_io) = acp::ClientSideConnection::new(
gw_tx,
outgoing_write.compat_write(),
incoming,
|fut| { tokio::task::spawn_local(fut); },
);
let gw_rx = AcpGatewayReceiver::new(agent_channel.rx, conn).with_tracing(true);
tokio::task::spawn_local(handle_io);
tokio::task::spawn_local(gw_rx.run());
tokio::task::yield_now().await;
bridge_cancel.cancelled().await;
reader_task.abort();
writer_task.abort();
Ok(())
})
})?;
Ok(LeaderBridge {
channel: client_channel,
cancel,
thread_handle,
})
}
#[cfg(test)]
mod tests {
use super::*;
use kigi_acp_lib::acp_send;
#[tokio::test]
async fn forward_outbound_line_delivers_on_live_channel() {
let (tx, mut rx) = mpsc::unbounded_channel::<String>();
let shared = TokioMutex::new(tx);
let cancel = CancellationToken::new();
assert_eq!(
forward_outbound_line(&shared, &cancel, "hello".into()).await,
ForwardOutcome::Sent
);
assert_eq!(rx.recv().await.as_deref(), Some("hello"));
}
/// Connection-scoped redelivery semantics: a line whose send failed is
/// HELD (blocking later lines) until the reader swaps in a fresh tx, then
/// DROPPED — neither discarded at first failure (silent outbound loss)
/// nor replayed onto the new connection (a stale `session/load` would
/// double-replay the transcript). Lines queued behind it flow onto the
/// new connection.
#[tokio::test]
async fn forward_outbound_line_drops_stale_line_after_swap_and_sends_next() {
let (dead_tx, dead_rx) = mpsc::unbounded_channel::<String>();
drop(dead_rx);
let shared = Arc::new(TokioMutex::new(dead_tx));
let cancel = CancellationToken::new();
let (new_tx, mut new_rx) = mpsc::unbounded_channel::<String>();
let swapped = Arc::new(std::sync::atomic::AtomicBool::new(false));
let swapper_shared = shared.clone();
let swapper_swapped = swapped.clone();
let swapper = tokio::spawn(async move {
tokio::time::sleep(std::time::Duration::from_millis(250)).await;
// Flag BEFORE the swap: a `DroppedStale` return implies the new
// tx was observed, which happens-after this store.
swapper_swapped.store(true, std::sync::atomic::Ordering::SeqCst);
*swapper_shared.lock().await = new_tx;
});
assert_eq!(
forward_outbound_line(&shared, &cancel, "stale request".into()).await,
ForwardOutcome::DroppedStale
);
assert!(
swapped.load(std::sync::atomic::Ordering::SeqCst),
"the line must be HELD until the swap, not dropped on first failure"
);
swapper.await.unwrap();
assert_eq!(
forward_outbound_line(&shared, &cancel, "fresh request".into()).await,
ForwardOutcome::Sent
);
assert_eq!(
new_rx.recv().await.as_deref(),
Some("fresh request"),
"the first line on the new connection is the post-swap one"
);
assert!(
new_rx.try_recv().is_err(),
"the stale line must not be re-delivered onto the new connection"
);
}
#[tokio::test]
async fn forward_outbound_line_cancellation_exits_retry() {
let (dead_tx, dead_rx) = mpsc::unbounded_channel::<String>();
drop(dead_rx);
let shared = TokioMutex::new(dead_tx);
let cancel = CancellationToken::new();
cancel.cancel();
assert_eq!(
forward_outbound_line(&shared, &cancel, "x".into()).await,
ForwardOutcome::Cancelled
);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bridge_passes_initialize_round_trip() {
let cancel = CancellationToken::new();
let (fake_leader_tx, bridge_leader_rx) = mpsc::unbounded_channel::<String>();
let (bridge_leader_tx, mut fake_leader_rx) = mpsc::unbounded_channel::<String>();
let bridge = bridge_channels(
bridge_leader_tx,
bridge_leader_rx,
cancel.clone(),
None,
ReconnectPolicy::bounded(),
)
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
let fake_leader = tokio::spawn(async move {
let msg = fake_leader_rx.recv().await.expect("expected a message");
let req: serde_json::Value =
serde_json::from_str(&msg).expect("invalid JSON from bridge");
let id = req.get("id").expect("missing id").clone();
let response = serde_json::json!({
"jsonrpc": "2.0",
"id": id,
"result": {
"protocolVersion": "1",
"serverCapabilities": {},
"authMethods": []
}
});
fake_leader_tx
.send(serde_json::to_string(&response).unwrap())
.unwrap();
});
let _resp: acp::InitializeResponse = acp_send(
acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(
acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false),
),
&bridge.channel.tx,
)
.await
.expect("initialize should succeed through bridge");
fake_leader.await.unwrap();
cancel.cancel();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn bridge_cancels_on_leader_disconnect_without_reconnector() {
let cancel = CancellationToken::new();
let (leader_inbound_tx, bridge_leader_rx) = mpsc::unbounded_channel::<String>();
let (bridge_leader_tx, _leader_outbound_rx) = mpsc::unbounded_channel::<String>();
let bridge = bridge_channels(
bridge_leader_tx,
bridge_leader_rx,
cancel.clone(),
None,
ReconnectPolicy::bounded(),
)
.unwrap();
drop(leader_inbound_tx);
tokio::time::timeout(std::time::Duration::from_secs(5), bridge.cancel.cancelled())
.await
.expect("bridge should cancel after leader disconnect");
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn writer_does_not_break_pipe_on_send_failure() {
let cancel = CancellationToken::new();
let (leader_inbound_tx, bridge_leader_rx) = mpsc::unbounded_channel::<String>();
let (bridge_leader_tx, leader_outbound_rx) = mpsc::unbounded_channel::<String>();
let bridge = bridge_channels(
bridge_leader_tx,
bridge_leader_rx,
cancel.clone(),
None,
ReconnectPolicy::bounded(),
)
.unwrap();
tokio::time::sleep(std::time::Duration::from_millis(50)).await;
// Drop the outbound receiver so leader_tx.send() will fail in the writer.
drop(leader_outbound_rx);
// Spawn a background task that pushes an outbound ACP request.
// acp_send blocks for a response that will never arrive (the outbound
// receiver is dropped), but the writer task should survive the failed
// send rather than breaking the simplex pipe.
let tx = bridge.channel.tx.clone();
tokio::spawn(async move {
let _ = acp_send(
acp::InitializeRequest::new(acp::ProtocolVersion::V1).client_capabilities(
acp::ClientCapabilities::new()
.fs(acp::FileSystemCapabilities::new())
.terminal(false),
),
&tx,
)
.await;
});
// Give the writer time to hit the send failure and sleep/retry.
tokio::time::sleep(std::time::Duration::from_millis(300)).await;
// The bridge should still be alive — the writer must not have broken
// the pipe by exiting.
assert!(
!cancel.is_cancelled(),
"writer should survive send failures (reader not disconnected yet)"
);
// Now disconnect the leader fully so the reader triggers cancel.
drop(leader_inbound_tx);
tokio::time::timeout(std::time::Duration::from_secs(5), bridge.cancel.cancelled())
.await
.expect("bridge should eventually cancel after full disconnect");
}
}
+231
View File
@@ -0,0 +1,231 @@
//! Strongly-typed notification metadata.
//!
//! Parses the `_meta` JSON from `SessionNotification` into a struct with
//! typed fields. All fields are `Option` — gracefully degrades when
//! grok-shell hasn't been updated or meta is absent.
use serde::{Deserialize, Serialize};
/// Parsed fields from `SessionNotification._meta`.
///
/// Extracted once in [`acp_handler`](super::super::app::acp_handler) and passed
/// downstream to the tracker and agent state.
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
pub struct NotificationMeta {
/// Accumulated token count across the session (`totalTokens`).
pub total_tokens: Option<u64>,
/// UTC ms when this notification was sent (`agentTimestampMs`).
pub agent_timestamp_ms: Option<i64>,
/// UTC ms when the current LLM streaming response started (`streamStartMs`).
/// Resets each tool-use loop iteration.
pub stream_start_ms: Option<i64>,
/// UTC ms when the current turn started (`turnStartMs`).
/// Constant for the entire turn.
pub turn_start_ms: Option<i64>,
/// Stable id for the prompt this notification belongs to (`promptId`).
/// The client passes a UUID in `PromptRequest._meta.promptId`; the agent
/// echoes it on every notification it emits while processing that
/// prompt. Used to drop chunks for cancelled / rewound turns.
pub prompt_id: Option<String>,
/// Whether this notification is historical replay from `session/load`.
pub is_replay: bool,
/// Raw `eventId` string (`"{sessionId}-{counter}"`). Tracked per session
/// as the reconnect cursor (`_meta.cursor` on `session/load`): the agent
/// resolves it by exact string match against persisted lines, so the full
/// id is kept — the numeric suffix alone is ambiguous across the
/// non-monotonic counter runs of a multi-resume history.
pub event_id: Option<String>,
/// Monotonic per-process sequence parsed from `eventId`
/// (`"{sessionId}-{counter}"`, see `kigi-shell util::event_id`). The
/// agent stamps the SAME `eventId` on the live emission and on the persisted
/// line that is later replayed, so a client can dedup an event it receives
/// twice (replay/live overlap, a re-emit after the reconnect gate, or
/// duplicate routing). Per-session events arrive in increasing order, so the
/// pager keeps a highwater and drops anything `<=` it. `None` when the agent
/// didn't stamp an `eventId` (older shell) — such updates always apply.
pub event_seq: Option<u64>,
}
/// Serializable counterpart of the replay stamp the agent injects on
/// replayed notifications (`_meta.isReplay`, stamped by kigi-shell's
/// `forward_raw_replay_line` during `session/load`).
///
/// [`NotificationMeta::from_json`] is the parse side; this is the build
/// side, so code that constructs a replay-stamped `_meta` (test fixtures,
/// playgrounds) shares the wire key with the parser instead of hand-writing
/// `json!` literals.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ReplayMetaStamp {
pub is_replay: bool,
}
impl ReplayMetaStamp {
/// `_meta` value for a replayed (`session/load`) notification.
pub fn replayed() -> serde_json::Value {
serde_json::to_value(Self { is_replay: true }).expect("serialize replay meta stamp")
}
}
/// User-prompt content-block `_meta` keys (`TextContent.meta`), shared by the
/// producers (`dispatch/queue.rs` drain, `effects.rs` prompt send) and the
/// replay consumer (`acp/tracker.rs` `handle_user_message`) so the wire keys
/// cannot drift. Tests keep raw literals — they pin the wire values.
pub mod user_prompt_meta {
/// Clean display text shown in scrollback instead of the wire text.
pub const DISPLAY_TEXT: &str = "displayText";
/// Render the display text as a skill invocation (teal leading token).
pub const DISPLAY_AS_SKILL: &str = "displayAsSkill";
/// Render the display text as a scheduled (cron) prompt.
pub const DISPLAY_AS_CRON: &str = "displayAsCron";
/// `[[start, end], …]` byte ranges of recognized slash tokens into the
/// block's `text`; only meaningful when that text is displayed verbatim
/// (never stamped alongside `displayText`).
pub const SKILL_TOKEN_RANGES: &str = "skillTokenRanges";
}
/// `UserMessageChunk` / `ContentChunk._meta` keys stamped by the shell and
/// read by the pager (live and on replay).
pub mod user_message_chunk_meta {
/// Prompt index for rewind / attribution.
pub const PROMPT_INDEX: &str = "promptIndex";
/// When true, the chunk must not become a scrollback user prompt
/// ([`kigi_shell::session::PromptOrigin::hide_user_echo_from_scrollback`]).
pub const HIDE_FROM_SCROLLBACK: &str = "hideFromScrollback";
}
impl NotificationMeta {
/// Parse from the `_meta` JSON map on a `SessionNotification`.
pub fn from_json(meta: Option<&serde_json::Map<String, serde_json::Value>>) -> Self {
let Some(m) = meta else {
return Self::default();
};
let event_id = m
.get("eventId")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
// `eventId` is `"{sessionId}-{counter}"`; the counter is the part
// after the LAST '-' (session ids themselves contain '-').
let event_seq = event_id
.as_deref()
.and_then(|s| s.rsplit('-').next())
.and_then(|c| c.parse::<u64>().ok());
Self {
total_tokens: m.get("totalTokens").and_then(|v| v.as_u64()),
agent_timestamp_ms: m.get("agentTimestampMs").and_then(|v| v.as_i64()),
stream_start_ms: m.get("streamStartMs").and_then(|v| v.as_i64()),
turn_start_ms: m.get("turnStartMs").and_then(|v| v.as_i64()),
prompt_id: m
.get("promptId")
.and_then(|v| v.as_str())
.map(|s| s.to_string()),
is_replay: m.get("isReplay").and_then(|v| v.as_bool()).unwrap_or(false),
event_id,
event_seq,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn parse_full_meta() {
let meta_json = json!({
"totalTokens": 5000u64,
"agentTimestampMs": 1700000000000i64,
"streamStartMs": 1700000000000i64 - 3200,
"turnStartMs": 1700000000000i64 - 5000,
"eventId": "sess-1-7",
});
let map = meta_json.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert_eq!(meta.total_tokens, Some(5000));
assert_eq!(meta.agent_timestamp_ms, Some(1700000000000));
assert_eq!(meta.stream_start_ms, Some(1700000000000 - 3200));
assert_eq!(meta.turn_start_ms, Some(1700000000000 - 5000));
assert!(!meta.is_replay);
assert_eq!(meta.event_id.as_deref(), Some("sess-1-7"));
assert_eq!(meta.event_seq, Some(7));
}
#[test]
fn parse_missing_new_fields() {
// Simulate old grok-shell that doesn't send streamStartMs/turnStartMs
let meta_json = json!({
"totalTokens": 1000u64,
"agentTimestampMs": 1700000000000i64,
});
let map = meta_json.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert_eq!(meta.total_tokens, Some(1000));
assert_eq!(meta.agent_timestamp_ms, Some(1700000000000));
assert_eq!(meta.stream_start_ms, None);
assert_eq!(meta.turn_start_ms, None);
assert!(!meta.is_replay);
}
#[test]
fn parse_replay_flag() {
let meta_json = json!({
"isReplay": true,
});
let map = meta_json.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert!(meta.is_replay);
}
/// The build side ([`ReplayMetaStamp::replayed`]) and the parse side
/// ([`NotificationMeta::from_json`]) must agree on the wire key — a
/// rename on either side breaks replay detection silently otherwise.
#[test]
fn replay_meta_stamp_round_trips_through_parser() {
let stamp = ReplayMetaStamp::replayed();
let map = stamp.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert!(meta.is_replay, "stamped meta must parse as a replay");
}
#[test]
fn parse_event_id_keeps_raw_string_and_seq() {
let meta_json = json!({
"eventId": "sess-ab-12-42",
});
let map = meta_json.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert_eq!(meta.event_id.as_deref(), Some("sess-ab-12-42"));
assert_eq!(meta.event_seq, Some(42));
}
/// A non-numeric suffix yields no seq (dedup disabled) but the raw id is
/// still kept for the reconnect cursor (string-matched agent-side).
#[test]
fn parse_event_id_without_numeric_suffix() {
let meta_json = json!({
"eventId": "weird-id-zzz",
});
let map = meta_json.as_object().unwrap();
let meta = NotificationMeta::from_json(Some(map));
assert_eq!(meta.event_id.as_deref(), Some("weird-id-zzz"));
assert_eq!(meta.event_seq, None);
}
#[test]
fn parse_none_meta() {
let meta = NotificationMeta::from_json(None);
assert_eq!(meta.total_tokens, None);
assert_eq!(meta.agent_timestamp_ms, None);
assert_eq!(meta.stream_start_ms, None);
assert_eq!(meta.turn_start_ms, None);
assert!(!meta.is_replay);
assert_eq!(meta.event_id, None);
assert_eq!(meta.event_seq, None);
}
}
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,649 @@
//! Model state — tracks available models and current selection.
use agent_client_protocol as acp;
use indexmap::IndexMap;
use kigi_shell::sampling::types::{
ReasoningEffort, ReasoningEffortOption, parse_reasoning_effort_meta,
parse_reasoning_efforts_meta, supports_reasoning_effort_meta,
};
use crate::slash::commands::effort_levels::legacy_effort_options;
/// Why an effort token could not be applied to a model. Shared by every effort
/// surface (`/effort`, the CLI deferred switch, and headless) so they classify
/// the same input identically and differ only in how they surface the error.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum EffortTokenError {
/// The target model does not advertise `supportsReasoningEffort`.
Unsupported,
/// The token is neither a menu id nor a canonical value offered by this
/// model's menu. `offered` is the model-specific list of option ids the
/// user can type (never a hardcoded global set — so we do not advertise
/// `none`/`minimal` when the model does not offer them).
UnknownToken { token: String, offered: Vec<String> },
/// No active model to resolve the effort against.
NoActiveModel,
}
impl EffortTokenError {
pub(crate) fn message(&self) -> String {
match self {
Self::Unsupported => "current model does not support reasoning effort".to_string(),
Self::UnknownToken { token, offered } => {
if offered.is_empty() {
format!(
"unknown effort level '{token}'; this model has no selectable effort levels"
)
} else {
format!(
"unknown effort level '{token}'; use one of: {}",
offered.join(", ")
)
}
}
Self::NoActiveModel => "no active model to apply effort to".to_string(),
}
}
}
/// Per-agent model state.
#[derive(Debug, Clone, Default)]
pub struct ModelState {
pub available: IndexMap<acp::ModelId, acp::ModelInfo>,
pub current: Option<acp::ModelId>,
pub reasoning_effort: Option<ReasoningEffort>,
/// External override for the context window size (tokens).
/// When set, `get_context_window()` returns this instead of
/// reading from the current model's metadata. Used for subagent
/// views where SubagentProgress reports the actual window size.
context_window_override: Option<u64>,
}
impl ModelState {
pub fn is_empty(&self) -> bool {
self.available.is_empty()
}
/// Display name for the current model.
pub fn current_model_name(&self) -> Option<String> {
let current = self.current.as_ref()?;
if let Some(model_info) = self.available.get(current) {
Some(model_info.name.clone())
} else {
Some(current.0.to_string())
}
}
/// Machine-readable model ID string for the current model (e.g. "grok-4.5").
pub fn current_model_id_str(&self) -> Option<&str> {
Some(self.current.as_ref()?.0.as_ref())
}
/// Total context window tokens for the current model (if available).
fn current_context_window_tokens(&self) -> Option<u64> {
let meta = self.available.get(self.current.as_ref()?)?.meta.as_ref()?;
meta.get("totalContextTokens")
.and_then(|value| match value {
serde_json::Value::Number(number) => number.as_u64(),
_ => None,
})
}
/// Whether the current model accepts image input, read from the model's
/// `meta` (the ACP extension point — same source as `totalContextTokens`).
///
/// Honors an explicit `acceptsImages` bool, else an `inputModalities` array
/// containing `"image"`. DEFAULTS TO `true` when neither key is present:
/// correct today (all current Grok models accept images, so nothing is
/// suppressed) and forward-compatible (suppresses non-vision models once the
/// ACP server populates the key). Populating that key server-side is a
/// separate change.
pub fn current_model_accepts_images(&self) -> bool {
let Some(meta) = self
.current
.as_ref()
.and_then(|id| self.available.get(id))
.and_then(|info| info.meta.as_ref())
else {
return true;
};
if let Some(accepts) = meta.get("acceptsImages").and_then(|v| v.as_bool()) {
return accepts;
}
if let Some(modalities) = meta.get("inputModalities").and_then(|v| v.as_array()) {
return modalities
.iter()
.any(|m| m.as_str().is_some_and(|s| s.eq_ignore_ascii_case("image")));
}
true
}
/// Get the effective context window size (tokens).
///
/// Returns the override if set, otherwise reads from the current model's
/// metadata. The override is set by `override_context_window()` when an
/// external source (e.g., SubagentProgress) reports the actual window size.
pub fn get_context_window(&self) -> Option<u64> {
self.context_window_override
.or_else(|| self.current_context_window_tokens())
}
/// Override the context window size.
///
/// Used for subagent views where the actual context window is reported
/// via SubagentProgress and may differ from the inherited model's metadata.
pub fn override_context_window(&mut self, tokens: u64) {
self.context_window_override = Some(tokens);
}
/// Replace the available models, preserving current selection if still valid.
pub fn update_catalog(
&mut self,
new_available: IndexMap<acp::ModelId, acp::ModelInfo>,
fallback_current: Option<acp::ModelId>,
) {
let previous_current_model = self.current.clone();
self.available = new_available;
if let Some(ref id) = self.current {
if !self.available.contains_key(id) {
self.current = fallback_current;
}
} else {
self.current = fallback_current;
}
// The models/update broadcast carries each model's static default effort,
// not this session's choice; only re-derive when the model changed so a
// catalog refresh can't clobber a user-set effort.
if self.current != previous_current_model {
self.reasoning_effort = self
.current
.as_ref()
.and_then(|id| self.available.get(id))
.and_then(|info| parse_reasoning_effort_meta(info.meta.as_ref()));
}
}
/// Set the current model and resolve reasoning effort from catalog meta.
pub fn set_current(
&mut self,
model_id: acp::ModelId,
effort_override: Option<ReasoningEffort>,
) {
self.current = Some(model_id.clone());
self.reasoning_effort = effort_override.or_else(|| {
self.available
.get(&model_id)
.and_then(|info| parse_reasoning_effort_meta(info.meta.as_ref()))
});
}
/// The reasoning-effort menu for the current model. Gate-first: an unset or
/// unsupported model yields no menu; a supported model uses the server list
/// when present, else the built-in fallback.
pub fn reasoning_effort_options(&self) -> Vec<ReasoningEffortOption> {
match self.current.as_ref() {
Some(id) => self.reasoning_effort_options_for(id),
None => Vec::new(),
}
}
/// Menu for a specific catalog model id (used by `/model`'s effort phase).
/// `parse_reasoning_efforts_meta` returns `None` for absent, non-array, or
/// present-but-unusable lists, so all of those fall back to the built-in menu
/// exactly as the shell's session picker does.
pub(crate) fn reasoning_effort_options_for(
&self,
id: &acp::ModelId,
) -> Vec<ReasoningEffortOption> {
let Some(info) = self.available.get(id) else {
return Vec::new();
};
if !supports_reasoning_effort_meta(info.meta.as_ref()) {
return Vec::new();
}
parse_reasoning_efforts_meta(info.meta.as_ref()).unwrap_or_else(legacy_effort_options)
}
/// Map a typed/selected effort token to its canonical value for the current
/// model. Accepts a menu option id (case-insensitive) or a canonical level
/// that appears as a **value** in that model's menu. Levels the model does
/// not offer (e.g. `none` on grok-4.5) are rejected so we fail in the TUI
/// instead of sending a blocked effort to the API.
pub fn resolve_effort_token(&self, token: &str) -> Option<ReasoningEffort> {
match self.current.as_ref() {
Some(id) => self.resolve_effort_token_for(id, token),
// No model yet: still parse so deferred CLI can hold a token; it is
// re-validated with `resolve_effort_for_model` once a model is active.
None => token.parse::<ReasoningEffort>().ok(),
}
}
/// [`Self::resolve_effort_token`] scoped to a specific catalog model id.
pub(crate) fn resolve_effort_token_for(
&self,
id: &acp::ModelId,
token: &str,
) -> Option<ReasoningEffort> {
let options = self.reasoning_effort_options_for(id);
if let Some(option) = options
.iter()
.find(|opt| opt.id.eq_ignore_ascii_case(token))
{
return Some(option.value);
}
// Canonical level (e.g. "high", "max"→xhigh) only if the model menu
// actually offers that value — not free-form power-user aliases that
// would 400 on the server (e.g. `none` on grok-4.5).
let parsed = token.parse::<ReasoningEffort>().ok()?;
options
.iter()
.find(|opt| opt.value == parsed)
.map(|o| o.value)
}
/// Canonical effort-token policy: gate on the model's support flag first,
/// then resolve the token (menu id or canonical level). This is the single
/// decision shared by `/effort`, the CLI deferred switch, and headless —
/// each caller only maps the [`EffortTokenError`] to its own surface.
pub(crate) fn resolve_effort_for_model(
&self,
id: &acp::ModelId,
token: &str,
) -> Result<ReasoningEffort, EffortTokenError> {
let supports = self
.available
.get(id)
.map(|info| supports_reasoning_effort_meta(info.meta.as_ref()))
.unwrap_or(false);
if !supports {
return Err(EffortTokenError::Unsupported);
}
self.resolve_effort_token_for(id, token)
.ok_or_else(|| EffortTokenError::UnknownToken {
token: token.to_string(),
// Menu option ids only — matches `/effort` autocomplete and
// never invents levels (none/minimal/…) the model does not offer.
offered: self
.reasoning_effort_options_for(id)
.into_iter()
.map(|opt| opt.id)
.collect(),
})
}
/// Resolve a user-supplied name to a `ModelId` via case-insensitive
/// ASCII match against the catalog.
pub fn resolve_by_name_or_id(&self, query: &str) -> Option<acp::ModelId> {
self.available.iter().find_map(|(id, info)| {
if info.name.eq_ignore_ascii_case(query) || id.0.as_ref().eq_ignore_ascii_case(query) {
Some(id.clone())
} else {
None
}
})
}
/// Look up the display name for a `ModelId` in the catalog.
pub fn display_name_for(&self, id: &acp::ModelId) -> String {
self.available
.get(id)
.map(|info| info.name.clone())
.unwrap_or_else(|| id.0.to_string())
}
/// Cycle to the next model.
pub fn next_model(&self) -> Option<acp::ModelId> {
if self.available.is_empty() {
None
} else if let Some(ref current) = self.current {
let idx = self.available.get_index_of(current)?;
let idx = (idx + 1) % self.available.len();
Some(self.available.get_index(idx)?.0.clone())
} else {
Some(self.available.first()?.0.clone())
}
}
}
impl From<Option<acp::SessionModelState>> for ModelState {
fn from(state: Option<acp::SessionModelState>) -> Self {
state
.map(|state| {
let mut models = IndexMap::new();
for model in state.available_models {
models.insert(model.model_id.clone(), model);
}
let current_model = models
.contains_key(&state.current_model_id)
.then_some(state.current_model_id);
let reasoning_effort = current_model
.as_ref()
.and_then(|id| models.get(id))
.and_then(|info| parse_reasoning_effort_meta(info.meta.as_ref()));
Self {
available: models,
current: current_model,
reasoning_effort,
context_window_override: None,
}
})
.unwrap_or_default()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Arc;
fn sample_models() -> ModelState {
let mut state = ModelState::default();
let id_a = acp::ModelId::new(Arc::from("model-a"));
let id_b = acp::ModelId::new(Arc::from("model-b"));
state.available.insert(
id_a.clone(),
acp::ModelInfo::new(id_a.clone(), "Model A".to_string()),
);
state.available.insert(
id_b.clone(),
acp::ModelInfo::new(id_b.clone(), "Model B".to_string()),
);
state.current = Some(id_a);
state
}
#[test]
fn test_current_model_name() {
let state = sample_models();
assert_eq!(state.current_model_name(), Some("Model A".to_string()));
}
#[test]
fn test_next_model_cycles() {
let state = sample_models();
let next = state.next_model().unwrap();
assert_eq!(next.0.as_ref(), "model-b");
}
#[test]
fn test_next_model_wraps() {
let mut state = sample_models();
state.current = Some(acp::ModelId::new(Arc::from("model-b")));
let next = state.next_model().unwrap();
assert_eq!(next.0.as_ref(), "model-a");
}
#[test]
fn test_empty_state() {
let state = ModelState::default();
assert!(state.is_empty());
assert!(state.current_model_name().is_none());
assert!(state.next_model().is_none());
}
fn model_with_effort(id: &str, name: &str, effort: &str) -> acp::ModelInfo {
acp::ModelInfo::new(acp::ModelId::new(Arc::from(id)), name.to_string()).meta(
serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEffort": effort,
})
.as_object()
.cloned(),
)
}
#[test]
fn update_catalog_preserves_user_effort_when_model_unchanged() {
let id = acp::ModelId::new(Arc::from("grok-build"));
let mut state = ModelState::default();
state.available.insert(
id.clone(),
model_with_effort("grok-build", "Grok Build", "high"),
);
state.set_current(id.clone(), Some(ReasoningEffort::Xhigh));
assert_eq!(state.reasoning_effort, Some(ReasoningEffort::Xhigh));
// The broadcast carries the model's static default (high) for the same model.
let mut refreshed = IndexMap::new();
refreshed.insert(
id.clone(),
model_with_effort("grok-build", "Grok Build", "high"),
);
state.update_catalog(refreshed, Some(id.clone()));
assert_eq!(
state.reasoning_effort,
Some(ReasoningEffort::Xhigh),
"catalog refresh must not clobber a user-set per-session effort"
);
}
#[test]
fn update_catalog_rederives_effort_when_current_model_changes() {
let id_a = acp::ModelId::new(Arc::from("model-a"));
let mut state = ModelState::default();
state.available.insert(
id_a.clone(),
model_with_effort("model-a", "Model A", "high"),
);
state.set_current(id_a.clone(), Some(ReasoningEffort::Xhigh));
// Refresh drops model-a; fall back to model-b whose default is low.
let id_b = acp::ModelId::new(Arc::from("model-b"));
let mut refreshed = IndexMap::new();
refreshed.insert(id_b.clone(), model_with_effort("model-b", "Model B", "low"));
state.update_catalog(refreshed, Some(id_b.clone()));
assert_eq!(state.current, Some(id_b));
assert_eq!(state.reasoning_effort, Some(ReasoningEffort::Low));
}
fn state_with_meta(meta: Option<serde_json::Value>) -> ModelState {
let id = acp::ModelId::new(Arc::from("m"));
let mut state = ModelState::default();
state.available.insert(
id.clone(),
acp::ModelInfo::new(id.clone(), "M".to_string())
.meta(meta.and_then(|v| v.as_object().cloned())),
);
state.current = Some(id);
state
}
#[test]
fn accepts_images_defaults_true_when_meta_absent() {
// No current model, empty meta, and a meta without the key all default
// permissive — correct today and a no-op until the server populates it.
assert!(ModelState::default().current_model_accepts_images());
assert!(state_with_meta(None).current_model_accepts_images());
assert!(
state_with_meta(Some(serde_json::json!({ "totalContextTokens": 256000 })))
.current_model_accepts_images()
);
}
#[test]
fn reasoning_effort_options_renders_server_list() {
let state = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEfforts": [
{ "id": "balanced", "value": "medium", "label": "Balanced" },
{ "id": "deep", "value": "xhigh", "label": "Deep", "description": "Max" },
],
})));
let opts = state.reasoning_effort_options();
assert_eq!(opts.len(), 2);
assert_eq!(opts[0].label, "Balanced");
assert_eq!(opts[0].value, ReasoningEffort::Medium);
assert_eq!(opts[1].id, "deep");
assert_eq!(opts[1].description.as_deref(), Some("Max"));
}
#[test]
fn reasoning_effort_options_gate_first_empty_when_unsupported() {
// No current model → empty.
assert!(ModelState::default().reasoning_effort_options().is_empty());
// Current model that does not support effort → empty (even with a list).
let state = state_with_meta(Some(serde_json::json!({
"reasoningEfforts": [{ "value": "high" }],
})));
assert!(state.reasoning_effort_options().is_empty());
}
#[test]
fn reasoning_effort_options_falls_back_to_builtin_menu() {
// Supported but no server list → today's four-row built-in menu.
let state = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
})));
let ids: Vec<_> = state
.reasoning_effort_options()
.into_iter()
.map(|o| o.id)
.collect();
assert_eq!(ids, ["xhigh", "high", "medium", "low"]);
}
#[test]
fn reasoning_effort_options_falls_back_when_list_present_but_unusable() {
// Matches the shell picker: an explicit empty list, and a list where every
// entry skip-invalidated under version skew, both fall back to the built-in
// menu rather than silently vanishing.
for meta in [
serde_json::json!({ "supportsReasoningEffort": true, "reasoningEfforts": [] }),
serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEfforts": [{ "value": "quantum" }],
}),
] {
let ids: Vec<_> = state_with_meta(Some(meta.clone()))
.reasoning_effort_options()
.into_iter()
.map(|o| o.id)
.collect();
assert_eq!(ids, ["xhigh", "high", "medium", "low"], "for meta {meta}");
}
}
#[test]
fn resolve_effort_token_maps_remap_id_to_canonical_value() {
let state = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEfforts": [
{ "id": "deep", "value": "xhigh", "label": "Deep" },
{ "id": "high", "value": "high", "label": "High" },
],
})));
// Design-2 remap: the typed id resolves to its canonical wire value.
assert_eq!(
state.resolve_effort_token("deep"),
Some(ReasoningEffort::Xhigh)
);
assert_eq!(
state.resolve_effort_token("DEEP"),
Some(ReasoningEffort::Xhigh)
);
// Canonical level offered by the menu is accepted by value.
assert_eq!(
state.resolve_effort_token("high"),
Some(ReasoningEffort::High)
);
// Levels the model does not offer (none/minimal on 4.5-style menus)
// are rejected — better than a server-side 400.
assert!(state.resolve_effort_token("minimal").is_none());
assert!(state.resolve_effort_token("none").is_none());
assert!(state.resolve_effort_token("bogus").is_none());
}
#[test]
fn resolve_effort_token_accepts_none_only_when_menu_offers_it() {
let with_none = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEfforts": [
{ "value": "none", "label": "None", "default": true },
{ "value": "high", "label": "High" },
],
})));
assert_eq!(
with_none.resolve_effort_token("none"),
Some(ReasoningEffort::None)
);
let without_none = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
"reasoningEfforts": [
{ "value": "high", "label": "High", "default": true },
{ "value": "low", "label": "Low" },
],
})));
assert!(without_none.resolve_effort_token("none").is_none());
let err = without_none
.resolve_effort_for_model(without_none.current.as_ref().unwrap(), "none")
.unwrap_err();
assert_eq!(
err,
EffortTokenError::UnknownToken {
token: "none".to_string(),
offered: vec!["high".to_string(), "low".to_string()],
}
);
// Error copy must list only this model's options — never hardcode
// none/minimal/… as offered values (the rejected token may still appear
// quoted in "unknown effort level '…'").
let msg = err.message();
assert!(msg.contains("use one of: high, low"), "msg={msg}");
let offered_half = msg
.split_once("; ")
.map(|(_, rest)| rest)
.expect("message should have '; ' separator");
assert!(
!offered_half.contains("none"),
"must not advertise blocked level: {msg}"
);
assert!(
!offered_half.contains("minimal"),
"must not advertise blocked level: {msg}"
);
assert!(
!msg.contains("unset"),
"unset is log-only, not a user token: {msg}"
);
}
#[test]
fn resolve_effort_token_legacy_menu_rejects_none() {
// supportsReasoningEffort without a server list → built-in low..xhigh.
let state = state_with_meta(Some(serde_json::json!({
"supportsReasoningEffort": true,
})));
assert!(state.resolve_effort_token("none").is_none());
assert!(state.resolve_effort_token("minimal").is_none());
assert_eq!(
state.resolve_effort_token("low"),
Some(ReasoningEffort::Low)
);
}
#[test]
fn accepts_images_honors_explicit_meta() {
assert!(
!state_with_meta(Some(serde_json::json!({ "acceptsImages": false })))
.current_model_accepts_images()
);
assert!(
state_with_meta(Some(serde_json::json!({ "acceptsImages": true })))
.current_model_accepts_images()
);
// inputModalities array form.
assert!(
state_with_meta(Some(
serde_json::json!({ "inputModalities": ["text", "image"] })
))
.current_model_accepts_images()
);
assert!(
!state_with_meta(Some(serde_json::json!({ "inputModalities": ["text"] })))
.current_model_accepts_images()
);
}
}
+126
View File
@@ -0,0 +1,126 @@
//! Agent spawning — creates the agent process and ACP channels.
//!
//! Simplified to only support GrokShell (in-process) mode.
//! Subprocess and remote modes can be added later if needed.
use std::rc::Rc;
use std::thread;
use anyhow::Result;
use tokio_util::sync::CancellationToken;
use kigi_acp_lib::{
AcpAgentChannel, AcpClientChannel, AcpClientTx, AcpGatewayReceiver, AcpGatewaySender,
acp_channels,
};
use kigi_shell::{
agent::{MvpAgent, config::Config as AgentConfig, models::RefreshStrategy},
auth::AuthManager,
util::kigi_home::kigi_home,
};
/// Result of spawning a child agent.
pub struct SpawnedAgent {
/// Kept alive so the thread isn't detached. Will be used for graceful shutdown.
pub _thread_handle: thread::JoinHandle<Result<()>>,
pub channel: AcpClientChannel,
pub cancel: CancellationToken,
/// The agent's `AuthManager`, shared so pager-side consumers
/// channel) resolve the same refreshing bearer as chat traffic.
pub auth_manager: std::sync::Arc<AuthManager>,
}
/// Spawn a GrokShell agent in a background thread.
///
/// Returns the ACP client channel for communication and a cancellation token.
pub async fn spawn_grok_shell(
agent_config: AgentConfig,
cancel: &CancellationToken,
memory_config: Option<kigi_shell::config::MemoryConfig>,
) -> Result<SpawnedAgent> {
let auth_manager = std::sync::Arc::new(AuthManager::new(
&kigi_home(),
agent_config.grok_com_config.clone(),
));
auth_manager.configure_refresher(agent_config.grok_com_config.auth_provider_command.clone());
// Pause token refreshes across system sleep so an OIDC refresh can't
// straddle a suspend (which can revoke the refresh token and force
// re-login). No-op where the OS listener is unavailable.
auth_manager.start_system_power_listener();
// Best-effort refresh of managed policy before bootstrap reads it (repairs a wrong-identity/missing
// cache). Never errors — the OS-protected system/MDM layers still apply.
kigi_shell::managed_config::ensure_managed_policy_present(&auth_manager).await;
// Run the full bootstrap sequence: config resolution, process-level
// singletons (including `extract_bundled_files` which writes compiled-in
// skills to ~/.kigi/skills/), and model catalog construction.
let (agent_config, models_manager) =
kigi_shell::agent::init::bootstrap(&agent_config, &auth_manager, None)
.map_err(|e| anyhow::anyhow!(e))?;
models_manager
.list_models(RefreshStrategy::OnlineIfUncached)
.await;
let agent_cancel = cancel.child_token();
let (acp_client, acp_agent) = acp_channels();
// Clone before `auth_manager` is moved into the agent closure below, so the
// pager can share the same refreshing bearer.
let auth_manager_for_pager = auth_manager.clone();
let spawn_fn: Box<dyn FnOnce(AcpClientTx) -> Result<Rc<MvpAgent>> + Send + 'static> = {
Box::new(move |client_tx| {
let gateway = AcpGatewaySender::new(client_tx);
let mut agent =
MvpAgent::with_models(gateway, &agent_config, auth_manager, models_manager);
if let Some(mc) = memory_config {
agent.set_memory_config(mc);
}
Ok(Rc::new(agent))
})
};
// Spawn the agent thread with direct dispatch
let handle = spawn_agent_thread_direct(spawn_fn, acp_agent, agent_cancel.clone())?;
Ok(SpawnedAgent {
_thread_handle: handle,
channel: acp_client,
cancel: agent_cancel,
auth_manager: auth_manager_for_pager,
})
}
/// Spawn an agent in a dedicated thread with direct RPC dispatch.
///
/// The agent runs on a single-threaded tokio LocalSet runtime.
/// RPC requests go directly to the agent via Rc, bypassing simplex pipes.
fn spawn_agent_thread_direct(
spawn_agent: Box<dyn FnOnce(AcpClientTx) -> Result<Rc<MvpAgent>> + Send + 'static>,
channel: AcpAgentChannel,
cancel: CancellationToken,
) -> Result<thread::JoinHandle<Result<()>>> {
Ok(thread::Builder::new()
.name("acp-agent-worker".into())
.spawn(move || -> Result<()> {
let rt = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let local = tokio::task::LocalSet::new();
local.block_on(&rt, async move {
let client_tx = channel.tx.clone();
let agent_rc = spawn_agent(client_tx)?;
// Direct dispatch: RPC requests go straight to the agent
let gw_rx = AcpGatewayReceiver::new(channel.rx, agent_rc).with_tracing(true);
tokio::task::spawn_local(gw_rx.run());
tokio::task::yield_now().await;
// Keep running until cancelled
cancel.cancelled().await;
anyhow::Result::Ok(())
})
})?)
}
File diff suppressed because it is too large Load Diff