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:
@@ -0,0 +1,225 @@
|
||||
//! `x.ai/auth/*` and legacy `x.ai/{get,set}ApiKey` extension handlers.
|
||||
//!
|
||||
//! These methods let the client read/write the API key via the agent and
|
||||
//! drive the OAuth login flow. The agent is the single source of truth for
|
||||
//! `auth.json`.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::ExtMethodResult;
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/auth/getBearerToken" => handle_get_bearer_token(agent).await,
|
||||
"x.ai/getApiKey" => handle_get_api_key(),
|
||||
"x.ai/setApiKey" => handle_set_api_key(args),
|
||||
"x.ai/auth/submit_code" => handle_submit_code(agent, args),
|
||||
"x.ai/auth/get_url" => handle_get_url(agent).await,
|
||||
"x.ai/auth/logout" => handle_logout(agent, args).await,
|
||||
"x.ai/auth/info" => handle_info(agent),
|
||||
"x.ai/auth/check_subscription" => handle_check_subscription(agent).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_get_bearer_token(agent: &MvpAgent) -> ExtResult {
|
||||
let token = match agent.auth_manager.get_valid_token().await {
|
||||
Ok(token) => Some(token),
|
||||
Err(_) => agent
|
||||
.sampling_config
|
||||
.borrow()
|
||||
.api_key
|
||||
.clone()
|
||||
.or_else(|| agent.auth_manager.current().map(|a| a.key)),
|
||||
};
|
||||
ExtMethodResult::success(serde_json::json!({ "token": token }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
fn handle_get_api_key() -> ExtResult {
|
||||
let key = crate::agent::auth_method::read_xai_api_key_env().ok();
|
||||
ExtMethodResult::success(serde_json::json!({ "key": key }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
fn handle_set_api_key(args: &acp::ExtRequest) -> ExtResult {
|
||||
let params: serde_json::Value = parse_params(args)?;
|
||||
let key = params.get("key").and_then(|v| v.as_str());
|
||||
let kigi_home = crate::util::kigi_home::kigi_home();
|
||||
if let Some(k) = key {
|
||||
if k.is_empty() {
|
||||
crate::auth::clear_api_key(&kigi_home)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::remove_var("XAI_API_KEY") };
|
||||
} else {
|
||||
crate::auth::store_api_key(&kigi_home, k)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::set_var("XAI_API_KEY", k) };
|
||||
}
|
||||
} else {
|
||||
crate::auth::clear_api_key(&kigi_home)
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// SAFETY: ext_method is single-threaded per agent
|
||||
unsafe { std::env::remove_var("XAI_API_KEY") };
|
||||
}
|
||||
ExtMethodResult::success(serde_json::json!({ "ok": true }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
/// Handle auth code submission from TUI.
|
||||
fn handle_submit_code(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
struct SubmitCodeParams {
|
||||
code: String,
|
||||
}
|
||||
|
||||
let params: SubmitCodeParams = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
|
||||
let auth_code_tx = agent.auth_code_tx.borrow();
|
||||
if let Some(ref tx) = *auth_code_tx {
|
||||
tx.try_send(params.code).map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("failed to submit auth code: {e}"))
|
||||
})?;
|
||||
to_raw_response(&serde_json::json!({ "submitted": true }))
|
||||
} else {
|
||||
Err(acp::Error::invalid_params().data("no pending auth session"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Awaits the auth URL from the oneshot channel (blocks until ready).
|
||||
async fn handle_get_url(agent: &MvpAgent) -> ExtResult {
|
||||
let rx = agent.auth_url_rx.borrow_mut().take();
|
||||
// `None` when no URL was sent (cached creds, early error, second poll):
|
||||
// report mode as `null` rather than mislabeling it `loopback`.
|
||||
let (auth_url, mode) = match rx {
|
||||
Some(rx) => match rx.await {
|
||||
Ok(info) => (Some(info.url), Some(info.mode)),
|
||||
Err(_) => (None, None),
|
||||
},
|
||||
None => (None, None),
|
||||
};
|
||||
to_raw_response(&serde_json::json!({
|
||||
"auth_url": auth_url,
|
||||
// `external_provider` kept for older clients; `mode` is authoritative.
|
||||
"external_provider": mode.is_some_and(|m| m.is_external_provider()),
|
||||
"mode": mode.map(|m| m.as_wire_str()),
|
||||
}))
|
||||
}
|
||||
|
||||
async fn handle_logout(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
struct LogoutParams {
|
||||
scope: Option<String>,
|
||||
}
|
||||
|
||||
let params: LogoutParams = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
|
||||
let result = crate::auth::perform_logout(&agent.auth_manager, params.scope.as_deref())
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("failed to logout: {e}")))?;
|
||||
// `auth.lifecycle` (not `auth`) avoids colliding with the pre-existing
|
||||
// per-request `AuthManager::auth()` `#[instrument]` span.
|
||||
tracing::info_span!("auth.lifecycle", action = "logout", success = true).in_scope(|| {});
|
||||
|
||||
agent.models_manager.on_auth_changed().await;
|
||||
|
||||
to_raw_response(&serde_json::json!({
|
||||
"ok": true,
|
||||
"was_logged_in": result.was_logged_in,
|
||||
"email": result.email,
|
||||
"api_key_still_set": result.api_key_still_set,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Single-shot subscription re-check (retry button on paywall screen).
|
||||
///
|
||||
/// Calls `retry_subscription_check()`, then returns the updated auth
|
||||
/// response with gate info so the pager can refresh the gate state.
|
||||
async fn handle_check_subscription(agent: &MvpAgent) -> ExtResult {
|
||||
agent.retry_subscription_check().await;
|
||||
let response = agent.auth_response_with_meta();
|
||||
to_raw_response(&serde_json::json!({
|
||||
"authenticated": response.meta.is_some(),
|
||||
"meta": response.meta,
|
||||
}))
|
||||
}
|
||||
|
||||
/// Returns current auth method ID, user profile fields, and team/principal
|
||||
/// metadata.
|
||||
fn handle_info(agent: &MvpAgent) -> ExtResult {
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct AuthInfoResponse {
|
||||
method_id: Option<String>,
|
||||
email: Option<String>,
|
||||
first_name: Option<String>,
|
||||
last_name: Option<String>,
|
||||
/// `grok-asset://` URL resolved by the Electron protocol handler,
|
||||
/// or a full `http(s)://` URL passed through unchanged.
|
||||
profile_image_url: Option<String>,
|
||||
team_id: Option<String>,
|
||||
team_name: Option<String>,
|
||||
team_role: Option<String>,
|
||||
organization_id: Option<String>,
|
||||
organization_name: Option<String>,
|
||||
organization_role: Option<String>,
|
||||
principal_type: Option<String>,
|
||||
principal_id: Option<String>,
|
||||
user_blocked_reason: Option<String>,
|
||||
team_blocked_reasons: Vec<String>,
|
||||
coding_data_retention_opt_out: bool,
|
||||
}
|
||||
|
||||
let method_id = agent
|
||||
.auth_method_id
|
||||
.load()
|
||||
.as_ref()
|
||||
.map(|m| m.0.to_string());
|
||||
let auth = agent.auth_manager.current();
|
||||
let raw_asset_id = auth.as_ref().and_then(|a| a.profile_image_asset_id.clone());
|
||||
|
||||
// Return a grok-asset:// URL that the Electron renderer resolves at
|
||||
// display time via a custom protocol handler. The handler proxies
|
||||
// through cli-chat-proxy's /asset endpoint; Electron's HTTP cache
|
||||
// handles reuse. No disk-cache or network call needed here.
|
||||
let profile_image_url = match raw_asset_id.as_deref().filter(|k| !k.is_empty()) {
|
||||
Some(key) if key.starts_with("http://") || key.starts_with("https://") => {
|
||||
Some(key.to_owned())
|
||||
}
|
||||
Some(key) => Some(format!("grok-asset:///{key}")),
|
||||
None => None,
|
||||
};
|
||||
to_raw_response(&AuthInfoResponse {
|
||||
method_id,
|
||||
email: auth.as_ref().and_then(|a| a.email.clone()),
|
||||
first_name: auth.as_ref().and_then(|a| a.first_name.clone()),
|
||||
last_name: auth.as_ref().and_then(|a| a.last_name.clone()),
|
||||
profile_image_url,
|
||||
team_id: auth.as_ref().and_then(|a| a.team_id.clone()),
|
||||
team_name: auth.as_ref().and_then(|a| a.team_name.clone()),
|
||||
team_role: auth.as_ref().and_then(|a| a.team_role.clone()),
|
||||
organization_id: auth.as_ref().and_then(|a| a.organization_id.clone()),
|
||||
organization_name: auth.as_ref().and_then(|a| a.organization_name.clone()),
|
||||
organization_role: auth.as_ref().and_then(|a| a.organization_role.clone()),
|
||||
principal_type: auth.as_ref().and_then(|a| a.principal_type.clone()),
|
||||
principal_id: auth.as_ref().and_then(|a| a.principal_id.clone()),
|
||||
user_blocked_reason: auth.as_ref().and_then(|a| a.user_blocked_reason.clone()),
|
||||
team_blocked_reasons: auth
|
||||
.as_ref()
|
||||
.map(|a| a.team_blocked_reasons.clone())
|
||||
.unwrap_or_default(),
|
||||
coding_data_retention_opt_out: auth
|
||||
.as_ref()
|
||||
.is_some_and(|a| a.coding_data_retention_opt_out),
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,18 @@
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use crate::auth::{AuthManager, GrokAuth};
|
||||
|
||||
/// Require xAI auth from a sync context, accepting tokens in the client-side buffer window.
|
||||
pub(crate) fn require_xai_auth(
|
||||
auth_manager: &AuthManager,
|
||||
missing_message: &'static str,
|
||||
non_xai_message: &'static str,
|
||||
) -> Result<GrokAuth, acp::Error> {
|
||||
let auth = auth_manager
|
||||
.current_or_expired()
|
||||
.ok_or_else(|| acp::Error::auth_required().data(missing_message))?;
|
||||
if !auth.is_xai_auth() {
|
||||
return Err(acp::Error::auth_required().data(non_xai_message));
|
||||
}
|
||||
Ok(auth)
|
||||
}
|
||||
@@ -0,0 +1,610 @@
|
||||
//! `x.ai/billing` extension handler.
|
||||
//!
|
||||
//! Fetches the authenticated user's Grok Build billing configuration
|
||||
//! (credit limit, usage, on-demand cap, billing period, history) from
|
||||
//! the backend. Used by the pager/desktop to display credits and usage.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{ExtResult, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
/// Billing period cycle identifier.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BillingCycle {
|
||||
pub year: i32,
|
||||
pub month: i32,
|
||||
}
|
||||
|
||||
/// Cent value from the billing API (USD cents).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Cent {
|
||||
/// proto3 JSON omits zero-valued scalars, so a `$0` Cent arrives as `{}`;
|
||||
/// default to 0 rather than failing the whole parse.
|
||||
#[serde(default)]
|
||||
pub val: i64,
|
||||
}
|
||||
|
||||
/// A usage period (weekly or monthly) from the newer credits config.
|
||||
///
|
||||
/// `start`/`end` are RFC 3339 timestamps. `period_type` is the proto enum name
|
||||
/// (e.g. `USAGE_PERIOD_TYPE_WEEKLY`); kept so callers can distinguish weekly
|
||||
/// vs monthly cycles.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UsagePeriod {
|
||||
#[serde(rename = "type", default, skip_serializing_if = "Option::is_none")]
|
||||
pub period_type: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub start: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub end: Option<String>,
|
||||
}
|
||||
|
||||
/// Usage summary for one past billing period.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BillingPeriodUsage {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub billing_cycle: Option<BillingCycle>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub included_used: Option<Cent>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub on_demand_used: Option<Cent>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub total_used: Option<Cent>,
|
||||
}
|
||||
|
||||
/// Current billing configuration for Grok Build coding credits.
|
||||
///
|
||||
/// Carries both the newer credits-config fields (`credit_usage_percent`,
|
||||
/// `current_period`) and the deprecated `GrokBuildBillingConfig` fields
|
||||
/// (`monthly_limit`, `used`, `billing_period_*`). Consumers should prefer the
|
||||
/// new fields and fall back to the deprecated ones, so the same struct works
|
||||
/// against both the new `GetGrokCreditsConfig` and the legacy
|
||||
/// `GetGrokBuildBillingConfig` backend responses.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct BillingConfig {
|
||||
/// Included credit usage as a percentage of the allowance (0.0–100.0).
|
||||
/// Preferred over deriving from `monthly_limit`/`used`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub credit_usage_percent: Option<f64>,
|
||||
/// Current usage period (weekly or monthly). Preferred over
|
||||
/// `billing_period_start`/`billing_period_end`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub current_period: Option<UsagePeriod>,
|
||||
/// Deprecated: included monthly credit budget. Use `credit_usage_percent`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub monthly_limit: Option<Cent>,
|
||||
/// Deprecated: credits used this period. Use `credit_usage_percent`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub used: Option<Cent>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub on_demand_cap: Option<Cent>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub on_demand_used: Option<Cent>,
|
||||
/// Remaining prepaid (purchased) credit balance, positive — the "bought
|
||||
/// credits" the user has topped up. Populated from the credits config
|
||||
/// (`GetGrokCreditsConfig.prepaid_balance`); absent in the legacy billing
|
||||
/// shape.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub prepaid_balance: Option<Cent>,
|
||||
/// Whether this user is on unified usage billing (shared weekly/monthly
|
||||
/// pool). From `GrokCreditsConfig.is_unified_billing_user`, which billing
|
||||
/// sets from remote settings `unified_consumer_billing_enabled`. `None` when
|
||||
/// absent (legacy `GetGrokBuildBillingConfig` shape or older servers).
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub is_unified_billing_user: Option<bool>,
|
||||
/// Deprecated: use `current_period.start`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub billing_period_start: Option<String>,
|
||||
/// Deprecated: use `current_period.end`.
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub billing_period_end: Option<String>,
|
||||
#[serde(default, skip_serializing_if = "Vec::is_empty")]
|
||||
pub history: Vec<BillingPeriodUsage>,
|
||||
}
|
||||
|
||||
/// Top-level response (primarily from `GET /rest/grok/credits` + auto-topup-rule).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BillingConfigResponse {
|
||||
pub config: Option<BillingConfig>,
|
||||
/// Whether on-demand credit usage is enabled. When `false`, the pager
|
||||
/// should hide on-demand controls. Populated from `RemoteSettings`.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub on_demand_enabled: Option<bool>,
|
||||
/// User-friendly subscription tier name (e.g. "SuperGrok Heavy").
|
||||
/// Populated from `RemoteSettings` so the pager can update its cached
|
||||
/// tier on every billing fetch without an extra request.
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub subscription_tier: Option<String>,
|
||||
}
|
||||
|
||||
/// Auto top-up configuration (from GetAutoTopupRule).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct AutoTopupRule {
|
||||
/// proto3 JSON omits `false`, so a disabled rule arrives without this field;
|
||||
/// default to `false` rather than failing the parse (which would otherwise
|
||||
/// keep a stale cached rule in the pager).
|
||||
#[serde(default)]
|
||||
pub enabled: bool,
|
||||
pub min_before_hitting_sl: Option<Cent>,
|
||||
pub topup_amount: Option<Cent>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub max_amount_per_month: Option<Cent>,
|
||||
}
|
||||
|
||||
/// Wrapper for the auto top-up rule response.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GetAutoTopupRuleResponse {
|
||||
#[serde(default)]
|
||||
pub rule: Option<AutoTopupRule>,
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/billing" => {
|
||||
tracing::info!("handling billing config request");
|
||||
handle_get_billing(agent).await
|
||||
}
|
||||
"x.ai/auto-topup-rule" => {
|
||||
tracing::info!("handling auto top-up rule request");
|
||||
handle_get_auto_topup_rule(agent).await
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Structured context for unified-log entries from a successful billing fetch.
|
||||
///
|
||||
/// Keeps history to a count + the most recent period so `~/.kigi/logs/unified.jsonl`
|
||||
/// stays useful without dumping unbounded period arrays.
|
||||
fn billing_unified_log_ctx(billing: &BillingConfigResponse) -> serde_json::Value {
|
||||
let history_len = billing
|
||||
.config
|
||||
.as_ref()
|
||||
.map(|c| c.history.len())
|
||||
.unwrap_or(0);
|
||||
let latest_history = billing
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|c| c.history.last())
|
||||
.and_then(|p| serde_json::to_value(p).ok());
|
||||
|
||||
let mut config_value = billing
|
||||
.config
|
||||
.as_ref()
|
||||
.and_then(|c| serde_json::to_value(c).ok())
|
||||
.unwrap_or(serde_json::Value::Null);
|
||||
if let Some(obj) = config_value.as_object_mut() {
|
||||
// Drop full history array; surface length + latest entry instead.
|
||||
obj.remove("history");
|
||||
obj.insert("historyLen".into(), serde_json::json!(history_len));
|
||||
if let Some(latest) = latest_history {
|
||||
obj.insert("latestHistory".into(), latest);
|
||||
}
|
||||
}
|
||||
|
||||
serde_json::json!({
|
||||
"config": config_value,
|
||||
"onDemandEnabled": billing.on_demand_enabled,
|
||||
"subscriptionTier": billing.subscription_tier,
|
||||
})
|
||||
}
|
||||
|
||||
async fn handle_get_billing(agent: &MvpAgent) -> ExtResult {
|
||||
let auth = super::auth_gate::require_xai_auth(
|
||||
&agent.auth_manager,
|
||||
"Authentication required to fetch billing data",
|
||||
"Billing data requires auth with grok.com. Run `grok login` to authenticate.",
|
||||
)?;
|
||||
|
||||
let proxy_base = agent.cli_chat_proxy_base_url();
|
||||
let base = proxy_base.trim_end_matches('/');
|
||||
|
||||
// Credits balance / usage (new billing system) via the CLI proxy, which
|
||||
// forwards to the backend `GetGrokCreditsConfig`.
|
||||
let credits_url = format!("{}/billing?format=credits", base);
|
||||
let credits_resp = crate::http::shared_client()
|
||||
.get(&credits_url)
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header(
|
||||
"X-XAI-Token-Auth",
|
||||
crate::auth::GrokComConfig::default().token_header,
|
||||
)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.timeout(std::time::Duration::from_secs(15))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!(error = %e, "billing: upstream request failed");
|
||||
kigi_log::unified_log::warn(
|
||||
"billing: upstream request failed",
|
||||
None,
|
||||
Some(serde_json::json!({ "error": e.to_string() })),
|
||||
);
|
||||
acp::Error::internal_error().data(format!("Failed to fetch billing data: {e}"))
|
||||
})?;
|
||||
|
||||
if !credits_resp.status().is_success() {
|
||||
let status = credits_resp.status().as_u16();
|
||||
let body = credits_resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(status, url = %credits_url, "billing: upstream error");
|
||||
|
||||
let detail = serde_json::from_str::<serde_json::Value>(&body)
|
||||
.ok()
|
||||
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from))
|
||||
.unwrap_or_else(|| format!("HTTP {status}"));
|
||||
|
||||
kigi_log::unified_log::warn(
|
||||
"billing: upstream error",
|
||||
None,
|
||||
Some(serde_json::json!({
|
||||
"status": status,
|
||||
"detail": detail,
|
||||
})),
|
||||
);
|
||||
|
||||
return Err(acp::Error::internal_error().data(format!("Billing service error: {detail}")));
|
||||
}
|
||||
|
||||
let mut billing: BillingConfigResponse = credits_resp.json().await.map_err(|e| {
|
||||
tracing::error!(error = %e, "billing: failed to parse response");
|
||||
kigi_log::unified_log::warn(
|
||||
"billing: failed to parse response",
|
||||
None,
|
||||
Some(serde_json::json!({ "error": e.to_string() })),
|
||||
);
|
||||
acp::Error::internal_error().data(format!("Failed to parse billing data: {e}"))
|
||||
})?;
|
||||
|
||||
// Enrich with fields from remote settings.
|
||||
let rs = agent.cfg.borrow().remote_settings.clone();
|
||||
billing.on_demand_enabled = rs.as_ref().and_then(|rs| rs.on_demand_enabled);
|
||||
billing.subscription_tier = rs.as_ref().and_then(|rs| {
|
||||
rs.subscription_tier_display
|
||||
.clone()
|
||||
.or_else(|| rs.subscription_tier.clone())
|
||||
});
|
||||
|
||||
// Every prompt / /usage / poll path hits `x.ai/billing`; log the fetched
|
||||
// credits snapshot so support can correlate limit UX with real balances.
|
||||
kigi_log::unified_log::info(
|
||||
"billing: fetched credits config",
|
||||
None,
|
||||
Some(billing_unified_log_ctx(&billing)),
|
||||
);
|
||||
|
||||
to_raw_response(&billing)
|
||||
}
|
||||
|
||||
async fn handle_get_auto_topup_rule(agent: &MvpAgent) -> ExtResult {
|
||||
let auth = super::auth_gate::require_xai_auth(
|
||||
&agent.auth_manager,
|
||||
"Authentication required to fetch auto top-up rule",
|
||||
"Auto top-up data requires auth with grok.com. Run `grok login` to authenticate.",
|
||||
)?;
|
||||
|
||||
let proxy_base = agent.cli_chat_proxy_base_url();
|
||||
let base = proxy_base.trim_end_matches('/');
|
||||
|
||||
// Auto top-up rule via the CLI proxy, which forwards to the backend
|
||||
// `GetAutoTopupRule`.
|
||||
let url = format!("{}/auto-topup-rule", base);
|
||||
let response = crate::http::shared_client()
|
||||
.get(&url)
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header(
|
||||
"X-XAI-Token-Auth",
|
||||
crate::auth::GrokComConfig::default().token_header,
|
||||
)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.timeout(std::time::Duration::from_secs(10))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!(error = %e, "auto-topup: upstream request failed");
|
||||
acp::Error::internal_error().data(format!("Failed to fetch auto top-up rule: {e}"))
|
||||
})?;
|
||||
|
||||
if !response.status().is_success() {
|
||||
let status = response.status().as_u16();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
tracing::warn!(status, url = %url, "auto-topup: upstream error");
|
||||
|
||||
let detail = serde_json::from_str::<serde_json::Value>(&body)
|
||||
.ok()
|
||||
.and_then(|v| v.get("error").and_then(|e| e.as_str()).map(String::from))
|
||||
.unwrap_or_else(|| format!("HTTP {status}"));
|
||||
|
||||
return Err(
|
||||
acp::Error::internal_error().data(format!("Auto top-up service error: {detail}"))
|
||||
);
|
||||
}
|
||||
|
||||
// Return the upstream response body verbatim (as a JSON value) so /usage
|
||||
// can print the exact data from this request unformatted.
|
||||
let body_text = response.text().await.unwrap_or_default();
|
||||
let value: serde_json::Value =
|
||||
serde_json::from_str(&body_text).unwrap_or(serde_json::json!({"raw": body_text}));
|
||||
to_raw_response(&value)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn auto_topup_disabled_rule_omits_enabled_field() {
|
||||
// proto3 JSON omits `false` / `0`, so a disabled rule arrives without
|
||||
// `enabled` (and zero Cents as `{}`). It must still deserialize (as
|
||||
// disabled) rather than erroring — otherwise the pager keeps a stale
|
||||
// cached rule.
|
||||
let json = serde_json::json!({
|
||||
"rule": { "topupAmount": {"val": 500}, "minBeforeHittingSl": {} }
|
||||
});
|
||||
let resp: GetAutoTopupRuleResponse = serde_json::from_value(json).unwrap();
|
||||
let rule = resp.rule.expect("rule present");
|
||||
assert!(!rule.enabled);
|
||||
assert_eq!(rule.topup_amount.unwrap().val, 500);
|
||||
assert_eq!(rule.min_before_hitting_sl.unwrap().val, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_response_deserializes_from_backend_json() {
|
||||
let json = serde_json::json!({
|
||||
"config": {
|
||||
"monthlyLimit": {"val": 2000},
|
||||
"used": {"val": 1234},
|
||||
"onDemandCap": {"val": 500},
|
||||
"billingPeriodStart": "2025-04-01T00:00:00Z",
|
||||
"billingPeriodEnd": "2025-05-01T00:00:00Z",
|
||||
"history": [
|
||||
{
|
||||
"billingCycle": {"year": 2025, "month": 3},
|
||||
"includedUsed": {"val": 1800},
|
||||
"onDemandUsed": {"val": 0},
|
||||
"totalUsed": {"val": 1800}
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
|
||||
let config = resp.config.unwrap();
|
||||
assert_eq!(config.monthly_limit.unwrap().val, 2000);
|
||||
assert_eq!(config.used.unwrap().val, 1234);
|
||||
assert_eq!(config.on_demand_cap.unwrap().val, 500);
|
||||
assert_eq!(
|
||||
config.billing_period_start.as_deref(),
|
||||
Some("2025-04-01T00:00:00Z")
|
||||
);
|
||||
assert_eq!(config.history.len(), 1);
|
||||
let period = &config.history[0];
|
||||
let cycle = period.billing_cycle.as_ref().unwrap();
|
||||
assert_eq!(cycle.year, 2025);
|
||||
assert_eq!(cycle.month, 3);
|
||||
assert_eq!(period.included_used.as_ref().unwrap().val, 1800);
|
||||
assert_eq!(period.total_used.as_ref().unwrap().val, 1800);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_unified_log_ctx_includes_credits_and_collapses_history() {
|
||||
let resp = BillingConfigResponse {
|
||||
config: Some(BillingConfig {
|
||||
credit_usage_percent: Some(42.5),
|
||||
current_period: Some(UsagePeriod {
|
||||
period_type: Some("USAGE_PERIOD_TYPE_WEEKLY".into()),
|
||||
start: Some("2025-04-01T00:00:00Z".into()),
|
||||
end: Some("2025-04-08T00:00:00Z".into()),
|
||||
}),
|
||||
monthly_limit: Some(Cent { val: 2000 }),
|
||||
used: Some(Cent { val: 850 }),
|
||||
on_demand_cap: Some(Cent { val: 500 }),
|
||||
on_demand_used: Some(Cent { val: 0 }),
|
||||
prepaid_balance: Some(Cent { val: 100 }),
|
||||
is_unified_billing_user: Some(true),
|
||||
billing_period_start: None,
|
||||
billing_period_end: None,
|
||||
history: vec![
|
||||
BillingPeriodUsage {
|
||||
billing_cycle: Some(BillingCycle {
|
||||
year: 2025,
|
||||
month: 2,
|
||||
}),
|
||||
included_used: Some(Cent { val: 1000 }),
|
||||
on_demand_used: Some(Cent { val: 0 }),
|
||||
total_used: Some(Cent { val: 1000 }),
|
||||
},
|
||||
BillingPeriodUsage {
|
||||
billing_cycle: Some(BillingCycle {
|
||||
year: 2025,
|
||||
month: 3,
|
||||
}),
|
||||
included_used: Some(Cent { val: 1800 }),
|
||||
on_demand_used: Some(Cent { val: 0 }),
|
||||
total_used: Some(Cent { val: 1800 }),
|
||||
},
|
||||
],
|
||||
}),
|
||||
on_demand_enabled: Some(true),
|
||||
subscription_tier: Some("SuperGrok".into()),
|
||||
};
|
||||
let ctx = billing_unified_log_ctx(&resp);
|
||||
assert_eq!(ctx["onDemandEnabled"], true);
|
||||
assert_eq!(ctx["subscriptionTier"], "SuperGrok");
|
||||
let config = ctx["config"].as_object().expect("config object");
|
||||
assert!(
|
||||
config.get("history").is_none(),
|
||||
"full history must be collapsed"
|
||||
);
|
||||
assert_eq!(config["historyLen"], 2);
|
||||
assert_eq!(
|
||||
config["latestHistory"]["billingCycle"]["month"], 3,
|
||||
"latest history period retained"
|
||||
);
|
||||
assert_eq!(config["creditUsagePercent"], 42.5);
|
||||
assert_eq!(config["prepaidBalance"]["val"], 100);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_response_roundtrips_through_json() {
|
||||
let config = BillingConfig {
|
||||
credit_usage_percent: None,
|
||||
current_period: None,
|
||||
monthly_limit: Some(Cent { val: 5000 }),
|
||||
used: Some(Cent { val: 123 }),
|
||||
on_demand_cap: Some(Cent { val: 0 }),
|
||||
on_demand_used: Some(Cent { val: 50 }),
|
||||
prepaid_balance: Some(Cent { val: 750 }),
|
||||
is_unified_billing_user: None,
|
||||
billing_period_start: Some("2025-04-01T00:00:00Z".to_string()),
|
||||
billing_period_end: Some("2025-05-01T00:00:00Z".to_string()),
|
||||
history: vec![BillingPeriodUsage {
|
||||
billing_cycle: Some(BillingCycle {
|
||||
year: 2025,
|
||||
month: 3,
|
||||
}),
|
||||
included_used: Some(Cent { val: 4500 }),
|
||||
on_demand_used: Some(Cent { val: 100 }),
|
||||
total_used: Some(Cent { val: 4600 }),
|
||||
}],
|
||||
};
|
||||
let resp = BillingConfigResponse {
|
||||
config: Some(config),
|
||||
on_demand_enabled: None,
|
||||
subscription_tier: None,
|
||||
};
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
let roundtripped: BillingConfigResponse = serde_json::from_value(json).unwrap();
|
||||
let rt_config = roundtripped.config.unwrap();
|
||||
assert_eq!(rt_config.monthly_limit.unwrap().val, 5000);
|
||||
assert_eq!(rt_config.used.unwrap().val, 123);
|
||||
assert_eq!(rt_config.prepaid_balance.unwrap().val, 750);
|
||||
assert_eq!(rt_config.history.len(), 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_response_handles_null_config() {
|
||||
let json = serde_json::json!({"config": null});
|
||||
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
|
||||
assert!(resp.config.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_response_handles_empty_history() {
|
||||
let json = serde_json::json!({
|
||||
"config": {
|
||||
"monthlyLimit": {"val": 1000},
|
||||
"used": {"val": 0}
|
||||
}
|
||||
});
|
||||
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
|
||||
let config = resp.config.unwrap();
|
||||
assert_eq!(config.monthly_limit.unwrap().val, 1000);
|
||||
assert!(config.history.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_serializes_camel_case() {
|
||||
let config = BillingConfig {
|
||||
credit_usage_percent: None,
|
||||
current_period: None,
|
||||
monthly_limit: Some(Cent { val: 100 }),
|
||||
used: None,
|
||||
on_demand_cap: None,
|
||||
on_demand_used: None,
|
||||
prepaid_balance: None,
|
||||
is_unified_billing_user: None,
|
||||
billing_period_start: None,
|
||||
billing_period_end: None,
|
||||
history: vec![],
|
||||
};
|
||||
let json = serde_json::to_value(&config).unwrap();
|
||||
assert!(json.get("monthlyLimit").is_some());
|
||||
// Fields with None are skipped
|
||||
assert!(json.get("creditUsagePercent").is_none());
|
||||
assert!(json.get("currentPeriod").is_none());
|
||||
assert!(json.get("used").is_none());
|
||||
assert!(json.get("onDemandCap").is_none());
|
||||
assert!(json.get("onDemandUsed").is_none());
|
||||
assert!(json.get("prepaidBalance").is_none());
|
||||
assert!(json.get("billingPeriodStart").is_none());
|
||||
// Empty history is skipped
|
||||
assert!(json.get("history").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn billing_config_deserializes_credits_config_shape() {
|
||||
// Newer `GetGrokCreditsConfig` response: percentage-based usage,
|
||||
// a typed current period, and history keyed by `period`.
|
||||
let json = serde_json::json!({
|
||||
"config": {
|
||||
"creditUsagePercent": 42.5,
|
||||
"currentPeriod": {
|
||||
"type": "USAGE_PERIOD_TYPE_WEEKLY",
|
||||
"start": "2026-06-01T00:00:00Z",
|
||||
"end": "2026-06-08T00:00:00Z"
|
||||
},
|
||||
"onDemandCap": {"val": 5000},
|
||||
"onDemandUsed": {"val": 300},
|
||||
"prepaidBalance": {"val": 1250},
|
||||
"isUnifiedBillingUser": true,
|
||||
"productUsage": [
|
||||
{"product": "PRODUCT_GROK_BUILD", "usagePercent": 61.2}
|
||||
],
|
||||
"history": [
|
||||
{
|
||||
"period": {
|
||||
"type": "USAGE_PERIOD_TYPE_WEEKLY",
|
||||
"start": "2026-05-25T00:00:00Z",
|
||||
"end": "2026-06-01T00:00:00Z"
|
||||
},
|
||||
"onDemandUsed": {"val": 120}
|
||||
}
|
||||
]
|
||||
}
|
||||
});
|
||||
let resp: BillingConfigResponse = serde_json::from_value(json).unwrap();
|
||||
let config = resp.config.unwrap();
|
||||
assert_eq!(config.credit_usage_percent, Some(42.5));
|
||||
let period = config.current_period.as_ref().unwrap();
|
||||
assert_eq!(
|
||||
period.period_type.as_deref(),
|
||||
Some("USAGE_PERIOD_TYPE_WEEKLY")
|
||||
);
|
||||
assert_eq!(period.end.as_deref(), Some("2026-06-08T00:00:00Z"));
|
||||
// Deprecated fields are absent in the credits shape.
|
||||
assert!(config.monthly_limit.is_none());
|
||||
assert!(config.billing_period_end.is_none());
|
||||
assert_eq!(config.on_demand_cap.unwrap().val, 5000);
|
||||
assert_eq!(config.on_demand_used.unwrap().val, 300);
|
||||
// Bought (prepaid) credit balance is parsed from the credits config.
|
||||
assert_eq!(config.prepaid_balance.unwrap().val, 1250);
|
||||
assert_eq!(config.is_unified_billing_user, Some(true));
|
||||
// productUsage is still unused by the CLI billing surface.
|
||||
assert_eq!(config.history.len(), 1);
|
||||
assert_eq!(config.history[0].on_demand_used.as_ref().unwrap().val, 120);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cent_serializes_as_val_field() {
|
||||
let c = Cent { val: 4299 };
|
||||
let json = serde_json::to_value(&c).unwrap();
|
||||
assert_eq!(json, serde_json::json!({"val": 4299}));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,13 @@
|
||||
//! `x.ai/session/load_history`: fetch one older page of a gateway-backed
|
||||
//! conversation by client-owned cursor (`beforeId` → `nextBeforeId`).
|
||||
use super::ExtResult;
|
||||
use crate::agent::MvpAgent;
|
||||
use agent_client_protocol as acp;
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
if true {
|
||||
let _ = (agent, args);
|
||||
return Err(acp::Error::method_not_found());
|
||||
}
|
||||
Err(acp::Error::method_not_found())
|
||||
}
|
||||
@@ -0,0 +1,439 @@
|
||||
//! Code Navigation Extension Methods
|
||||
//!
|
||||
//! Provides go-to-definition, go-to-references, and symbol lookup functionality
|
||||
//! using the kigi-codebase-graph index.
|
||||
//!
|
||||
//! ## Extension Methods
|
||||
//!
|
||||
//! | Method | Description |
|
||||
//! |--------|-------------|
|
||||
//! | `x.ai/code/goto-definition` | Definition location(s) for symbol at position |
|
||||
//! | `x.ai/code/goto-references` | Reference location(s) for symbol at position |
|
||||
//! | `x.ai/code/find-definitions` | All definitions of a symbol by name |
|
||||
//! | `x.ai/code/find-references` | All references to a symbol by name |
|
||||
//! | `x.ai/code/status` | Indexing status |
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use crate::agent::mvp_agent::{CodeNavEligibility, MvpAgent};
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Record a structured telemetry event at the end of a code-nav handler call.
|
||||
///
|
||||
/// This is called once per request with the method name, triggering session,
|
||||
/// cwd, whether the index was newly spawned or reused, and total elapsed time.
|
||||
/// These fields make it possible to:
|
||||
/// - identify first-use latency (newly spawned + high elapsed_ms)
|
||||
/// - identify reuse latency (reused + low elapsed_ms)
|
||||
/// - attribute slowness to index startup vs query processing
|
||||
fn log_code_nav_telemetry(
|
||||
method: &str,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
cwd: &Path,
|
||||
was_newly_started: bool,
|
||||
elapsed_ms: u128,
|
||||
) {
|
||||
tracing::info!(
|
||||
method,
|
||||
session_id = session_id.map(|s| s.0.as_ref()).unwrap_or(""),
|
||||
cwd = %cwd.display(),
|
||||
index_newly_started = was_newly_started,
|
||||
elapsed_ms,
|
||||
"code-nav request completed"
|
||||
);
|
||||
}
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
// ========== Request Types ==========
|
||||
|
||||
/// Position-based query request (for goto-definition, goto-references).
|
||||
/// Position parameters are 1-indexed (matching editor display).
|
||||
///
|
||||
/// **`sessionId` is required** for all code-nav requests. Per-client
|
||||
/// capability gating requires a valid session so eligibility is resolved
|
||||
/// correctly in both simple and leader modes. Requests without `sessionId`
|
||||
/// receive `reason: sessionRequired` in the error response.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GotoRequest {
|
||||
/// Session ID — required for code navigation.
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
/// Working directory (optional when session_id is provided).
|
||||
pub cwd: Option<String>,
|
||||
/// Relative path to the file within the cwd
|
||||
pub path: String,
|
||||
/// 1-indexed line number
|
||||
pub row: usize,
|
||||
/// 1-indexed column number
|
||||
pub column: usize,
|
||||
}
|
||||
|
||||
/// Symbol name query request (for find-definitions, find-references).
|
||||
///
|
||||
/// **`sessionId` is required** — same contract as [`GotoRequest`].
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FindSymbolRequest {
|
||||
/// Session ID — required for code navigation.
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
/// Working directory (optional when session_id is provided).
|
||||
pub cwd: Option<String>,
|
||||
/// Symbol name to search for
|
||||
pub symbol: String,
|
||||
/// Optional context file path for ranking results
|
||||
pub context_path: Option<String>,
|
||||
}
|
||||
|
||||
/// Status request — check indexing status.
|
||||
///
|
||||
/// **`sessionId` is required** — same contract as [`GotoRequest`].
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct StatusRequest {
|
||||
/// Session ID — required for code navigation.
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
/// Working directory (optional when session_id is provided).
|
||||
pub cwd: Option<String>,
|
||||
}
|
||||
|
||||
// ========== Response Types ==========
|
||||
|
||||
/// Response for goto-definition and goto-references queries.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CodeNavResponse {
|
||||
/// The symbol that was queried
|
||||
pub symbol: String,
|
||||
/// List of locations where the symbol was found
|
||||
pub locations: Vec<SymbolLocation>,
|
||||
}
|
||||
|
||||
/// A symbol location in a file.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SymbolLocation {
|
||||
/// Absolute path to the file
|
||||
pub path: String,
|
||||
/// 1-indexed line number
|
||||
pub line: usize,
|
||||
/// 1-indexed column (start of symbol, if available)
|
||||
pub column: usize,
|
||||
/// 1-indexed end line
|
||||
pub end_line: usize,
|
||||
/// 1-indexed end column
|
||||
pub end_column: usize,
|
||||
/// The matched symbol name (useful for aliases/imports)
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub matched_symbol: Option<String>,
|
||||
}
|
||||
|
||||
/// Reason string for the `x.ai/code/status` response.
|
||||
///
|
||||
/// Serialised as a camelCase string so clients can pattern-match on it.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub enum IndexStatusReason {
|
||||
/// Index is running and ready.
|
||||
Active,
|
||||
/// Index is eligible but has not been started yet (first code-nav request
|
||||
/// will trigger lazy startup).
|
||||
NotStarted,
|
||||
/// Client type is not web (web-only for initial rollout).
|
||||
ClientNotWeb,
|
||||
/// Client did not advertise `x.ai/codeNavigation.enabled`.
|
||||
CapabilityNotAdvertised,
|
||||
/// `codebase_indexing` feature is disabled in config.
|
||||
DisabledByConfig,
|
||||
/// The cwd is not inside a git repository.
|
||||
NotGitRepo,
|
||||
/// `sessionId` is required but was absent or refers to an unknown session.
|
||||
SessionRequired,
|
||||
}
|
||||
|
||||
/// Response for status query.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct StatusResponse {
|
||||
/// Whether an index is currently active for this cwd.
|
||||
pub indexed: bool,
|
||||
/// Whether this client is eligible to use codebase indexing.
|
||||
pub eligible: bool,
|
||||
/// Reason code describing the current status.
|
||||
pub reason: IndexStatusReason,
|
||||
/// Number of files in the index (present only when `indexed` is true).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub file_count: Option<usize>,
|
||||
}
|
||||
|
||||
// ========== Handler ==========
|
||||
|
||||
/// Handle code navigation extension methods.
|
||||
///
|
||||
/// Routes through [`WorkspaceOps`]. Eligibility checks still run in shell since
|
||||
/// they depend on agent-level config (client type, feature flags).
|
||||
pub async fn handle(
|
||||
agent: &MvpAgent,
|
||||
ops: &kigi_workspace::WorkspaceOps,
|
||||
args: &acp::ExtRequest,
|
||||
) -> ExtResult {
|
||||
use kigi_workspace::workspace_ops::*;
|
||||
|
||||
match args.method.as_ref() {
|
||||
"x.ai/code/goto-definition" => {
|
||||
let req: GotoRequest = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
let was_newly_started =
|
||||
ensure_eligible_and_started(agent, req.session_id.as_ref(), &cwd)?;
|
||||
let start = std::time::Instant::now();
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&CodeGotoDefinitionReq {
|
||||
root: Some(cwd.clone()),
|
||||
file: cwd.join(&req.path).to_string_lossy().to_string(),
|
||||
line: req.row,
|
||||
col: req.column,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("code nav error: {e}")))?;
|
||||
log_code_nav_telemetry(
|
||||
"goto-definition",
|
||||
req.session_id.as_ref(),
|
||||
&cwd,
|
||||
was_newly_started,
|
||||
start.elapsed().as_millis(),
|
||||
);
|
||||
to_code_nav_ext_response(result)
|
||||
}
|
||||
"x.ai/code/goto-references" => {
|
||||
let req: GotoRequest = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
let was_newly_started =
|
||||
ensure_eligible_and_started(agent, req.session_id.as_ref(), &cwd)?;
|
||||
let start = std::time::Instant::now();
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&CodeGotoReferencesReq {
|
||||
root: Some(cwd.clone()),
|
||||
file: cwd.join(&req.path).to_string_lossy().to_string(),
|
||||
line: req.row,
|
||||
col: req.column,
|
||||
include_definition: true,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("code nav error: {e}")))?;
|
||||
log_code_nav_telemetry(
|
||||
"goto-references",
|
||||
req.session_id.as_ref(),
|
||||
&cwd,
|
||||
was_newly_started,
|
||||
start.elapsed().as_millis(),
|
||||
);
|
||||
to_code_nav_ext_response(result)
|
||||
}
|
||||
"x.ai/code/find-definitions" => {
|
||||
let req: FindSymbolRequest = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
let was_newly_started =
|
||||
ensure_eligible_and_started(agent, req.session_id.as_ref(), &cwd)?;
|
||||
let start = std::time::Instant::now();
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&CodeFindDefinitionsReq {
|
||||
root: Some(cwd.clone()),
|
||||
symbol: req.symbol.clone(),
|
||||
context_file: req
|
||||
.context_path
|
||||
.as_ref()
|
||||
.map(|p| cwd.join(p).to_string_lossy().to_string()),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("code nav error: {e}")))?;
|
||||
log_code_nav_telemetry(
|
||||
"find-definitions",
|
||||
req.session_id.as_ref(),
|
||||
&cwd,
|
||||
was_newly_started,
|
||||
start.elapsed().as_millis(),
|
||||
);
|
||||
to_code_nav_ext_response(result)
|
||||
}
|
||||
"x.ai/code/find-references" => {
|
||||
let req: FindSymbolRequest = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
let was_newly_started =
|
||||
ensure_eligible_and_started(agent, req.session_id.as_ref(), &cwd)?;
|
||||
let start = std::time::Instant::now();
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&CodeFindReferencesReq {
|
||||
root: Some(cwd.clone()),
|
||||
symbol: req.symbol.clone(),
|
||||
context_file: req
|
||||
.context_path
|
||||
.as_ref()
|
||||
.map(|p| cwd.join(p).to_string_lossy().to_string()),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("code nav error: {e}")))?;
|
||||
log_code_nav_telemetry(
|
||||
"find-references",
|
||||
req.session_id.as_ref(),
|
||||
&cwd,
|
||||
was_newly_started,
|
||||
start.elapsed().as_millis(),
|
||||
);
|
||||
to_code_nav_ext_response(result)
|
||||
}
|
||||
"x.ai/code/status" => {
|
||||
let req: StatusRequest = serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
|
||||
// Check eligibility for the status response.
|
||||
let (eligible, reason, indexed, file_count) = match agent
|
||||
.code_nav_eligibility_for_request(req.session_id.as_ref(), &cwd)
|
||||
{
|
||||
Ok(()) => {
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&CodeIndexStatusReq {
|
||||
root: Some(cwd.clone()),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("code nav error: {e}"))
|
||||
})?;
|
||||
if result.active {
|
||||
(true, IndexStatusReason::Active, true, result.file_count)
|
||||
} else {
|
||||
(true, IndexStatusReason::NotStarted, false, None)
|
||||
}
|
||||
}
|
||||
Err(ineligible) => {
|
||||
let reason = match ineligible {
|
||||
CodeNavEligibility::ClientNotWeb => IndexStatusReason::ClientNotWeb,
|
||||
CodeNavEligibility::CapabilityNotAdvertised => {
|
||||
IndexStatusReason::CapabilityNotAdvertised
|
||||
}
|
||||
CodeNavEligibility::DisabledByConfig => IndexStatusReason::DisabledByConfig,
|
||||
CodeNavEligibility::NotGitRepo => IndexStatusReason::NotGitRepo,
|
||||
CodeNavEligibility::SessionRequired => IndexStatusReason::SessionRequired,
|
||||
};
|
||||
(false, reason, false, None)
|
||||
}
|
||||
};
|
||||
|
||||
let status = StatusResponse {
|
||||
indexed,
|
||||
eligible,
|
||||
reason,
|
||||
file_count,
|
||||
};
|
||||
super::to_ext_response(Ok(status))
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert workspace CodeNavResponse to the shell's CodeNavResponse format
|
||||
/// and wrap in the `ExtMethodResult` envelope that clients expect.
|
||||
fn to_code_nav_ext_response(resp: kigi_workspace::workspace_ops::CodeNavResponse) -> ExtResult {
|
||||
let symbol = resp
|
||||
.locations
|
||||
.first()
|
||||
.and_then(|l| l.symbol.clone())
|
||||
.unwrap_or_default();
|
||||
let shell_resp = CodeNavResponse {
|
||||
symbol,
|
||||
locations: resp
|
||||
.locations
|
||||
.into_iter()
|
||||
.map(|loc| SymbolLocation {
|
||||
path: loc.path,
|
||||
line: loc.line,
|
||||
column: 0,
|
||||
end_line: loc.line,
|
||||
end_column: 0,
|
||||
matched_symbol: loc.symbol,
|
||||
})
|
||||
.collect(),
|
||||
};
|
||||
super::to_ext_response(Ok(shell_resp))
|
||||
}
|
||||
|
||||
/// Check eligibility, ensure the codebase index is started, and return
|
||||
/// whether the index was newly created (for telemetry).
|
||||
fn ensure_eligible_and_started(
|
||||
agent: &MvpAgent,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
cwd: &Path,
|
||||
) -> Result<bool, acp::Error> {
|
||||
if let Err(reason) = agent.code_nav_eligibility_for_request(session_id, cwd) {
|
||||
return Err(eligibility_error(reason));
|
||||
}
|
||||
// Start the index if not already running (lazy creation).
|
||||
let was_newly_started = agent
|
||||
.start_codebase_index_for_code_nav(session_id, cwd)
|
||||
.map(|(_, was_new)| was_new)
|
||||
.unwrap_or(false);
|
||||
Ok(was_newly_started)
|
||||
}
|
||||
|
||||
// ========== Helper Functions ==========
|
||||
|
||||
/// Resolve cwd from session_id or direct cwd parameter.
|
||||
fn resolve_cwd(
|
||||
agent: &MvpAgent,
|
||||
cwd: Option<String>,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
) -> Result<PathBuf, acp::Error> {
|
||||
// Prefer direct cwd if provided
|
||||
if let Some(cwd_str) = cwd {
|
||||
return Ok(PathBuf::from(cwd_str));
|
||||
}
|
||||
|
||||
// Fall back to session's cwd
|
||||
if let Some(sid) = session_id
|
||||
&& let Some(session_cwd) = agent.get_session_cwd(sid)
|
||||
{
|
||||
return Ok(session_cwd);
|
||||
}
|
||||
|
||||
Err(acp::Error::invalid_params().data("either cwd or valid sessionId must be provided"))
|
||||
}
|
||||
|
||||
/// Map a `CodeNavEligibility` error to a human-readable ACP error.
|
||||
fn eligibility_error(reason: CodeNavEligibility) -> acp::Error {
|
||||
let msg = match reason {
|
||||
CodeNavEligibility::ClientNotWeb => {
|
||||
"code navigation is currently only enabled for grok-web clients"
|
||||
}
|
||||
CodeNavEligibility::CapabilityNotAdvertised => {
|
||||
"client must advertise x.ai/codeNavigation.enabled to use code navigation"
|
||||
}
|
||||
CodeNavEligibility::DisabledByConfig => "code navigation is disabled by configuration",
|
||||
CodeNavEligibility::NotGitRepo => {
|
||||
"code navigation requires the workspace to be inside a git repository"
|
||||
}
|
||||
CodeNavEligibility::SessionRequired => {
|
||||
"sessionId is required for code navigation and must refer to a valid active session"
|
||||
}
|
||||
};
|
||||
acp::Error::invalid_params().data(msg)
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
//! `x.ai/debug/*` extension handlers for local client testing.
|
||||
//!
|
||||
//! These methods bypass heuristics, sampling, cooldowns, and enabled checks
|
||||
//! so client engineers can exercise notification → response flows without
|
||||
//! needing real experiments, real sessions, or real model inference.
|
||||
//!
|
||||
//! - `trigger_feedback`: fire a synthetic `FeedbackRequestNotification`.
|
||||
//! - `arm_auto_compact`: arm the next turn to unconditionally trigger
|
||||
//! auto-compaction, regardless of context window usage.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{ExtResult, parse_params};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::{ExtMethodResult, SessionCommand};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/debug/trigger_feedback" => {
|
||||
tracing::info!("debug: triggering test feedback request");
|
||||
handle_trigger_feedback(agent, args).await
|
||||
}
|
||||
"x.ai/debug/arm_auto_compact" => handle_arm_auto_compact(agent, args),
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_trigger_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
use crate::session::feedback::{FeedbackMode, FeedbackTier};
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DebugTriggerParams {
|
||||
#[serde(alias = "session_id")]
|
||||
session_id: String,
|
||||
/// "tier1" | "tier2" | "tier3" (default: "tier1")
|
||||
#[serde(default)]
|
||||
tier: Option<String>,
|
||||
/// "thumbs" | "stars" | "text" | "thumbs_text" | "stars_text" (default: "thumbs_text")
|
||||
#[serde(default)]
|
||||
mode: Option<String>,
|
||||
}
|
||||
|
||||
let params: DebugTriggerParams = parse_params(args)?;
|
||||
|
||||
let tier = match params.tier.as_deref() {
|
||||
Some("tier2") => FeedbackTier::Tier2,
|
||||
Some("tier3") => FeedbackTier::Tier3,
|
||||
Some("tier1") | None => FeedbackTier::Tier1,
|
||||
Some(other) => {
|
||||
return Err(acp::Error::invalid_params().data(format!(
|
||||
"unknown tier: {other:?} (expected tier1/tier2/tier3)"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let mode = match params.mode.as_deref() {
|
||||
Some("thumbs") => FeedbackMode::Thumbs,
|
||||
Some("stars") => FeedbackMode::Stars,
|
||||
Some("text") => FeedbackMode::Text,
|
||||
Some("stars_text") => FeedbackMode::StarsText,
|
||||
Some("thumbs_text") | None => FeedbackMode::ThumbsText,
|
||||
Some(other) => {
|
||||
return Err(acp::Error::invalid_params().data(format!(
|
||||
"unknown mode: {other:?} (expected thumbs/stars/text/thumbs_text/stars_text)"
|
||||
)));
|
||||
}
|
||||
};
|
||||
|
||||
let session_id = acp::SessionId::new(params.session_id.clone());
|
||||
let handle = agent
|
||||
.sessions
|
||||
.borrow()
|
||||
.get(&session_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| {
|
||||
acp::Error::invalid_params().data(format!("session not found: {}", params.session_id))
|
||||
})?;
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::TriggerTestFeedback {
|
||||
tier,
|
||||
mode,
|
||||
respond_to: tx,
|
||||
})
|
||||
.map_err(|_| {
|
||||
acp::Error::internal_error().data("failed to dispatch debug trigger to session")
|
||||
})?;
|
||||
|
||||
rx.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("Internal error: {e:?}")))
|
||||
}
|
||||
|
||||
fn handle_arm_auto_compact(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let params: serde_json::Value = parse_params(args)?;
|
||||
|
||||
let session_id_str = params["sessionId"]
|
||||
.as_str()
|
||||
.or_else(|| params["session_id"].as_str())
|
||||
.ok_or_else(|| acp::Error::invalid_params().data("sessionId required"))?;
|
||||
let session_id = acp::SessionId::new(session_id_str);
|
||||
|
||||
let handle = agent
|
||||
.sessions
|
||||
.borrow()
|
||||
.get(&session_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
|
||||
|
||||
handle
|
||||
.force_compact
|
||||
.store(true, std::sync::atomic::Ordering::Relaxed);
|
||||
|
||||
tracing::info!(
|
||||
session_id = %session_id_str,
|
||||
"debug: armed auto-compact for next turn"
|
||||
);
|
||||
|
||||
ExtMethodResult::success(serde_json::json!({ "armed": true }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
//! `x.ai/feedback`, `x.ai/feedback/dismiss`, `x.ai/btw`, and `x.ai/review/*`
|
||||
//! extension handlers.
|
||||
//!
|
||||
//! - `feedback`/`feedback/dismiss`: persist user ratings/text locally and
|
||||
//! forward to cli-chat-proxy.
|
||||
//! - `btw`: dispatch a side question to the active session via
|
||||
//! `SessionCommand::SideQuestion` and return the answer.
|
||||
//! - `review/comment` and `review/comment/delete`: record inline code review
|
||||
//! events to cloud storage.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use super::{ExtResult, parse_params};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::persistence::{LocalFeedbackEntry, UserFeedbackEntry};
|
||||
use crate::session::{
|
||||
ClientFeedbackInput, CommentDeleteRequest, CommentDeleteResponse, CommentRequest,
|
||||
CommentResponse, FeedbackRequestDismiss, FeedbackResponse, SessionCommand,
|
||||
};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/btw" => {
|
||||
tracing::info!("handling /btw side question");
|
||||
handle_btw(agent, args).await
|
||||
}
|
||||
"x.ai/feedback" | "x.ai/feedback/dismiss" => {
|
||||
tracing::info!("handling user feedback");
|
||||
handle_feedback(agent, args).await
|
||||
}
|
||||
m if m.starts_with("x.ai/review") => {
|
||||
tracing::info!("handling review comment");
|
||||
handle_review(agent, args).await
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Handle `x.ai/btw` -- a side question that doesn't interrupt the current turn.
|
||||
async fn handle_btw(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct BtwRequest {
|
||||
session_id: String,
|
||||
question: String,
|
||||
}
|
||||
|
||||
let req: BtwRequest = parse_params(args)?;
|
||||
let sid: acp::SessionId = req.session_id.clone().into();
|
||||
let session_handle = {
|
||||
let sessions = agent.sessions.borrow();
|
||||
sessions.get(&sid).cloned()
|
||||
};
|
||||
let Some(session) = session_handle else {
|
||||
return Err(
|
||||
acp::Error::invalid_params().data(format!("session not found: {}", req.session_id))
|
||||
);
|
||||
};
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = session.cmd_tx.send(SessionCommand::SideQuestion {
|
||||
question: req.question,
|
||||
respond_to: tx,
|
||||
});
|
||||
let result = rx
|
||||
.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?;
|
||||
match result {
|
||||
Ok(answer) => super::to_ext_response(Ok(serde_json::json!({
|
||||
"answer": answer,
|
||||
}))),
|
||||
Err(e) => Err(acp::Error::internal_error().data(e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_feedback(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
if !agent.cfg.borrow().is_feedback_enabled() {
|
||||
return Err(acp::Error::internal_error().data(
|
||||
"Feedback is disabled. To enable, set KIGI_FEEDBACK_ENABLED=true or \
|
||||
[features] feedback = true in config.toml.",
|
||||
));
|
||||
}
|
||||
|
||||
match args.method.as_ref() {
|
||||
"x.ai/feedback" => {
|
||||
// Parse the input -- try the full ClientFeedbackInput first,
|
||||
// then fall back to the simple FeedbackRequest (from /feedback slash command)
|
||||
// which only has {session_id, feedback_text} and no client_type.
|
||||
let feedback_input: ClientFeedbackInput =
|
||||
match serde_json::from_str::<ClientFeedbackInput>(args.params.get()) {
|
||||
Ok(input) => input,
|
||||
Err(_) => {
|
||||
// Fallback: parse simple FeedbackRequest from /feedback command
|
||||
let simple: crate::session::FeedbackRequest = parse_params(args)?;
|
||||
ClientFeedbackInput {
|
||||
session_id: simple.session_id,
|
||||
client_type:
|
||||
prod_mc_cli_chat_proxy_types::feedback_types::ClientType::Tui,
|
||||
rating_type: None,
|
||||
rating_value: None,
|
||||
feedback_text: Some(simple.feedback_text),
|
||||
feedback_categories: vec![],
|
||||
context_type: None,
|
||||
turn_number: None,
|
||||
request_id: None,
|
||||
client_version: None,
|
||||
metadata: None,
|
||||
terminal_info: None,
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let session_id = acp::SessionId::new(feedback_input.session_id.clone());
|
||||
let session_handle = agent.sessions.borrow().get(&session_id).cloned();
|
||||
|
||||
let (model_id, model_metadata) = if let Some(ref session) = session_handle {
|
||||
let (tx1, rx1) = tokio::sync::oneshot::channel();
|
||||
let _ = session
|
||||
.cmd_tx
|
||||
.send(SessionCommand::GetCurrentModel { responds_to: tx1 });
|
||||
let model_id = rx1.await.ok();
|
||||
|
||||
let model_metadata = session.get_model_metadata().await;
|
||||
|
||||
(model_id, model_metadata)
|
||||
} else {
|
||||
let sampling_config = agent.sampling_config.borrow().clone();
|
||||
(Some(sampling_config.model.clone()), Default::default())
|
||||
};
|
||||
|
||||
let turn_number = feedback_input.turn_number.or_else(|| {
|
||||
agent
|
||||
.session_turn_number(&session_id)
|
||||
.map(|t| t.saturating_sub(1) as i64)
|
||||
});
|
||||
|
||||
let mut submission = feedback_input.to_submission(
|
||||
model_id.clone(),
|
||||
model_metadata.resolved_model_id,
|
||||
model_metadata.model_fingerprint,
|
||||
turn_number,
|
||||
);
|
||||
let turn_number = submission.turn_number;
|
||||
|
||||
if let Some(user_meta) =
|
||||
crate::agent::mvp_agent::parse_json_object_env("KIGI_USER_METADATA")
|
||||
{
|
||||
submission.merge_metadata(user_meta);
|
||||
}
|
||||
|
||||
// Enrich with session context for Slack notifications (best-effort).
|
||||
if let Some(ref session_handle) = session_handle {
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let _ = session_handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::GetFeedbackContext {
|
||||
turn_number,
|
||||
responds_to: tx,
|
||||
});
|
||||
if let Ok(ctx) = rx.await {
|
||||
submission.tool_outcomes = ctx.tool_outcomes;
|
||||
submission.session_cwd = Some(ctx.session_cwd);
|
||||
submission.compaction_count = Some(ctx.compaction_count);
|
||||
submission.context_window_usage = Some(ctx.context_window_usage);
|
||||
submission.context_tokens_used = Some(ctx.context_tokens_used);
|
||||
submission.context_window_tokens = Some(ctx.context_window_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
// Track rating in session signals
|
||||
if let (Some(session_handle), Some(rating_value)) =
|
||||
(&session_handle, feedback_input.rating_value)
|
||||
{
|
||||
use prod_mc_cli_chat_proxy_types::feedback_types::RatingType;
|
||||
let (is_positive, is_negative) = match feedback_input.rating_type {
|
||||
// Thumbs: -1 = down, 0 = neutral, 1 = up
|
||||
Some(RatingType::Thumbs) | None => (rating_value > 0, rating_value < 0),
|
||||
// Stars (1-5): >= 4 positive, <= 2 negative, 3 neutral
|
||||
Some(RatingType::Stars) => (rating_value >= 4, rating_value <= 2),
|
||||
// NPS (0-10): 9-10 promoter, 0-6 detractor, 7-8 passive
|
||||
Some(RatingType::Nps) => (rating_value >= 9, rating_value <= 6),
|
||||
};
|
||||
if is_positive {
|
||||
session_handle.signals_handle.record_positive_rating();
|
||||
} else if is_negative {
|
||||
session_handle.signals_handle.record_negative_rating();
|
||||
}
|
||||
}
|
||||
|
||||
// Log feedback type for debugging
|
||||
if feedback_input.is_solicited() {
|
||||
tracing::info!(
|
||||
session_id = %feedback_input.session_id,
|
||||
request_id = ?feedback_input.request_id(),
|
||||
turn_number = ?turn_number,
|
||||
"Solicited feedback received (response to feedback request)"
|
||||
);
|
||||
} else {
|
||||
tracing::info!(
|
||||
session_id = %feedback_input.session_id,
|
||||
turn_number = ?turn_number,
|
||||
"Spontaneous user feedback received"
|
||||
);
|
||||
}
|
||||
|
||||
let client = agent.feedback_client();
|
||||
if client.is_none() {
|
||||
tracing::warn!(
|
||||
"no feedback client available (missing proxy credentials); feedback saved locally only"
|
||||
);
|
||||
}
|
||||
let outcome = crate::session::feedback_manager::submit_feedback_workflow(
|
||||
&mut submission,
|
||||
client.as_ref(),
|
||||
session_handle.as_ref().map(|h| &h.persistence_tx),
|
||||
feedback_input.is_solicited(),
|
||||
)
|
||||
.await;
|
||||
|
||||
match &outcome {
|
||||
crate::session::feedback_manager::SubmitOutcome::Submitted => {
|
||||
tracing::info!("feedback submitted to proxy successfully");
|
||||
}
|
||||
crate::session::feedback_manager::SubmitOutcome::LocalOnly => {
|
||||
tracing::warn!("feedback saved locally only (no proxy client)");
|
||||
}
|
||||
crate::session::feedback_manager::SubmitOutcome::Failed(e) => {
|
||||
tracing::error!(error = %e, "feedback submission to proxy failed");
|
||||
return Err(acp::Error::internal_error()
|
||||
.data(format!("Feedback submission failed: {e}")));
|
||||
}
|
||||
}
|
||||
|
||||
let value = serde_json::to_value(FeedbackResponse { success: true })
|
||||
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
|
||||
.expect("to work")
|
||||
.expect("to work");
|
||||
Ok(acp::ExtResponse::new(value))
|
||||
}
|
||||
"x.ai/feedback/dismiss" => {
|
||||
let dismiss_input: FeedbackRequestDismiss = parse_params(args)?;
|
||||
|
||||
tracing::info!(
|
||||
session_id = %dismiss_input.session_id,
|
||||
request_id = %dismiss_input.request_id,
|
||||
"Feedback request dismissed by user"
|
||||
);
|
||||
|
||||
// Count dismissals too (else event_type is always "responded" and
|
||||
// response-rate is unknowable).
|
||||
{
|
||||
tracing::info_span!(
|
||||
"feedback.survey",
|
||||
survey_type = "session",
|
||||
event_type = "dismissed",
|
||||
appearance_id = %dismiss_input.request_id,
|
||||
has_feedback_text = false,
|
||||
is_solicited = true,
|
||||
)
|
||||
.in_scope(|| {});
|
||||
}
|
||||
|
||||
// Persist dismiss locally; flushed before storage CopyFile by the persistence actor.
|
||||
{
|
||||
let session_id = acp::SessionId::new(dismiss_input.session_id.clone());
|
||||
if let Some(session_handle) = agent.sessions.borrow().get(&session_id) {
|
||||
session_handle.persist_feedback(LocalFeedbackEntry::UserFeedback(
|
||||
UserFeedbackEntry {
|
||||
submitted_at: chrono::Utc::now(),
|
||||
session_id: dismiss_input.session_id.clone(),
|
||||
turn_number: None,
|
||||
solicited: true,
|
||||
request_id: Some(dismiss_input.request_id.clone()),
|
||||
dismissed: true,
|
||||
submission: None,
|
||||
},
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
let request_id = dismiss_input.request_id.clone();
|
||||
let client = agent
|
||||
.feedback_client()
|
||||
.ok_or_else(|| acp::Error::internal_error().data("No credentials for feedback"))?;
|
||||
let feedback_base_url = agent.cfg.borrow().endpoints.resolve_feedback_base_url();
|
||||
match client.dismiss_request(&request_id).await {
|
||||
Ok(response) => {
|
||||
tracing::info!(
|
||||
request_id = %response.request_id,
|
||||
status = %response.status,
|
||||
feedback_url = %feedback_base_url,
|
||||
"Feedback request dismissed"
|
||||
);
|
||||
let value = serde_json::to_value(&response)
|
||||
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
|
||||
.expect("to work")
|
||||
.expect("to work");
|
||||
Ok(acp::ExtResponse::new(value))
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
error = %e,
|
||||
request_id = %request_id,
|
||||
feedback_url = %feedback_base_url,
|
||||
"Failed to dismiss feedback request"
|
||||
);
|
||||
Err(acp::Error::internal_error()
|
||||
.data(format!("Failed to dismiss feedback request: {e}")))
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Record inline code review events.
|
||||
///
|
||||
/// Methods:
|
||||
/// - `x.ai/review/comment`: record a new inline code comment to cloud storage
|
||||
/// - `x.ai/review/comment/delete`: record a tombstone event for a deleted comment
|
||||
async fn handle_review(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/review/comment" => {
|
||||
let request: CommentRequest = parse_params(args)?;
|
||||
|
||||
let comment_id = uuid::Uuid::now_v7().to_string();
|
||||
|
||||
tracing::info!(
|
||||
comment_id = %comment_id,
|
||||
session_id = %request.session_id,
|
||||
prompt_index = request.prompt_index,
|
||||
path = %request.citation.path,
|
||||
lines = %format!("{}-{}", request.citation.start_line, request.citation.end_line),
|
||||
"Comment received"
|
||||
);
|
||||
|
||||
let value = serde_json::to_value(CommentResponse {
|
||||
comment_id,
|
||||
recorded: true,
|
||||
})
|
||||
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
|
||||
.expect("to work")
|
||||
.expect("to work");
|
||||
Ok(acp::ExtResponse::new(value))
|
||||
}
|
||||
"x.ai/review/comment/delete" => {
|
||||
let request: CommentDeleteRequest = parse_params(args)?;
|
||||
|
||||
tracing::info!(
|
||||
comment_id = %request.comment_id,
|
||||
session_id = %request.session_id,
|
||||
"Comment delete received"
|
||||
);
|
||||
|
||||
let value = serde_json::to_value(CommentDeleteResponse {
|
||||
comment_id: request.comment_id,
|
||||
deleted: true,
|
||||
})
|
||||
.map(|value| serde_json::value::to_raw_value(&value).map(Arc::from))
|
||||
.expect("to work")
|
||||
.expect("to work");
|
||||
Ok(acp::ExtResponse::new(value))
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,252 @@
|
||||
//! Filesystem extension API layer.
|
||||
//!
|
||||
//! Routing: absolute paths work directly; relative paths require sessionId for lookup.
|
||||
//! Business logic delegated to `session::file_system::*` pure functions.
|
||||
use super::{Empty, ExtResult, parse_params, to_ext_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::ExtMethodResult;
|
||||
use crate::session::file_system::{
|
||||
self as fs, FsListParams, FsReadFileData, check_file_size_limits,
|
||||
};
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_workspace::file_system::FsReadEncoding;
|
||||
use serde::Deserialize;
|
||||
use std::path::{Path, PathBuf};
|
||||
fn default_depth() -> usize {
|
||||
1
|
||||
}
|
||||
fn default_limit() -> usize {
|
||||
1000
|
||||
}
|
||||
fn default_follow_symlinks() -> bool {
|
||||
true
|
||||
}
|
||||
fn default_respect_git_ignore() -> bool {
|
||||
true
|
||||
}
|
||||
fn default_max_bytes() -> usize {
|
||||
1_048_576
|
||||
}
|
||||
fn default_create_dirs() -> bool {
|
||||
true
|
||||
}
|
||||
fn default_include_hidden() -> bool {
|
||||
true
|
||||
}
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsListRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
pub path: String,
|
||||
#[serde(default = "default_depth")]
|
||||
pub depth: usize,
|
||||
#[serde(default = "default_include_hidden")]
|
||||
pub include_hidden: bool,
|
||||
#[serde(default = "default_limit")]
|
||||
pub limit: usize,
|
||||
/// Pagination offset applied after the dirs-first sort (default 0).
|
||||
#[serde(default)]
|
||||
pub offset: u64,
|
||||
#[serde(default = "default_follow_symlinks")]
|
||||
pub follow_symlinks: bool,
|
||||
#[serde(default = "default_respect_git_ignore")]
|
||||
pub respect_git_ignore: bool,
|
||||
#[serde(default)]
|
||||
pub include_globs: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub exclude_globs: Vec<String>,
|
||||
}
|
||||
impl FsListRequest {
|
||||
fn to_params(&self) -> FsListParams {
|
||||
FsListParams {
|
||||
path: self.path.clone(),
|
||||
depth: self.depth,
|
||||
limit: self.limit,
|
||||
offset: self.offset,
|
||||
follow_symlinks: self.follow_symlinks,
|
||||
respect_git_ignore: self.respect_git_ignore,
|
||||
include_hidden: self.include_hidden,
|
||||
include_globs: self.include_globs.clone(),
|
||||
exclude_globs: self.exclude_globs.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsExistsRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
pub path: String,
|
||||
}
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsReadFileRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
pub path: String,
|
||||
#[serde(default = "default_max_bytes")]
|
||||
pub max_bytes: usize,
|
||||
#[serde(default)]
|
||||
pub max_lines: Option<usize>,
|
||||
/// Byte offset for a binary-safe ranged read. When `offset`/`length`
|
||||
/// is set (or `encoding` is `base64`) the read returns the chunk
|
||||
/// `[offset, offset + length)`; otherwise the whole file is read
|
||||
/// (legacy behavior).
|
||||
#[serde(default)]
|
||||
pub offset: Option<u64>,
|
||||
/// Bytes to read for a ranged read. Absent means "to EOF", but the
|
||||
/// effective read is always capped at `max_bytes` (default 1 MiB) and the
|
||||
/// server's hard limit, so an unset `length` still yields at most
|
||||
/// `max_bytes`. Detect "more data" by comparing the returned bytes (from
|
||||
/// `offset`) against the response `size`.
|
||||
#[serde(default)]
|
||||
pub length: Option<u64>,
|
||||
/// Transfer encoding for ranged reads (default `utf8`; non-UTF-8
|
||||
/// ranges fall back to base64 regardless).
|
||||
#[serde(default)]
|
||||
pub encoding: FsReadEncoding,
|
||||
}
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsWriteFileRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
pub path: String,
|
||||
pub content: String,
|
||||
#[serde(default = "default_create_dirs")]
|
||||
pub create_dirs: bool,
|
||||
}
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FsDeleteFileRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
pub path: String,
|
||||
}
|
||||
/// Resolve path from explicit value or session lookup.
|
||||
/// For absolute paths, use directly. For relative paths, resolve from session cwd.
|
||||
fn resolve_path(
|
||||
agent: &MvpAgent,
|
||||
path: &str,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
) -> Result<PathBuf, acp::Error> {
|
||||
let p = Path::new(path);
|
||||
if p.is_absolute() {
|
||||
return Ok(p.to_path_buf());
|
||||
}
|
||||
if let Some(sid) = session_id {
|
||||
if let Some(cwd) = agent.get_session_cwd(sid) {
|
||||
return Ok(cwd.join(p));
|
||||
}
|
||||
return Err(acp::Error::invalid_params().data(format!("session not found: {}", sid.0)));
|
||||
}
|
||||
Err(acp::Error::invalid_params().data("sessionId is required for relative paths"))
|
||||
}
|
||||
/// Confine `path` to the workspace root, falling back to the session cwd for
|
||||
/// worktree sessions (rooted outside it). Returns the resolved path and an
|
||||
/// optional confining walk root (`None` when confinement is off — the default,
|
||||
/// so the fallback and error paths only apply on a confining sandbox workspace).
|
||||
async fn confine_local(
|
||||
agent: &MvpAgent,
|
||||
path: &Path,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
) -> Result<(PathBuf, Option<PathBuf>), acp::Error> {
|
||||
let ops = agent.resolve_workspace_ops()?;
|
||||
let handle = ops.workspace_handle().ok_or_else(|| {
|
||||
acp::Error::internal_error().data("no local workspace handle for fs confinement")
|
||||
})?;
|
||||
let workspace_err = match handle.confine_to_workspace_root(path).await {
|
||||
Ok(confined) => return Ok(confined),
|
||||
Err(e) => e,
|
||||
};
|
||||
if let Some(sid) = session_id
|
||||
&& let Some(session_cwd) = agent.get_session_cwd(sid)
|
||||
&& let Ok(confined) = handle.confine_to_root(path, &session_cwd).await
|
||||
{
|
||||
return Ok(confined);
|
||||
}
|
||||
Err(acp::Error::invalid_params().data(workspace_err.to_string()))
|
||||
}
|
||||
pub(crate) fn is_fs_method(method: &str) -> bool {
|
||||
method.starts_with("x.ai/fs/")
|
||||
}
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/fs/list" => {
|
||||
let req = parse_params::<FsListRequest>(args)?;
|
||||
let path = resolve_path(agent, &req.path, req.session_id.as_ref())?;
|
||||
let (path, confine_root) = confine_local(agent, &path, req.session_id.as_ref()).await?;
|
||||
let params = req.to_params();
|
||||
let result = fs::list(&path, ¶ms, confine_root).await;
|
||||
to_ext_response(result)
|
||||
}
|
||||
"x.ai/fs/exists" => {
|
||||
let req = parse_params::<FsExistsRequest>(args)?;
|
||||
let path = resolve_path(agent, &req.path, req.session_id.as_ref())?;
|
||||
let (path, _) = match confine_local(agent, &path, req.session_id.as_ref()).await {
|
||||
Ok(confined) => confined,
|
||||
Err(_) => return to_ext_response(Ok(fs::FsExistsData { exists: false })),
|
||||
};
|
||||
let result = fs::exists(&path).await;
|
||||
to_ext_response(result)
|
||||
}
|
||||
"x.ai/fs/read_file" => {
|
||||
let req = parse_params::<FsReadFileRequest>(args)?;
|
||||
let max_lines = req.max_lines;
|
||||
let path_str = req.path.clone();
|
||||
let ranged = req.offset.is_some()
|
||||
|| req.length.is_some()
|
||||
|| req.encoding == FsReadEncoding::Base64;
|
||||
let path = resolve_path(agent, &req.path, req.session_id.as_ref())?;
|
||||
let (path, _) = confine_local(agent, &path, req.session_id.as_ref()).await?;
|
||||
let read_result: anyhow::Result<FsReadFileData> = if ranged {
|
||||
fs::read_file_ranged(
|
||||
&path,
|
||||
req.offset.unwrap_or(0),
|
||||
req.length.unwrap_or(u64::MAX),
|
||||
req.max_bytes as u64,
|
||||
req.encoding,
|
||||
)
|
||||
.await
|
||||
} else {
|
||||
fs::read_file(&path).await
|
||||
};
|
||||
match read_result {
|
||||
Ok(data) => {
|
||||
let size_check = if ranged {
|
||||
Ok(())
|
||||
} else {
|
||||
check_file_size_limits(&data, &path_str, None, max_lines)
|
||||
};
|
||||
if let Err(err) = size_check {
|
||||
let ext_result: ExtMethodResult<FsReadFileData> = err.into();
|
||||
ext_result
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
} else {
|
||||
to_ext_response(Ok(data))
|
||||
}
|
||||
}
|
||||
Err(e) => to_ext_response(Err::<FsReadFileData, _>(e)),
|
||||
}
|
||||
}
|
||||
"x.ai/fs/write_file" => {
|
||||
let req = parse_params::<FsWriteFileRequest>(args)?;
|
||||
let path = resolve_path(agent, &req.path, req.session_id.as_ref())?;
|
||||
let (path, _) = confine_local(agent, &path, req.session_id.as_ref()).await?;
|
||||
let result = fs::write_file(&path, &req.content, req.create_dirs)
|
||||
.await
|
||||
.map(|_| Empty {});
|
||||
to_ext_response(result)
|
||||
}
|
||||
"x.ai/fs/delete_file" => {
|
||||
let req = parse_params::<FsDeleteFileRequest>(args)?;
|
||||
let path = resolve_path(agent, &req.path, req.session_id.as_ref())?;
|
||||
let (path, _) = confine_local(agent, &path, req.session_id.as_ref()).await?;
|
||||
let result = fs::delete_file(&path).await.map(|_| Empty {});
|
||||
to_ext_response(result)
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,702 @@
|
||||
//! Git extension API layer.
|
||||
//!
|
||||
//! Routing: prefers explicit `gitRoot`, falls back to session lookup via `sessionId`.
|
||||
//! Business logic delegated to `session::git::*` pure functions.
|
||||
//!
|
||||
//! **Phase 4 design note**: Git/JJ functions (`git_cli`, `status`,
|
||||
//! `detect_vcs_kind`, `find_git_root_from_path`, etc.) are stateless
|
||||
//! utilities that take a `&Path` and shell out to `git`/`jj`. They do
|
||||
//! not access workspace state and therefore remain direct calls rather
|
||||
//! than routing through `WorkspaceChannel`. The channel's VCS stubs
|
||||
//! (`git_status`, `git_diff`, etc.) are reserved for future stateful
|
||||
//! operations (e.g. cached VCS state, cross-session conflict detection).
|
||||
use super::{Empty, ExtResult, parse_params, to_ext_response, to_ext_response_partial};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::ExtMethodResult;
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_workspace::session::git::{
|
||||
self, DiscardScope, GIT_STATUS_CACHE_TTL, GitDiffsData, GitStatusData, check_diff_size_limits,
|
||||
};
|
||||
use kigi_workspace::workspace_ops::{
|
||||
GitBranchesReq, GitCheckoutCommitReq, GitCheckoutReq, GitCommitReq, GitCurrentCommitReq,
|
||||
GitDiffReq, GitDiscardReq, GitFilesReq, GitInfoReq, GitStageContentReq, GitStageReq,
|
||||
GitStashReq, GitStatusExtReq, GitStatusFormat, GitUnstageReq,
|
||||
};
|
||||
use parking_lot::Mutex;
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
/// Global cache for git status results, keyed by git_root path.
|
||||
/// This provides caching at the extension API layer while keeping git::status pure.
|
||||
static GIT_STATUS_CACHE: std::sync::LazyLock<Mutex<HashMap<PathBuf, GitStatusCacheEntry>>> =
|
||||
std::sync::LazyLock::new(|| Mutex::new(HashMap::new()));
|
||||
struct GitStatusCacheEntry {
|
||||
result: GitStatusData,
|
||||
commit: String,
|
||||
cached_at: Instant,
|
||||
include_untracked: bool,
|
||||
include_stats: bool,
|
||||
}
|
||||
impl GitStatusCacheEntry {
|
||||
fn is_valid(&self, commit: &str, include_untracked: bool, include_stats: bool) -> bool {
|
||||
self.commit == commit
|
||||
&& self.include_untracked == include_untracked
|
||||
&& self.include_stats == include_stats
|
||||
&& self.cached_at.elapsed() < GIT_STATUS_CACHE_TTL
|
||||
}
|
||||
}
|
||||
/// Invalidate the git status cache for a given git_root.
|
||||
/// Should be called after any mutation operation (stage, unstage, discard, commit).
|
||||
fn invalidate_status_cache(git_root: &PathBuf) {
|
||||
let mut cache = GIT_STATUS_CACHE.lock();
|
||||
cache.remove(git_root);
|
||||
}
|
||||
fn default_head() -> String {
|
||||
"HEAD".to_string()
|
||||
}
|
||||
fn default_working() -> String {
|
||||
"working".to_string()
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitStatusRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub include_untracked: Option<bool>,
|
||||
pub include_stats: Option<bool>,
|
||||
pub ignore_submodules: Option<bool>,
|
||||
pub include_patches: Option<bool>,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitFilesRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub paths: Vec<String>,
|
||||
#[serde(default = "default_head")]
|
||||
pub version: String,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitDiffsRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
#[serde(default)]
|
||||
pub paths: Option<Vec<String>>,
|
||||
#[serde(default = "default_head")]
|
||||
pub from: String,
|
||||
#[serde(default = "default_working")]
|
||||
pub to: String,
|
||||
#[serde(default)]
|
||||
pub include_patch: bool,
|
||||
#[serde(default)]
|
||||
pub include_content: bool,
|
||||
#[serde(default)]
|
||||
pub max_patch_bytes: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub max_patch_lines: Option<usize>,
|
||||
#[serde(default)]
|
||||
pub merge_base: bool,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitStageRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub paths: Option<Vec<String>>,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitStageContentRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub path: String,
|
||||
pub content: String,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitUnstageRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub paths: Option<Vec<String>>,
|
||||
}
|
||||
#[derive(Clone, Copy, Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum GitDiscardScope {
|
||||
Working,
|
||||
Staged,
|
||||
#[default]
|
||||
Both,
|
||||
}
|
||||
impl From<GitDiscardScope> for DiscardScope {
|
||||
fn from(s: GitDiscardScope) -> Self {
|
||||
match s {
|
||||
GitDiscardScope::Working => DiscardScope::Working,
|
||||
GitDiscardScope::Staged => DiscardScope::Staged,
|
||||
GitDiscardScope::Both => DiscardScope::Both,
|
||||
}
|
||||
}
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitDiscardRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub paths: Option<Vec<String>>,
|
||||
#[serde(default)]
|
||||
pub include_untracked: bool,
|
||||
#[serde(default)]
|
||||
scope: GitDiscardScope,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitCommitRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub message: String,
|
||||
#[serde(default)]
|
||||
pub amend: bool,
|
||||
#[serde(default)]
|
||||
pub signoff: bool,
|
||||
#[serde(default)]
|
||||
pub push: bool,
|
||||
#[serde(default)]
|
||||
pub sync: bool,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitStashRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
#[serde(default)]
|
||||
pub include_untracked: bool,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitCheckoutRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
pub branch: String,
|
||||
#[serde(default)]
|
||||
pub create: bool,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct CheckoutSessionHeadRequest {
|
||||
pub session_id: acp::SessionId,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
#[serde(default)]
|
||||
pub stash_if_dirty: bool,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitInfoRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitBranchesRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
}
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitCurrentCommitRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
}
|
||||
/// Request for x.ai/git/checkout_commit extension method.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GitCheckoutCommitRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub git_root: Option<String>,
|
||||
/// Commit hash or ref to checkout.
|
||||
pub commit: String,
|
||||
#[serde(default)]
|
||||
pub stash_if_dirty: bool,
|
||||
}
|
||||
/// Resolve git_root from explicit value or session lookup via [`WorkspaceOps`].
|
||||
async fn resolve_git_root(
|
||||
agent: &MvpAgent,
|
||||
ops: &kigi_workspace::WorkspaceOps,
|
||||
git_root: Option<String>,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
) -> Result<PathBuf, acp::Error> {
|
||||
if let Some(root) = git_root {
|
||||
return Ok(PathBuf::from(root));
|
||||
}
|
||||
if let Some(sid) = session_id {
|
||||
if let Some(cwd) = agent.get_session_cwd(sid) {
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::GitResolveRootReq { cwd },
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::invalid_params()
|
||||
.data(format!("cannot find git root from session cwd: {}", e))
|
||||
})?;
|
||||
return result.ok_or_else(|| {
|
||||
acp::Error::invalid_params()
|
||||
.data("cannot find git root from session cwd: not a git repository")
|
||||
});
|
||||
}
|
||||
return Err(acp::Error::invalid_params().data(format!("session not found: {}", sid.0)));
|
||||
}
|
||||
Err(acp::Error::invalid_params().data("either gitRoot or sessionId is required"))
|
||||
}
|
||||
/// Try to extract a git_root from the request params (best-effort, for jj routing).
|
||||
async fn try_resolve_git_root(
|
||||
agent: &MvpAgent,
|
||||
ops: &kigi_workspace::WorkspaceOps,
|
||||
args: &acp::ExtRequest,
|
||||
) -> Option<PathBuf> {
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct Probe {
|
||||
git_root: Option<String>,
|
||||
session_id: Option<agent_client_protocol::SessionId>,
|
||||
}
|
||||
let probe: Probe = serde_json::from_str(args.params.get()).ok()?;
|
||||
if let Some(root) = probe.git_root {
|
||||
return Some(PathBuf::from(root));
|
||||
}
|
||||
if let Some(sid) = &probe.session_id
|
||||
&& let Some(cwd) = agent.get_session_cwd(sid)
|
||||
{
|
||||
return ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::GitResolveRootReq { cwd },
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
}
|
||||
None
|
||||
}
|
||||
pub async fn handle(
|
||||
agent: &MvpAgent,
|
||||
ops: &kigi_workspace::WorkspaceOps,
|
||||
args: &acp::ExtRequest,
|
||||
) -> ExtResult {
|
||||
if let Some(git_root) = try_resolve_git_root(agent, ops, args).await {
|
||||
let vcs_kind = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::DetectVcsKindReq {
|
||||
path: git_root.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(kigi_workspace::session::git::VcsKind::Git);
|
||||
if vcs_kind.is_jj()
|
||||
&& let Some(result) =
|
||||
super::jj::try_handle(args.method.as_ref(), &git_root, &args.params).await
|
||||
{
|
||||
return result;
|
||||
}
|
||||
}
|
||||
match args.method.as_ref() {
|
||||
"x.ai/git/git_repo_root" => {
|
||||
let req: git::GitRepoRequest = parse_params(args)?;
|
||||
let response = git::is_git_repo(&req).await?;
|
||||
super::to_raw_response(&response)
|
||||
}
|
||||
"x.ai/git/serialize_changes" => {
|
||||
let _ = (args, ops);
|
||||
to_ext_response::<()>(Err(anyhow::anyhow!(
|
||||
"git serialize_changes is unavailable in this build"
|
||||
)))
|
||||
}
|
||||
"x.ai/git/status" => {
|
||||
let req = parse_params::<GitStatusRequest>(args)?;
|
||||
let include_untracked = req.include_untracked.unwrap_or(true);
|
||||
let include_stats = req.include_stats.unwrap_or(false);
|
||||
let ignore_submodules = req.ignore_submodules.unwrap_or(true);
|
||||
let include_patches = req.include_patches.unwrap_or(false);
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
if let Some(ref git_root) = git_root {
|
||||
let current_commit = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::GitCurrentCommitReq {
|
||||
git_root: git_root.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(None);
|
||||
if let Some(commit) = ¤t_commit {
|
||||
let cached_result = {
|
||||
let cache = GIT_STATUS_CACHE.lock();
|
||||
cache.get(git_root).and_then(|entry| {
|
||||
if entry.is_valid(commit, include_untracked, include_stats) {
|
||||
tracing::debug!("git.status (cached)");
|
||||
Some(entry.result.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
};
|
||||
if let Some(result) = cached_result {
|
||||
return to_ext_response(Ok(result));
|
||||
}
|
||||
}
|
||||
}
|
||||
let op = GitStatusExtReq {
|
||||
git_root: git_root.clone(),
|
||||
include_untracked,
|
||||
include_stats,
|
||||
ignore_submodules,
|
||||
include_patches,
|
||||
format: GitStatusFormat::Structured,
|
||||
};
|
||||
let response = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
let result = response.data.ok_or_else(|| {
|
||||
acp::Error::internal_error().data("git_status_ext returned no structured data")
|
||||
})?;
|
||||
if let Some(git_root) = git_root
|
||||
&& let Some(ref commit) = result.commit
|
||||
{
|
||||
let mut cache = GIT_STATUS_CACHE.lock();
|
||||
cache.insert(
|
||||
git_root,
|
||||
GitStatusCacheEntry {
|
||||
result: result.clone(),
|
||||
commit: commit.clone(),
|
||||
cached_at: Instant::now(),
|
||||
include_untracked,
|
||||
include_stats,
|
||||
},
|
||||
);
|
||||
}
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/files" => {
|
||||
let req = parse_params::<GitFilesRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitFilesReq {
|
||||
git_root,
|
||||
paths: req.paths.clone(),
|
||||
version: req.version.clone(),
|
||||
};
|
||||
let result = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/diffs" => {
|
||||
let req = parse_params::<GitDiffsRequest>(args)?;
|
||||
let max_bytes = req.max_patch_bytes;
|
||||
let max_lines = req.max_patch_lines;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitDiffReq {
|
||||
git_root,
|
||||
paths: req.paths.clone(),
|
||||
from: req.from.clone(),
|
||||
to: req.to.clone(),
|
||||
include_patch: req.include_patch,
|
||||
include_content: req.include_content,
|
||||
merge_base: req.merge_base,
|
||||
};
|
||||
let data = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Err(err) = check_diff_size_limits(&data, max_bytes, max_lines) {
|
||||
let ext_result = ExtMethodResult::<GitDiffsData>::failure(err.message());
|
||||
ext_result
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
} else {
|
||||
to_ext_response(Ok(data))
|
||||
}
|
||||
}
|
||||
"x.ai/git/stage" => {
|
||||
let req = parse_params::<GitStageRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitStageReq {
|
||||
git_root: git_root.clone(),
|
||||
paths: req.paths,
|
||||
};
|
||||
let result = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/stage/content" => {
|
||||
let req = parse_params::<GitStageContentRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitStageContentReq {
|
||||
git_root: git_root.clone(),
|
||||
path: req.path.clone(),
|
||||
content: req.content.clone(),
|
||||
};
|
||||
ops.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
"x.ai/git/unstage" => {
|
||||
let req = parse_params::<GitUnstageRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitUnstageReq {
|
||||
git_root: git_root.clone(),
|
||||
paths: req.paths,
|
||||
};
|
||||
ops.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
"x.ai/git/discard" => {
|
||||
let req = parse_params::<GitDiscardRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitDiscardReq {
|
||||
git_root: git_root.clone(),
|
||||
paths: req.paths,
|
||||
scope: req.scope.into(),
|
||||
include_untracked: req.include_untracked,
|
||||
};
|
||||
ops.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
"x.ai/git/commit" => {
|
||||
let req = parse_params::<GitCommitRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitCommitReq {
|
||||
git_root: git_root.clone(),
|
||||
message: req.message.clone(),
|
||||
amend: req.amend,
|
||||
signoff: req.signoff,
|
||||
push: req.push,
|
||||
sync: req.sync,
|
||||
};
|
||||
let commit_result = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response_partial(Ok(commit_result.data), commit_result.warning)
|
||||
}
|
||||
"x.ai/git/checkout" => {
|
||||
let req = parse_params::<GitCheckoutRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitCheckoutReq {
|
||||
git_root: git_root.clone(),
|
||||
branch: req.branch.clone(),
|
||||
create: req.create,
|
||||
};
|
||||
ops.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
"x.ai/git/stash" => {
|
||||
let req = parse_params::<GitStashRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let op = GitStashReq {
|
||||
git_root: git_root.clone(),
|
||||
include_untracked: req.include_untracked,
|
||||
};
|
||||
ops.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
if let Some(ref git_root) = git_root {
|
||||
invalidate_status_cache(git_root);
|
||||
}
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
"x.ai/git/info" => {
|
||||
let req = parse_params::<GitInfoRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let result = ops
|
||||
.dispatch(&GitInfoReq { git_root }, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/branches" => {
|
||||
let req = parse_params::<GitBranchesRequest>(args)?;
|
||||
let git_root = resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok();
|
||||
let result = ops
|
||||
.dispatch(&GitBranchesReq { git_root }, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/current_commit" => {
|
||||
let req = parse_params::<GitCurrentCommitRequest>(args)?;
|
||||
let result = match resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref())
|
||||
.await
|
||||
.ok()
|
||||
{
|
||||
Some(git_root) => ops
|
||||
.dispatch(&GitCurrentCommitReq { git_root }, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?,
|
||||
None => None,
|
||||
};
|
||||
to_ext_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/checkout_session_head" => {
|
||||
let req = parse_params::<CheckoutSessionHeadRequest>(args)?;
|
||||
let git_root =
|
||||
resolve_git_root(agent, ops, req.git_root, Some(&req.session_id)).await?;
|
||||
let vcs_kind = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::DetectVcsKindReq {
|
||||
path: git_root.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(kigi_workspace::session::git::VcsKind::Git);
|
||||
if vcs_kind.is_jj() {
|
||||
return Err(acp::Error::invalid_request()
|
||||
.data("checkout_session_head is not supported in jj repositories"));
|
||||
}
|
||||
let summary =
|
||||
crate::session::persistence::find_summary_by_session_id(&req.session_id.0)
|
||||
.ok_or_else(|| {
|
||||
acp::Error::invalid_params()
|
||||
.data(format!("session {} not found", req.session_id.0))
|
||||
})?;
|
||||
let head_commit = summary.head_commit.ok_or_else(|| {
|
||||
acp::Error::invalid_params().data(format!(
|
||||
"session {} has no persisted HEAD commit",
|
||||
req.session_id.0
|
||||
))
|
||||
})?;
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::GitCheckoutCommitReq {
|
||||
git_root: git_root.clone(),
|
||||
head_commit,
|
||||
head_branch: summary.head_branch,
|
||||
stash_if_dirty: req.stash_if_dirty,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("checkout failed: {e}")))?;
|
||||
invalidate_status_cache(&git_root);
|
||||
super::to_raw_response(&result)
|
||||
}
|
||||
"x.ai/git/checkout_commit" => {
|
||||
let req = parse_params::<GitCheckoutCommitRequest>(args)?;
|
||||
let git_root =
|
||||
resolve_git_root(agent, ops, req.git_root, req.session_id.as_ref()).await?;
|
||||
let vcs_kind = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::DetectVcsKindReq {
|
||||
path: git_root.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(kigi_workspace::session::git::VcsKind::Git);
|
||||
if vcs_kind.is_jj() {
|
||||
return Err(acp::Error::invalid_request().data(
|
||||
"checkout_commit is not supported in jj repos; use `jj new` or `jj edit`",
|
||||
));
|
||||
}
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&GitCheckoutCommitReq {
|
||||
git_root: git_root.clone(),
|
||||
head_commit: req.commit,
|
||||
head_branch: None,
|
||||
stash_if_dirty: req.stash_if_dirty,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("checkout failed: {e}")))?;
|
||||
invalidate_status_cache(&git_root);
|
||||
super::to_raw_response(&result)
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
//! `x.ai/hooks/*` extension handlers.
|
||||
//!
|
||||
//! The file-hook list/action endpoints for the pager's hooks modal, plus the
|
||||
//! client-registered hook wire types and `parse_client_hooks`.
|
||||
|
||||
use std::collections::HashMap;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_hooks::event::{HookEventEnvelope, HookEventName};
|
||||
use kigi_hooks::matcher::HookMatcher;
|
||||
use kigi_hooks_plugins_types::{HookEvent, HookHandlerType, HookInfo};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ListRequest {
|
||||
session_id: String,
|
||||
}
|
||||
|
||||
pub fn hook_spec_to_info(spec: &kigi_hooks::config::HookSpec) -> HookInfo {
|
||||
use kigi_hooks::event::HookEventName;
|
||||
|
||||
let event = match spec.event {
|
||||
// Session lifecycle
|
||||
HookEventName::SessionStart => HookEvent::SessionStart,
|
||||
HookEventName::SessionEnd => HookEvent::SessionEnd,
|
||||
HookEventName::Stop => HookEvent::Stop,
|
||||
HookEventName::StopFailure => HookEvent::StopFailure,
|
||||
// Tool events
|
||||
HookEventName::PreToolUse => HookEvent::PreToolUse,
|
||||
HookEventName::PostToolUse => HookEvent::PostToolUse,
|
||||
HookEventName::PostToolUseFailure => HookEvent::PostToolUseFailure,
|
||||
HookEventName::PermissionDenied => HookEvent::PermissionDenied,
|
||||
// User / notification
|
||||
HookEventName::UserPromptSubmit => HookEvent::UserPromptSubmit,
|
||||
HookEventName::Notification => HookEvent::Notification,
|
||||
// Subagent
|
||||
HookEventName::SubagentStart => HookEvent::SubagentStart,
|
||||
HookEventName::SubagentStop | HookEventName::SubagentEnd => HookEvent::SubagentStop,
|
||||
// Compaction
|
||||
HookEventName::PreCompact => HookEvent::PreCompact,
|
||||
HookEventName::PostCompact => HookEvent::PostCompact,
|
||||
};
|
||||
|
||||
let handler_type = if spec.url.is_some() {
|
||||
HookHandlerType::Http
|
||||
} else {
|
||||
HookHandlerType::Command
|
||||
};
|
||||
|
||||
// Display the pre-expansion source string when available so the
|
||||
// pager UI / ACP DTO never leaks values resolved from the user
|
||||
// `env` map (which may contain secrets like API tokens). Fall back
|
||||
// to the post-expansion form for any future code path that builds
|
||||
// a `HookSpec` without populating the raw source.
|
||||
let command_display = spec
|
||||
.command_raw
|
||||
.clone()
|
||||
.or_else(|| spec.command.as_ref().map(|p| p.display().to_string()));
|
||||
let url_display = spec.url_raw.clone().or_else(|| spec.url.clone());
|
||||
|
||||
HookInfo {
|
||||
name: spec.name.clone(),
|
||||
event,
|
||||
handler_type,
|
||||
matcher: spec.configured_matcher.clone(),
|
||||
command: command_display,
|
||||
url: url_display,
|
||||
timeout_ms: spec.timeout_ms,
|
||||
source_dir: spec.source_dir.display().to_string(),
|
||||
disabled: kigi_hooks::trust::is_hook_disabled(&spec.name),
|
||||
}
|
||||
}
|
||||
|
||||
// Wire types for client-registered hooks (`x.ai/hooks/run`); the gate that uses
|
||||
// them lives in `session::acp_session::hooks`.
|
||||
|
||||
/// A matcher group from the client's registration: `{ matcher, hookCallbackIds, timeout }`.
|
||||
///
|
||||
/// `pub` (not `pub(crate)`) because [`ClientHooks`] flows through the public
|
||||
/// `SessionCommand::SnapshotClientHooks` so subagents can inherit the parent's hooks.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ClientHookGroup {
|
||||
/// `None` (wire `null`, `""`, or `"*"`) matches every tool.
|
||||
pub matcher: Option<HookMatcher>,
|
||||
pub callback_ids: Vec<String>,
|
||||
/// Per-group reply deadline for the `PreToolUse` gate (wire value in seconds). `None`
|
||||
/// falls back to the default gate timeout.
|
||||
pub timeout: Option<std::time::Duration>,
|
||||
}
|
||||
|
||||
pub type ClientHooks = HashMap<HookEventName, Vec<ClientHookGroup>>;
|
||||
|
||||
/// One hook dispatched to a client callback: the shared [`HookEventEnvelope`]
|
||||
/// (flattened, camelCase) plus the `hookCallbackId` it targets. The same shape is sent
|
||||
/// for both the `x.ai/hooks/run` request (gate) and the `x.ai/hooks/event` notification
|
||||
/// (observe-only), so the client decodes one payload for every hook.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct ClientHookDispatch<'a> {
|
||||
pub hook_callback_id: &'a str,
|
||||
#[serde(flatten)]
|
||||
pub envelope: &'a HookEventEnvelope,
|
||||
}
|
||||
|
||||
/// Only `Deny` blocks the tool; every other value proceeds (fail-open).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, serde::Deserialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub(crate) enum ClientHookDecision {
|
||||
#[default]
|
||||
Continue,
|
||||
Deny,
|
||||
#[serde(other)]
|
||||
Other,
|
||||
}
|
||||
|
||||
/// Response payload for `x.ai/hooks/run` (client to agent). `Default` (used on
|
||||
/// timeout, transport error, or a malformed reply) proceeds.
|
||||
#[derive(Debug, Clone, Default, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub(crate) struct ClientHookResponse {
|
||||
#[serde(default)]
|
||||
pub decision: ClientHookDecision,
|
||||
/// Deny reason surfaced to the model/user; consumed only when `decision` is `Deny`.
|
||||
#[serde(default)]
|
||||
pub system_message: Option<String>,
|
||||
}
|
||||
|
||||
/// Parse client hooks from `session/new` `_meta["x.ai/hooks"]`, shaped
|
||||
/// `{ "<Event>": [{ matcher, hookCallbackIds }] }` (PascalCase or snake_case
|
||||
/// events). Each `matcher` is compiled with the agent's [`HookMatcher`] so client
|
||||
/// and file hooks match identically. Unknown events, malformed groups, invalid
|
||||
/// matchers, and callback-less groups are skipped; absent meta yields no hooks.
|
||||
pub(crate) fn parse_client_hooks(meta: Option<&acp::Meta>) -> ClientHooks {
|
||||
let mut hooks = ClientHooks::new();
|
||||
let Some(map) = meta
|
||||
.and_then(|m| m.get("x.ai/hooks"))
|
||||
.and_then(|h| h.as_object())
|
||||
else {
|
||||
return hooks;
|
||||
};
|
||||
for (event_name, value) in map {
|
||||
let de = serde::de::value::StrDeserializer::<serde::de::value::Error>::new(event_name);
|
||||
let Ok(event) = HookEventName::deserialize(de) else {
|
||||
tracing::warn!(event = %event_name, "ignoring unknown x.ai/hooks event");
|
||||
continue;
|
||||
};
|
||||
let Some(array) = value.as_array() else {
|
||||
tracing::warn!(event = %event_name, "x.ai/hooks event value is not an array; skipping");
|
||||
continue;
|
||||
};
|
||||
let groups: Vec<ClientHookGroup> = array
|
||||
.iter()
|
||||
.filter_map(|group| parse_hook_group(event, group))
|
||||
.collect();
|
||||
if !groups.is_empty() {
|
||||
// Key by the canonical event so a registration under an alias (e.g.
|
||||
// `SubagentEnd`) still matches the event the agent fires (`SubagentStop`).
|
||||
hooks.entry(event.canonical()).or_default().extend(groups);
|
||||
}
|
||||
}
|
||||
hooks
|
||||
}
|
||||
|
||||
/// Hooks to apply on a `load_session` reconnect: `Some` (possibly empty, an explicit
|
||||
/// clear) when the request meta carries `x.ai/hooks`, else `None` so a reconnect that
|
||||
/// omits the key leaves the live registrations from `session/new` untouched.
|
||||
pub(crate) fn reconnect_client_hooks(meta: Option<&acp::Meta>) -> Option<ClientHooks> {
|
||||
meta.and_then(|m| m.get("x.ai/hooks"))
|
||||
.map(|_| parse_client_hooks(meta))
|
||||
}
|
||||
|
||||
/// Parse one `{ matcher, hookCallbackIds }` registration entry. Returns `None`
|
||||
/// (with a warning) when the entry is malformed, carries no callback ids, or its
|
||||
/// matcher fails to compile.
|
||||
fn parse_hook_group(event: HookEventName, value: &serde_json::Value) -> Option<ClientHookGroup> {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct WireGroup {
|
||||
#[serde(default)]
|
||||
matcher: Option<String>,
|
||||
#[serde(default)]
|
||||
hook_callback_ids: Vec<String>,
|
||||
/// Per-group gate timeout in seconds.
|
||||
#[serde(default)]
|
||||
timeout: Option<f64>,
|
||||
}
|
||||
|
||||
let group = WireGroup::deserialize(value)
|
||||
.inspect_err(|err| tracing::warn!(%event, %err, "ignoring malformed x.ai/hooks group"))
|
||||
.ok()?;
|
||||
if group.hook_callback_ids.is_empty() {
|
||||
tracing::warn!(%event, "ignoring x.ai/hooks group with no hookCallbackIds");
|
||||
return None;
|
||||
}
|
||||
// Drop a non-finite/non-positive timeout (fall back to the default gate timeout) and
|
||||
// cap it so a client can't make a tool hang on the gate for an unbounded time.
|
||||
const MAX_HOOK_TIMEOUT_SECS: f64 = 300.0;
|
||||
let timeout = group
|
||||
.timeout
|
||||
.filter(|s| s.is_finite() && *s > 0.0)
|
||||
.map(|s| std::time::Duration::from_secs_f64(s.min(MAX_HOOK_TIMEOUT_SECS)));
|
||||
let matcher = match group.matcher.as_deref() {
|
||||
// Match-all tokens map to no matcher (group always fires). `HookMatcher::new`
|
||||
// also treats these as match-all; short-circuiting here keeps the intent explicit.
|
||||
None | Some("") | Some("*") => None,
|
||||
Some(pattern) => match HookMatcher::new(pattern) {
|
||||
Ok(matcher) => Some(matcher),
|
||||
Err(err) => {
|
||||
tracing::warn!(%event, pattern, %err, "ignoring x.ai/hooks group with invalid matcher");
|
||||
return None;
|
||||
}
|
||||
},
|
||||
};
|
||||
Some(ClientHookGroup {
|
||||
matcher,
|
||||
callback_ids: group.hook_callback_ids,
|
||||
timeout,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/hooks/list" => {
|
||||
let req: ListRequest = super::parse_params(args)?;
|
||||
let sid = acp::SessionId::new(req.session_id);
|
||||
|
||||
let result = agent
|
||||
.list_hooks(&sid)
|
||||
.await
|
||||
.ok_or_else(|| anyhow::anyhow!("session not found"));
|
||||
super::to_ext_response(result)
|
||||
}
|
||||
"x.ai/hooks/action" => {
|
||||
let req: kigi_hooks_plugins_types::HooksActionRequest = super::parse_params(args)?;
|
||||
let sid = acp::SessionId::new(req.session_id);
|
||||
|
||||
let result = agent
|
||||
.execute_hooks_action(&sid, req.action)
|
||||
.await
|
||||
.ok_or_else(|| anyhow::anyhow!("session not found"));
|
||||
super::to_ext_response(result)
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_hooks::config::HookSpec;
|
||||
use kigi_hooks::event::HookEventName;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Minimal `HookSpec` for `hook_spec_to_info` tests (`handler_type` is unused;
|
||||
/// the DTO derives it from `url`).
|
||||
fn make_spec(
|
||||
command_raw: Option<&str>,
|
||||
command: Option<&str>,
|
||||
url_raw: Option<&str>,
|
||||
url: Option<&str>,
|
||||
) -> HookSpec {
|
||||
HookSpec {
|
||||
name: "test:pre_tool_use[0].hooks[0]".to_string(),
|
||||
event: HookEventName::PreToolUse,
|
||||
handler_type: "command".to_string(),
|
||||
configured_matcher: None,
|
||||
matcher: None,
|
||||
enabled: true,
|
||||
command: command.map(PathBuf::from),
|
||||
command_raw: command_raw.map(str::to_string),
|
||||
url: url.map(str::to_string),
|
||||
url_raw: url_raw.map(str::to_string),
|
||||
timeout_ms: 5000,
|
||||
source_dir: PathBuf::from("/tmp"),
|
||||
extra_env: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// `*_raw` (pre-expansion) wins over the resolved value so secrets never reach
|
||||
/// the DTO; then the resolved value, else `None`. Same for `command` and `url`.
|
||||
#[test]
|
||||
fn hook_spec_to_info_display_precedence() {
|
||||
let command =
|
||||
|raw, resolved| hook_spec_to_info(&make_spec(raw, resolved, None, None)).command;
|
||||
assert_eq!(
|
||||
command(Some("${VAR}/x"), Some("/resolved/x")).as_deref(),
|
||||
Some("${VAR}/x")
|
||||
);
|
||||
assert_eq!(
|
||||
command(None, Some("/legacy/x")).as_deref(),
|
||||
Some("/legacy/x")
|
||||
);
|
||||
assert!(command(None, None).is_none());
|
||||
|
||||
let url = |raw, resolved| hook_spec_to_info(&make_spec(None, None, raw, resolved)).url;
|
||||
assert_eq!(
|
||||
url(
|
||||
Some("https://${HOST}/p?token=${TOKEN}"),
|
||||
Some("https://api/p?token=ghp_X")
|
||||
)
|
||||
.as_deref(),
|
||||
Some("https://${HOST}/p?token=${TOKEN}"),
|
||||
);
|
||||
assert_eq!(
|
||||
url(None, Some("https://h/c")).as_deref(),
|
||||
Some("https://h/c")
|
||||
);
|
||||
assert!(url(None, None).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_client_hooks_parses_valid_groups() {
|
||||
let meta = serde_json::json!({
|
||||
"x.ai/hooks": {
|
||||
"PreToolUse": [
|
||||
{ "matcher": "run_terminal_command", "hookCallbackIds": ["cb_0"] },
|
||||
{ "matcher": null, "hookCallbackIds": ["cb_1"] },
|
||||
{ "matcher": "*", "hookCallbackIds": ["cb_2"] }
|
||||
],
|
||||
"post_tool_use": [{ "hookCallbackIds": ["cb_3"] }]
|
||||
}
|
||||
});
|
||||
let hooks = parse_client_hooks(meta.as_object());
|
||||
|
||||
let pre = &hooks[&HookEventName::PreToolUse];
|
||||
assert_eq!(pre.len(), 3);
|
||||
assert_eq!(pre[0].callback_ids, ["cb_0"]);
|
||||
let matcher = pre[0].matcher.as_ref().unwrap();
|
||||
assert!(matcher.is_match("run_terminal_command"));
|
||||
assert!(!matcher.is_match("read_file"));
|
||||
assert!(pre[1].matcher.is_none()); // null / "*" = match-all
|
||||
assert!(pre[2].matcher.is_none());
|
||||
assert!(hooks.contains_key(&HookEventName::PostToolUse)); // snake_case resolves
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_client_hooks_drops_invalid_and_absent() {
|
||||
assert!(parse_client_hooks(None).is_empty());
|
||||
assert!(
|
||||
parse_client_hooks(serde_json::json!({ "askUserQuestion": true }).as_object())
|
||||
.is_empty()
|
||||
);
|
||||
|
||||
let meta = serde_json::json!({
|
||||
"NotARealEvent": [{ "hookCallbackIds": ["x"] }],
|
||||
"x.ai/hooks": {
|
||||
"PreToolUse": [
|
||||
{ "matcher": "[invalid", "hookCallbackIds": ["bad_regex"] },
|
||||
{ "matcher": "run_terminal_command", "hookCallbackIds": [] },
|
||||
{ "matcher": "read_file", "hookCallbackIds": ["good"] }
|
||||
]
|
||||
}
|
||||
});
|
||||
let groups = &parse_client_hooks(meta.as_object())[&HookEventName::PreToolUse];
|
||||
assert_eq!(groups.len(), 1);
|
||||
assert_eq!(groups[0].callback_ids, ["good"]);
|
||||
}
|
||||
|
||||
/// A group's `timeout` (seconds) parses to a `Duration`; absent or non-positive falls
|
||||
/// back to the default gate timeout (`None`).
|
||||
#[test]
|
||||
fn parse_client_hooks_reads_group_timeout() {
|
||||
let meta = serde_json::json!({
|
||||
"x.ai/hooks": {
|
||||
"PreToolUse": [
|
||||
{ "hookCallbackIds": ["a"], "timeout": 5.0 },
|
||||
{ "hookCallbackIds": ["b"], "timeout": 0 },
|
||||
{ "hookCallbackIds": ["c"] },
|
||||
{ "hookCallbackIds": ["d"], "timeout": 100000 }
|
||||
]
|
||||
}
|
||||
});
|
||||
let groups = &parse_client_hooks(meta.as_object())[&HookEventName::PreToolUse];
|
||||
assert_eq!(groups[0].timeout, Some(std::time::Duration::from_secs(5)));
|
||||
assert_eq!(groups[1].timeout, None); // non-positive -> default
|
||||
assert_eq!(groups[2].timeout, None); // absent -> default
|
||||
assert_eq!(groups[3].timeout, Some(std::time::Duration::from_secs(300))); // capped
|
||||
}
|
||||
|
||||
/// A registration under the `SubagentEnd` alias must land on the canonical
|
||||
/// `SubagentStop` key the agent fires.
|
||||
#[test]
|
||||
fn parse_client_hooks_canonicalizes_subagent_alias() {
|
||||
let meta = serde_json::json!({
|
||||
"x.ai/hooks": { "SubagentEnd": [{ "hookCallbackIds": ["cb"] }] }
|
||||
});
|
||||
let hooks = parse_client_hooks(meta.as_object());
|
||||
assert!(hooks.contains_key(&HookEventName::SubagentStop));
|
||||
assert!(!hooks.contains_key(&HookEventName::SubagentEnd));
|
||||
}
|
||||
|
||||
/// Reconnect refresh applies hooks only when the load meta carries `x.ai/hooks`:
|
||||
/// an absent key returns `None` (don't wipe `session/new` registrations); a present
|
||||
/// key returns `Some` (an empty object is an explicit clear).
|
||||
#[test]
|
||||
fn reconnect_client_hooks_only_when_key_present() {
|
||||
assert!(reconnect_client_hooks(None).is_none());
|
||||
assert!(reconnect_client_hooks(serde_json::json!({ "other": true }).as_object()).is_none());
|
||||
|
||||
let cleared = reconnect_client_hooks(serde_json::json!({ "x.ai/hooks": {} }).as_object());
|
||||
assert!(cleared.is_some_and(|h| h.is_empty()));
|
||||
|
||||
let set = reconnect_client_hooks(
|
||||
serde_json::json!({
|
||||
"x.ai/hooks": { "PreToolUse": [{ "hookCallbackIds": ["cb"] }] }
|
||||
})
|
||||
.as_object(),
|
||||
);
|
||||
assert!(set.is_some_and(|h| h.contains_key(&HookEventName::PreToolUse)));
|
||||
}
|
||||
|
||||
/// `deny` parses to `Deny` (+ optional message); everything else fails open:
|
||||
/// unknown values to `Other`, missing/empty/default to `Continue`.
|
||||
#[test]
|
||||
fn client_hook_response_deserialization() {
|
||||
let deny: ClientHookResponse =
|
||||
serde_json::from_str(r#"{"decision":"deny","systemMessage":"blocked"}"#).unwrap();
|
||||
assert_eq!(deny.decision, ClientHookDecision::Deny);
|
||||
assert_eq!(deny.system_message.as_deref(), Some("blocked"));
|
||||
|
||||
let unknown: ClientHookResponse =
|
||||
serde_json::from_str(r#"{"decision":"maybe_later"}"#).unwrap();
|
||||
assert_eq!(unknown.decision, ClientHookDecision::Other);
|
||||
|
||||
let empty: ClientHookResponse = serde_json::from_str("{}").unwrap();
|
||||
assert_eq!(empty.decision, ClientHookDecision::Continue);
|
||||
assert!(empty.system_message.is_none());
|
||||
assert_eq!(
|
||||
ClientHookResponse::default().decision,
|
||||
ClientHookDecision::Continue
|
||||
);
|
||||
}
|
||||
|
||||
/// The callback id sits beside the flattened envelope (camelCase keys,
|
||||
/// `hookEventName` snake_case); the one shape sent for both run and event.
|
||||
#[test]
|
||||
fn client_hook_dispatch_serializes_envelope() {
|
||||
use kigi_hooks::event::{HookEventEnvelope, HookPayload};
|
||||
|
||||
let envelope = HookEventEnvelope {
|
||||
hook_event_name: HookEventName::PreToolUse,
|
||||
session_id: "s1".into(),
|
||||
cwd: "/work".into(),
|
||||
workspace_root: "/work".into(),
|
||||
timestamp: "t".into(),
|
||||
transcript_path: None,
|
||||
client_identifier: None,
|
||||
prompt_id: None,
|
||||
payload: HookPayload::PreToolUse {
|
||||
tool_name: "run_terminal_command".into(),
|
||||
tool_use_id: "call_1".into(),
|
||||
tool_input: serde_json::json!({ "command": "ls" }),
|
||||
tool_input_truncated: true,
|
||||
permission_mode: None,
|
||||
subagent_type: None,
|
||||
},
|
||||
};
|
||||
let dispatch = ClientHookDispatch {
|
||||
hook_callback_id: "cb_0",
|
||||
envelope: &envelope,
|
||||
};
|
||||
let value = serde_json::to_value(&dispatch).unwrap();
|
||||
assert_eq!(value["hookCallbackId"], "cb_0");
|
||||
assert_eq!(value["hookEventName"], "pre_tool_use");
|
||||
assert_eq!(value["sessionId"], "s1");
|
||||
assert_eq!(value["cwd"], "/work");
|
||||
assert_eq!(value["toolUseId"], "call_1");
|
||||
assert_eq!(value["toolName"], "run_terminal_command");
|
||||
assert_eq!(value["toolInput"]["command"], "ls");
|
||||
assert_eq!(value["toolInputTruncated"], true);
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,122 @@
|
||||
//! `x.ai/interject` extension handler.
|
||||
//!
|
||||
//! Queues a mid-turn interjection into the active session's pending
|
||||
//! interjection buffer. The session actor drains it at the next safe
|
||||
//! point in `process_conversation_turn`.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{ExtResult, parse_params};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::SessionCommand;
|
||||
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct InterjectRequest {
|
||||
session_id: String,
|
||||
text: String,
|
||||
#[serde(default)]
|
||||
interjection_id: Option<String>,
|
||||
/// Optional structured blocks (text + images) from image-capable
|
||||
/// clients; absent = legacy text-only wire shape (empty after default).
|
||||
#[serde(default)]
|
||||
content: Vec<acp::ContentBlock>,
|
||||
}
|
||||
|
||||
/// Split a `content` array into the model-safe text and the image blocks.
|
||||
///
|
||||
/// The Text block (when present and non-empty) is the client's REWRITTEN
|
||||
/// text — failed-orphan placeholders stripped, `[Image #N: <path>]` paths
|
||||
/// dropped — and must win over the raw `text` param, which exists for
|
||||
/// legacy clients and display. Returns `(text_override, images)`.
|
||||
fn split_content(content: Vec<acp::ContentBlock>) -> (Option<String>, Vec<acp::ImageContent>) {
|
||||
let text_override = content.iter().find_map(|block| match block {
|
||||
acp::ContentBlock::Text(tb) if !tb.text.trim().is_empty() => Some(tb.text.clone()),
|
||||
_ => None,
|
||||
});
|
||||
(text_override, crate::session::image_blocks(content))
|
||||
}
|
||||
|
||||
/// Handle `x.ai/interject` — queue a mid-turn user interjection.
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: InterjectRequest = parse_params(args)?;
|
||||
let sid: acp::SessionId = req.session_id.clone().into();
|
||||
// Load-race-tolerant: an interjection racing a reconnect-replayed
|
||||
// `session/load` (leader restart) waits for the load instead of failing.
|
||||
let session_handle = agent.session_handle_waiting_for_load(&sid).await;
|
||||
let Some(session) = session_handle else {
|
||||
return Err(
|
||||
acp::Error::invalid_params().data(format!("session not found: {}", req.session_id))
|
||||
);
|
||||
};
|
||||
|
||||
let (text_override, images) = split_content(req.content);
|
||||
let _ = session.cmd_tx.send(SessionCommand::Interject {
|
||||
text: text_override.unwrap_or(req.text),
|
||||
id: req.interjection_id,
|
||||
images,
|
||||
});
|
||||
|
||||
super::to_ext_response(Ok(serde_json::json!({
|
||||
"status": "queued",
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
/// Legacy wire shape (no `content`) parses byte-identically: text-only,
|
||||
/// zero images, no text override.
|
||||
#[test]
|
||||
fn parse_without_content_is_legacy_text_only() {
|
||||
let req: InterjectRequest = serde_json::from_value(serde_json::json!({
|
||||
"sessionId": "s1",
|
||||
"text": "steer left",
|
||||
"interjectionId": "i1",
|
||||
}))
|
||||
.expect("legacy params must parse");
|
||||
assert_eq!(req.text, "steer left");
|
||||
assert_eq!(req.interjection_id.as_deref(), Some("i1"));
|
||||
let (text_override, images) = split_content(req.content);
|
||||
assert_eq!(text_override, None);
|
||||
assert!(images.is_empty());
|
||||
}
|
||||
|
||||
/// `content` with text + image blocks parses; the images are extracted
|
||||
/// and the Text block (the client's rewritten, path-stripped text)
|
||||
/// overrides the raw `text` param.
|
||||
#[test]
|
||||
fn parse_with_content_extracts_images_and_prefers_block_text() {
|
||||
let req: InterjectRequest = serde_json::from_value(serde_json::json!({
|
||||
"sessionId": "s1",
|
||||
"text": "look at [Image #1: /tmp/x.png]",
|
||||
"content": [
|
||||
{ "type": "text", "text": "look at [Image #1]" },
|
||||
{ "type": "image", "data": "aGVsbG8=", "mimeType": "image/png" },
|
||||
],
|
||||
}))
|
||||
.expect("content params must parse");
|
||||
let (text_override, images) = split_content(req.content);
|
||||
assert_eq!(
|
||||
text_override.as_deref(),
|
||||
Some("look at [Image #1]"),
|
||||
"rewritten block text must win over the raw text param"
|
||||
);
|
||||
assert_eq!(images.len(), 1);
|
||||
assert_eq!(images[0].mime_type, "image/png");
|
||||
assert_eq!(images[0].data, "aGVsbG8=");
|
||||
}
|
||||
|
||||
/// Garbage `content` fails the whole parse (strict, like other params)
|
||||
/// instead of silently dropping attachments.
|
||||
#[test]
|
||||
fn parse_with_garbage_content_is_an_error() {
|
||||
let result: Result<InterjectRequest, _> = serde_json::from_value(serde_json::json!({
|
||||
"sessionId": "s1",
|
||||
"text": "steer",
|
||||
"content": "not an array",
|
||||
}));
|
||||
assert!(result.is_err(), "garbage content must be rejected");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
//! Jujutsu extension handlers — delegates to [`kigi_workspace::session::jj`].
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{Empty, ExtResult, to_ext_response, to_ext_response_partial};
|
||||
use kigi_workspace::session::git::{CommitData, StageData};
|
||||
use kigi_workspace::session::jj;
|
||||
|
||||
/// Handle a `x.ai/git/*` method for a jj-colocated repo.
|
||||
///
|
||||
/// Returns `Some(result)` if handled, `None` to fall through to git.
|
||||
pub async fn try_handle(
|
||||
method: &str,
|
||||
git_root: &std::path::Path,
|
||||
raw_params: &serde_json::value::RawValue,
|
||||
) -> Option<ExtResult> {
|
||||
match method {
|
||||
"x.ai/git/status" => Some(to_ext_response(jj::status(git_root).await)),
|
||||
"x.ai/git/info" => Some(to_ext_response(jj::info(git_root).await)),
|
||||
// git HEAD points at `@-` in a colocated repo; route to jj so we report
|
||||
// the working-copy commit (`@`), consistent with `status`/`info`.
|
||||
"x.ai/git/current_commit" => Some(to_ext_response(jj::current_commit(git_root).await)),
|
||||
"x.ai/git/branches" => Some(to_ext_response(jj::list_bookmarks(git_root).await)),
|
||||
|
||||
// jj has no staging area — stage/unstage are no-ops
|
||||
"x.ai/git/stage" => Some(to_ext_response(Ok(StageData { paths: Vec::new() }))),
|
||||
"x.ai/git/stage/content" | "x.ai/git/unstage" => Some(to_ext_response(Ok(Empty {}))),
|
||||
|
||||
"x.ai/git/discard" => {
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct Req {
|
||||
#[serde(default)]
|
||||
paths: Option<Vec<String>>,
|
||||
}
|
||||
let req: Req = serde_json::from_str(raw_params.get()).ok()?;
|
||||
Some(to_ext_response(
|
||||
jj::discard(git_root, req.paths).await.map(|_| Empty {}),
|
||||
))
|
||||
}
|
||||
|
||||
"x.ai/git/commit" => {
|
||||
#[derive(serde::Deserialize)]
|
||||
struct Req {
|
||||
message: String,
|
||||
}
|
||||
let req: Req = serde_json::from_str(raw_params.get()).ok()?;
|
||||
let result = jj::commit(git_root, &req.message).await;
|
||||
Some(match result {
|
||||
Ok(r) => to_ext_response_partial(Ok(r.data), r.warning),
|
||||
Err(e) => to_ext_response(Err::<CommitData, _>(e)),
|
||||
})
|
||||
}
|
||||
|
||||
// Operations that don't apply to jj
|
||||
"x.ai/git/checkout" => Some(Err(acp::Error::invalid_params()
|
||||
.data("checkout is not supported in jj repos; use `jj new` or `jj edit`"))),
|
||||
"x.ai/git/stash" => Some(Err(acp::Error::invalid_params()
|
||||
.data("stash is not supported in jj repos; changes are always committed"))),
|
||||
|
||||
// Everything else (diffs, files, serialize_changes) falls through to git
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,101 @@
|
||||
//! `x.ai/memory/flush`, `x.ai/memory/rewrite`, and `x.ai/compact_conversation`
|
||||
//! extension handlers.
|
||||
//!
|
||||
//! - `compact_conversation`: trigger an on-demand compaction for a session.
|
||||
//! - `memory/flush`: trigger an on-demand memory flush for a session.
|
||||
//! - `memory/rewrite`: rewrite a raw memory note into structured markdown via
|
||||
//! a one-shot LLM call.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use super::{Empty, ExtResult, parse_params, to_ext_response, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::{CompactConversationRequest, CompactConversationResponse, SessionCommand};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
m if m.starts_with("x.ai/compact_conversation") => handle_compact(agent, args).await,
|
||||
"x.ai/memory/flush" => handle_flush(agent, args).await,
|
||||
"x.ai/memory/rewrite" => handle_rewrite(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_compact(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: CompactConversationRequest = parse_params(args)?;
|
||||
// send over the compact query here properly
|
||||
let session_handle = {
|
||||
let sessions = agent.sessions.borrow();
|
||||
sessions.get(&req.session_id.into()).cloned()
|
||||
};
|
||||
let (tx, rx) = oneshot::channel();
|
||||
if let Some(session) = session_handle {
|
||||
let _ = session.cmd_tx.send(SessionCommand::CompactSession {
|
||||
user_context: req.user_context,
|
||||
respond_to: tx,
|
||||
});
|
||||
}
|
||||
rx.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("Internal error: {:?}", e)))?;
|
||||
to_raw_response(&CompactConversationResponse {})
|
||||
}
|
||||
|
||||
async fn handle_flush(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
struct MemoryFlushRequest {
|
||||
session_id: String,
|
||||
}
|
||||
|
||||
let req: MemoryFlushRequest = parse_params(args)?;
|
||||
let not_found_err = format!("session not found: {}", req.session_id);
|
||||
let session_handle = {
|
||||
let sessions = agent.sessions.borrow();
|
||||
sessions.get(&req.session_id.into()).cloned()
|
||||
};
|
||||
let Some(session) = session_handle else {
|
||||
return Err(acp::Error::invalid_params().data(not_found_err));
|
||||
};
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = session
|
||||
.cmd_tx
|
||||
.send(SessionCommand::FlushMemory { respond_to: tx });
|
||||
rx.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("{:?}", e)))?;
|
||||
to_ext_response(Ok(Empty {}))
|
||||
}
|
||||
|
||||
async fn handle_rewrite(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct RewriteRequest {
|
||||
session_id: String,
|
||||
raw_text: String,
|
||||
context_summary: String,
|
||||
}
|
||||
|
||||
let req: RewriteRequest = parse_params(args)?;
|
||||
let not_found_err = format!("session not found: {}", req.session_id);
|
||||
let session_handle = {
|
||||
let sessions = agent.sessions.borrow();
|
||||
sessions.get(&req.session_id.into()).cloned()
|
||||
};
|
||||
let Some(session) = session_handle else {
|
||||
return Err(acp::Error::invalid_params().data(not_found_err));
|
||||
};
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = session.cmd_tx.send(SessionCommand::RewriteMemoryNote {
|
||||
raw_text: req.raw_text,
|
||||
context_summary: req.context_summary,
|
||||
respond_to: tx,
|
||||
});
|
||||
let rewritten = rx
|
||||
.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(e))?;
|
||||
to_raw_response(&serde_json::json!({ "rewritten": rewritten }))
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
pub mod auth;
|
||||
pub(crate) mod auth_gate;
|
||||
pub mod billing;
|
||||
pub mod bundle;
|
||||
pub mod chat_conversation_history;
|
||||
pub mod code_nav;
|
||||
pub mod debug;
|
||||
pub mod feedback;
|
||||
pub mod fs;
|
||||
pub mod git;
|
||||
pub mod hooks;
|
||||
pub mod hunk_tracker;
|
||||
pub mod interject;
|
||||
pub mod jj;
|
||||
pub mod mcp;
|
||||
pub mod memory;
|
||||
pub mod notification;
|
||||
pub mod plugins;
|
||||
pub mod pr;
|
||||
pub mod privacy;
|
||||
pub mod prompt_history;
|
||||
pub mod prompt_meta;
|
||||
pub mod recap;
|
||||
pub mod repair;
|
||||
pub mod rewind;
|
||||
pub mod rollout;
|
||||
pub mod routing;
|
||||
pub mod search;
|
||||
pub mod session_admin;
|
||||
pub mod session_search;
|
||||
pub mod session_updates;
|
||||
pub mod share;
|
||||
pub mod skills;
|
||||
pub mod suggest;
|
||||
pub mod task;
|
||||
pub mod terminal;
|
||||
pub mod worktree;
|
||||
use crate::session::ExtMethodResult;
|
||||
use agent_client_protocol as acp;
|
||||
use serde::Serialize;
|
||||
use serde::de::DeserializeOwned;
|
||||
use std::sync::Arc;
|
||||
pub type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
pub fn parse_params<T: DeserializeOwned>(args: &acp::ExtRequest) -> Result<T, acp::Error> {
|
||||
parse_params_str(args.params.get())
|
||||
}
|
||||
/// Deserialize ACP params from their raw JSON string, mapping a parse failure
|
||||
/// to `invalid_params`. Used by [`parse_params`] and the bridge `encode` hooks,
|
||||
/// which hold the params `RawValue` directly.
|
||||
pub fn parse_params_str<T: DeserializeOwned>(raw: &str) -> Result<T, acp::Error> {
|
||||
serde_json::from_str(raw)
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {}", e)))
|
||||
}
|
||||
/// Extract the session ID from an extension request's params.
|
||||
pub fn parse_session_id(args: &acp::ExtRequest) -> Option<acp::SessionId> {
|
||||
let v: serde_json::Value = serde_json::from_str(args.params.get()).ok()?;
|
||||
let sid = v.get("sessionId")?.as_str()?;
|
||||
Some(acp::SessionId::new(sid))
|
||||
}
|
||||
pub fn to_ext_response<T: Serialize>(result: anyhow::Result<T>) -> ExtResult {
|
||||
ExtMethodResult::from_result(result)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
/// Wrap a serializable value as an `ExtResponse` without the `ExtMethodResult` envelope.
|
||||
pub fn to_raw_response<T: Serialize>(v: &T) -> ExtResult {
|
||||
serde_json::value::to_raw_value(v)
|
||||
.map(|raw| acp::ExtResponse::new(Arc::from(raw)))
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
/// Convert a result with optional warning to an ExtResponse.
|
||||
pub fn to_ext_response_partial<T: Serialize>(
|
||||
result: anyhow::Result<T>,
|
||||
warning: Option<String>,
|
||||
) -> ExtResult {
|
||||
let ext_result = match (result, warning) {
|
||||
(Ok(data), Some(warn)) => ExtMethodResult::partial(data, warn),
|
||||
(Ok(data), None) => ExtMethodResult::success(data),
|
||||
(Err(e), _) => ExtMethodResult::failure(e),
|
||||
};
|
||||
ext_result
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
/// Empty response for operations that return no data.
|
||||
#[derive(Debug, Serialize)]
|
||||
pub struct Empty {}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,284 @@
|
||||
//! `x.ai/plugins/*` extension handlers.
|
||||
//!
|
||||
//! Provides the plugins list endpoint for the pager's hooks/plugins modal.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_hooks_plugins_types::{
|
||||
HookStatus, McpStatus, PluginInfo, PluginOrigin, PluginScope, PluginsListResponse,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ListRequest {
|
||||
session_id: String,
|
||||
}
|
||||
|
||||
/// Convert a `LoadedPlugin` to a `PluginInfo` DTO.
|
||||
pub fn loaded_plugin_to_info(plugin: &kigi_agent::plugins::LoadedPlugin) -> PluginInfo {
|
||||
use kigi_agent::plugins::discovery::PluginScope as AgentScope;
|
||||
|
||||
let scope = match plugin.scope {
|
||||
AgentScope::CliOverride => PluginScope::Cli,
|
||||
AgentScope::Project => PluginScope::Project,
|
||||
AgentScope::User => PluginScope::User,
|
||||
AgentScope::ConfigPath => PluginScope::Config,
|
||||
};
|
||||
|
||||
let origin = origin_to_dto(&plugin.origin);
|
||||
|
||||
let hook_status = if !plugin.has_hooks {
|
||||
HookStatus::None
|
||||
} else if !plugin.trusted {
|
||||
HookStatus::Blocked
|
||||
} else if plugin.has_inline_hooks_only {
|
||||
HookStatus::ActiveInline
|
||||
} else {
|
||||
HookStatus::Active
|
||||
};
|
||||
|
||||
let mcp_status = if plugin.mcp_server_count == 0 {
|
||||
McpStatus::None
|
||||
} else if !plugin.trusted {
|
||||
McpStatus::Blocked
|
||||
} else if plugin.has_inline_mcp_only {
|
||||
McpStatus::ActiveInline
|
||||
} else {
|
||||
McpStatus::Active
|
||||
};
|
||||
|
||||
PluginInfo {
|
||||
name: plugin.name.clone(),
|
||||
id: plugin.id.0.clone(),
|
||||
root: plugin.root.display().to_string(),
|
||||
scope,
|
||||
trusted: plugin.trusted,
|
||||
enabled: plugin.enabled,
|
||||
version: plugin.version.clone(),
|
||||
description: plugin.description.clone(),
|
||||
skill_count: plugin.skill_count,
|
||||
skill_names: plugin.skill_names.clone(),
|
||||
agent_count: plugin.agent_count,
|
||||
agent_names: plugin.agent_names.clone(),
|
||||
hook_status,
|
||||
hook_count: plugin.hook_count,
|
||||
mcp_server_count: plugin.mcp_server_count,
|
||||
mcp_status,
|
||||
marketplace_source: marketplace_source_label(&origin),
|
||||
origin: Some(origin),
|
||||
conflict: plugin.conflict.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Map the agent-side origin to the wire DTO.
|
||||
fn origin_to_dto(origin: &kigi_agent::plugins::PluginOrigin) -> PluginOrigin {
|
||||
use kigi_agent::plugins::PluginOrigin as AgentOrigin;
|
||||
match origin {
|
||||
AgentOrigin::CliOverride => PluginOrigin::CliOverride,
|
||||
AgentOrigin::ProjectGrok => PluginOrigin::ProjectGrok,
|
||||
AgentOrigin::ProjectClaude => PluginOrigin::ProjectClaude,
|
||||
AgentOrigin::UserGrok => PluginOrigin::UserGrok,
|
||||
AgentOrigin::UserClaude => PluginOrigin::UserClaude,
|
||||
AgentOrigin::ClaudeMarketplace { marketplace } => PluginOrigin::ClaudeMarketplace {
|
||||
marketplace: marketplace.clone(),
|
||||
},
|
||||
AgentOrigin::ClaudeInstalled { marketplace } => PluginOrigin::ClaudeInstalled {
|
||||
marketplace: marketplace.clone(),
|
||||
},
|
||||
AgentOrigin::MarketplaceInstall { git_url } => PluginOrigin::MarketplaceInstall {
|
||||
source_name: None,
|
||||
git_url: git_url.clone(),
|
||||
},
|
||||
AgentOrigin::ConfigPath => PluginOrigin::ConfigPath,
|
||||
}
|
||||
}
|
||||
|
||||
/// Derive the legacy `marketplace_source` label (older-pager compat) from the
|
||||
/// origin: a `git: owner/repo` label for direct git installs.
|
||||
fn marketplace_source_label(origin: &PluginOrigin) -> Option<String> {
|
||||
match origin {
|
||||
PluginOrigin::MarketplaceInstall {
|
||||
git_url: Some(url), ..
|
||||
} => {
|
||||
// Derive short name from URL: "https://github.com/obra/superpowers.git" → "obra/superpowers"
|
||||
let label = url
|
||||
.trim_end_matches(".git")
|
||||
.rsplit("://")
|
||||
.next()
|
||||
.and_then(|s| {
|
||||
s.strip_prefix("github.com/")
|
||||
.or_else(|| s.strip_prefix("gitlab.com/"))
|
||||
.or(Some(s))
|
||||
})
|
||||
.unwrap_or(url);
|
||||
Some(format!("git: {label}"))
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/plugins/list" => {
|
||||
let req: ListRequest = super::parse_params(args)?;
|
||||
|
||||
// A known session answers from its own registry, which includes
|
||||
// `_meta.pluginDirs` plugins. Only an unknown session (a pull
|
||||
// before any session exists) falls back to the shared snapshot.
|
||||
let sid = acp::SessionId::new(req.session_id);
|
||||
let registry = match agent.session_handle_waiting_for_load(&sid).await {
|
||||
Some(handle) => handle.plugins_list().await,
|
||||
None => agent.plugin_registry_snapshot(),
|
||||
};
|
||||
let response = match registry {
|
||||
Some(registry) => {
|
||||
let plugins = registry
|
||||
.list()
|
||||
.iter()
|
||||
.map(|p| loaded_plugin_to_info(p))
|
||||
.collect();
|
||||
PluginsListResponse { plugins }
|
||||
}
|
||||
None => PluginsListResponse {
|
||||
plugins: Vec::new(),
|
||||
},
|
||||
};
|
||||
super::to_ext_response(Ok::<_, anyhow::Error>(response))
|
||||
}
|
||||
"x.ai/plugins/action" => {
|
||||
let req: kigi_hooks_plugins_types::PluginsActionRequest = super::parse_params(args)?;
|
||||
let sid = acp::SessionId::new(req.session_id);
|
||||
|
||||
let result = agent
|
||||
.execute_plugins_action(&sid, req.action)
|
||||
.await
|
||||
.ok_or_else(|| anyhow::anyhow!("session not found"));
|
||||
super::to_ext_response(result)
|
||||
}
|
||||
"x.ai/plugins/notify-updates" => {
|
||||
// Broadcast a PluginUpdatesInstalled notification to the session.
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct NotifyUpdatesRequest {
|
||||
session_id: String,
|
||||
updates: Vec<(String, String, String)>, // (name, old_ver, new_ver)
|
||||
}
|
||||
let req: NotifyUpdatesRequest = super::parse_params(args)?;
|
||||
let sid = acp::SessionId::new(req.session_id);
|
||||
if let Some(handle) = agent.get_session_handle(&sid) {
|
||||
handle.notify_plugin_updates(req.updates).await;
|
||||
}
|
||||
super::to_ext_response(Ok::<_, anyhow::Error>(serde_json::json!({ "ok": true })))
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_agent::plugins::PluginOrigin as AgentOrigin;
|
||||
use kigi_agent::plugins::discovery::{PluginId, PluginScope as AgentScope};
|
||||
|
||||
fn make_loaded_plugin(origin: AgentOrigin) -> kigi_agent::plugins::LoadedPlugin {
|
||||
let root = std::path::PathBuf::from("/tmp/test-plugin");
|
||||
kigi_agent::plugins::LoadedPlugin {
|
||||
name: "test-plugin".to_string(),
|
||||
id: PluginId::new(AgentScope::User, &root, "test-plugin"),
|
||||
root: root.clone(),
|
||||
canonical_root: root,
|
||||
scope: AgentScope::User,
|
||||
origin,
|
||||
trusted: true,
|
||||
enabled: true,
|
||||
version: Some("1.0.0".to_string()),
|
||||
description: None,
|
||||
skill_dirs: vec![],
|
||||
command_dirs: vec![],
|
||||
agent_dirs: vec![],
|
||||
hooks_path: None,
|
||||
mcp_config_path: None,
|
||||
lsp_config_path: None,
|
||||
skill_count: 0,
|
||||
agent_count: 0,
|
||||
skill_names: vec![],
|
||||
agent_names: vec![],
|
||||
has_hooks: false,
|
||||
hook_count: 0,
|
||||
has_inline_hooks_only: false,
|
||||
mcp_server_count: 0,
|
||||
has_inline_mcp_only: false,
|
||||
lsp_server_count: 0,
|
||||
has_inline_lsp_only: false,
|
||||
inline_hooks: None,
|
||||
inline_mcp_servers: None,
|
||||
inline_lsp_servers: None,
|
||||
conflict: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_git_install_gets_git_label() {
|
||||
let plugin = make_loaded_plugin(AgentOrigin::MarketplaceInstall {
|
||||
git_url: Some("https://github.com/obra/superpowers.git".to_string()),
|
||||
});
|
||||
let info = loaded_plugin_to_info(&plugin);
|
||||
assert_eq!(
|
||||
info.marketplace_source.as_deref(),
|
||||
Some("git: obra/superpowers")
|
||||
);
|
||||
assert_eq!(
|
||||
info.origin,
|
||||
Some(PluginOrigin::MarketplaceInstall {
|
||||
source_name: None,
|
||||
git_url: Some("https://github.com/obra/superpowers.git".to_string()),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn direct_local_install_has_no_marketplace_source() {
|
||||
let plugin = make_loaded_plugin(AgentOrigin::MarketplaceInstall { git_url: None });
|
||||
let info = loaded_plugin_to_info(&plugin);
|
||||
assert_eq!(info.marketplace_source, None);
|
||||
assert_eq!(
|
||||
info.origin,
|
||||
Some(PluginOrigin::MarketplaceInstall {
|
||||
source_name: None,
|
||||
git_url: None,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn claude_origins_map_to_dto_without_marketplace_source() {
|
||||
for (agent_origin, expected) in [
|
||||
(
|
||||
AgentOrigin::ClaudeMarketplace {
|
||||
marketplace: "mp".to_string(),
|
||||
},
|
||||
PluginOrigin::ClaudeMarketplace {
|
||||
marketplace: "mp".to_string(),
|
||||
},
|
||||
),
|
||||
(
|
||||
AgentOrigin::ClaudeInstalled {
|
||||
marketplace: Some("mp".to_string()),
|
||||
},
|
||||
PluginOrigin::ClaudeInstalled {
|
||||
marketplace: Some("mp".to_string()),
|
||||
},
|
||||
),
|
||||
(AgentOrigin::UserClaude, PluginOrigin::UserClaude),
|
||||
(AgentOrigin::ProjectClaude, PluginOrigin::ProjectClaude),
|
||||
] {
|
||||
let info = loaded_plugin_to_info(&make_loaded_plugin(agent_origin));
|
||||
assert_eq!(info.origin, Some(expected));
|
||||
assert_eq!(info.marketplace_source, None);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,230 @@
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{ExtResult, parse_params, to_ext_response};
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PrStatusRequest {
|
||||
pub cwd: String,
|
||||
pub branch: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PrStatusResponse {
|
||||
pub pr: Option<PrData>,
|
||||
pub updated_session_ids: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PrData {
|
||||
pub url: String,
|
||||
pub state: String,
|
||||
pub is_in_merge_queue: bool,
|
||||
pub number: Option<u64>,
|
||||
pub title: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GhPrViewResponse {
|
||||
state: Option<String>,
|
||||
url: Option<String>,
|
||||
is_draft: Option<bool>,
|
||||
number: Option<u64>,
|
||||
title: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GhGraphqlResponse {
|
||||
data: Option<GhGraphqlData>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct GhGraphqlData {
|
||||
resource: Option<GhGraphqlPullRequest>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GhGraphqlPullRequest {
|
||||
is_in_merge_queue: Option<bool>,
|
||||
}
|
||||
|
||||
pub async fn handle(_agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/pr/status" => {
|
||||
let req = parse_params::<PrStatusRequest>(args)?;
|
||||
to_ext_response(handle_pr_status(&req.cwd, &req.branch).await)
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_pr_status(cwd: &str, branch: &str) -> anyhow::Result<PrStatusResponse> {
|
||||
Ok(PrStatusResponse {
|
||||
pr: gh_pr_view_by_branch(cwd, branch).await,
|
||||
updated_session_ids: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn gh_pr_view_by_branch(cwd: &str, branch: &str) -> Option<PrData> {
|
||||
let mut cmd = tokio::process::Command::new("gh");
|
||||
cmd.args([
|
||||
"pr",
|
||||
"view",
|
||||
branch,
|
||||
"--json",
|
||||
"state,url,isDraft,number,title",
|
||||
])
|
||||
.current_dir(cwd)
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_command(&mut cmd);
|
||||
cmd.envs(kigi_tools::util::pager_env());
|
||||
let output = cmd.output().await.ok()?;
|
||||
|
||||
if !output.status.success() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let parsed = serde_json::from_slice::<GhPrViewResponse>(&output.stdout).ok()?;
|
||||
let url = parsed.url?;
|
||||
let state = match parsed
|
||||
.state
|
||||
.as_deref()
|
||||
.map(str::to_ascii_lowercase)
|
||||
.as_deref()
|
||||
{
|
||||
Some("merged") => "merged",
|
||||
Some("closed") => "closed",
|
||||
_ if parsed.is_draft.unwrap_or(false) => "draft",
|
||||
_ => "open",
|
||||
};
|
||||
let is_in_merge_queue = state == "open" && gh_pr_is_in_merge_queue(cwd, &url).await;
|
||||
|
||||
Some(PrData {
|
||||
url,
|
||||
state: state.to_string(),
|
||||
is_in_merge_queue,
|
||||
number: parsed.number,
|
||||
title: parsed.title,
|
||||
})
|
||||
}
|
||||
|
||||
/// `gh pr view --json` does not expose `isInMergeQueue`; query GraphQL via `gh api`.
|
||||
async fn gh_pr_is_in_merge_queue(cwd: &str, pr_url: &str) -> bool {
|
||||
const QUERY: &str =
|
||||
"query($url: URI!) { resource(url: $url) { ... on PullRequest { isInMergeQueue } } }";
|
||||
let mut cmd = tokio::process::Command::new("gh");
|
||||
cmd.args([
|
||||
"api",
|
||||
"graphql",
|
||||
"-f",
|
||||
&format!("query={QUERY}"),
|
||||
"-f",
|
||||
&format!("url={pr_url}"),
|
||||
])
|
||||
.current_dir(cwd)
|
||||
.stdin(std::process::Stdio::null());
|
||||
kigi_tools::util::detach_command(&mut cmd);
|
||||
cmd.envs(kigi_tools::util::pager_env());
|
||||
cmd.env("NO_COLOR", "1");
|
||||
let output = match cmd.output().await {
|
||||
Ok(output) => output,
|
||||
Err(_) => return false,
|
||||
};
|
||||
if !output.status.success() {
|
||||
let stderr_snippet: String = String::from_utf8_lossy(&output.stderr)
|
||||
.chars()
|
||||
.take(200)
|
||||
.collect();
|
||||
tracing::warn!(
|
||||
status = %output.status,
|
||||
stderr = %stderr_snippet,
|
||||
"gh api graphql isInMergeQueue lookup failed"
|
||||
);
|
||||
return false;
|
||||
}
|
||||
parse_is_in_merge_queue(&output.stdout).unwrap_or(false)
|
||||
}
|
||||
|
||||
fn parse_is_in_merge_queue(stdout: &[u8]) -> Option<bool> {
|
||||
let stripped = strip_ansi_csi(stdout);
|
||||
let parsed = match serde_json::from_slice::<GhGraphqlResponse>(&stripped) {
|
||||
Ok(parsed) => parsed,
|
||||
Err(error) => {
|
||||
tracing::warn!(error = %error, "failed to parse gh api graphql isInMergeQueue response");
|
||||
return None;
|
||||
}
|
||||
};
|
||||
parsed.data?.resource?.is_in_merge_queue
|
||||
}
|
||||
|
||||
/// `gh` can colorize stdout even when piped (e.g. `GH_FORCE_TTY`, `--color always`
|
||||
/// in config), which would break serde parsing of the JSON payload.
|
||||
fn strip_ansi_csi(bytes: &[u8]) -> Vec<u8> {
|
||||
let mut out = Vec::with_capacity(bytes.len());
|
||||
let mut i = 0;
|
||||
while i < bytes.len() {
|
||||
if bytes[i] == 0x1b && bytes.get(i + 1) == Some(&b'[') {
|
||||
i += 2;
|
||||
while i < bytes.len() && !(0x40..=0x7e).contains(&bytes[i]) {
|
||||
i += 1;
|
||||
}
|
||||
i += 1;
|
||||
} else {
|
||||
out.push(bytes[i]);
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_true() {
|
||||
let stdout = br#"{"data":{"resource":{"isInMergeQueue":true}}}"#;
|
||||
assert_eq!(parse_is_in_merge_queue(stdout), Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_false() {
|
||||
let stdout = br#"{"data":{"resource":{"isInMergeQueue":false}}}"#;
|
||||
assert_eq!(parse_is_in_merge_queue(stdout), Some(false));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_missing_resource() {
|
||||
let stdout = br#"{"data":{"resource":null}}"#;
|
||||
assert_eq!(parse_is_in_merge_queue(stdout), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_missing_data() {
|
||||
assert_eq!(parse_is_in_merge_queue(b"{}"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_malformed_json() {
|
||||
assert_eq!(parse_is_in_merge_queue(b"not json"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_ansi_wrapped_json() {
|
||||
let stdout =
|
||||
b"\x1b[1;32m{\"data\":{\"resource\":{\"isInMergeQueue\":\x1b[0;36mtrue\x1b[0m}}}\x1b[0m";
|
||||
assert_eq!(parse_is_in_merge_queue(stdout), Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_is_in_merge_queue_ansi_wrapped_false() {
|
||||
let stdout = b"\x1b[38;5;208m{\"data\":{\"resource\":{\"isInMergeQueue\":false}}}\x1b[0m\n";
|
||||
assert_eq!(parse_is_in_merge_queue(stdout), Some(false));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
//! `x.ai/privacy/setCodingDataRetention` extension handler.
|
||||
//!
|
||||
//! PUTs the new opt-out flag to cli-chat-proxy and updates local auth state
|
||||
//! to match. The local update is fire-and-forget (best-effort cache refresh).
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/privacy/setCodingDataRetention" => handle_set(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_set(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct Params {
|
||||
coding_data_retention_opt_out: bool,
|
||||
}
|
||||
|
||||
let params: Params = parse_params(args)?;
|
||||
|
||||
let auth = agent.auth_manager.auth().await.map_err(|e| {
|
||||
tracing::warn!(error = %e, "privacy: auth resolution failed");
|
||||
acp::Error::auth_required()
|
||||
.data("Authentication required. Run `grok login` to re-authenticate.")
|
||||
})?;
|
||||
|
||||
let proxy_url = agent.cfg.borrow().endpoints.proxy_url();
|
||||
let url = format!("{proxy_url}/privacy/coding-data-retention");
|
||||
let token_header = agent.auth_manager.grok_com_config().token_header.clone();
|
||||
|
||||
let body = serde_json::json!({
|
||||
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
|
||||
});
|
||||
|
||||
let provider: std::sync::Arc<dyn kigi_auth::AuthCredentialProvider> = std::sync::Arc::new(
|
||||
crate::auth::credential_provider::ShellAuthCredentialProvider::new(
|
||||
agent.auth_manager.clone(),
|
||||
None,
|
||||
None,
|
||||
),
|
||||
);
|
||||
let client = crate::http::with_auth_retry(crate::http::shared_client(), provider);
|
||||
|
||||
let resp = client
|
||||
.put(&url)
|
||||
.header("X-XAI-Token-Auth", &token_header)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("HTTP request failed: {e}")))?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status().as_u16();
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
tracing::warn!(status, "setCodingDataRetention request failed");
|
||||
let friendly = serde_json::from_str::<serde_json::Value>(&body)
|
||||
.ok()
|
||||
.and_then(|v| {
|
||||
v.get("error")
|
||||
.or_else(|| v.get("message"))
|
||||
.and_then(|e| e.as_str().map(String::from))
|
||||
})
|
||||
.unwrap_or_else(|| format!("server returned HTTP {status}"));
|
||||
return Err(acp::Error::internal_error().data(friendly));
|
||||
}
|
||||
|
||||
// Update local auth state to reflect the change.
|
||||
// Use save_without_enrichment to avoid a race: update() spawns a
|
||||
// background GET /user enrichment that may read stale ACL state
|
||||
// and overwrite the opt-out flag back to its previous value.
|
||||
let mut updated = auth.clone();
|
||||
updated.coding_data_retention_opt_out = params.coding_data_retention_opt_out;
|
||||
let _ = agent.auth_manager.save_without_enrichment(updated).await;
|
||||
|
||||
to_raw_response(&serde_json::json!({
|
||||
"codingDataRetentionOptOut": params.coding_data_retention_opt_out,
|
||||
}))
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
//! `x.ai/prompt_history` extension handler.
|
||||
//!
|
||||
//! Returns the user-prompt history for a given cwd. Three paths:
|
||||
//! - **fast path** (no ids): reads the per-CWD `prompt_history.jsonl` file
|
||||
//! directly so Ctrl+R is instant; returns all sessions, most-recent-first.
|
||||
//! - **fast scoped path** (`filter_session_id`): the same file filtered to a
|
||||
//! single session, most-recent-first. This is what the pager's up-arrow /
|
||||
//! Ctrl+R overlay uses to scope history to the current session.
|
||||
//! - **slow path** (`session_id`): rebuilds prompts from session storage in
|
||||
//! chronological order with stable per-session indices. Not used by the
|
||||
//! pager; retained for clients that request session-scoped history this way.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::persistence::list_summaries;
|
||||
use crate::session::prompt_history;
|
||||
use crate::session::storage::StorageAdapter;
|
||||
use crate::session::storage::jsonl::JsonlStorageAdapter;
|
||||
use crate::timed;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct PromptHistoryRequest {
|
||||
cwd: String,
|
||||
/// Optional session ID to filter to a specific session. Routes to the
|
||||
/// session-storage "slow path" (chronological order, stable per-session
|
||||
/// indices). Not used by the pager — see `filter_session_id`.
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
/// Optional session ID to restrict the **fast** per-CWD history file to a
|
||||
/// single session, keeping most-recent-first ordering. Used by the pager's
|
||||
/// up-arrow / Ctrl+R overlay to scope history to the current session.
|
||||
/// Takes precedence over `session_id` when both are set.
|
||||
#[serde(default)]
|
||||
filter_session_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct PromptHistoryResponse {
|
||||
prompts: Vec<String>,
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(_agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/prompt_history" => handle_prompt_history(args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_prompt_history(args: &acp::ExtRequest) -> ExtResult {
|
||||
let request: PromptHistoryRequest = parse_params(args)?;
|
||||
|
||||
// If session_id is specified, use slow path (needed for rewind feature).
|
||||
// Use timed!(try: ...) so we still log timing even when returning early on error.
|
||||
let all_prompts = timed!(try: "prompt_history: load prompts", async {
|
||||
tracing::debug!(
|
||||
"Loading prompt history for cwd: {}, session_id: {:?}, filter_session_id: {:?}",
|
||||
request.cwd,
|
||||
request.session_id,
|
||||
request.filter_session_id
|
||||
);
|
||||
|
||||
if let Some(filter_session_id) = request.filter_session_id.as_deref() {
|
||||
// Fast path, scoped to a single session: filter the per-CWD history
|
||||
// file by session id, preserving most-recent-first ordering.
|
||||
prompt_history::load_prompts_for_session_async(
|
||||
request.cwd.clone(),
|
||||
filter_session_id.to_string(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error()
|
||||
.data(format!("failed to load prompt history: {e}"))
|
||||
})
|
||||
} else if request.session_id.is_some() {
|
||||
// Slow path: load from session storage for per-session queries
|
||||
load_session_prompts(&request.cwd, request.session_id.as_deref()).await
|
||||
} else {
|
||||
// Fast path: use per-CWD prompt history file
|
||||
prompt_history::load_prompts_async(request.cwd.clone())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error()
|
||||
.data(format!("failed to load prompt history: {e}"))
|
||||
})
|
||||
}
|
||||
})?;
|
||||
|
||||
tracing::debug!(
|
||||
"Found {} prompts for cwd {}",
|
||||
all_prompts.len(),
|
||||
request.cwd
|
||||
);
|
||||
|
||||
to_raw_response(&PromptHistoryResponse {
|
||||
prompts: all_prompts,
|
||||
})
|
||||
}
|
||||
|
||||
/// Load prompts using the slow path (session-based loading).
|
||||
/// Used when `session_id` is specified: rebuilds prompts from session storage
|
||||
/// in chronological order with stable per-session indices.
|
||||
async fn load_session_prompts(
|
||||
cwd: &str,
|
||||
session_id: Option<&str>,
|
||||
) -> Result<Vec<String>, acp::Error> {
|
||||
// Load session summaries - either all for the cwd or just the specific session
|
||||
let mut summaries = list_summaries(Some(cwd)).await.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("failed to load session history: {e}"))
|
||||
})?;
|
||||
|
||||
// Filter to specific session if session_id is provided
|
||||
if let Some(target_session_id) = session_id {
|
||||
summaries.retain(|s| s.info.id.0.as_ref() == target_session_id);
|
||||
}
|
||||
|
||||
// Sort sessions by updated_at ascending (oldest first)
|
||||
// so that when we reverse the final list, most recent prompts are first
|
||||
summaries.sort_by_key(|a| a.updated_at);
|
||||
|
||||
// Load only user prompts using the optimized method (avoids loading full session data)
|
||||
let root_dir = crate::util::kigi_home::kigi_home();
|
||||
let storage = JsonlStorageAdapter::with_root(root_dir);
|
||||
|
||||
// Load prompts from sessions with bounded concurrency using stream
|
||||
// Using `buffered` (not `buffer_unordered`) to preserve session order
|
||||
use futures::stream::{self, StreamExt};
|
||||
|
||||
// Limit concurrent file reads to avoid overwhelming the blocking thread pool
|
||||
const MAX_CONCURRENT_READS: usize = 32;
|
||||
|
||||
let mut all_prompts: Vec<String> = stream::iter(summaries)
|
||||
.map(|summary| {
|
||||
let storage = storage.clone();
|
||||
async move {
|
||||
storage
|
||||
.load_prompts_only(&summary.info)
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
}
|
||||
})
|
||||
.buffered(MAX_CONCURRENT_READS)
|
||||
.flat_map(stream::iter)
|
||||
.collect()
|
||||
.await;
|
||||
|
||||
// Deduplicate consecutive identical prompts
|
||||
all_prompts.dedup();
|
||||
|
||||
// DON'T reverse when filtering to a single session - keep chronological
|
||||
// order so per-session prompt indices stay stable (0-indexed from the first
|
||||
// prompt). Only reverse when showing all sessions (history search, most
|
||||
// recent first).
|
||||
if session_id.is_none() {
|
||||
all_prompts.reverse();
|
||||
}
|
||||
|
||||
Ok(all_prompts)
|
||||
}
|
||||
@@ -0,0 +1,71 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
/// Typed metadata for a prompt `TextContent._meta` field.
|
||||
///
|
||||
/// Replaces ad-hoc `serde_json::json!()` construction on the sender side
|
||||
/// and manual `.get()` parsing on the receiver side.
|
||||
///
|
||||
/// Wire-compatible with the existing format: `{"bash_command": "ls -la"}`
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct PromptBlockMeta {
|
||||
/// Direct bash command to execute (bypasses agent loop).
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub bash_command: Option<String>,
|
||||
}
|
||||
|
||||
impl PromptBlockMeta {
|
||||
/// Create meta for a direct bash command.
|
||||
pub fn bash(command: impl Into<String>) -> Self {
|
||||
Self {
|
||||
bash_command: Some(command.into()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Try to parse from a freeform `_meta` map.
|
||||
pub fn from_value(value: &agent_client_protocol::Meta) -> Option<Self> {
|
||||
serde_json::from_value(serde_json::Value::Object(value.clone())).ok()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn bash_roundtrip_serde() {
|
||||
let meta = PromptBlockMeta::bash("ls -la");
|
||||
let json = serde_json::to_value(&meta).unwrap();
|
||||
let parsed: PromptBlockMeta = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(parsed.bash_command, Some("ls -la".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_value_legacy_compat() {
|
||||
let val = serde_json::json!({"bash_command": "ls"});
|
||||
let meta = PromptBlockMeta::from_value(val.as_object().unwrap()).unwrap();
|
||||
assert_eq!(meta.bash_command, Some("ls".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_value_unrelated_meta() {
|
||||
let val = serde_json::json!({"other": 1});
|
||||
let meta = PromptBlockMeta::from_value(val.as_object().unwrap());
|
||||
assert!(meta.is_some());
|
||||
assert_eq!(meta.unwrap().bash_command, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn from_value_empty_object() {
|
||||
let val = serde_json::json!({});
|
||||
let meta = PromptBlockMeta::from_value(val.as_object().unwrap());
|
||||
assert!(meta.is_some());
|
||||
assert_eq!(meta.unwrap().bash_command, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skip_serializing_none() {
|
||||
let meta = PromptBlockMeta { bash_command: None };
|
||||
let json = serde_json::to_value(&meta).unwrap();
|
||||
assert!(!json.as_object().unwrap().contains_key("bash_command"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
//! `x.ai/recap` extension handler.
|
||||
//!
|
||||
//! Triggers generation of a session recap — a short "where was I" summary of
|
||||
//! the session so far — via [`SessionCommand::Recap`]. This is fire-and-forget:
|
||||
//! the recap is delivered asynchronously to every attached client as a
|
||||
//! [`SessionUpdate::SessionRecap`](crate::extensions::notification::SessionUpdate::SessionRecap)
|
||||
//! notification, so the handler returns as soon as the command is queued rather
|
||||
//! than blocking on the model call.
|
||||
//!
|
||||
//! Invoked on demand via the `/recap` slash command (`auto = false`) and
|
||||
//! automatically when the user returns to the terminal after being away
|
||||
//! (`auto = true`).
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{ExtResult, parse_params, to_ext_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::SessionCommand;
|
||||
|
||||
#[tracing::instrument(skip_all)]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct RecapRequest {
|
||||
session_id: String,
|
||||
#[serde(default)]
|
||||
auto: bool,
|
||||
}
|
||||
|
||||
let req: RecapRequest = parse_params(args)?;
|
||||
tracing::info!(auto = req.auto, "handling /recap request");
|
||||
|
||||
// Feature gate: remote setting / `[features] session_recap`
|
||||
// config.toml key / `KIGI_SESSION_RECAP` env (default ON). Gates both the
|
||||
// manual `/recap` and the automatic recap.
|
||||
if !agent.cfg.borrow().is_session_recap_enabled() {
|
||||
tracing::debug!("session recap disabled by config/feature flag; ignoring request");
|
||||
return to_ext_response(Ok(serde_json::json!({ "ok": true, "disabled": true })));
|
||||
}
|
||||
|
||||
let sid: acp::SessionId = req.session_id.clone().into();
|
||||
// Load-race-tolerant: an automatic recap fires on return-from-away, when a
|
||||
// leader restart may have a reconnect-replayed `session/load` still in
|
||||
// flight. Wait for that load instead of failing with "session not found".
|
||||
let Some(session) = agent.session_handle_waiting_for_load(&sid).await else {
|
||||
return Err(
|
||||
acp::Error::invalid_params().data(format!("session not found: {}", req.session_id))
|
||||
);
|
||||
};
|
||||
|
||||
// Fire-and-forget: the recap is emitted later as a SessionRecap
|
||||
// notification. We only ack that the request was accepted.
|
||||
let _ = session
|
||||
.cmd_tx
|
||||
.send(SessionCommand::Recap { auto: req.auto });
|
||||
|
||||
to_ext_response(Ok(serde_json::json!({ "ok": true })))
|
||||
}
|
||||
@@ -0,0 +1,295 @@
|
||||
//! `x.ai/session/repair` — out-of-band recovery for sessions bricked by
|
||||
//! corrupted tool-pairing history.
|
||||
//!
|
||||
//! A `ToolResult` whose owning assistant `tool_call` is missing (e.g. a
|
||||
//! torn/merged `chat_history.jsonl` line skipped on load) makes every request
|
||||
//! 400 with "unexpected `tool_use_id` found in `tool_result` blocks". No
|
||||
//! in-band path can recover — compaction's sanitizer needs a model call that
|
||||
//! itself 400s — so the client invokes this method against the session.
|
||||
//!
|
||||
//! Repairs via [`kigi_chat_state::compaction_utils::repair_history`]. Resident
|
||||
//! sessions go through `SessionCommand::RepairHistory` (serialized with
|
||||
//! session activity, rejected mid-turn); non-resident sessions are repaired
|
||||
//! on disk via the atomic `replace_chat_history`.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_chat_state::compaction_utils::HistoryRepairReport;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::sync::oneshot;
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::SessionCommand;
|
||||
use crate::session::storage::StorageAdapter;
|
||||
use crate::session::storage::jsonl::JsonlStorageAdapter;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct RepairSessionRequest {
|
||||
session_id: String,
|
||||
/// Report what would change without mutating memory or disk.
|
||||
#[serde(default)]
|
||||
dry_run: bool,
|
||||
}
|
||||
|
||||
/// Response payload for `x.ai/session/repair`.
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RepairSessionResponse {
|
||||
/// Whether the repair modified (or, for `dryRun`, would modify) the history.
|
||||
pub repaired: bool,
|
||||
/// Echo of the request's `dryRun` flag.
|
||||
pub dry_run: bool,
|
||||
/// Whether the session was resident (repaired via the live actor) or
|
||||
/// repaired directly on disk.
|
||||
pub resident: bool,
|
||||
/// Duplicate `ToolResult` entries removed.
|
||||
pub duplicates_removed: usize,
|
||||
/// `tool_call_id`s of orphaned/displaced `ToolResult`s stripped.
|
||||
pub stripped_tool_result_ids: Vec<String>,
|
||||
/// Synthetic `ToolResult`s inserted for unanswered tool calls.
|
||||
pub synthetic_results_inserted: usize,
|
||||
}
|
||||
|
||||
impl RepairSessionResponse {
|
||||
fn new(report: HistoryRepairReport, dry_run: bool, resident: bool) -> Self {
|
||||
Self {
|
||||
repaired: report.changed(),
|
||||
dry_run,
|
||||
resident,
|
||||
duplicates_removed: report.duplicates_removed,
|
||||
stripped_tool_result_ids: report.stripped_tool_result_ids,
|
||||
synthetic_results_inserted: report.synthetic_results_inserted,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/session/repair" => handle_session_repair(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_session_repair(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: RepairSessionRequest = parse_params(args)?;
|
||||
let session_id = acp::SessionId::new(req.session_id.as_str());
|
||||
|
||||
// Resident rail. The load-waiting lookup keeps a repair racing a
|
||||
// reconnect replay from falling through to the disk rail.
|
||||
if let Some(handle) = agent.session_handle_waiting_for_load(&session_id).await {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::RepairHistory {
|
||||
dry_run: req.dry_run,
|
||||
respond_to: tx,
|
||||
})
|
||||
.map_err(|_| acp::Error::internal_error().data("failed to send repair command"))?;
|
||||
let report = rx
|
||||
.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("repair failed: {e}")))?;
|
||||
return to_raw_response(&RepairSessionResponse::new(report, req.dry_run, true));
|
||||
}
|
||||
|
||||
// Disk rail: session not resident — repair `chat_history.jsonl` in place.
|
||||
repair_on_disk(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
&req.session_id,
|
||||
req.dry_run,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
/// Repair a non-resident session's history on disk: load via the resume
|
||||
/// path's corruption-tolerant reader (legacy upgrades apply), repair, write
|
||||
/// back atomically. `grok_root` is injectable for tests.
|
||||
async fn repair_on_disk(grok_root: &std::path::Path, session_id: &str, dry_run: bool) -> ExtResult {
|
||||
let summary = crate::session::persistence::find_summary_by_session_id_in_root(
|
||||
session_id,
|
||||
&grok_root.join("sessions"),
|
||||
)
|
||||
.ok_or_else(|| {
|
||||
acp::Error::resource_not_found(Some(format!("session not found: {session_id}")))
|
||||
})?;
|
||||
let info = summary.info.clone();
|
||||
|
||||
let storage = JsonlStorageAdapter::with_root(grok_root.to_path_buf());
|
||||
let mut chat_history = storage
|
||||
.load_session_without_updates(&info)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("failed to load session history: {e}"))
|
||||
})?
|
||||
.chat_history;
|
||||
|
||||
let report = kigi_chat_state::compaction_utils::repair_history(&mut chat_history);
|
||||
|
||||
if report.changed() && !dry_run {
|
||||
storage
|
||||
.replace_chat_history(&info, &chat_history)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("failed to write repaired history: {e}"))
|
||||
})?;
|
||||
tracing::warn!(
|
||||
session_id,
|
||||
duplicates_removed = report.duplicates_removed,
|
||||
stripped_tool_result_ids = ?report.stripped_tool_result_ids,
|
||||
synthetic_results_inserted = report.synthetic_results_inserted,
|
||||
"session history repaired on disk"
|
||||
);
|
||||
}
|
||||
|
||||
to_raw_response(&RepairSessionResponse::new(report, dry_run, false))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::sampling::ConversationItem;
|
||||
use crate::session::info::Info;
|
||||
use crate::session::persistence::default_model_id;
|
||||
use kigi_sampling_types::ToolCall;
|
||||
use tempfile::TempDir;
|
||||
|
||||
const SESSION_ID: &str = "019f3df7-3d70-7f60-8ca0-a38d2d005670";
|
||||
|
||||
/// Seed `{root}/sessions/{cwd}/{id}/` with a summary and the given chat
|
||||
/// history, returning the adapter + info for follow-up reads.
|
||||
async fn seed_session(
|
||||
root: &std::path::Path,
|
||||
items: &[ConversationItem],
|
||||
) -> (JsonlStorageAdapter, Info) {
|
||||
let adapter = JsonlStorageAdapter::with_root(root.to_path_buf());
|
||||
let info = Info {
|
||||
id: acp::SessionId::new(SESSION_ID),
|
||||
cwd: "/work".to_string(),
|
||||
};
|
||||
adapter
|
||||
.init_session(&info, default_model_id())
|
||||
.await
|
||||
.expect("init session");
|
||||
for item in items {
|
||||
adapter
|
||||
.append_chat_message(&info, item)
|
||||
.await
|
||||
.expect("append chat message");
|
||||
}
|
||||
(adapter, info)
|
||||
}
|
||||
|
||||
/// The bricked-session shape: the assistant line owning `call_LOST` is
|
||||
/// gone (torn/merged JSONL line skipped on load), leaving an orphaned
|
||||
/// tool result that 400s on every request.
|
||||
fn corrupted_history() -> Vec<ConversationItem> {
|
||||
vec![
|
||||
ConversationItem::system("sys"),
|
||||
ConversationItem::user("prompt"),
|
||||
ConversationItem::tool_result("call_LOST", "orphaned result"),
|
||||
ConversationItem::assistant_tool_calls(vec![ToolCall {
|
||||
id: "call_OK".into(),
|
||||
name: "read_file".to_string(),
|
||||
arguments: "{}".into(),
|
||||
}]),
|
||||
ConversationItem::tool_result("call_OK", "fine"),
|
||||
]
|
||||
}
|
||||
|
||||
fn parse(resp: &acp::ExtResponse) -> serde_json::Value {
|
||||
serde_json::from_str(resp.0.get()).expect("repair response json")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disk_repair_strips_orphaned_result_and_rewrites_file() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let (adapter, info) = seed_session(tmp.path(), &corrupted_history()).await;
|
||||
|
||||
let resp = repair_on_disk(tmp.path(), SESSION_ID, false)
|
||||
.await
|
||||
.expect("repair ok");
|
||||
let v = parse(&resp);
|
||||
assert_eq!(v["repaired"], true);
|
||||
assert_eq!(v["resident"], false);
|
||||
assert_eq!(v["dryRun"], false);
|
||||
assert_eq!(v["strippedToolResultIds"], serde_json::json!(["call_LOST"]));
|
||||
assert_eq!(v["duplicatesRemoved"], 0);
|
||||
assert_eq!(v["syntheticResultsInserted"], 0);
|
||||
|
||||
// The rewritten file must reload as a valid conversation with the
|
||||
// orphan gone and the intact pair preserved.
|
||||
let reloaded = adapter
|
||||
.load_session_without_updates(&info)
|
||||
.await
|
||||
.expect("reload")
|
||||
.chat_history;
|
||||
assert_eq!(reloaded.len(), 4);
|
||||
assert!(!reloaded.iter().any(|i| matches!(
|
||||
i,
|
||||
ConversationItem::ToolResult(tr) if tr.tool_call_id == "call_LOST"
|
||||
)));
|
||||
|
||||
// A second repair is a no-op: the corruption is really gone.
|
||||
let v2 = parse(
|
||||
&repair_on_disk(tmp.path(), SESSION_ID, false)
|
||||
.await
|
||||
.expect("second repair ok"),
|
||||
);
|
||||
assert_eq!(v2["repaired"], false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disk_repair_dry_run_reports_without_writing() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let (adapter, info) = seed_session(tmp.path(), &corrupted_history()).await;
|
||||
|
||||
let v = parse(
|
||||
&repair_on_disk(tmp.path(), SESSION_ID, true)
|
||||
.await
|
||||
.expect("dry run ok"),
|
||||
);
|
||||
assert_eq!(v["repaired"], true);
|
||||
assert_eq!(v["dryRun"], true);
|
||||
assert_eq!(v["strippedToolResultIds"], serde_json::json!(["call_LOST"]));
|
||||
|
||||
// Disk untouched: the orphan is still there.
|
||||
let reloaded = adapter
|
||||
.load_session_without_updates(&info)
|
||||
.await
|
||||
.expect("reload")
|
||||
.chat_history;
|
||||
assert_eq!(reloaded.len(), 5);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disk_repair_noop_on_valid_history() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let valid = vec![
|
||||
ConversationItem::system("sys"),
|
||||
ConversationItem::user("prompt"),
|
||||
ConversationItem::assistant("done"),
|
||||
];
|
||||
seed_session(tmp.path(), &valid).await;
|
||||
|
||||
let v = parse(
|
||||
&repair_on_disk(tmp.path(), SESSION_ID, false)
|
||||
.await
|
||||
.expect("repair ok"),
|
||||
);
|
||||
assert_eq!(v["repaired"], false);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn disk_repair_unknown_session_is_resource_not_found() {
|
||||
let tmp = TempDir::new().unwrap();
|
||||
let err = repair_on_disk(tmp.path(), "no-such-session", false)
|
||||
.await
|
||||
.expect_err("must fail");
|
||||
assert_eq!(
|
||||
err.code,
|
||||
acp::Error::resource_not_found(None::<String>).code
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,86 @@
|
||||
//! `x.ai/rewind/*` extension handlers.
|
||||
//!
|
||||
//! - `rewind/execute`: rewind a session to a target prompt index, optionally
|
||||
//! forcing past in-flight prompts and choosing a `RewindMode`.
|
||||
//! - `rewind/points`: list the prompt indices that can be rewound to.
|
||||
//!
|
||||
//! Local mode dispatches [`handle`]. In gateway-bridge mode the agent's
|
||||
//! routing hook calls [`handle_bridge`], which composes the server's
|
||||
//! conversation rewind with the local file half so the pager-facing ACP
|
||||
//! response is identical either way.
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::handle::SessionHandle;
|
||||
use crate::session::{RewindMode, RewindRequest, SessionCommand};
|
||||
use agent_client_protocol as acp;
|
||||
use serde::Deserialize;
|
||||
use tokio::sync::oneshot;
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
tracing::info!("handling rewind request: {}", args.method);
|
||||
match args.method.as_ref() {
|
||||
"x.ai/rewind/execute" => handle_execute(agent, args).await,
|
||||
"x.ai/rewind/points" => handle_points(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct RewindSessionRequest {
|
||||
#[serde(alias = "sessionId")]
|
||||
session_id: String,
|
||||
#[serde(alias = "targetPromptIndex")]
|
||||
target_prompt_index: usize,
|
||||
#[serde(default)]
|
||||
force: bool,
|
||||
#[serde(default)]
|
||||
mode: Option<RewindMode>,
|
||||
}
|
||||
#[derive(Deserialize)]
|
||||
struct RewindPointsRequest {
|
||||
#[serde(alias = "sessionId")]
|
||||
session_id: String,
|
||||
}
|
||||
/// Look up a `SessionHandle` by id string, or return a `resource_not_found`
|
||||
/// `acp::Error`. Used by both arms below.
|
||||
fn lookup_session(agent: &MvpAgent, session_id: String) -> Result<SessionHandle, acp::Error> {
|
||||
agent
|
||||
.sessions
|
||||
.borrow()
|
||||
.get(&acp::SessionId::new(session_id))
|
||||
.cloned()
|
||||
.ok_or_else(|| acp::Error::resource_not_found(Some("session not found".into())))
|
||||
}
|
||||
async fn handle_execute(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let request: RewindSessionRequest = parse_params(args)?;
|
||||
let handle = lookup_session(agent, request.session_id)?;
|
||||
let (tx, rx) = oneshot::channel();
|
||||
handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::Rewind {
|
||||
request: RewindRequest {
|
||||
target_prompt_index: request.target_prompt_index,
|
||||
force: request.force,
|
||||
mode: request.mode.unwrap_or(RewindMode::All),
|
||||
},
|
||||
respond_to: tx,
|
||||
})
|
||||
.map_err(|_| acp::Error::internal_error().data("failed to send rewind command"))?;
|
||||
let result = rx
|
||||
.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("Rewind failed: {:?}", e)))?;
|
||||
to_raw_response(&result)
|
||||
}
|
||||
async fn handle_points(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let request: RewindPointsRequest = parse_params(args)?;
|
||||
let handle = lookup_session(agent, request.session_id)?;
|
||||
let (tx, rx) = oneshot::channel();
|
||||
handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::GetRewindPoints { respond_to: tx })
|
||||
.map_err(|_| acp::Error::internal_error().data("failed to send command"))?;
|
||||
let result = rx
|
||||
.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session failed to respond"))?;
|
||||
to_raw_response(&result)
|
||||
}
|
||||
@@ -0,0 +1,39 @@
|
||||
//! `x.ai/rollout/survey` extension handler.
|
||||
//!
|
||||
//! Logs a rollout-survey submission via telemetry (Mixpanel + BigQuery).
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::{RolloutSurveyRequest, RolloutSurveyResponse};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(_agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/rollout/survey" => {
|
||||
let req: RolloutSurveyRequest = parse_params(args)?;
|
||||
|
||||
tracing::info_span!(
|
||||
"feedback.survey",
|
||||
survey_type = "rollout",
|
||||
event_type = "responded",
|
||||
has_feedback_text = !req.feedback.is_empty(),
|
||||
preference_count = req.preferences.len() as i64,
|
||||
)
|
||||
.in_scope(|| {});
|
||||
|
||||
// Log the survey via telemetry (this will go to Mixpanel and BigQuery)
|
||||
|
||||
tracing::info!(
|
||||
"Rollout survey received for session {}: preferences={:?}, feedback={}",
|
||||
req.session_id,
|
||||
req.preferences,
|
||||
req.feedback,
|
||||
);
|
||||
|
||||
to_raw_response(&RolloutSurveyResponse { success: true })
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,143 @@
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
// Re-export from workspace crate (canonical home for fuzzy search).
|
||||
pub use kigi_workspace::file_system::{ClientId, TargetClientId};
|
||||
|
||||
/// Metadata from the request, used for routing notifications back to the
|
||||
/// correct client.
|
||||
#[derive(Clone, Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct RequestMeta {
|
||||
#[serde(default)]
|
||||
pub client_id: TargetClientId,
|
||||
}
|
||||
|
||||
/// Notification-side routing metadata. Embeddable in any outgoing
|
||||
/// notification struct via `#[serde(rename = "_meta")]`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct NotificationMeta {
|
||||
#[serde(skip_serializing_if = "TargetClientId::is_none")]
|
||||
pub target_client_id: TargetClientId,
|
||||
}
|
||||
|
||||
/// Inject `targetClientId` into the `_meta` field of a JSON params object.
|
||||
/// Merges with any existing `_meta` fields rather than replacing them.
|
||||
pub fn inject_routing_meta(params: &mut serde_json::Value, target_client_id: &TargetClientId) {
|
||||
if target_client_id.is_none() {
|
||||
return;
|
||||
}
|
||||
let meta = params.as_object_mut().and_then(|obj| {
|
||||
obj.entry("_meta")
|
||||
.or_insert_with(|| serde_json::json!({}))
|
||||
.as_object_mut()
|
||||
});
|
||||
if let Some(meta) = meta
|
||||
&& let Ok(val) = serde_json::to_value(target_client_id)
|
||||
{
|
||||
meta.insert("targetClientId".to_string(), val);
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a fire-and-forget ext notification with optional client routing.
|
||||
///
|
||||
/// If `target_client_id` is set, injects `_meta.targetClientId` into `params`
|
||||
/// so the gateway can route the notification to the correct client.
|
||||
pub fn send_routed_notification(
|
||||
gateway: &GatewaySender,
|
||||
method: &str,
|
||||
mut params: serde_json::Value,
|
||||
target_client_id: &TargetClientId,
|
||||
) {
|
||||
inject_routing_meta(&mut params, target_client_id);
|
||||
if let Ok(raw) = serde_json::value::to_raw_value(¶ms) {
|
||||
gateway.forward_fire_and_forget(acp::ExtNotification::new(method, raw.into()));
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn target_client_id_serialization() {
|
||||
let none = TargetClientId::None;
|
||||
assert_eq!(serde_json::to_string(&none).unwrap(), "null");
|
||||
|
||||
let client_id = TargetClientId::ClientId(ClientId {
|
||||
instance_id: "inst-123".to_string(),
|
||||
conn_id: "conn-456".to_string(),
|
||||
});
|
||||
let json = serde_json::to_string(&client_id).unwrap();
|
||||
assert_eq!(json, r#"{"instanceId":"inst-123","connId":"conn-456"}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_meta_deserialization() {
|
||||
let json = r#"{"clientId": {"instanceId": "inst-1", "connId": "conn-2"}}"#;
|
||||
let meta: RequestMeta = serde_json::from_str(json).unwrap();
|
||||
assert!(matches!(meta.client_id, TargetClientId::ClientId(_)));
|
||||
|
||||
let json = r#"{}"#;
|
||||
let meta: RequestMeta = serde_json::from_str(json).unwrap();
|
||||
assert!(meta.client_id.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn notification_meta_serialization() {
|
||||
let meta = NotificationMeta {
|
||||
target_client_id: TargetClientId::ClientId(ClientId {
|
||||
instance_id: "inst-abc".to_string(),
|
||||
conn_id: "conn-xyz".to_string(),
|
||||
}),
|
||||
};
|
||||
let json = serde_json::to_string(&meta).unwrap();
|
||||
assert_eq!(
|
||||
json,
|
||||
r#"{"targetClientId":{"instanceId":"inst-abc","connId":"conn-xyz"}}"#
|
||||
);
|
||||
|
||||
let meta_none = NotificationMeta {
|
||||
target_client_id: TargetClientId::None,
|
||||
};
|
||||
let json = serde_json::to_string(&meta_none).unwrap();
|
||||
assert_eq!(json, r#"{}"#);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_routing_meta_inserts_into_empty_params() {
|
||||
let mut params = serde_json::json!({"terminalId": "abc"});
|
||||
let target = TargetClientId::ClientId(ClientId {
|
||||
instance_id: "inst-1".to_string(),
|
||||
conn_id: "conn-2".to_string(),
|
||||
});
|
||||
inject_routing_meta(&mut params, &target);
|
||||
assert_eq!(params["_meta"]["targetClientId"]["instanceId"], "inst-1");
|
||||
assert_eq!(params["_meta"]["targetClientId"]["connId"], "conn-2");
|
||||
assert_eq!(params["terminalId"], "abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_routing_meta_merges_with_existing_meta() {
|
||||
let mut params = serde_json::json!({
|
||||
"data": "x",
|
||||
"_meta": {"eventId": "evt-1"}
|
||||
});
|
||||
let target = TargetClientId::ClientId(ClientId {
|
||||
instance_id: "inst-1".to_string(),
|
||||
conn_id: "conn-2".to_string(),
|
||||
});
|
||||
inject_routing_meta(&mut params, &target);
|
||||
assert_eq!(params["_meta"]["eventId"], "evt-1");
|
||||
assert_eq!(params["_meta"]["targetClientId"]["instanceId"], "inst-1");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn inject_routing_meta_skips_when_none() {
|
||||
let mut params = serde_json::json!({"terminalId": "abc"});
|
||||
inject_routing_meta(&mut params, &TargetClientId::None);
|
||||
assert!(params.get("_meta").is_none());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,354 @@
|
||||
//! Search extension API layer (fuzzy file search, content search).
|
||||
//!
|
||||
//! Routing: prefers explicit `cwd`, falls back to session lookup via `sessionId`.
|
||||
|
||||
use crate::agent::mvp_agent::MvpAgent;
|
||||
use crate::session::ExtMethodResult;
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_workspace::file_system::ContentSearchRequest as ContentSearchRequestParams;
|
||||
use kigi_workspace::workspace_ops::{FuzzyChangeReq, FuzzyCloseReq, FuzzyOpenReq};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::path::PathBuf;
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyOpenResponse {
|
||||
pub session_id: String,
|
||||
pub search_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyChangeResponse {
|
||||
pub session_id: String,
|
||||
pub search_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyCloseResponse {
|
||||
pub session_id: String,
|
||||
pub search_id: String,
|
||||
pub closed: bool,
|
||||
}
|
||||
|
||||
pub use crate::extensions::routing::{ClientId, NotificationMeta, RequestMeta, TargetClientId};
|
||||
|
||||
fn parse<T: for<'de> Deserialize<'de>>(s: &str) -> Result<T, acp::Error> {
|
||||
serde_json::from_str::<T>(s)
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {}", e)))
|
||||
}
|
||||
|
||||
/// Resolve the search root, preferring an explicit `cwd` over a `sessionId` lookup.
|
||||
fn resolve_cwd(
|
||||
agent: &MvpAgent,
|
||||
cwd: Option<String>,
|
||||
session_id: Option<&acp::SessionId>,
|
||||
) -> Result<PathBuf, acp::Error> {
|
||||
if let Some(cwd) = cwd {
|
||||
return Ok(PathBuf::from(cwd));
|
||||
}
|
||||
|
||||
if let Some(session_id) = session_id {
|
||||
if let Some(cwd) = agent.get_session_cwd(session_id) {
|
||||
return Ok(cwd);
|
||||
}
|
||||
return Err(
|
||||
acp::Error::invalid_params().data(format!("session not found: {}", session_id.0))
|
||||
);
|
||||
}
|
||||
|
||||
Err(acp::Error::invalid_params().data("either cwd or sessionId is required"))
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyOpenRequest {
|
||||
/// Optional session ID - used to lookup cwd if cwd not provided directly
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
/// Optional absolute cwd path - preferred over session_id lookup
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
/// Optional relative path within the resolved cwd
|
||||
#[serde(default)]
|
||||
pub root: Option<String>,
|
||||
#[serde(default)]
|
||||
pub request_id: Option<String>,
|
||||
#[serde(default)]
|
||||
pub hidden: bool,
|
||||
/// Metadata for routing (contains client_id from relay).
|
||||
#[serde(default, rename = "_meta")]
|
||||
pub meta: Option<RequestMeta>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyChangeRequest {
|
||||
pub search_id: String,
|
||||
pub query: String,
|
||||
#[serde(default)]
|
||||
pub dirs_only: bool,
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct FuzzyCloseRequest {
|
||||
pub search_id: String,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ContentSearchRequest {
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub params: ContentSearchRequestParams,
|
||||
}
|
||||
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/search/fuzzy/open" => {
|
||||
let req: FuzzyOpenRequest = parse(args.params.get())?;
|
||||
let cwd = resolve_cwd(agent, req.cwd, req.session_id.as_ref())?;
|
||||
let search_root = match &req.root {
|
||||
Some(r) => cwd.join(r),
|
||||
None => cwd,
|
||||
};
|
||||
let session_id = req.session_id.map(|s| s.0.to_string());
|
||||
let target_client_id = req.meta.map(|m| m.client_id).unwrap_or_default();
|
||||
|
||||
let ops = agent
|
||||
.resolve_workspace_ops()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
let search_id = ops
|
||||
.dispatch(
|
||||
&FuzzyOpenReq {
|
||||
root: Some(search_root),
|
||||
request_id: req.request_id,
|
||||
hidden: req.hidden,
|
||||
session_id: session_id.clone(),
|
||||
target_client_id,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
let response = FuzzyOpenResponse {
|
||||
session_id: session_id.unwrap_or_else(|| "agent".to_string()),
|
||||
search_id,
|
||||
};
|
||||
ExtMethodResult::success(response)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
"x.ai/search/fuzzy/change" => {
|
||||
let req: FuzzyChangeRequest = parse(args.params.get())?;
|
||||
let ops = agent
|
||||
.resolve_workspace_ops()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// The workspace owns the manager and spawns the status driver, which
|
||||
// streams `x.ai/search/fuzzy/status` through the client sink.
|
||||
let found = ops
|
||||
.dispatch(
|
||||
&FuzzyChangeReq {
|
||||
search_id: req.search_id.clone(),
|
||||
query: req.query.clone(),
|
||||
dirs_only: req.dirs_only,
|
||||
limit: req.limit,
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
if !found {
|
||||
return Err(acp::Error::invalid_params()
|
||||
.data(format!("search not found: {}", req.search_id)));
|
||||
}
|
||||
|
||||
let response = FuzzyChangeResponse {
|
||||
session_id: "agent".to_string(),
|
||||
search_id: req.search_id,
|
||||
};
|
||||
ExtMethodResult::success(response)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
"x.ai/search/fuzzy/close" => {
|
||||
let req: FuzzyCloseRequest = parse(args.params.get())?;
|
||||
let ops = agent
|
||||
.resolve_workspace_ops()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
let closed = ops
|
||||
.dispatch(
|
||||
&FuzzyCloseReq {
|
||||
search_id: req.search_id.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
let response = FuzzyCloseResponse {
|
||||
session_id: "agent".to_string(),
|
||||
search_id: req.search_id,
|
||||
closed,
|
||||
};
|
||||
ExtMethodResult::success(response)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
"x.ai/search/content" => {
|
||||
let req: ContentSearchRequest = parse(args.params.get())?;
|
||||
let cwd = resolve_cwd(agent, req.cwd.clone(), req.session_id.as_ref())?;
|
||||
let context_id = req
|
||||
.session_id
|
||||
.as_ref()
|
||||
.map(|s| s.0.to_string())
|
||||
.unwrap_or_else(|| "agent".to_string());
|
||||
let ops = agent
|
||||
.resolve_workspace_ops()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
// The workspace runs the streaming search and emits
|
||||
// `x.ai/search/content/status` batches through the client sink.
|
||||
let mut op = req.params;
|
||||
op.cwd = Some(cwd);
|
||||
op.context_id = Some(context_id);
|
||||
let data = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
ExtMethodResult::success(data)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_with_cwd() {
|
||||
let json = r#"{"cwd": "/path/to/project", "hidden": false}"#;
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.cwd, Some("/path/to/project".to_string()));
|
||||
assert_eq!(req.session_id, None);
|
||||
assert!(!req.hidden);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_with_session_id() {
|
||||
let json = r#"{"sessionId": "session-123", "hidden": true}"#;
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert!(req.session_id.is_some());
|
||||
assert_eq!(req.session_id.unwrap().0.as_ref(), "session-123");
|
||||
assert_eq!(req.cwd, None);
|
||||
assert!(req.hidden);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_with_both_cwd_and_session_id() {
|
||||
let json = r#"{"sessionId": "session-123", "cwd": "/path/to/project"}"#;
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
// Both should be present - cwd takes precedence in resolve_cwd
|
||||
assert!(req.session_id.is_some());
|
||||
assert_eq!(req.cwd, Some("/path/to/project".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_with_root() {
|
||||
let json = r#"{"cwd": "/home/user", "root": "src"}"#;
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.cwd, Some("/home/user".to_string()));
|
||||
assert_eq!(req.root, Some("src".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_change_request() {
|
||||
let json = r#"{"searchId": "search-456", "query": "main.rs", "limit": 10}"#;
|
||||
let req: FuzzyChangeRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.search_id, "search-456");
|
||||
assert_eq!(req.query, "main.rs");
|
||||
assert_eq!(req.limit, Some(10));
|
||||
assert!(!req.dirs_only);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_close_request() {
|
||||
let json = r#"{"searchId": "search-456"}"#;
|
||||
let req: FuzzyCloseRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.search_id, "search-456");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_defaults() {
|
||||
let json = r#"{"cwd": "/path"}"#;
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.session_id, None);
|
||||
assert_eq!(req.root, None);
|
||||
assert_eq!(req.request_id, None);
|
||||
assert!(!req.hidden);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_change_request_defaults() {
|
||||
let json = r#"{"searchId": "s1", "query": "q"}"#;
|
||||
let req: FuzzyChangeRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert!(!req.dirs_only);
|
||||
assert_eq!(req.limit, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_fuzzy_open_request_with_meta() {
|
||||
// Test that FuzzyOpenRequest correctly deserializes _meta.clientId
|
||||
// This is what the relay injects into the request
|
||||
let json = r#"{
|
||||
"cwd": "/path/to/project",
|
||||
"requestId": "req-123",
|
||||
"hidden": false,
|
||||
"_meta": {
|
||||
"clientId": {
|
||||
"instanceId": "relay-instance-1",
|
||||
"connId": "client-conn-abc"
|
||||
}
|
||||
}
|
||||
}"#;
|
||||
|
||||
let req: FuzzyOpenRequest = serde_json::from_str(json).unwrap();
|
||||
|
||||
assert_eq!(req.cwd, Some("/path/to/project".to_string()));
|
||||
assert!(req.meta.is_some());
|
||||
|
||||
let meta = req.meta.unwrap();
|
||||
match &meta.client_id {
|
||||
TargetClientId::ClientId(client_id) => {
|
||||
assert_eq!(client_id.instance_id, "relay-instance-1");
|
||||
assert_eq!(client_id.conn_id, "client-conn-abc");
|
||||
}
|
||||
TargetClientId::None => {
|
||||
panic!("Expected ClientId, got None");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
//! Session-administration extension handlers.
|
||||
//!
|
||||
//! Methods grouped here are operational/admin endpoints that mutate
|
||||
//! persistent or shared agent state but are not part of the per-turn prompt
|
||||
//! lifecycle:
|
||||
//!
|
||||
//! - `x.ai/session/rename` rename a session locally + remote
|
||||
//! - `x.ai/session/delete` delete a session locally + remote
|
||||
//! - `x.ai/session/update_mcp_servers` mid-session MCP server swap
|
||||
//! - `x.ai/session/fork` fork a session into a new one
|
||||
//! - `x.ai/internal/reload_all_mcp_servers` config hot-reload, all sessions
|
||||
//! - `x.ai/internal/reload_project_mcp_servers` config hot-reload, cwd-scoped
|
||||
//! - `x.ai/internal/reload_skills` skills file watcher fan-out
|
||||
//! - `x.ai/internal/reload_models` model list hot-reload from config.toml
|
||||
//! - `x.ai/internal/reload_models_cache` model catalog hot-reload from disk cache
|
||||
//! - `x.ai/internal/auth_cleared` auth hot-clear cleanup
|
||||
//! - `x.ai/plugins/reload` rebuild shared plugin registry
|
||||
//! - `x.ai/commands/list` list slash commands
|
||||
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use agent_client_protocol::Client as _;
|
||||
use serde::Deserialize;
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::session::persistence::list_summaries;
|
||||
use crate::session::storage::StorageAdapter;
|
||||
use crate::session::storage::jsonl::JsonlStorageAdapter;
|
||||
use crate::session::unified_list::SessionKind;
|
||||
use crate::session::{ExtMethodResult, SessionCommand};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/session/rename" => handle_session_rename(agent, args).await,
|
||||
"x.ai/session/delete" => handle_session_delete(agent, args).await,
|
||||
"x.ai/session/update_mcp_servers" => handle_update_mcp_servers(agent, args).await,
|
||||
"x.ai/session/fork" => handle_session_fork(agent, args).await,
|
||||
"x.ai/internal/reload_all_mcp_servers" => handle_reload_all_mcp_servers(agent).await,
|
||||
"x.ai/internal/reload_project_mcp_servers" => {
|
||||
handle_reload_project_mcp_servers(agent, args).await
|
||||
}
|
||||
"x.ai/internal/reload_skills" => handle_reload_skills(agent),
|
||||
"x.ai/internal/reload_models" => handle_reload_models(agent),
|
||||
"x.ai/internal/reload_models_cache" => handle_reload_models_cache(agent),
|
||||
"x.ai/internal/auth_cleared" => handle_auth_cleared(agent),
|
||||
"x.ai/plugins/reload" => handle_plugins_reload(agent).await,
|
||||
"x.ai/commands/list" => handle_commands_list(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
// session/rename
|
||||
|
||||
/// Handles renaming a session.
|
||||
async fn handle_session_rename(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct RenameRequest {
|
||||
session_id: String,
|
||||
title: String,
|
||||
#[serde(default)]
|
||||
cwd: Option<String>,
|
||||
#[serde(default)]
|
||||
kind: SessionKind,
|
||||
}
|
||||
|
||||
let mut req: RenameRequest = parse_params(args)?;
|
||||
// Manual titles must be non-blank: `Summary.title_is_manual` binds to a
|
||||
// real `generated_title`, so reject whitespace-only input at the boundary.
|
||||
req.title = req.title.trim().to_string();
|
||||
if req.title.is_empty() {
|
||||
return Err(acp::Error::invalid_request().data("title must not be blank"));
|
||||
}
|
||||
|
||||
if req.kind == SessionKind::Chat {
|
||||
return rename_chat_conversation(agent, &req.session_id, &req.title).await;
|
||||
}
|
||||
|
||||
let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));
|
||||
|
||||
// Find the session info, scoping to cwd if provided
|
||||
let summaries = list_summaries(req.cwd.as_deref())
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("failed to list sessions: {e}")))?;
|
||||
|
||||
let summary = summaries
|
||||
.iter()
|
||||
.find(|s| s.info.id == session_id)
|
||||
.ok_or_else(|| {
|
||||
acp::Error::invalid_request().data(format!("session not found: {}", req.session_id))
|
||||
})?;
|
||||
|
||||
let info = summary.info.clone();
|
||||
|
||||
// Update the session title in local storage
|
||||
let storage = JsonlStorageAdapter::default();
|
||||
storage
|
||||
.update_session_title(&info, req.title.clone())
|
||||
.await
|
||||
.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("failed to update session title: {e}"))
|
||||
})?;
|
||||
|
||||
// Update session search index with new title
|
||||
crate::session::storage::search::notify_session_updated(&info.id.to_string(), &info.cwd);
|
||||
|
||||
// Send a SessionSummaryGenerated notification so the TUI updates its title
|
||||
notify_session_title(agent, session_id, &req.title).await;
|
||||
|
||||
if agent.is_writeback_storage()
|
||||
&& let Some(auth) = agent.current_auth()
|
||||
&& !auth.is_zdr_team()
|
||||
{
|
||||
use crate::remote::client::BackendClient;
|
||||
use crate::session::export::ExportedMetadata;
|
||||
|
||||
let mut metadata = ExportedMetadata::from_summary(summary);
|
||||
metadata.title = Some(req.title.clone());
|
||||
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
|
||||
if let Err(e) = BackendClient::new()
|
||||
.with_auth_manager(agent.auth_manager.clone())
|
||||
.save_session_data(&req.session_id, &[], Some(&metadata))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(?e, session_id = %req.session_id, "failed to sync renamed title to backend");
|
||||
}
|
||||
}
|
||||
|
||||
// Hook 2: update session replica with summary (fire-and-forget)
|
||||
if let Some(client) = agent.session_registry_client() {
|
||||
let sid = req.session_id.to_string();
|
||||
let title = if agent
|
||||
.auth_manager
|
||||
.current_or_expired()
|
||||
.is_some_and(|a| a.is_zdr_team())
|
||||
{
|
||||
None
|
||||
} else {
|
||||
Some(req.title.clone())
|
||||
};
|
||||
tokio::spawn(async move {
|
||||
let update = crate::agent::session_registry_client::UpdateRequest {
|
||||
summary: title,
|
||||
first_prompt: None,
|
||||
last_turn_number: None,
|
||||
repo_head_at_end: None,
|
||||
restorable_turn_number: None,
|
||||
};
|
||||
if let Err(e) = client.update(&sid, &update).await {
|
||||
tracing::warn!(error = %e, "session registry summary update failed (non-fatal)");
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
tracing::info!(session_id = %req.session_id, title = %req.title, "Session renamed");
|
||||
|
||||
to_raw_response(&serde_json::json!({ "success": true }))
|
||||
}
|
||||
|
||||
/// Notify connected clients of a session's new title via
|
||||
/// `SessionSummaryGenerated`.
|
||||
async fn notify_session_title(agent: &MvpAgent, session_id: acp::SessionId, title: &str) {
|
||||
use crate::extensions::notification::{SessionNotification, SessionUpdate};
|
||||
|
||||
let notification = SessionNotification {
|
||||
session_id,
|
||||
update: SessionUpdate::SessionSummaryGenerated {
|
||||
session_summary: title.to_owned(),
|
||||
},
|
||||
meta: None,
|
||||
};
|
||||
if let Ok(params) = serde_json::value::to_raw_value(¬ification) {
|
||||
let ext_notification =
|
||||
acp::ExtNotification::new("x.ai/session_notification", params.into());
|
||||
let _ = agent.gateway.ext_notification(ext_notification).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn rename_chat_conversation(
|
||||
agent: &MvpAgent,
|
||||
conversation_id: &str,
|
||||
title: &str,
|
||||
) -> ExtResult {
|
||||
use crate::remote::{ConvError, UpdateConversationBody};
|
||||
|
||||
let Some(client) = agent.conversations_client() else {
|
||||
return Err(acp::Error::invalid_request()
|
||||
.data("chat session rename requires the conversations lane (OIDC + chat feature)"));
|
||||
};
|
||||
|
||||
let body = UpdateConversationBody {
|
||||
title: Some(title.to_owned()),
|
||||
starred: None,
|
||||
};
|
||||
client
|
||||
.update_conversation(conversation_id, &body)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
ConvError::NoOauth => acp::Error::invalid_request()
|
||||
.data("chat session rename requires xAI OAuth credentials"),
|
||||
ConvError::Http { status: 404 } => acp::Error::invalid_request()
|
||||
.data(format!("conversation not found: {conversation_id}")),
|
||||
other => acp::Error::internal_error()
|
||||
.data(format!("chat conversation rename failed: {other}")),
|
||||
})?;
|
||||
|
||||
// If this conversation is open live, notify clients of the new title.
|
||||
let session_id = acp::SessionId::new(Arc::from(conversation_id));
|
||||
if agent.sessions.borrow().contains_key(&session_id) {
|
||||
notify_session_title(agent, session_id, title).await;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
session_id = %conversation_id,
|
||||
title = %title,
|
||||
"Chat conversation renamed"
|
||||
);
|
||||
|
||||
to_raw_response(&serde_json::json!({ "success": true }))
|
||||
}
|
||||
|
||||
// session/delete
|
||||
|
||||
/// Delete a session from history.
|
||||
async fn handle_session_delete(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DeleteRequest {
|
||||
session_id: String,
|
||||
#[serde(default)]
|
||||
cwd: Option<String>,
|
||||
#[serde(default)]
|
||||
kind: SessionKind,
|
||||
}
|
||||
|
||||
let req: DeleteRequest = parse_params(args)?;
|
||||
|
||||
if req.kind == SessionKind::Chat {
|
||||
return soft_delete_chat_conversation(agent, &req.session_id).await;
|
||||
}
|
||||
|
||||
let session_id = acp::SessionId::new(Arc::from(req.session_id.as_str()));
|
||||
|
||||
// For writeback storage (non-ZDR): remote delete is authoritative for
|
||||
// the cloud history and runs first; on failure no local bits are
|
||||
// touched so the pager does not remove the row or toast success.
|
||||
let needs_remote =
|
||||
agent.is_writeback_storage() && agent.current_auth().is_some_and(|a| !a.is_zdr_team());
|
||||
|
||||
// Shared delete: remote-first, then local disk + FTS eviction.
|
||||
// Mirrored by the `grok sessions delete <id>` CLI path.
|
||||
crate::session::persistence::delete_session_history(
|
||||
&req.session_id,
|
||||
req.cwd.as_deref(),
|
||||
needs_remote,
|
||||
agent.auth_manager.clone(),
|
||||
)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
if let crate::session::persistence::DeleteSessionError::Remote(_) = &e {
|
||||
tracing::warn!(?e, session_id = %req.session_id, "failed to delete remote session data");
|
||||
}
|
||||
acp::Error::internal_error().data(e.to_string())
|
||||
})?;
|
||||
|
||||
// If an in-memory live session exists for this id (e.g. the user
|
||||
// deleted history for a session that is still open in another agent
|
||||
// or the current one), shut it down and drop the MvpAgent bookkeeping
|
||||
// so we don't leave a live actor whose on-disk/FTS state is gone.
|
||||
if agent.sessions.borrow().contains_key(&session_id) {
|
||||
agent.request_session_shutdown(&session_id);
|
||||
agent.remove_session(&session_id);
|
||||
}
|
||||
|
||||
tracing::info!(session_id = %req.session_id, "Session deleted");
|
||||
|
||||
to_raw_response(&serde_json::json!({ "success": true }))
|
||||
}
|
||||
|
||||
async fn soft_delete_chat_conversation(agent: &MvpAgent, conversation_id: &str) -> ExtResult {
|
||||
use crate::remote::ConvError;
|
||||
|
||||
let Some(client) = agent.conversations_client() else {
|
||||
return Err(acp::Error::invalid_request()
|
||||
.data("chat session delete requires the conversations lane (OIDC + chat feature)"));
|
||||
};
|
||||
|
||||
client
|
||||
.soft_delete_conversation(conversation_id)
|
||||
.await
|
||||
.map_err(|e| match e {
|
||||
ConvError::NoOauth => acp::Error::invalid_request()
|
||||
.data("chat session delete requires xAI OAuth credentials"),
|
||||
other => acp::Error::internal_error()
|
||||
.data(format!("chat conversation soft-delete failed: {other}")),
|
||||
})?;
|
||||
|
||||
let session_id = acp::SessionId::new(Arc::from(conversation_id));
|
||||
if agent.sessions.borrow().contains_key(&session_id) {
|
||||
agent.request_session_shutdown(&session_id);
|
||||
agent.remove_session(&session_id);
|
||||
}
|
||||
|
||||
tracing::info!(session_id = %conversation_id, "Chat conversation soft-deleted");
|
||||
|
||||
to_raw_response(&serde_json::json!({ "success": true }))
|
||||
}
|
||||
|
||||
// session/update_mcp_servers
|
||||
|
||||
async fn handle_update_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct Params {
|
||||
session_id: acp::SessionId,
|
||||
mcp_servers: Vec<acp::McpServer>,
|
||||
}
|
||||
|
||||
let params: Params = parse_params(args)?;
|
||||
|
||||
let (handle, cwd) = {
|
||||
let sessions = agent.sessions.borrow();
|
||||
let h = sessions
|
||||
.get(¶ms.session_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| acp::Error::invalid_params().data("unknown session id"))?;
|
||||
let cwd = std::path::PathBuf::from(&h.info.cwd);
|
||||
(h, cwd)
|
||||
};
|
||||
|
||||
let managed = agent.get_managed_mcp_configs().await;
|
||||
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
|
||||
params.mcp_servers.clone(),
|
||||
&cwd,
|
||||
&managed,
|
||||
agent.plugin_registry_handle().snapshot().as_deref(),
|
||||
&agent.cfg.borrow().compat_resolved,
|
||||
);
|
||||
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::UpdateMcpServers {
|
||||
mcp_servers: merged,
|
||||
respond_to: tx,
|
||||
})
|
||||
.map_err(|_| acp::Error::internal_error().data("session closed"))?;
|
||||
|
||||
// Wait for the session actor to finish MCP re-initialization.
|
||||
rx.await
|
||||
.map_err(|_| acp::Error::internal_error().data("session closed"))?
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
// Persist the new client set on the handle so config hot-reloads
|
||||
// (`reload_all_mcp_servers` / `reload_project_mcp_servers`) re-merge from
|
||||
// the client's latest intent rather than the `session/new` snapshot —
|
||||
// otherwise a reload would resurrect servers the client just removed
|
||||
// (or drop ones it just added).
|
||||
if let Some(h) = agent.sessions.borrow_mut().get_mut(¶ms.session_id) {
|
||||
h.initial_client_mcp_servers = params.mcp_servers;
|
||||
}
|
||||
|
||||
ExtMethodResult::success(serde_json::json!({ "ok": true }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
// internal/reload_skills
|
||||
|
||||
/// Reload skills for ALL active sessions. Called by the skills file watcher
|
||||
/// (via ACP injection from `app.rs`) when `SKILL.md` files change.
|
||||
fn handle_reload_skills(agent: &MvpAgent) -> ExtResult {
|
||||
let session_ids: Vec<acp::SessionId> = agent.sessions.borrow().keys().cloned().collect();
|
||||
for sid in &session_ids {
|
||||
if let Some(handle) = agent.sessions.borrow().get(sid).cloned() {
|
||||
let _ = handle.cmd_tx.send(SessionCommand::ReloadSkills);
|
||||
}
|
||||
}
|
||||
ExtMethodResult::success(serde_json::json!({ "reloaded": session_ids.len() }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
// internal/reload_all_mcp_servers
|
||||
|
||||
/// Reload MCP servers for ALL active sessions. Called by the config
|
||||
/// hot-reload watcher when `[mcp_servers]` changes in config.toml.
|
||||
async fn handle_reload_all_mcp_servers(agent: &MvpAgent) -> ExtResult {
|
||||
let session_ids: Vec<acp::SessionId> = agent.sessions.borrow().keys().cloned().collect();
|
||||
|
||||
if session_ids.is_empty() {
|
||||
return ExtMethodResult::success(serde_json::json!({ "updated": 0 }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()));
|
||||
}
|
||||
|
||||
let managed = agent.get_managed_mcp_configs().await;
|
||||
let mut updated = 0u32;
|
||||
for session_id in &session_ids {
|
||||
let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
|
||||
continue;
|
||||
};
|
||||
let cwd = std::path::PathBuf::from(&handle.info.cwd);
|
||||
let compat = agent.cfg.borrow().compat_resolved;
|
||||
// Re-seed the merge with the session's original client-provided MCP
|
||||
// servers (e.g. a managed connector injected at `session/new` by a
|
||||
// client session binding). `merge_managed_mcp_servers` already
|
||||
// re-reads every disk source (config.toml, plugins, ~/.claude.json,
|
||||
// ~/.cursor/mcp.json, .mcp.json) internally, so passing
|
||||
// `load_mcp_servers()` output here was redundant — and silently
|
||||
// dropped client servers that exist in no on-disk config, tearing
|
||||
// them down on every config hot-reload.
|
||||
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
|
||||
handle.initial_client_mcp_servers.clone(),
|
||||
&cwd,
|
||||
&managed,
|
||||
agent.plugin_registry_handle().snapshot().as_deref(),
|
||||
&compat,
|
||||
);
|
||||
|
||||
let (tx, _rx) = tokio::sync::oneshot::channel();
|
||||
if handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::UpdateMcpServers {
|
||||
mcp_servers: merged,
|
||||
respond_to: tx,
|
||||
})
|
||||
.is_ok()
|
||||
{
|
||||
updated += 1;
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
updated,
|
||||
total = session_ids.len(),
|
||||
"reloaded MCP servers for active sessions"
|
||||
);
|
||||
ExtMethodResult::success(serde_json::json!({ "updated": updated }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
// internal/reload_project_mcp_servers
|
||||
|
||||
/// Reload MCP servers for sessions whose `cwd` matches (or sits beneath)
|
||||
/// the project root passed in `params.cwd`. Called by the config
|
||||
/// hot-reload watcher when `<cwd>/.kigi/config.toml`,
|
||||
/// `<cwd>/.mcp.json`, or `<cwd>/.claude.json` changes.
|
||||
///
|
||||
/// Sessions in unrelated cwds are intentionally NOT touched — that is
|
||||
/// the whole point of [`crate::config::reloader::ConfigUpdate::
|
||||
/// ProjectMcpServersChanged`] being a per-cwd variant. The legacy
|
||||
/// [`handle_reload_all_mcp_servers`] is still the fan-out for global
|
||||
/// `~/.kigi/config.toml` edits.
|
||||
async fn handle_reload_project_mcp_servers(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
#[derive(Deserialize)]
|
||||
struct Params {
|
||||
cwd: String,
|
||||
}
|
||||
|
||||
let params: Params = parse_params(args)?;
|
||||
let target_cwd = std::path::PathBuf::from(¶ms.cwd);
|
||||
|
||||
// Collect (session_id, cwd) pairs once so we don't hold the
|
||||
// `sessions` RefCell borrow across `.await` points.
|
||||
let session_ids: Vec<(acp::SessionId, std::path::PathBuf)> = agent
|
||||
.sessions
|
||||
.borrow()
|
||||
.iter()
|
||||
.map(|(sid, h)| (sid.clone(), std::path::PathBuf::from(&h.info.cwd)))
|
||||
.filter(|(_, cwd)| cwd_matches(cwd, &target_cwd))
|
||||
.collect();
|
||||
|
||||
if session_ids.is_empty() {
|
||||
return ExtMethodResult::success(serde_json::json!({ "updated": 0 }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()));
|
||||
}
|
||||
|
||||
let managed = agent.get_managed_mcp_configs().await;
|
||||
let mut updated = 0u32;
|
||||
for (session_id, cwd) in &session_ids {
|
||||
let Some(handle) = agent.sessions.borrow().get(session_id).cloned() else {
|
||||
continue;
|
||||
};
|
||||
// See `handle_reload_all_mcp_servers`: seed with the session's
|
||||
// client-provided servers, not `load_mcp_servers()` — the merge
|
||||
// re-reads all disk sources itself, and client-provided servers
|
||||
// (session bindings) must survive config hot-reloads.
|
||||
let merged = crate::session::managed_mcp::merge_managed_mcp_servers(
|
||||
handle.initial_client_mcp_servers.clone(),
|
||||
cwd,
|
||||
&managed,
|
||||
agent.plugin_registry_handle().snapshot().as_deref(),
|
||||
&agent.cfg.borrow().compat_resolved,
|
||||
);
|
||||
|
||||
let (tx, _rx) = tokio::sync::oneshot::channel();
|
||||
if handle
|
||||
.cmd_tx
|
||||
.send(SessionCommand::UpdateMcpServers {
|
||||
mcp_servers: merged,
|
||||
respond_to: tx,
|
||||
})
|
||||
.is_ok()
|
||||
{
|
||||
updated += 1;
|
||||
}
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
updated,
|
||||
total = session_ids.len(),
|
||||
cwd = %target_cwd.display(),
|
||||
"reloaded project MCP servers for matching sessions"
|
||||
);
|
||||
ExtMethodResult::success(serde_json::json!({ "updated": updated }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
/// Returns `true` iff `session_cwd` equals `target_cwd` or sits
|
||||
/// beneath it (so a `<repo>/` edit reloads `<repo>/subdir/` sessions
|
||||
/// too).
|
||||
///
|
||||
/// This uses `Path::starts_with`, which is
|
||||
/// **component-aware** — `/repo-test` does NOT match `/repo` even
|
||||
/// though the byte prefix matches. That is the desired behavior
|
||||
/// (component-aware avoids the `/foo-bar` ⊂ `/foo` foot-gun). Paths
|
||||
/// come from `SessionInfo::cwd` (always absolute) and the watcher's
|
||||
/// emitted path (also absolute), so no canonicalization is needed
|
||||
/// here. The `==` short-circuit is redundant (`Path::starts_with` is
|
||||
/// reflexive) but kept for an explicit zero-allocation fast path.
|
||||
fn cwd_matches(session_cwd: &std::path::Path, target_cwd: &std::path::Path) -> bool {
|
||||
session_cwd == target_cwd || session_cwd.starts_with(target_cwd)
|
||||
}
|
||||
|
||||
// internal/reload_models
|
||||
|
||||
/// Re-resolve the agent model list from config.toml. Called by the config
|
||||
/// hot-reload watcher when `[model.*]` or `[models]` changes.
|
||||
///
|
||||
/// Re-reads config from disk, re-runs the same resolution logic as
|
||||
/// `new_with_models()` for user TOML config entries, and swaps the model list
|
||||
/// in-place. Prefetched (API) and default models are NOT re-fetched -- only
|
||||
/// BYOK entries from config are updated.
|
||||
fn handle_reload_models(agent: &MvpAgent) -> ExtResult {
|
||||
let disk_config = crate::config::load_effective_config()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
let toml_config = crate::agent::config::Config::new_from_toml_cfg(&disk_config)
|
||||
.map_err(|e| acp::Error::internal_error().data(e))?;
|
||||
|
||||
// Merge TOML-derived model fields into the agent's in-memory config so
|
||||
// runtime-only fields (#[serde(skip)]: remote_settings, endpoints, CLI
|
||||
// flags) are preserved. Only model-related TOML fields are refreshed.
|
||||
{
|
||||
let agent_config = agent.cfg.borrow();
|
||||
let overrides = crate::config::ModelOverrideConfig::resolve(
|
||||
agent_config.web_search_model_override.as_deref(),
|
||||
agent_config.session_summary_model_override.as_deref(),
|
||||
&disk_config,
|
||||
agent_config.remote_settings.as_ref(),
|
||||
);
|
||||
drop(agent_config);
|
||||
let mut agent_config = agent.cfg.borrow_mut();
|
||||
agent_config.models = toml_config.models.clone();
|
||||
agent_config.config_models = toml_config.config_models.clone();
|
||||
agent_config.web_search_model = overrides.web_search;
|
||||
agent_config.session_summary_model = overrides.session_summary;
|
||||
agent_config.image_description_model = overrides.image_description;
|
||||
agent_config.prompt_suggest_model_pin = overrides.prompt_suggestion;
|
||||
}
|
||||
// Recompute the campaign overlay + `pre_campaign_default` (the catalog-miss
|
||||
// fallback) so reload matches spawn; `new_from_toml_cfg` reset it to None.
|
||||
{
|
||||
let mut agent_config = agent.cfg.borrow_mut();
|
||||
crate::util::config::sync_campaign_fields(&mut agent_config);
|
||||
}
|
||||
let merged_config = agent.cfg.borrow().clone();
|
||||
|
||||
agent.models_manager.apply_config(merged_config);
|
||||
|
||||
let count = agent.models_manager.models().len();
|
||||
tracing::info!(count, "model list reloaded from config.toml");
|
||||
ExtMethodResult::success(serde_json::json!({ "models": count }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
// internal/reload_models_cache
|
||||
|
||||
/// Hot-reload the model catalog from `~/.kigi/models_cache.json` after an
|
||||
/// external write detected by the config watcher.
|
||||
///
|
||||
/// Routed through the agent's ACP stream (injected by the
|
||||
/// `ConfigUpdate::ModelsCacheChanged` arm in `agent/app.rs`) instead of being
|
||||
/// applied directly on the manager from the config-update task: stream
|
||||
/// requests are processed in order, so when `config.toml` and
|
||||
/// `models_cache.json` change in the same watcher batch this runs strictly
|
||||
/// after `reload_models`' `apply_config` accepted or rejected the new config,
|
||||
/// rather than rebuilding the catalog and notifying clients mid-flight.
|
||||
fn handle_reload_models_cache(agent: &MvpAgent) -> ExtResult {
|
||||
agent.models_manager.reload_from_disk_cache();
|
||||
ExtMethodResult::success(serde_json::json!({ "reloaded": true }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
fn handle_auth_cleared(agent: &MvpAgent) -> ExtResult {
|
||||
agent.disable_managed_gateway_tools_and_refresh_sessions();
|
||||
ExtMethodResult::success(serde_json::json!({ "ok": true }))
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
// plugins/reload
|
||||
|
||||
async fn handle_plugins_reload(agent: &MvpAgent) -> ExtResult {
|
||||
// Rebuild the shared registry so future/new sessions clone the latest.
|
||||
let session_cwd = agent
|
||||
.sessions
|
||||
.borrow()
|
||||
.values()
|
||||
.next()
|
||||
.map(|h| std::path::PathBuf::from(&h.info.cwd));
|
||||
let mut plugins = agent.cfg.borrow().plugins.clone();
|
||||
plugins.merge_claude_enabled_plugins(session_cwd.as_deref());
|
||||
let disk_cfg = plugins.to_discovery_config();
|
||||
// Folder-trust gates repo-local project plugins (hooks/MCP). Resolve and
|
||||
// record the verdict for this cwd (honoring the real remote), then gate
|
||||
// plugins on it.
|
||||
let project_trusted = session_cwd.as_deref().is_some_and(|c| {
|
||||
let remote_settings = agent.cfg.borrow().remote_settings.clone();
|
||||
crate::agent::folder_trust::resolve_and_record(c, remote_settings.as_ref(), false)
|
||||
});
|
||||
// Explicit desktop `x.ai/plugins/reload`: force a full local-install re-copy.
|
||||
agent
|
||||
.plugin_registry_handle()
|
||||
.reload(session_cwd.as_deref(), &disk_cfg, project_trusted, true);
|
||||
|
||||
// Eagerly fan out the new registry to every live session: each adopts a
|
||||
// cwd-correct snapshot (hooks + MCP + skills + client slash-command
|
||||
// catalog), the same refresh the originating session of a reload gets.
|
||||
agent.broadcast_plugin_registry_to_sessions(None);
|
||||
|
||||
super::to_ext_response(Ok(serde_json::json!({"ok": true})))
|
||||
}
|
||||
|
||||
// commands/list
|
||||
|
||||
async fn handle_commands_list(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: crate::session::slash_commands::ListCommandsRequest = parse_params(args)?;
|
||||
|
||||
let skills_config = agent.cfg.borrow().skills.clone();
|
||||
let compat = agent.cfg.borrow().compat_resolved;
|
||||
let availability = agent.command_availability();
|
||||
|
||||
// For a given cwd, compute the plugin registry the same way a session would
|
||||
// at spawn time (via build_for_cwd) and the same way reload_plugins_impl does
|
||||
// (ancestor project config walk + vendor compat merge). This is required so
|
||||
// that `x.ai/commands/list` (the pull used by grok-desktop after session
|
||||
// start) returns plugin-provided slash commands for the target cwd.
|
||||
//
|
||||
// The shared snapshot is only populated at agent boot (using process CWD)
|
||||
// and by explicit reloads. In desktop<->docker (and ssh) setups the agent's
|
||||
// launch CWD is unrelated to the user's chosen workspace dir, so relying on
|
||||
// snapshot() alone meant the post-start pull returned no project plugin
|
||||
// skills until the user manually reloaded.
|
||||
let plugin_reg = if let Some(cwd_str) = &req.cwd {
|
||||
let cwd = Path::new(cwd_str);
|
||||
|
||||
// Folder-trust gates repo-local project plugins (hooks/MCP). Resolve and
|
||||
// record the verdict for this cwd (honoring the real remote) BEFORE the
|
||||
// plugins-config read below: that read gates its project-paths merge on
|
||||
// the recorded verdict, and a cold cwd (client-supplied, no session
|
||||
// resolve yet) must not first take the gate's remote-less backstop —
|
||||
// that would record a kill-switch-blind deny no later resolve can lift.
|
||||
let remote_settings = agent.cfg.borrow().remote_settings.clone();
|
||||
let project_trusted =
|
||||
crate::agent::folder_trust::resolve_and_record(cwd, remote_settings.as_ref(), false);
|
||||
|
||||
// Effective [plugins] config (global + ancestor project configs +
|
||||
// vendor compat merge), shared with reload_plugins_impl and the eager
|
||||
// fan-out so the menu agrees with each session's registry for this cwd.
|
||||
let disk_cfg = crate::config::resolve_effective_plugins_config(cwd).to_discovery_config();
|
||||
|
||||
// Fresh discovery for *this* cwd (includes .kigi/plugins under it, plus
|
||||
// the cli --plugin-dir dirs). Does not mutate the shared snapshot.
|
||||
agent
|
||||
.plugin_registry_handle()
|
||||
.build_for_cwd(cwd, &disk_cfg, &[], project_trusted)
|
||||
} else {
|
||||
// No cwd: global/user skills only (pre-session case). Use the boot snapshot.
|
||||
agent.plugin_registry_handle().snapshot()
|
||||
};
|
||||
|
||||
let response = crate::session::slash_commands::list_commands(
|
||||
req.cwd.as_deref(),
|
||||
&skills_config,
|
||||
plugin_reg.as_deref(),
|
||||
availability,
|
||||
compat,
|
||||
)
|
||||
.await;
|
||||
Ok(acp::ExtResponse::new(Arc::from(
|
||||
serde_json::value::to_raw_value(&response)?,
|
||||
)))
|
||||
}
|
||||
|
||||
// session/fork
|
||||
|
||||
async fn handle_session_fork(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
use crate::session::fork::{ForkSessionRequest, fork_session};
|
||||
|
||||
let request: ForkSessionRequest = parse_params(args)?;
|
||||
|
||||
let agent_id = crate::util::agent_id::agent_id();
|
||||
let response = fork_session(request, &agent_id, Some(agent.auth_manager.clone()))
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
to_raw_response(&response)
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
//! ACP extension handler for session search (`x.ai/session/search`).
|
||||
//!
|
||||
//! Exposes session full-text search as an ACP extension method.
|
||||
//! The client sends a query and receives ranked results across all
|
||||
//! (or workspace-filtered) past sessions.
|
||||
//!
|
||||
//! ```text
|
||||
//! JSON-RPC -> mvp_agent.ext_method()
|
||||
//! -> session_search::handle()
|
||||
//! -> storage::search::execute_search()
|
||||
//! -> search_fts::SessionSearchIndex (SQLite FTS5)
|
||||
//! ```
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::session::storage::search::{SessionSearchRequest, SessionSearchResponse};
|
||||
use crate::session::storage::search_fts::SessionSearchRow;
|
||||
|
||||
use super::ExtResult;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SearchSessionsRequest {
|
||||
/// The search query string.
|
||||
pub query: String,
|
||||
/// Optional workspace directory to scope results to.
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
/// Maximum number of results to return. Defaults to 20.
|
||||
#[serde(default = "default_limit")]
|
||||
pub limit: usize,
|
||||
/// Offset for pagination. Defaults to 0.
|
||||
#[serde(default)]
|
||||
pub offset: usize,
|
||||
/// Whether to include content snippets in results.
|
||||
#[serde(default)]
|
||||
pub include_content: bool,
|
||||
}
|
||||
|
||||
fn default_limit() -> usize {
|
||||
20
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SearchSessionsResponse {
|
||||
pub results: Vec<SearchSessionHit>,
|
||||
pub next_offset: Option<usize>,
|
||||
pub total_estimate: Option<usize>,
|
||||
/// True when the FTS5 index is still being bootstrapped.
|
||||
pub bootstrapping: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SearchSessionHit {
|
||||
pub session_id: String,
|
||||
pub cwd: String,
|
||||
/// Session title/summary for display
|
||||
pub summary: String,
|
||||
/// RFC 3339 formatted updated_at
|
||||
pub updated_at: String,
|
||||
pub score: f32,
|
||||
pub matched_fields: Vec<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub snippet: Option<String>,
|
||||
}
|
||||
|
||||
/// Route `x.ai/session/search` extension method calls.
|
||||
pub async fn handle(args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/session/search" => {
|
||||
let req: SearchSessionsRequest = super::parse_params(args)?;
|
||||
let internal_req = SessionSearchRequest {
|
||||
query: req.query,
|
||||
cwd: req.cwd,
|
||||
limit: req.limit,
|
||||
offset: req.offset,
|
||||
include_content: req.include_content,
|
||||
};
|
||||
|
||||
let root_dir = crate::util::kigi_home::kigi_home();
|
||||
let result = crate::session::storage::search::execute_search(&root_dir, &internal_req)
|
||||
.await
|
||||
.map(to_response)
|
||||
.map_err(|e| anyhow::anyhow!(e));
|
||||
|
||||
super::to_ext_response(result)
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Convert the internal response to the ACP-facing response.
|
||||
fn to_response(resp: SessionSearchResponse) -> SearchSessionsResponse {
|
||||
SearchSessionsResponse {
|
||||
results: resp
|
||||
.results
|
||||
.into_iter()
|
||||
.map(|row: SessionSearchRow| SearchSessionHit {
|
||||
session_id: row.session_id,
|
||||
cwd: row.cwd,
|
||||
summary: row.title,
|
||||
updated_at: chrono::DateTime::from_timestamp(row.updated_at_unix, 0)
|
||||
.map(|dt| dt.to_rfc3339())
|
||||
.unwrap_or_default(),
|
||||
score: row.score,
|
||||
matched_fields: row.matched_fields,
|
||||
snippet: row.snippet,
|
||||
})
|
||||
.collect(),
|
||||
next_offset: resp.next_offset,
|
||||
total_estimate: resp.total_estimate,
|
||||
bootstrapping: resp.bootstrapping,
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,212 @@
|
||||
//! `x.ai/share_session` extension handler.
|
||||
//!
|
||||
//! Loads a local session, exports it, and asks the backend for a public
|
||||
//! share URL.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
use crate::remote::client::BackendClient;
|
||||
use crate::session::export::ExportedSession;
|
||||
use crate::session::info::Info as SessionInfo;
|
||||
use crate::session::persistence::list_summaries;
|
||||
use crate::session::share::{ShareSessionRequest, ShareSessionResponse};
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/share_session" => {
|
||||
tracing::info!("handling share session request");
|
||||
handle_share_session(agent, args).await
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_share_session(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let request: ShareSessionRequest = parse_params(args)?;
|
||||
|
||||
// Get auth - required for sharing.
|
||||
let auth = require_xai_auth_for_share(&agent.auth_manager)?;
|
||||
|
||||
// Remote settings / feature-flag gate: sharing_enabled defaults to false
|
||||
// and is only enabled for eligible accounts.
|
||||
let sharing_enabled = agent
|
||||
.cfg
|
||||
.borrow()
|
||||
.remote_settings
|
||||
.as_ref()
|
||||
.and_then(|rs| rs.sharing_enabled)
|
||||
.unwrap_or(false);
|
||||
if !sharing_enabled {
|
||||
return Err(
|
||||
acp::Error::invalid_params().data("Session sharing is not available for your account.")
|
||||
);
|
||||
}
|
||||
|
||||
// Only block for ZDR teams (hard data-retention policy), not for
|
||||
// coding-data-retention opt-out — sharing is user-initiated.
|
||||
if auth.is_zdr_team() {
|
||||
return Err(acp::Error::invalid_params()
|
||||
.data("Session sharing is disabled for your team's data retention policy"));
|
||||
}
|
||||
|
||||
// Find session info by searching through summaries
|
||||
let summaries = list_summaries(None).await.map_err(|e| {
|
||||
acp::Error::internal_error().data(format!("Failed to list sessions: {}", e))
|
||||
})?;
|
||||
|
||||
let summary = summaries
|
||||
.iter()
|
||||
.find(|s| s.info.id.0.as_ref() == request.session_id.as_str())
|
||||
.ok_or_else(|| acp::Error::resource_not_found(Some("Session not found".into())))?;
|
||||
|
||||
let info = SessionInfo {
|
||||
id: acp::SessionId::new(request.session_id.clone()),
|
||||
cwd: summary.info.cwd.clone(),
|
||||
};
|
||||
|
||||
// Load and export session
|
||||
let exported = ExportedSession::from_local_session(&info)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(format!("Failed to load session: {}", e)))?;
|
||||
|
||||
// Check for empty session
|
||||
if exported.messages.is_empty() {
|
||||
return Err(acp::Error::invalid_params().data("No messages to share yet"));
|
||||
}
|
||||
|
||||
// Upload to backend and get share URL.
|
||||
let client = BackendClient::new().with_auth_manager(agent.auth_manager.clone());
|
||||
let agent_id = crate::util::agent_id::agent_id();
|
||||
let share_url = client
|
||||
.share_session(&exported, &agent_id)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
tracing::error!(error = %e, "Failed to share session with backend");
|
||||
acp::Error::internal_error().data(format!("Failed to share session: {}", e))
|
||||
})?;
|
||||
|
||||
let response = ShareSessionResponse { share_url };
|
||||
to_raw_response(&response)
|
||||
}
|
||||
|
||||
fn require_xai_auth_for_share(
|
||||
auth_manager: &crate::auth::AuthManager,
|
||||
) -> Result<crate::auth::GrokAuth, acp::Error> {
|
||||
super::auth_gate::require_xai_auth(
|
||||
auth_manager,
|
||||
"Authentication required to share session",
|
||||
"Share session is disabled. Run `grok login` to authenticate.",
|
||||
)
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::auth::GrokComConfig;
|
||||
use crate::auth::{AuthMode, GrokAuth};
|
||||
use chrono::{Duration, Utc};
|
||||
use std::sync::Arc;
|
||||
use tempfile::tempdir;
|
||||
|
||||
fn make_auth_manager_with_token_expiring_in(
|
||||
ttl: Duration,
|
||||
) -> (Arc<crate::auth::AuthManager>, tempfile::TempDir) {
|
||||
let dir = tempdir().expect("tempdir for share auth test");
|
||||
let mgr = Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
GrokComConfig::default(),
|
||||
));
|
||||
|
||||
let expires_at = Utc::now() + ttl;
|
||||
|
||||
// We must explicitly set oidc_issuer to a first-party xAI issuer.
|
||||
// Only OIDC tokens against https://auth.x.ai (or the local-dev equivalent)
|
||||
// return true from is_xai_auth(). This is required for the share tests to
|
||||
// exercise the happy path through require_xai_auth_for_share.
|
||||
let auth = GrokAuth {
|
||||
auth_mode: AuthMode::Oidc,
|
||||
oidc_issuer: Some("https://auth.x.ai".to_string()),
|
||||
key: "test-key".into(),
|
||||
expires_at: Some(expires_at),
|
||||
create_time: Utc::now() - Duration::hours(1),
|
||||
..Default::default()
|
||||
};
|
||||
mgr.hot_swap(auth);
|
||||
(mgr, dir)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_works_outside_the_5m_early_invalidation_window() {
|
||||
let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::minutes(10));
|
||||
assert!(mgr.current().is_some());
|
||||
assert!(require_xai_auth_for_share(&mgr).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_succeeds_inside_the_5m_early_invalidation_window() {
|
||||
let (mgr, _dir) = make_auth_manager_with_token_expiring_in(Duration::seconds(1));
|
||||
// This is exactly the state that triggered the user bug:
|
||||
assert!(
|
||||
mgr.current().is_none(),
|
||||
"current() drops the token inside the buffer"
|
||||
);
|
||||
assert!(mgr.expired_auth().is_some());
|
||||
|
||||
// Now that we use current_or_expired(), this passes.
|
||||
let res = require_xai_auth_for_share(&mgr);
|
||||
assert!(
|
||||
res.is_ok(),
|
||||
"require_xai_auth_for_share must succeed for a still-valid buffered xAI token"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_fails_with_no_auth_at_all() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let mgr = Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
GrokComConfig::default(),
|
||||
));
|
||||
assert!(require_xai_auth_for_share(&mgr).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn share_rejects_non_xai_auth_with_actionable_grok_login_message() {
|
||||
let dir = tempdir().expect("tempdir");
|
||||
let mgr = Arc::new(crate::auth::AuthManager::new(
|
||||
dir.path(),
|
||||
GrokComConfig::default(),
|
||||
));
|
||||
|
||||
// API key is the simplest non-xAI credential (External and enterprise OIDC
|
||||
// are also rejected the same way).
|
||||
let non_xai = GrokAuth {
|
||||
auth_mode: AuthMode::ApiKey,
|
||||
key: "xai-test-key".into(),
|
||||
create_time: Utc::now(),
|
||||
..Default::default()
|
||||
};
|
||||
mgr.hot_swap(non_xai);
|
||||
|
||||
let err = require_xai_auth_for_share(&mgr)
|
||||
.expect_err("non-xAI accounts (API key, External, enterprise IdP) must be rejected");
|
||||
|
||||
// This is the key assertion the review asked for: we must test the *exact*
|
||||
// actionable data string for the non-xAI path (distinct from the generic
|
||||
// "Authentication required to share session" path).
|
||||
let serialized =
|
||||
serde_json::to_value(&err).expect("acp::Error serializes to JSON-RPC shape");
|
||||
let data = serialized
|
||||
.get("data")
|
||||
.and_then(|v| v.as_str())
|
||||
.expect("auth_required error carries a data string");
|
||||
|
||||
assert_eq!(
|
||||
data,
|
||||
"Share session is disabled. Run `grok login` to authenticate."
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,630 @@
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::util::config as cli_config;
|
||||
use kigi_agent::prompt::skills::{CompatConfig, SkillInfo, SkillsConfig, list_skills_with_plugins};
|
||||
|
||||
use super::ExtResult;
|
||||
|
||||
/// Generic params for methods that only need an optional `cwd`.
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct CwdParams {
|
||||
#[serde(default)]
|
||||
cwd: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsAddRequest {
|
||||
/// Path to add (directory or SKILL.md file). Supports `~` expansion.
|
||||
pub path: String,
|
||||
/// Working directory for skill discovery context.
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsAddResponse {
|
||||
/// Number of skills discovered at the added path.
|
||||
pub added_count: usize,
|
||||
/// Total number of skills loaded across all sources.
|
||||
pub total: usize,
|
||||
/// The path that was added to config.
|
||||
pub path: String,
|
||||
/// Full updated skill list after reload.
|
||||
pub skills: Vec<SkillInfo>,
|
||||
/// Human-readable message.
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsRemoveRequest {
|
||||
/// Path to remove from config paths.
|
||||
pub path: String,
|
||||
/// Working directory for skill discovery context.
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsRemoveResponse {
|
||||
/// The path that was removed.
|
||||
pub path: String,
|
||||
/// Full updated skill list after reload.
|
||||
pub skills: Vec<SkillInfo>,
|
||||
/// Human-readable message.
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsResetResponse {
|
||||
/// Full updated skill list after reload.
|
||||
pub skills: Vec<SkillInfo>,
|
||||
/// Human-readable message.
|
||||
pub message: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsToggleRequest {
|
||||
/// Skill name to toggle.
|
||||
pub name: String,
|
||||
/// Whether to enable (`true`) or disable (`false`) the skill.
|
||||
pub enabled: bool,
|
||||
/// Working directory for skill discovery context.
|
||||
#[serde(default)]
|
||||
pub cwd: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsListRequest {
|
||||
/// Working directory for skill discovery context.
|
||||
pub cwd: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsListResponse {
|
||||
/// All discovered skills.
|
||||
pub skills: Vec<SkillInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct SkillsConfigResponse {
|
||||
/// Configured paths from `[skills].paths`.
|
||||
pub paths: Vec<String>,
|
||||
/// Ignored paths from `[skills].ignore`.
|
||||
pub ignore: Vec<String>,
|
||||
/// Total loaded skill count.
|
||||
pub total_skills: usize,
|
||||
/// Human-readable summary.
|
||||
pub message: String,
|
||||
/// Full updated skill list.
|
||||
pub skills: Vec<SkillInfo>,
|
||||
}
|
||||
|
||||
/// Reload skills using the current config for the given working directory.
|
||||
#[tracing::instrument(skip_all, fields(cwd))]
|
||||
async fn reload_skills(
|
||||
cwd: &str,
|
||||
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
|
||||
compat: CompatConfig,
|
||||
) -> Vec<SkillInfo> {
|
||||
let config = cli_config::load_config().await.skills;
|
||||
let discovery = list_skills_with_plugins(Some(cwd), &config, plugin_registry, compat);
|
||||
match tokio::time::timeout(std::time::Duration::from_secs(5), discovery).await {
|
||||
Ok(skills) => skills,
|
||||
Err(_) => {
|
||||
tracing::warn!("Skills reload timed out");
|
||||
vec![]
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Count how many skills have paths starting with the given prefix.
|
||||
fn count_skills_from(skills: &[SkillInfo], dir: &std::path::Path) -> usize {
|
||||
let prefix = dir.to_str().unwrap_or("");
|
||||
skills.iter().filter(|s| s.path.starts_with(prefix)).count()
|
||||
}
|
||||
|
||||
/// Resolve a skill path to an absolute path.
|
||||
///
|
||||
/// Handles `~` expansion and relative path resolution against `cwd`.
|
||||
/// Falls back to the original string if canonicalization fails.
|
||||
fn resolve_skill_path(raw: &str, cwd: &str) -> String {
|
||||
use std::path::PathBuf;
|
||||
|
||||
// Expand ~ to $HOME
|
||||
let expanded = if let Some(rest) = raw.strip_prefix("~/") {
|
||||
std::env::var_os("HOME")
|
||||
.or_else(|| std::env::var_os("USERPROFILE"))
|
||||
.map(|home| PathBuf::from(home).join(rest))
|
||||
.unwrap_or_else(|| PathBuf::from(raw))
|
||||
} else if raw == "~" {
|
||||
std::env::var_os("HOME")
|
||||
.or_else(|| std::env::var_os("USERPROFILE"))
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|| PathBuf::from(raw))
|
||||
} else {
|
||||
PathBuf::from(raw)
|
||||
};
|
||||
|
||||
// If already absolute, canonicalize to resolve `..` etc.
|
||||
// If relative, join with cwd first.
|
||||
let absolute = if expanded.is_absolute() {
|
||||
expanded
|
||||
} else {
|
||||
PathBuf::from(cwd).join(&expanded)
|
||||
};
|
||||
|
||||
// canonicalize resolves symlinks and `..` — fall back to the joined path if it fails
|
||||
// (e.g. path doesn't exist yet)
|
||||
dunce::canonicalize(&absolute)
|
||||
.unwrap_or(absolute)
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Collect auto-discovered skill source directories and their counts.
|
||||
fn discover_auto_sources(cwd: &str, skills: &[SkillInfo]) -> Vec<(String, usize)> {
|
||||
let cwd_path = std::path::PathBuf::from(cwd);
|
||||
let kigi_home = kigi_tools::util::kigi_home::kigi_home();
|
||||
let git_root = git2::Repository::discover(&cwd_path)
|
||||
.ok()
|
||||
.and_then(|repo| repo.workdir().map(|p| p.to_path_buf()));
|
||||
|
||||
// Once the user has imported, stop scanning hardcoded
|
||||
// .claude/skills/ paths. Equivalent locations should be opted in via
|
||||
// [paths] extra_skill_dirs in config.toml (written by /import-claude).
|
||||
let imported = crate::claude_import::is_claude_import_marked();
|
||||
let local_dir_names: &[&str] = if imported {
|
||||
&[".kigi", ".agents"]
|
||||
} else {
|
||||
&[".kigi", ".agents", ".claude"]
|
||||
};
|
||||
|
||||
let mut sources: Vec<(String, usize)> = Vec::new();
|
||||
let subdirs = ["skills", "commands"];
|
||||
|
||||
let mut try_add_source = |dir: std::path::PathBuf, seen: Option<&[std::path::PathBuf]>| {
|
||||
if dir.is_dir() && !seen.is_some_and(|s| s.contains(&dir)) {
|
||||
let count = count_skills_from(skills, &dir);
|
||||
if count > 0 {
|
||||
sources.push((dir.to_string_lossy().to_string(), count));
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
let mut local_dirs: Vec<std::path::PathBuf> = Vec::new();
|
||||
for dir_name in local_dir_names {
|
||||
for subdir in &subdirs {
|
||||
let dir = cwd_path.join(dir_name).join(subdir);
|
||||
try_add_source(dir.clone(), None);
|
||||
local_dirs.push(dir);
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(ref root) = git_root {
|
||||
for dir_name in local_dir_names {
|
||||
for subdir in &subdirs {
|
||||
try_add_source(root.join(dir_name).join(subdir), Some(&local_dirs));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for subdir in &subdirs {
|
||||
try_add_source(kigi_home.join(subdir), None);
|
||||
}
|
||||
|
||||
let home = std::env::var_os("HOME").or_else(|| std::env::var_os("USERPROFILE"));
|
||||
if let Some(ref h) = home {
|
||||
let home_path = std::path::PathBuf::from(h);
|
||||
for subdir in &subdirs {
|
||||
try_add_source(home_path.join(".agents").join(subdir), None);
|
||||
}
|
||||
if !imported {
|
||||
for subdir in &subdirs {
|
||||
try_add_source(home_path.join(".claude").join(subdir), None);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// [paths] extra_skill_dirs from config.toml. These supplement the built-in
|
||||
// scan locations. Used both standalone and as the migration target after
|
||||
// /import-claude when the runtime .claude/skills/ scan is disabled.
|
||||
for dir in extra_skill_dirs_from_config() {
|
||||
let path = crate::claude_import::expand_home(&dir);
|
||||
if path.is_dir()
|
||||
&& !sources
|
||||
.iter()
|
||||
.any(|(s, _)| s.as_str() == path.to_string_lossy().as_ref())
|
||||
{
|
||||
sources.push((
|
||||
path.to_string_lossy().to_string(),
|
||||
count_skills_from(skills, &path),
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
sources
|
||||
}
|
||||
|
||||
/// Read `[paths] extra_skill_dirs` from the effective config. Returns empty
|
||||
/// on any read/parse failure so misconfiguration never breaks listing.
|
||||
fn extra_skill_dirs_from_config() -> Vec<String> {
|
||||
let Ok(root) = crate::config::load_effective_config() else {
|
||||
return Vec::new();
|
||||
};
|
||||
root.get("paths")
|
||||
.and_then(|v| v.get("extra_skill_dirs"))
|
||||
.and_then(|v| v.as_array())
|
||||
.map(|arr| {
|
||||
arr.iter()
|
||||
.filter_map(|v| v.as_str().map(|s| s.to_string()))
|
||||
.collect()
|
||||
})
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(
|
||||
args: &acp::ExtRequest,
|
||||
plugin_registry: Option<&kigi_agent::plugins::PluginRegistry>,
|
||||
compat: CompatConfig,
|
||||
) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/skills/add" => {
|
||||
let req: SkillsAddRequest = serde_json::from_str(args.params.get())?;
|
||||
let cwd = req.cwd.as_deref().unwrap_or(".");
|
||||
|
||||
// Resolve to absolute path so config entries work from any cwd.
|
||||
let resolved = resolve_skill_path(&req.path, cwd);
|
||||
|
||||
let p = resolved.clone();
|
||||
if let Err(e) = cli_config::update_config(|cfg| {
|
||||
cfg.skills.ignore.retain(|i| {
|
||||
!(i == &p || p.starts_with(i.as_str()) || i.starts_with(p.as_str()))
|
||||
});
|
||||
if !cfg.skills.paths.contains(&p) {
|
||||
cfg.skills.paths.push(p);
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
return super::to_ext_response(Err::<SkillsAddResponse, _>(anyhow::anyhow!(
|
||||
"Failed to save config: {e}"
|
||||
)));
|
||||
}
|
||||
|
||||
let skills = reload_skills(cwd, plugin_registry, compat).await;
|
||||
let added_count = skills
|
||||
.iter()
|
||||
.filter(|s| s.path.starts_with(&resolved))
|
||||
.count();
|
||||
let total = skills.len();
|
||||
let message = format!(
|
||||
"Added path {}. {} new skill{} found ({} total).",
|
||||
resolved,
|
||||
added_count,
|
||||
if added_count == 1 { "" } else { "s" },
|
||||
total,
|
||||
);
|
||||
|
||||
super::to_ext_response(Ok(SkillsAddResponse {
|
||||
added_count,
|
||||
total,
|
||||
path: resolved,
|
||||
skills,
|
||||
message,
|
||||
}))
|
||||
}
|
||||
|
||||
"x.ai/skills/remove" => {
|
||||
let req: SkillsRemoveRequest = serde_json::from_str(args.params.get())?;
|
||||
let cwd = req.cwd.as_deref().unwrap_or(".");
|
||||
|
||||
// Resolve so relative/tilde paths match what was saved by add.
|
||||
let resolved = resolve_skill_path(&req.path, cwd);
|
||||
|
||||
let p = resolved.clone();
|
||||
if let Err(e) = cli_config::update_config(|cfg| {
|
||||
cfg.skills.paths.retain(|i| i != &p);
|
||||
})
|
||||
.await
|
||||
{
|
||||
return super::to_ext_response(Err::<SkillsRemoveResponse, _>(anyhow::anyhow!(
|
||||
"Failed to save config: {e}"
|
||||
)));
|
||||
}
|
||||
|
||||
let skills = reload_skills(cwd, plugin_registry, compat).await;
|
||||
let total = skills.len();
|
||||
let message = format!(
|
||||
"Removed path {}. {} skill{} remaining.",
|
||||
resolved,
|
||||
total,
|
||||
if total == 1 { "" } else { "s" },
|
||||
);
|
||||
|
||||
super::to_ext_response(Ok(SkillsRemoveResponse {
|
||||
path: resolved,
|
||||
skills,
|
||||
message,
|
||||
}))
|
||||
}
|
||||
|
||||
"x.ai/skills/reset" => {
|
||||
let params: CwdParams =
|
||||
serde_json::from_str(args.params.get()).unwrap_or(CwdParams { cwd: None });
|
||||
let cwd = params.cwd.as_deref().unwrap_or(".");
|
||||
|
||||
if let Err(e) = cli_config::update_config(|cfg| {
|
||||
cfg.skills = SkillsConfig::default();
|
||||
})
|
||||
.await
|
||||
{
|
||||
return super::to_ext_response(Err::<SkillsResetResponse, _>(anyhow::anyhow!(
|
||||
"Failed to save config: {e}"
|
||||
)));
|
||||
}
|
||||
|
||||
let skills = reload_skills(cwd, plugin_registry, compat).await;
|
||||
let message = "Custom skills config reset".to_string();
|
||||
|
||||
super::to_ext_response(Ok(SkillsResetResponse { skills, message }))
|
||||
}
|
||||
|
||||
"x.ai/skills/list" => {
|
||||
let req: SkillsListRequest = serde_json::from_str(args.params.get())?;
|
||||
let skills = reload_skills(&req.cwd, plugin_registry, compat).await;
|
||||
super::to_ext_response(Ok(SkillsListResponse { skills }))
|
||||
}
|
||||
|
||||
"x.ai/skills/config" => {
|
||||
let params: CwdParams =
|
||||
serde_json::from_str(args.params.get()).unwrap_or(CwdParams { cwd: None });
|
||||
let cwd = params.cwd.as_deref().unwrap_or(".");
|
||||
|
||||
let config = cli_config::load_config().await.skills;
|
||||
let paths = config.paths.clone();
|
||||
let ignore = config.ignore.clone();
|
||||
|
||||
let skills = reload_skills(cwd, plugin_registry, compat).await;
|
||||
let total_skills = skills.len();
|
||||
|
||||
let auto_sources = discover_auto_sources(cwd, &skills);
|
||||
|
||||
let mut msg = String::new();
|
||||
|
||||
msg.push_str("Skill discovery sources:\n");
|
||||
for (source, count) in &auto_sources {
|
||||
msg.push_str(&format!(
|
||||
" • {} ({} skill{})\n",
|
||||
source,
|
||||
count,
|
||||
if *count == 1 { "" } else { "s" }
|
||||
));
|
||||
}
|
||||
if auto_sources.is_empty() {
|
||||
msg.push_str(" (no auto-discovered directories found)\n");
|
||||
}
|
||||
|
||||
if !paths.is_empty() {
|
||||
msg.push_str("\nCustom paths:\n");
|
||||
for p in &paths {
|
||||
let count = skills
|
||||
.iter()
|
||||
.filter(|s| s.path.starts_with(p.as_str()))
|
||||
.count();
|
||||
msg.push_str(&format!(
|
||||
" • {} ({} skill{})\n",
|
||||
p,
|
||||
count,
|
||||
if count == 1 { "" } else { "s" }
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
if !ignore.is_empty() {
|
||||
msg.push_str("\nIgnored:\n");
|
||||
for p in &ignore {
|
||||
msg.push_str(&format!(" • {}\n", p));
|
||||
}
|
||||
}
|
||||
|
||||
msg.push_str(&format!("\nTotal skills loaded: {}", total_skills));
|
||||
|
||||
super::to_ext_response(Ok(SkillsConfigResponse {
|
||||
paths,
|
||||
ignore,
|
||||
total_skills,
|
||||
message: msg,
|
||||
skills,
|
||||
}))
|
||||
}
|
||||
|
||||
"x.ai/skills/toggle" => {
|
||||
let req: SkillsToggleRequest = serde_json::from_str(args.params.get())?;
|
||||
let cwd = req.cwd.as_deref().unwrap_or(".");
|
||||
|
||||
// Validate the skill name exists before modifying config.
|
||||
let current_skills = reload_skills(cwd, plugin_registry, compat).await;
|
||||
if !current_skills.iter().any(|s| s.name == req.name) {
|
||||
return super::to_ext_response(Err::<SkillsListResponse, _>(anyhow::anyhow!(
|
||||
"Skill '{}' not found",
|
||||
req.name
|
||||
)));
|
||||
}
|
||||
|
||||
let name = req.name.clone();
|
||||
let enabled = req.enabled;
|
||||
if let Err(e) = cli_config::update_config(|cfg| {
|
||||
if enabled {
|
||||
cfg.skills.disabled.retain(|d| d != &name);
|
||||
} else if !cfg.skills.disabled.contains(&name) {
|
||||
cfg.skills.disabled.push(name.clone());
|
||||
}
|
||||
})
|
||||
.await
|
||||
{
|
||||
return super::to_ext_response(Err::<SkillsListResponse, _>(anyhow::anyhow!(
|
||||
"Failed to save config: {e}"
|
||||
)));
|
||||
}
|
||||
|
||||
// Re-apply disabled marking against the already-loaded skills
|
||||
// to reflect the config change without a second full discovery.
|
||||
let config = cli_config::load_config().await.skills;
|
||||
let disabled_set: std::collections::HashSet<&str> =
|
||||
config.disabled.iter().map(|s| s.as_str()).collect();
|
||||
let skills: Vec<SkillInfo> = current_skills
|
||||
.into_iter()
|
||||
.map(|mut s| {
|
||||
s.enabled = !disabled_set.contains(s.name.as_str());
|
||||
s
|
||||
})
|
||||
.collect();
|
||||
super::to_ext_response(Ok(SkillsListResponse { skills }))
|
||||
}
|
||||
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_add_request_with_cwd() {
|
||||
let json = r#"{"path": "/home/user/skills", "cwd": "/project"}"#;
|
||||
let req: SkillsAddRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.path, "/home/user/skills");
|
||||
assert_eq!(req.cwd, Some("/project".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_request_without_cwd() {
|
||||
let json = r#"{"path": "~/my-skills"}"#;
|
||||
let req: SkillsAddRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.path, "~/my-skills");
|
||||
assert_eq!(req.cwd, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_remove_request() {
|
||||
let json = r#"{"path": "/home/user/skills", "cwd": "/project"}"#;
|
||||
let req: SkillsRemoveRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.path, "/home/user/skills");
|
||||
assert_eq!(req.cwd, Some("/project".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_list_request() {
|
||||
let json = r#"{"cwd": "/project"}"#;
|
||||
let req: SkillsListRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.cwd, "/project");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_add_response_camel_case() {
|
||||
let resp = SkillsAddResponse {
|
||||
added_count: 3,
|
||||
total: 10,
|
||||
path: "/test".to_string(),
|
||||
skills: vec![],
|
||||
message: "ok".to_string(),
|
||||
};
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
assert_eq!(json["addedCount"], 3);
|
||||
assert_eq!(json["total"], 10);
|
||||
assert_eq!(json["path"], "/test");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_absolute_path_unchanged() {
|
||||
let resolved = resolve_skill_path("/absolute/path/to/skills", "/some/cwd");
|
||||
// Canonicalize will fail (path doesn't exist), so we get the joined absolute path
|
||||
assert_eq!(resolved, "/absolute/path/to/skills");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_relative_path_against_cwd() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let sub = tmp.path().join("sub");
|
||||
std::fs::create_dir(&sub).unwrap();
|
||||
|
||||
let resolved = resolve_skill_path("sub", &tmp.path().to_string_lossy());
|
||||
assert_eq!(
|
||||
resolved,
|
||||
dunce::canonicalize(&sub).unwrap().to_string_lossy()
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_resolve_dotdot_path() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let cwd = tmp.path().join("a").join("b");
|
||||
std::fs::create_dir_all(&cwd).unwrap();
|
||||
|
||||
let resolved = resolve_skill_path("../..", &cwd.to_string_lossy());
|
||||
assert_eq!(
|
||||
resolved,
|
||||
dunce::canonicalize(tmp.path()).unwrap().to_string_lossy()
|
||||
);
|
||||
}
|
||||
|
||||
/// Hermetic tilde expansion: pin HOME to a temp dir so remote sandboxes
|
||||
/// (missing HOME, symlink-resolved homes, pre-existing ~/my-skills) cannot
|
||||
/// make `starts_with($HOME)` fail spuriously. Serial because env mutation
|
||||
/// is process-global.
|
||||
#[test]
|
||||
#[serial_test::serial]
|
||||
fn test_resolve_tilde_path() {
|
||||
let tmp = tempfile::tempdir().unwrap();
|
||||
let home = tmp.path().to_path_buf();
|
||||
let prev_home = std::env::var_os("HOME");
|
||||
let prev_userprofile = std::env::var_os("USERPROFILE");
|
||||
// SAFETY: serial test; restored in the same scope below.
|
||||
unsafe {
|
||||
std::env::set_var("HOME", &home);
|
||||
std::env::remove_var("USERPROFILE");
|
||||
}
|
||||
let resolved = resolve_skill_path("~/my-skills", "/ignored");
|
||||
match prev_home {
|
||||
Some(v) => unsafe { std::env::set_var("HOME", v) },
|
||||
None => unsafe { std::env::remove_var("HOME") },
|
||||
}
|
||||
match prev_userprofile {
|
||||
Some(v) => unsafe { std::env::set_var("USERPROFILE", v) },
|
||||
None => unsafe { std::env::remove_var("USERPROFILE") },
|
||||
}
|
||||
let expected = home.join("my-skills");
|
||||
assert_eq!(
|
||||
std::path::PathBuf::from(&resolved),
|
||||
expected,
|
||||
"resolved={resolved}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_config_response_camel_case() {
|
||||
let resp = SkillsConfigResponse {
|
||||
paths: vec!["/a".into()],
|
||||
ignore: vec![],
|
||||
total_skills: 5,
|
||||
message: "ok".to_string(),
|
||||
skills: vec![],
|
||||
};
|
||||
let json = serde_json::to_value(&resp).unwrap();
|
||||
assert_eq!(json["totalSkills"], 5);
|
||||
assert!(json["paths"].is_array());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
use tokio::sync::{mpsc, oneshot};
|
||||
use tokio::time::Duration;
|
||||
|
||||
use super::{RankedSuggestion, SuggestionSource};
|
||||
use crate::session::commands::SessionCommand;
|
||||
|
||||
const AI_TIMEOUT: Duration = Duration::from_secs(2);
|
||||
const AI_PRIORITY: i32 = -10;
|
||||
|
||||
/// Request AI-powered shell command suggestions via the session actor.
|
||||
///
|
||||
/// Sends `SessionCommand::AISuggest` and awaits the response with a 2-second
|
||||
/// timeout. Returns at most one `RankedSuggestion` with `source: AI` and
|
||||
/// `priority: -10` (below history/path results).
|
||||
pub(crate) async fn suggest(
|
||||
cmd_tx: &mpsc::UnboundedSender<SessionCommand>,
|
||||
prefix: &str,
|
||||
cwd: &str,
|
||||
model_override: Option<String>,
|
||||
) -> Vec<RankedSuggestion> {
|
||||
if prefix.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let cmd = SessionCommand::AISuggest {
|
||||
prefix: prefix.to_owned(),
|
||||
cwd: cwd.to_owned(),
|
||||
model_override,
|
||||
respond_to: tx,
|
||||
};
|
||||
|
||||
if cmd_tx.send(cmd).is_err() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let result = match tokio::time::timeout(AI_TIMEOUT, rx).await {
|
||||
Ok(Ok(Some(text))) => text,
|
||||
_ => return Vec::new(),
|
||||
};
|
||||
|
||||
build_suggestion(prefix, &result)
|
||||
}
|
||||
|
||||
fn build_suggestion(prefix: &str, raw: &str) -> Vec<RankedSuggestion> {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.is_empty() || trimmed == prefix {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
// If the model returned the full command (including prefix), use it as-is.
|
||||
// Otherwise concatenate directly — the model output may start with a space
|
||||
// or continuation that should be appended verbatim after the prefix.
|
||||
let insert_text = if trimmed.starts_with(prefix) {
|
||||
trimmed.to_owned()
|
||||
} else if raw.starts_with(prefix) {
|
||||
raw.trim_end().to_owned()
|
||||
} else {
|
||||
format!("{prefix}{raw}").trim_end().to_owned()
|
||||
};
|
||||
|
||||
vec![RankedSuggestion {
|
||||
display: insert_text.clone(),
|
||||
description: String::new(),
|
||||
insert_text,
|
||||
source: SuggestionSource::AI,
|
||||
priority: AI_PRIORITY,
|
||||
// Whole-line; `handle_suggest` stamps the range (no full text here).
|
||||
replace_range: None,
|
||||
token_text: None,
|
||||
truncated: false,
|
||||
is_ghost_candidate: true,
|
||||
}]
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_with_prefix_continuation() {
|
||||
let result = build_suggestion("git", "git commit --amend");
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].insert_text, "git commit --amend");
|
||||
assert_eq!(result[0].source, SuggestionSource::AI);
|
||||
assert_eq!(result[0].priority, -10);
|
||||
assert!(result[0].is_ghost_candidate);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_prepends_prefix_when_missing() {
|
||||
let result = build_suggestion("git", " commit --amend");
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].insert_text, "git commit --amend");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_exact_match_returns_empty() {
|
||||
assert!(build_suggestion("git", "git").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_whitespace_only_returns_empty() {
|
||||
assert!(build_suggestion("git", " \n ").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_empty_returns_empty() {
|
||||
assert!(build_suggestion("git", "").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_no_separator_concatenates_directly() {
|
||||
// Model returned a continuation without leading space — result has no separator.
|
||||
// This is expected: the model should include the space if one is needed.
|
||||
let result = build_suggestion("git", "commit");
|
||||
assert_eq!(result[0].insert_text, "gitcommit");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_raw_starts_with_prefix_preserves_internal_whitespace() {
|
||||
let result = build_suggestion("git", "git commit \n");
|
||||
assert_eq!(result[0].insert_text, "git commit");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_suggestion_trims_surrounding_whitespace() {
|
||||
let result = build_suggestion("git", " git commit \n");
|
||||
assert_eq!(result[0].insert_text, "git commit");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_prefix_skips_channel() {
|
||||
let (tx, _rx) = mpsc::unbounded_channel();
|
||||
let result = suggest(&tx, "", "/tmp", None).await;
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn closed_channel_returns_empty() {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
drop(rx);
|
||||
let result = suggest(&tx, "git", "/tmp", None).await;
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn successful_response() {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Some(SessionCommand::AISuggest { respond_to, .. }) = rx.recv().await {
|
||||
let _ = respond_to.send(Some("git commit --amend".into()));
|
||||
}
|
||||
});
|
||||
|
||||
let result = suggest(&tx, "git", "/tmp", None).await;
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].insert_text, "git commit --amend");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn none_response_returns_empty() {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Some(SessionCommand::AISuggest { respond_to, .. }) = rx.recv().await {
|
||||
let _ = respond_to.send(None);
|
||||
}
|
||||
});
|
||||
|
||||
let result = suggest(&tx, "git", "/tmp", None).await;
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn slow_responder_times_out() {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Some(SessionCommand::AISuggest { respond_to, .. }) = rx.recv().await {
|
||||
// Respond well after the 2-second timeout.
|
||||
tokio::time::sleep(Duration::from_secs(10)).await;
|
||||
let _ = respond_to.send(Some("git commit --amend".into()));
|
||||
}
|
||||
});
|
||||
|
||||
let result = suggest(&tx, "git", "/tmp", None).await;
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sends_correct_fields_to_session() {
|
||||
let (tx, mut rx) = mpsc::unbounded_channel();
|
||||
|
||||
tokio::spawn(async move {
|
||||
if let Some(SessionCommand::AISuggest {
|
||||
prefix,
|
||||
cwd,
|
||||
model_override,
|
||||
respond_to,
|
||||
}) = rx.recv().await
|
||||
{
|
||||
assert_eq!(prefix, "docker");
|
||||
assert_eq!(cwd, "/home/user");
|
||||
assert_eq!(model_override.as_deref(), Some("custom-model"));
|
||||
let _ = respond_to.send(Some("docker compose up".into()));
|
||||
}
|
||||
});
|
||||
|
||||
let result = suggest(&tx, "docker", "/home/user", Some("custom-model".into())).await;
|
||||
assert_eq!(result.len(), 1);
|
||||
assert_eq!(result[0].insert_text, "docker compose up");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,727 @@
|
||||
use std::collections::HashSet;
|
||||
use std::io::{BufRead, BufReader};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
|
||||
use super::{RankedSuggestion, SuggestContext, SuggestionSource, stamp_whole_line_range};
|
||||
use crate::session::prompt_history;
|
||||
|
||||
const CACHE_TTL: Duration = Duration::from_secs(60);
|
||||
const MAX_CROSS_CWD_ENTRIES: usize = 200;
|
||||
const MAX_CROSS_CWD_DIRS: usize = 20;
|
||||
const MAX_SHELL_HISTORY_ENTRIES: usize = 200;
|
||||
const MAX_RESULTS: usize = 10;
|
||||
|
||||
pub(crate) struct HistoryProvider;
|
||||
|
||||
impl HistoryProvider {
|
||||
pub async fn suggest(&self, ctx: &SuggestContext) -> Vec<RankedSuggestion> {
|
||||
let prefix = ctx.prefix();
|
||||
if prefix.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let local = prompt_history::load_bash_prompts_async(ctx.cwd.clone())
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
let shell_history = get_or_refresh_shell_history_cache().await;
|
||||
let cross_cwd = get_or_refresh_cross_cwd_cache().await;
|
||||
|
||||
let mut results =
|
||||
rank_history_matches(prefix, &local, &shell_history.commands, &cross_cwd.prompts);
|
||||
// History carries the full command: it replaces the whole line.
|
||||
stamp_whole_line_range(&mut results, ctx.text.len());
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
/// Rank history matches from three tiers of history sources.
|
||||
///
|
||||
/// Priority order: local grok bash history > shell history > cross-CWD history.
|
||||
fn rank_history_matches(
|
||||
prefix: &str,
|
||||
local: &[String],
|
||||
shell_history: &[String],
|
||||
cross_cwd: &[String],
|
||||
) -> Vec<RankedSuggestion> {
|
||||
if prefix.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let mut seen: HashSet<&str> = HashSet::new();
|
||||
let mut results = Vec::new();
|
||||
|
||||
for prompt in local.iter().chain(shell_history).chain(cross_cwd) {
|
||||
if !prompt.starts_with(prefix) || !seen.insert(prompt.as_str()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
let base_priority = (10i32).saturating_sub(results.len() as i32).max(0);
|
||||
let priority = if *prompt == *prefix {
|
||||
base_priority + 30
|
||||
} else {
|
||||
base_priority
|
||||
};
|
||||
|
||||
let text = prompt.clone();
|
||||
results.push(RankedSuggestion {
|
||||
display: text.clone(),
|
||||
insert_text: text,
|
||||
description: String::new(),
|
||||
source: SuggestionSource::History,
|
||||
priority,
|
||||
is_ghost_candidate: results.is_empty(),
|
||||
replace_range: None,
|
||||
token_text: None,
|
||||
truncated: false,
|
||||
});
|
||||
|
||||
if results.len() >= MAX_RESULTS {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
results
|
||||
}
|
||||
|
||||
// --- Cross-CWD cache ---
|
||||
|
||||
struct CrossCwdCache {
|
||||
prompts: Vec<String>,
|
||||
updated_at: Instant,
|
||||
}
|
||||
|
||||
static CROSS_CWD_CACHE: OnceLock<ArcSwap<CrossCwdCache>> = OnceLock::new();
|
||||
static CROSS_CWD_REFRESHING: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
async fn get_or_refresh_cross_cwd_cache() -> Arc<CrossCwdCache> {
|
||||
let swap = CROSS_CWD_CACHE.get_or_init(|| {
|
||||
ArcSwap::from_pointee(CrossCwdCache {
|
||||
prompts: Vec::new(),
|
||||
updated_at: Instant::now() - CACHE_TTL - Duration::from_secs(1),
|
||||
})
|
||||
});
|
||||
|
||||
let current = swap.load_full();
|
||||
if current.updated_at.elapsed() < CACHE_TTL {
|
||||
return current;
|
||||
}
|
||||
|
||||
if CROSS_CWD_REFRESHING
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_err()
|
||||
{
|
||||
return current;
|
||||
}
|
||||
|
||||
let result = match tokio::task::spawn_blocking(scan_cross_cwd_prompts).await {
|
||||
Ok(prompts) => {
|
||||
let new = Arc::new(CrossCwdCache {
|
||||
prompts,
|
||||
updated_at: Instant::now(),
|
||||
});
|
||||
swap.store(Arc::clone(&new));
|
||||
new
|
||||
}
|
||||
Err(_) => current,
|
||||
};
|
||||
|
||||
CROSS_CWD_REFRESHING.store(false, Ordering::Release);
|
||||
result
|
||||
}
|
||||
|
||||
fn scan_cross_cwd_prompts() -> Vec<String> {
|
||||
let sessions_dir = crate::util::kigi_home::kigi_home().join("sessions");
|
||||
let entries = match std::fs::read_dir(&sessions_dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let mut dirs: Vec<(std::path::PathBuf, std::time::SystemTime)> = entries
|
||||
.filter_map(Result::ok)
|
||||
.filter(|e| e.file_type().map(|ft| ft.is_dir()).unwrap_or(false))
|
||||
.filter_map(|e| {
|
||||
let mtime = e.metadata().ok()?.modified().ok()?;
|
||||
Some((e.path(), mtime))
|
||||
})
|
||||
.collect();
|
||||
|
||||
dirs.sort_by_key(|entry| std::cmp::Reverse(entry.1));
|
||||
|
||||
let mut prompts = Vec::new();
|
||||
for (dir, _) in dirs.iter().take(MAX_CROSS_CWD_DIRS) {
|
||||
if prompts.len() >= MAX_CROSS_CWD_ENTRIES {
|
||||
break;
|
||||
}
|
||||
|
||||
let cwd = match crate::util::kigi_home::decode_cwd_from_dirname(dir) {
|
||||
Some(decoded) => decoded,
|
||||
None => continue,
|
||||
};
|
||||
|
||||
if let Ok(dir_prompts) = prompt_history::load_bash_prompts(&cwd) {
|
||||
let remaining = MAX_CROSS_CWD_ENTRIES - prompts.len();
|
||||
prompts.extend(dir_prompts.into_iter().take(remaining));
|
||||
}
|
||||
}
|
||||
|
||||
prompts
|
||||
}
|
||||
|
||||
// --- Shell history cache ---
|
||||
|
||||
struct ShellHistoryCache {
|
||||
commands: Vec<String>,
|
||||
updated_at: Instant,
|
||||
}
|
||||
|
||||
/// Longer TTL for shell history — the file rarely changes during a session.
|
||||
const SHELL_HISTORY_CACHE_TTL: Duration = Duration::from_secs(300);
|
||||
|
||||
static SHELL_HISTORY_CACHE: OnceLock<ArcSwap<ShellHistoryCache>> = OnceLock::new();
|
||||
static SHELL_HISTORY_REFRESHING: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
async fn get_or_refresh_shell_history_cache() -> Arc<ShellHistoryCache> {
|
||||
let swap = SHELL_HISTORY_CACHE.get_or_init(|| {
|
||||
ArcSwap::from_pointee(ShellHistoryCache {
|
||||
commands: Vec::new(),
|
||||
updated_at: Instant::now() - SHELL_HISTORY_CACHE_TTL - Duration::from_secs(1),
|
||||
})
|
||||
});
|
||||
|
||||
let current = swap.load_full();
|
||||
if current.updated_at.elapsed() < SHELL_HISTORY_CACHE_TTL {
|
||||
return current;
|
||||
}
|
||||
|
||||
if SHELL_HISTORY_REFRESHING
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_err()
|
||||
{
|
||||
return current;
|
||||
}
|
||||
|
||||
let result = match tokio::task::spawn_blocking(load_shell_history).await {
|
||||
Ok(commands) => {
|
||||
let new = Arc::new(ShellHistoryCache {
|
||||
commands,
|
||||
updated_at: Instant::now(),
|
||||
});
|
||||
swap.store(Arc::clone(&new));
|
||||
new
|
||||
}
|
||||
Err(_) => current,
|
||||
};
|
||||
|
||||
SHELL_HISTORY_REFRESHING.store(false, Ordering::Release);
|
||||
result
|
||||
}
|
||||
|
||||
/// Detect the user's shell and load history from the appropriate file.
|
||||
///
|
||||
/// Returns the most recent commands in reverse chronological order, capped
|
||||
/// at [`MAX_SHELL_HISTORY_ENTRIES`].
|
||||
fn load_shell_history() -> Vec<String> {
|
||||
let shell = std::env::var("SHELL").unwrap_or_default();
|
||||
let shell_name = std::path::Path::new(&shell)
|
||||
.file_name()
|
||||
.and_then(|n| n.to_str())
|
||||
.unwrap_or("");
|
||||
|
||||
let home = match home_dir() {
|
||||
Some(h) => h,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
|
||||
// Respect HISTFILE if set, otherwise use shell-specific defaults.
|
||||
let histfile = std::env::var("HISTFILE").ok().map(PathBuf::from);
|
||||
|
||||
match shell_name {
|
||||
"zsh" => {
|
||||
let path = histfile.unwrap_or_else(|| home.join(".zsh_history"));
|
||||
load_zsh_history(&path)
|
||||
}
|
||||
"fish" => load_fish_history(&home.join(".local/share/fish/fish_history")),
|
||||
// Default to bash (covers "bash" and unknown shells)
|
||||
_ => {
|
||||
let path = histfile.unwrap_or_else(|| home.join(".bash_history"));
|
||||
load_bash_history(&path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn home_dir() -> Option<PathBuf> {
|
||||
std::env::var_os("HOME").map(PathBuf::from)
|
||||
}
|
||||
|
||||
/// Keep only the most recent `max` entries, reverse to most-recent-first, and
|
||||
/// deduplicate consecutive identical entries.
|
||||
fn trim_to_recent(commands: &mut Vec<String>, max: usize) {
|
||||
let start = commands.len().saturating_sub(max);
|
||||
commands.drain(..start);
|
||||
commands.reverse();
|
||||
commands.dedup();
|
||||
}
|
||||
|
||||
/// Load bash history (one command per line).
|
||||
///
|
||||
/// Skips `#<timestamp>` lines emitted by `HISTTIMEFORMAT`.
|
||||
fn load_bash_history(path: &std::path::Path) -> Vec<String> {
|
||||
let file = match std::fs::File::open(path) {
|
||||
Ok(f) => f,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let reader = BufReader::new(file);
|
||||
let mut commands = Vec::new();
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = match line {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue,
|
||||
};
|
||||
let trimmed = line.trim();
|
||||
// Skip empty lines and HISTTIMEFORMAT timestamp markers (`#1700000000`)
|
||||
if trimmed.is_empty() || trimmed.starts_with('#') {
|
||||
continue;
|
||||
}
|
||||
commands.push(trimmed.to_owned());
|
||||
}
|
||||
|
||||
trim_to_recent(&mut commands, MAX_SHELL_HISTORY_ENTRIES);
|
||||
commands
|
||||
}
|
||||
|
||||
/// Load zsh history. Lines may be in extended format: `: timestamp:0;command`
|
||||
/// or plain format (one command per line).
|
||||
///
|
||||
/// TODO: zsh represents multiline commands with backslash-newline continuations
|
||||
/// in the history file. Currently each continuation line is treated as a
|
||||
/// separate command, yielding broken fragments for multiline entries.
|
||||
fn load_zsh_history(path: &std::path::Path) -> Vec<String> {
|
||||
let data = match std::fs::read(path) {
|
||||
Ok(d) => d,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
// zsh history may contain invalid UTF-8 (from metafied bytes); use lossy conversion
|
||||
let content = String::from_utf8_lossy(&data);
|
||||
let mut commands = Vec::new();
|
||||
|
||||
for line in content.lines() {
|
||||
let trimmed = line.trim();
|
||||
if trimmed.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Extended history format: `: 1234567890:0;actual command`
|
||||
let cmd = if let Some(rest) = trimmed.strip_prefix(": ") {
|
||||
// Find the `;` separator after the timestamp:duration part
|
||||
rest.find(';').map(|pos| &rest[pos + 1..]).unwrap_or(rest)
|
||||
} else {
|
||||
trimmed
|
||||
};
|
||||
|
||||
if !cmd.is_empty() {
|
||||
commands.push(cmd.to_owned());
|
||||
}
|
||||
}
|
||||
|
||||
trim_to_recent(&mut commands, MAX_SHELL_HISTORY_ENTRIES);
|
||||
commands
|
||||
}
|
||||
|
||||
/// Load fish history. The file uses a YAML-like format:
|
||||
/// ```text
|
||||
/// - cmd: some command
|
||||
/// when: 1234567890
|
||||
/// - cmd: another command
|
||||
/// when: 1234567891
|
||||
/// ```
|
||||
fn load_fish_history(path: &std::path::Path) -> Vec<String> {
|
||||
let file = match std::fs::File::open(path) {
|
||||
Ok(f) => f,
|
||||
Err(_) => return Vec::new(),
|
||||
};
|
||||
|
||||
let reader = BufReader::new(file);
|
||||
let mut commands = Vec::new();
|
||||
|
||||
for line in reader.lines() {
|
||||
let line = match line {
|
||||
Ok(l) => l,
|
||||
Err(_) => continue,
|
||||
};
|
||||
// Fish history entries start with "- cmd: "
|
||||
if let Some(cmd) = line.strip_prefix("- cmd: ") {
|
||||
let cmd = cmd.trim();
|
||||
if !cmd.is_empty() {
|
||||
commands.push(cmd.to_owned());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trim_to_recent(&mut commands, MAX_SHELL_HISTORY_ENTRIES);
|
||||
commands
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
use tempfile::NamedTempFile;
|
||||
|
||||
#[test]
|
||||
fn exact_prefix_match_gets_bonus() {
|
||||
let local = vec!["git commit".into(), "git checkout".into()];
|
||||
let results = rank_history_matches("git commit", &local, &[], &[]);
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].insert_text, "git commit");
|
||||
assert!(
|
||||
results[0].priority >= 30,
|
||||
"exact match priority {} should be >= 30",
|
||||
results[0].priority
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn partial_matches_decay_by_position() {
|
||||
let local = vec![
|
||||
"git commit -m fix".into(),
|
||||
"git checkout main".into(),
|
||||
"git cherry-pick abc".into(),
|
||||
];
|
||||
let results = rank_history_matches("git c", &local, &[], &[]);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(results[0].priority, 10);
|
||||
assert_eq!(results[1].priority, 9);
|
||||
assert_eq!(results[2].priority, 8);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn priorities_decrease_per_position() {
|
||||
let local: Vec<String> = (0..5).map(|i| format!("test_{i}")).collect();
|
||||
let results = rank_history_matches("test", &local, &[], &[]);
|
||||
assert_eq!(results.len(), 5);
|
||||
let priorities: Vec<i32> = results.iter().map(|r| r.priority).collect();
|
||||
assert_eq!(priorities, &[10, 9, 8, 7, 6]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn deduplicates_across_local_and_cross_cwd() {
|
||||
let local = vec!["git push".into(), "git pull".into()];
|
||||
let cross = vec!["git pull".into(), "git fetch".into()];
|
||||
let results = rank_history_matches("git p", &local, &[], &cross);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].insert_text, "git push");
|
||||
assert_eq!(results[1].insert_text, "git pull");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn local_takes_priority_over_cross_cwd() {
|
||||
let local = vec!["git push origin main".into()];
|
||||
let cross = vec!["git push origin dev".into()];
|
||||
let results = rank_history_matches("git push", &local, &[], &cross);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].insert_text, "git push origin main");
|
||||
assert!(results[0].priority > results[1].priority);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn first_match_is_ghost_candidate() {
|
||||
let local = vec!["ls -la".into(), "ls -lh".into()];
|
||||
let results = rank_history_matches("ls", &local, &[], &[]);
|
||||
assert!(results[0].is_ghost_candidate);
|
||||
assert!(!results[1].is_ghost_candidate);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_prefix_returns_nothing() {
|
||||
let local = vec!["git commit".into()];
|
||||
assert!(rank_history_matches("", &local, &[], &[]).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_matching_prefix_returns_empty() {
|
||||
let local = vec!["git commit".into()];
|
||||
assert!(rank_history_matches("docker", &local, &[], &[]).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn caps_at_max_results() {
|
||||
let local: Vec<String> = (0..20).map(|i| format!("test_cmd_{i}")).collect();
|
||||
let results = rank_history_matches("test", &local, &[], &[]);
|
||||
assert_eq!(results.len(), MAX_RESULTS);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn cross_cwd_duplicates_filtered() {
|
||||
let local = vec!["make build".into()];
|
||||
let cross = vec!["make build".into(), "make test".into()];
|
||||
let results = rank_history_matches("make", &local, &[], &cross);
|
||||
let texts: Vec<&str> = results.iter().map(|r| r.insert_text.as_str()).collect();
|
||||
assert_eq!(texts, &["make build", "make test"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn exact_match_among_partial_matches() {
|
||||
let local = vec![
|
||||
"cargo build --release".into(),
|
||||
"cargo build".into(),
|
||||
"cargo bench".into(),
|
||||
];
|
||||
let results = rank_history_matches("cargo build", &local, &[], &[]);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].priority, 10);
|
||||
assert_eq!(results[1].priority, 9 + 30);
|
||||
assert!(results[1].priority > results[0].priority);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_local_and_cross_cwd() {
|
||||
assert!(rank_history_matches("git", &[], &[], &[]).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn single_char_prefix() {
|
||||
let local = vec!["git commit".into(), "grep foo".into(), "ls".into()];
|
||||
let results = rank_history_matches("g", &local, &[], &[]);
|
||||
assert_eq!(results.len(), 2);
|
||||
assert_eq!(results[0].insert_text, "git commit");
|
||||
assert_eq!(results[1].insert_text, "grep foo");
|
||||
}
|
||||
|
||||
// --- Shell history priority ordering ---
|
||||
|
||||
#[test]
|
||||
fn shell_history_ranked_between_local_and_cross_cwd() {
|
||||
let local = vec!["git push origin main".into()];
|
||||
let shell = vec!["git push origin staging".into()];
|
||||
let cross = vec!["git push origin dev".into()];
|
||||
let results = rank_history_matches("git push", &local, &shell, &cross);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert_eq!(results[0].insert_text, "git push origin main");
|
||||
assert_eq!(results[1].insert_text, "git push origin staging");
|
||||
assert_eq!(results[2].insert_text, "git push origin dev");
|
||||
// Priority decreases: local > shell > cross
|
||||
assert!(results[0].priority > results[1].priority);
|
||||
assert!(results[1].priority > results[2].priority);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn shell_history_deduplicates_with_local() {
|
||||
let local = vec!["ls -la".into()];
|
||||
let shell = vec!["ls -la".into(), "ls -lh".into()];
|
||||
let results = rank_history_matches("ls", &local, &shell, &[]);
|
||||
let texts: Vec<&str> = results.iter().map(|r| r.insert_text.as_str()).collect();
|
||||
assert_eq!(texts, &["ls -la", "ls -lh"]);
|
||||
}
|
||||
|
||||
// --- Bash history parsing ---
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_basic() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "ls -la").unwrap();
|
||||
writeln!(f, "cd /tmp").unwrap();
|
||||
writeln!(f, "echo hello").unwrap();
|
||||
let commands = load_bash_history(f.path());
|
||||
// Reverse chrono order
|
||||
assert_eq!(commands, &["echo hello", "cd /tmp", "ls -la"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_skips_empty_lines() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "ls").unwrap();
|
||||
writeln!(f).unwrap();
|
||||
writeln!(f, " ").unwrap();
|
||||
writeln!(f, "pwd").unwrap();
|
||||
let commands = load_bash_history(f.path());
|
||||
assert_eq!(commands, &["pwd", "ls"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_caps_at_limit() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
for i in 0..300 {
|
||||
writeln!(f, "cmd_{i}").unwrap();
|
||||
}
|
||||
let commands = load_bash_history(f.path());
|
||||
assert_eq!(commands.len(), MAX_SHELL_HISTORY_ENTRIES);
|
||||
// Most recent first
|
||||
assert_eq!(commands[0], "cmd_299");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_deduplicates() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "ls").unwrap();
|
||||
writeln!(f, "pwd").unwrap();
|
||||
writeln!(f, "ls").unwrap();
|
||||
let commands = load_bash_history(f.path());
|
||||
// After reverse + dedup: ["ls", "pwd", "ls"] -> reversed = ["ls", "pwd", "ls"]
|
||||
// dedup removes consecutive dupes only. "ls", "pwd", "ls" has no consecutive dupes.
|
||||
assert_eq!(commands, &["ls", "pwd", "ls"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_skips_timestamp_markers() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "#1700000000").unwrap();
|
||||
writeln!(f, "ls -la").unwrap();
|
||||
writeln!(f, "#1700000001").unwrap();
|
||||
writeln!(f, "cd /tmp").unwrap();
|
||||
let commands = load_bash_history(f.path());
|
||||
assert_eq!(commands, &["cd /tmp", "ls -la"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_comment_only_lines() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "#1700000000").unwrap();
|
||||
writeln!(f, "# this is also a comment").unwrap();
|
||||
writeln!(f, "#not-a-timestamp-but-still-skipped").unwrap();
|
||||
writeln!(f, "echo hello").unwrap();
|
||||
let commands = load_bash_history(f.path());
|
||||
assert_eq!(commands, &["echo hello"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_bash_history_missing_file() {
|
||||
let commands = load_bash_history(std::path::Path::new("/nonexistent/.bash_history"));
|
||||
assert!(commands.is_empty());
|
||||
}
|
||||
|
||||
// --- Zsh history parsing ---
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_extended_format() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, ": 1700000000:0;ls -la").unwrap();
|
||||
writeln!(f, ": 1700000001:0;cd /tmp").unwrap();
|
||||
writeln!(f, ": 1700000002:0;echo hello world").unwrap();
|
||||
let commands = load_zsh_history(f.path());
|
||||
assert_eq!(commands, &["echo hello world", "cd /tmp", "ls -la"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_plain_format() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "ls -la").unwrap();
|
||||
writeln!(f, "cd /tmp").unwrap();
|
||||
let commands = load_zsh_history(f.path());
|
||||
assert_eq!(commands, &["cd /tmp", "ls -la"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_mixed_format() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "plain command").unwrap();
|
||||
writeln!(f, ": 1700000000:0;extended command").unwrap();
|
||||
let commands = load_zsh_history(f.path());
|
||||
assert_eq!(commands, &["extended command", "plain command"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_skips_empty_lines() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, ": 1700000000:0;ls").unwrap();
|
||||
writeln!(f).unwrap();
|
||||
writeln!(f, ": 1700000001:0;pwd").unwrap();
|
||||
let commands = load_zsh_history(f.path());
|
||||
assert_eq!(commands, &["pwd", "ls"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_caps_at_limit() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
for i in 0..300 {
|
||||
writeln!(f, ": {i}:0;cmd_{i}").unwrap();
|
||||
}
|
||||
let commands = load_zsh_history(f.path());
|
||||
assert_eq!(commands.len(), MAX_SHELL_HISTORY_ENTRIES);
|
||||
assert_eq!(commands[0], "cmd_299");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_missing_file() {
|
||||
let commands = load_zsh_history(std::path::Path::new("/nonexistent/.zsh_history"));
|
||||
assert!(commands.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_zsh_history_extended_with_semicolon_in_command() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, ": 1700000000:0;echo foo; echo bar").unwrap();
|
||||
let commands = load_zsh_history(f.path());
|
||||
// The command includes everything after the first `;`
|
||||
assert_eq!(commands, &["echo foo; echo bar"]);
|
||||
}
|
||||
|
||||
// --- Fish history parsing ---
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_basic() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "- cmd: ls -la").unwrap();
|
||||
writeln!(f, " when: 1700000000").unwrap();
|
||||
writeln!(f, "- cmd: cd /tmp").unwrap();
|
||||
writeln!(f, " when: 1700000001").unwrap();
|
||||
let commands = load_fish_history(f.path());
|
||||
assert_eq!(commands, &["cd /tmp", "ls -la"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_no_when() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "- cmd: ls").unwrap();
|
||||
writeln!(f, "- cmd: pwd").unwrap();
|
||||
let commands = load_fish_history(f.path());
|
||||
assert_eq!(commands, &["pwd", "ls"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_skips_non_cmd_lines() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "- cmd: ls").unwrap();
|
||||
writeln!(f, " when: 12345").unwrap();
|
||||
writeln!(f, " paths:").unwrap();
|
||||
writeln!(f, " - /some/path").unwrap();
|
||||
writeln!(f, "- cmd: pwd").unwrap();
|
||||
let commands = load_fish_history(f.path());
|
||||
assert_eq!(commands, &["pwd", "ls"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_caps_at_limit() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
for i in 0..300 {
|
||||
writeln!(f, "- cmd: cmd_{i}").unwrap();
|
||||
writeln!(f, " when: {i}").unwrap();
|
||||
}
|
||||
let commands = load_fish_history(f.path());
|
||||
assert_eq!(commands.len(), MAX_SHELL_HISTORY_ENTRIES);
|
||||
assert_eq!(commands[0], "cmd_299");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_missing_file() {
|
||||
let commands = load_fish_history(std::path::Path::new("/nonexistent/fish/fish_history"));
|
||||
assert!(commands.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_fish_history_empty_cmd_skipped() {
|
||||
let mut f = NamedTempFile::new().unwrap();
|
||||
writeln!(f, "- cmd: ").unwrap();
|
||||
writeln!(f, "- cmd: ls").unwrap();
|
||||
let commands = load_fish_history(f.path());
|
||||
assert_eq!(commands, &["ls"]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,680 @@
|
||||
mod ai_provider;
|
||||
mod file_provider;
|
||||
mod history_provider;
|
||||
mod path_provider;
|
||||
mod shell_token;
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use super::{ExtResult, parse_params, to_raw_response};
|
||||
use crate::agent::MvpAgent;
|
||||
|
||||
pub(crate) use file_provider::FilePathProvider;
|
||||
pub(crate) use history_provider::HistoryProvider;
|
||||
pub(crate) use path_provider::PathProvider;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SuggestRequest {
|
||||
text: String,
|
||||
cursor: usize,
|
||||
cwd: String,
|
||||
limit: usize,
|
||||
generation: u64,
|
||||
#[serde(default)]
|
||||
include_ai: bool,
|
||||
#[serde(default)]
|
||||
ai_model: Option<String>,
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
/// Deterministic Tab mode: run only the token providers (path/file).
|
||||
/// A history/AI row would make the set mixed — killing the pager's
|
||||
/// insta-accept/LCP semantics — and reparse history per keystroke.
|
||||
#[serde(default)]
|
||||
token_only: bool,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SuggestResponse {
|
||||
ghost: Option<GhostSuggestion>,
|
||||
completions: Vec<CompletionItem>,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct GhostSuggestion {
|
||||
full_text: String,
|
||||
suffix: String,
|
||||
source: String,
|
||||
}
|
||||
|
||||
/// One completion row. Wire-compat contract (leader mode and the cloud
|
||||
/// bridge mix shell/pager versions):
|
||||
/// - `insert_text` is ALWAYS a safe whole-line replacement — range-unaware
|
||||
/// pagers `set_text` it, so it must never be a bare token.
|
||||
/// - `replace_range` + `token_text` are the additive token-in-place upgrade:
|
||||
/// byte offsets `[start, end)` into the request `text` and the text that
|
||||
/// replaces that span. Range-aware pagers use them as an ATOMIC pair —
|
||||
/// a range without `token_text` (history/AI whole-line rows, where
|
||||
/// `insert_text` doubles as the span replacement) degrades to the
|
||||
/// equivalent whole-line accept.
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct CompletionItem {
|
||||
display: String,
|
||||
description: String,
|
||||
insert_text: String,
|
||||
source: String,
|
||||
priority: i32,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
replace_range: Option<(usize, usize)>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
token_text: Option<String>,
|
||||
/// The provider capped its scan/result set: the row set may be
|
||||
/// incomplete, so range-aware pagers keep dropdown-only semantics
|
||||
/// (absent = `false` for older shells).
|
||||
#[serde(skip_serializing_if = "std::ops::Not::not")]
|
||||
truncated: bool,
|
||||
}
|
||||
|
||||
pub(crate) struct SuggestContext {
|
||||
text: String,
|
||||
cursor: usize,
|
||||
cwd: String,
|
||||
}
|
||||
|
||||
impl SuggestContext {
|
||||
fn new(text: String, cursor: usize, cwd: String) -> Self {
|
||||
let mut cursor = cursor.min(text.len());
|
||||
while cursor > 0 && !text.is_char_boundary(cursor) {
|
||||
cursor -= 1;
|
||||
}
|
||||
Self { text, cursor, cwd }
|
||||
}
|
||||
|
||||
pub(crate) fn prefix(&self) -> &str {
|
||||
&self.text[..self.cursor]
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub(crate) struct RankedSuggestion {
|
||||
pub(crate) display: String,
|
||||
pub(crate) description: String,
|
||||
/// Whole-line replacement (see the [`CompletionItem`] compat contract).
|
||||
pub(crate) insert_text: String,
|
||||
pub(crate) source: SuggestionSource,
|
||||
pub(crate) priority: i32,
|
||||
pub(crate) is_ghost_candidate: bool,
|
||||
/// Request-text byte range the completion targets (token for path/file,
|
||||
/// whole line for history/AI); `None` keeps whole-line-only semantics.
|
||||
pub(crate) replace_range: Option<(usize, usize)>,
|
||||
/// Replacement for `replace_range` when it differs from `insert_text`.
|
||||
pub(crate) token_text: Option<String>,
|
||||
/// Provider capped its scan/results — a hidden row could disprove a
|
||||
/// sole match or an LCP.
|
||||
pub(crate) truncated: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(crate) enum SuggestionSource {
|
||||
History,
|
||||
Path,
|
||||
File,
|
||||
AI,
|
||||
}
|
||||
|
||||
impl SuggestionSource {
|
||||
fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::History => "history",
|
||||
Self::Path => "path",
|
||||
Self::File => "file",
|
||||
Self::AI => "ai",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<RankedSuggestion> for CompletionItem {
|
||||
fn from(s: RankedSuggestion) -> Self {
|
||||
Self {
|
||||
display: s.display,
|
||||
description: s.description,
|
||||
insert_text: s.insert_text,
|
||||
source: s.source.as_str().to_owned(),
|
||||
priority: s.priority,
|
||||
replace_range: s.replace_range,
|
||||
token_text: s.token_text,
|
||||
truncated: s.truncated,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark whole-line suggestions (history/AI carry the full command as
|
||||
/// `insert_text`) as replacing the entire request text.
|
||||
fn stamp_whole_line_range(results: &mut [RankedSuggestion], text_len: usize) {
|
||||
results
|
||||
.iter_mut()
|
||||
.for_each(|s| s.replace_range = Some((0, text_len)));
|
||||
}
|
||||
|
||||
/// Convert token-valued suggestions (path/file build `insert_text` as the
|
||||
/// token replacing `range`) into the wire pair: the token moves to
|
||||
/// `token_text` and `insert_text` becomes the full line with the token
|
||||
/// spliced in — the shape range-unaware pagers can safely `set_text`.
|
||||
fn splice_token_into_line(results: &mut [RankedSuggestion], text: &str, range: (usize, usize)) {
|
||||
for s in results {
|
||||
let token = std::mem::take(&mut s.insert_text);
|
||||
s.insert_text = format!("{}{}{}", &text[..range.0], token, &text[range.1..]);
|
||||
s.token_text = Some(token);
|
||||
}
|
||||
}
|
||||
|
||||
#[tracing::instrument(skip_all, fields(method = %args.method))]
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/suggest" => handle_suggest(agent, args).await,
|
||||
"x.ai/suggestPrompt" => handle_suggest_prompt(agent, args).await,
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Request/response for `x.ai/suggestPrompt` — predict the user's likely next
|
||||
/// prompt after a completed turn (tab-autocomplete ghost text).
|
||||
#[derive(Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SuggestPromptRequest {
|
||||
/// Client-side generation counter, echoed back so stale responses can be
|
||||
/// discarded (the client may have started a newer turn meanwhile).
|
||||
generation: u64,
|
||||
#[serde(default)]
|
||||
session_id: Option<String>,
|
||||
/// Client hint for the suggestion model (the pager sends its env
|
||||
/// override, or `grok-build-0.1` when its catalog offers it). One tier
|
||||
/// of the shell-side resolution in
|
||||
/// `prompt_suggest::effective_suggest_model`: env > config.toml > remote
|
||||
/// > this hint > `grok-build-0.1` default, catalog-guarded (a
|
||||
/// non-sampleable effective model skips the request; the session model
|
||||
/// is never used).
|
||||
#[serde(default)]
|
||||
model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SuggestPromptResponse {
|
||||
suggestion: Option<String>,
|
||||
generation: u64,
|
||||
}
|
||||
|
||||
/// Upper bound on the suggestion round-trip. Turn-end prediction is not
|
||||
/// latency-critical (the user is reading the agent's reply — the idle window
|
||||
/// after a turn is typically long), but a hung call must not pin the oneshot
|
||||
/// forever. Reasoning models (e.g. `grok-build`) can take ~30s on a cold
|
||||
/// cache; a late suggestion is still useful (the pager's generation guard
|
||||
/// and empty-prompt gating discard it if the user moved on).
|
||||
const SUGGEST_PROMPT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(45);
|
||||
|
||||
async fn handle_suggest_prompt(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: SuggestPromptRequest = parse_params(args)?;
|
||||
let generation = req.generation;
|
||||
|
||||
let suggestion = match find_session(agent, req.session_id.as_deref()) {
|
||||
Some(handle) => {
|
||||
let (tx, rx) = tokio::sync::oneshot::channel();
|
||||
let cmd = crate::session::commands::SessionCommand::SuggestPrompt {
|
||||
model_override: req.model,
|
||||
respond_to: tx,
|
||||
};
|
||||
if handle.cmd_tx.send(cmd).is_err() {
|
||||
tracing::debug!("suggestPrompt: session command channel closed");
|
||||
None
|
||||
} else {
|
||||
match tokio::time::timeout(SUGGEST_PROMPT_TIMEOUT, rx).await {
|
||||
Ok(Ok(suggestion)) => suggestion,
|
||||
Ok(Err(_)) => {
|
||||
tracing::debug!("suggestPrompt: responder dropped");
|
||||
None
|
||||
}
|
||||
Err(_) => {
|
||||
tracing::debug!("suggestPrompt: timed out");
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
None => {
|
||||
tracing::debug!(session_id = ?req.session_id, "suggestPrompt: session not found");
|
||||
None
|
||||
}
|
||||
};
|
||||
|
||||
to_raw_response(&SuggestPromptResponse {
|
||||
suggestion,
|
||||
generation,
|
||||
})
|
||||
}
|
||||
|
||||
async fn handle_suggest(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
let req: SuggestRequest = parse_params(args)?;
|
||||
let limit = req.limit;
|
||||
let generation = req.generation;
|
||||
let include_ai = req.include_ai;
|
||||
let token_only = req.token_only;
|
||||
let SuggestRequest {
|
||||
text,
|
||||
cursor,
|
||||
cwd,
|
||||
ai_model,
|
||||
session_id,
|
||||
..
|
||||
} = req;
|
||||
let ctx = SuggestContext::new(text, cursor, cwd);
|
||||
|
||||
let (history_results, path_results, file_results) = tokio::join!(
|
||||
async {
|
||||
if token_only {
|
||||
Vec::new()
|
||||
} else {
|
||||
HistoryProvider.suggest(&ctx).await
|
||||
}
|
||||
},
|
||||
PathProvider.suggest(&ctx),
|
||||
FilePathProvider.suggest(&ctx),
|
||||
);
|
||||
|
||||
let mut ai_results =
|
||||
if include_ai && !token_only && !should_skip_ai(&history_results, ctx.prefix()) {
|
||||
let session = find_session(agent, session_id.as_deref());
|
||||
match session {
|
||||
Some(handle) => {
|
||||
ai_provider::suggest(&handle.cmd_tx, ctx.prefix(), &ctx.cwd, ai_model).await
|
||||
}
|
||||
None => Vec::new(),
|
||||
}
|
||||
} else {
|
||||
Vec::new()
|
||||
};
|
||||
stamp_whole_line_range(&mut ai_results, ctx.text.len());
|
||||
|
||||
let (ghost, completions) = aggregate(
|
||||
history_results,
|
||||
path_results,
|
||||
file_results,
|
||||
ai_results,
|
||||
ctx.prefix(),
|
||||
limit,
|
||||
);
|
||||
|
||||
to_raw_response(&SuggestResponse {
|
||||
ghost,
|
||||
completions,
|
||||
generation,
|
||||
})
|
||||
}
|
||||
|
||||
fn find_session(
|
||||
agent: &MvpAgent,
|
||||
session_id: Option<&str>,
|
||||
) -> Option<crate::session::handle::SessionHandle> {
|
||||
let sessions = agent.sessions.borrow();
|
||||
if let Some(id) = session_id {
|
||||
sessions.get(&acp::SessionId::new(id)).cloned()
|
||||
} else {
|
||||
sessions.values().next().cloned()
|
||||
}
|
||||
}
|
||||
|
||||
fn aggregate(
|
||||
history: Vec<RankedSuggestion>,
|
||||
path: Vec<RankedSuggestion>,
|
||||
file: Vec<RankedSuggestion>,
|
||||
ai: Vec<RankedSuggestion>,
|
||||
prefix: &str,
|
||||
limit: usize,
|
||||
) -> (Option<GhostSuggestion>, Vec<CompletionItem>) {
|
||||
let mut all: Vec<RankedSuggestion> = history
|
||||
.into_iter()
|
||||
.chain(path)
|
||||
.chain(file)
|
||||
.chain(ai)
|
||||
.collect();
|
||||
// STABLE sort — load-bearing: providers pre-rank their items and ship
|
||||
// them at one shared priority per response (the file provider's fuzzy
|
||||
// tier/score/dirs-first order — see its `FILE_CMD_BOOST` doc), relying
|
||||
// on equal-priority order surviving to the wire. Do not "optimize"
|
||||
// into `sort_unstable_by`.
|
||||
all.sort_by_key(|s| std::cmp::Reverse(s.priority));
|
||||
|
||||
let ghost = all.iter().find(|s| s.is_ghost_candidate).map(|s| {
|
||||
let suffix = s.insert_text.strip_prefix(prefix).unwrap_or(&s.insert_text);
|
||||
GhostSuggestion {
|
||||
full_text: s.insert_text.clone(),
|
||||
suffix: suffix.to_owned(),
|
||||
source: s.source.as_str().to_owned(),
|
||||
}
|
||||
});
|
||||
|
||||
let completions = all
|
||||
.into_iter()
|
||||
.take(limit)
|
||||
.map(CompletionItem::from)
|
||||
.collect();
|
||||
|
||||
(ghost, completions)
|
||||
}
|
||||
|
||||
/// Determines whether AI suggestions can be skipped based on history quality.
|
||||
pub(crate) fn should_skip_ai(history_matches: &[RankedSuggestion], prefix: &str) -> bool {
|
||||
if history_matches.is_empty() {
|
||||
return false;
|
||||
}
|
||||
if history_matches[0].priority >= 30 {
|
||||
return true;
|
||||
}
|
||||
!prefix.is_empty() && history_matches.len() >= 3
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn ranked(
|
||||
priority: i32,
|
||||
source: SuggestionSource,
|
||||
ghost: bool,
|
||||
text: &str,
|
||||
) -> RankedSuggestion {
|
||||
RankedSuggestion {
|
||||
display: text.into(),
|
||||
description: String::new(),
|
||||
insert_text: text.into(),
|
||||
source,
|
||||
priority,
|
||||
is_ghost_candidate: ghost,
|
||||
replace_range: None,
|
||||
token_text: None,
|
||||
truncated: false,
|
||||
}
|
||||
}
|
||||
|
||||
// --- aggregate ---
|
||||
|
||||
#[test]
|
||||
fn aggregate_sorts_by_descending_priority() {
|
||||
let history = vec![
|
||||
ranked(10, SuggestionSource::History, true, "git commit"),
|
||||
ranked(5, SuggestionSource::History, false, "git checkout"),
|
||||
];
|
||||
let path = vec![ranked(0, SuggestionSource::Path, false, "git")];
|
||||
let (_, completions) = aggregate(history, path, vec![], vec![], "git", 10);
|
||||
assert_eq!(completions[0].priority, 10);
|
||||
assert_eq!(completions[1].priority, 5);
|
||||
assert_eq!(completions[2].priority, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_selects_ghost_with_correct_suffix() {
|
||||
let history = vec![ranked(10, SuggestionSource::History, true, "git commit")];
|
||||
let path = vec![ranked(0, SuggestionSource::Path, false, "git")];
|
||||
let (ghost, _) = aggregate(history, path, vec![], vec![], "git", 10);
|
||||
let ghost = ghost.unwrap();
|
||||
assert_eq!(ghost.full_text, "git commit");
|
||||
assert_eq!(ghost.suffix, " commit");
|
||||
assert_eq!(ghost.source, "history");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_no_ghost_without_candidate() {
|
||||
let path = vec![
|
||||
ranked(0, SuggestionSource::Path, false, "git"),
|
||||
ranked(0, SuggestionSource::Path, false, "grep"),
|
||||
];
|
||||
let (ghost, _) = aggregate(vec![], path, vec![], vec![], "g", 10);
|
||||
assert!(ghost.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_respects_limit() {
|
||||
let history: Vec<_> = (0..10)
|
||||
.map(|i| {
|
||||
ranked(
|
||||
10 - i,
|
||||
SuggestionSource::History,
|
||||
i == 0,
|
||||
&format!("cmd_{i}"),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let (_, completions) = aggregate(history, vec![], vec![], vec![], "cmd", 3);
|
||||
assert_eq!(completions.len(), 3);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_ghost_suffix_for_exact_prefix() {
|
||||
let history = vec![ranked(40, SuggestionSource::History, true, "ls")];
|
||||
let (ghost, _) = aggregate(history, vec![], vec![], vec![], "ls", 10);
|
||||
let ghost = ghost.unwrap();
|
||||
assert_eq!(ghost.suffix, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_includes_ai_results() {
|
||||
let history = vec![ranked(10, SuggestionSource::History, true, "git commit")];
|
||||
let ai = vec![ranked(
|
||||
-10,
|
||||
SuggestionSource::AI,
|
||||
true,
|
||||
"git commit --amend",
|
||||
)];
|
||||
let (_, completions) = aggregate(history, vec![], vec![], ai, "git", 10);
|
||||
assert_eq!(completions.len(), 2);
|
||||
assert_eq!(completions[0].source, "history");
|
||||
assert_eq!(completions[1].source, "ai");
|
||||
assert_eq!(completions[1].priority, -10);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_ai_ghost_used_when_no_history_ghost() {
|
||||
let ai = vec![ranked(
|
||||
-10,
|
||||
SuggestionSource::AI,
|
||||
true,
|
||||
"git commit --amend",
|
||||
)];
|
||||
let (ghost, _) = aggregate(vec![], vec![], vec![], ai, "git", 10);
|
||||
let ghost = ghost.unwrap();
|
||||
assert_eq!(ghost.source, "ai");
|
||||
assert_eq!(ghost.suffix, " commit --amend");
|
||||
}
|
||||
|
||||
// --- should_skip_ai ---
|
||||
|
||||
#[test]
|
||||
fn skip_ai_returns_false_for_empty_history() {
|
||||
assert!(!should_skip_ai(&[], "git"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skip_ai_on_exact_match() {
|
||||
let m = vec![ranked(40, SuggestionSource::History, true, "git commit")];
|
||||
assert!(should_skip_ai(&m, "git commit"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skip_ai_prefix_with_enough_matches() {
|
||||
let m = vec![
|
||||
ranked(10, SuggestionSource::History, true, "git commit"),
|
||||
ranked(9, SuggestionSource::History, false, "git checkout"),
|
||||
ranked(8, SuggestionSource::History, false, "git cherry-pick"),
|
||||
];
|
||||
assert!(should_skip_ai(&m, "git"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dont_skip_ai_prefix_with_few_matches() {
|
||||
let m = vec![ranked(5, SuggestionSource::History, true, "git commit")];
|
||||
assert!(!should_skip_ai(&m, "git"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn skip_ai_many_matches_empty_prefix() {
|
||||
let m = vec![
|
||||
ranked(5, SuggestionSource::History, true, "a"),
|
||||
ranked(4, SuggestionSource::History, false, "b"),
|
||||
ranked(3, SuggestionSource::History, false, "c"),
|
||||
];
|
||||
// empty prefix + 3 matches: !prefix.is_empty() is false, len >= 3 is true → false AND true → false
|
||||
assert!(!should_skip_ai(&m, ""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dont_skip_ai_few_matches_empty_prefix() {
|
||||
let m = vec![
|
||||
ranked(5, SuggestionSource::History, true, "a"),
|
||||
ranked(4, SuggestionSource::History, false, "b"),
|
||||
];
|
||||
assert!(!should_skip_ai(&m, ""));
|
||||
}
|
||||
|
||||
// --- context ---
|
||||
|
||||
#[test]
|
||||
fn context_clamps_cursor_to_len() {
|
||||
let ctx = SuggestContext::new("abc".into(), 100, "/tmp".into());
|
||||
assert_eq!(ctx.cursor, 3);
|
||||
assert_eq!(ctx.prefix(), "abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn context_adjusts_to_char_boundary() {
|
||||
let text = "caf\u{00e9}"; // "cafe" with e-acute (2 bytes for e-acute)
|
||||
assert_eq!(text.len(), 5);
|
||||
let ctx = SuggestContext::new(text.into(), 4, "/tmp".into()); // middle of 2-byte e-acute
|
||||
assert_eq!(ctx.prefix(), "caf");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_item_serializes_replace_range_and_token_as_camel_case() {
|
||||
let mut s = ranked(10, SuggestionSource::Path, false, "ls | grep");
|
||||
s.replace_range = Some((5, 7));
|
||||
s.token_text = Some("grep".into());
|
||||
let json = serde_json::to_value(CompletionItem::from(s)).unwrap();
|
||||
assert_eq!(json["replaceRange"], serde_json::json!([5, 7]));
|
||||
assert_eq!(json["tokenText"], "grep");
|
||||
// Whole-line compat field for range-unaware pagers.
|
||||
assert_eq!(json["insertText"], "ls | grep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_item_omits_absent_replace_range_and_token() {
|
||||
let json = serde_json::to_value(CompletionItem::from(ranked(
|
||||
0,
|
||||
SuggestionSource::Path,
|
||||
false,
|
||||
"grep",
|
||||
)))
|
||||
.unwrap();
|
||||
assert!(json.get("replaceRange").is_none());
|
||||
assert!(json.get("tokenText").is_none());
|
||||
// Additive: `truncated` only serializes when set.
|
||||
assert!(json.get("truncated").is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_item_serializes_truncated_when_set() {
|
||||
let mut s = ranked(0, SuggestionSource::File, false, "notes.md");
|
||||
s.truncated = true;
|
||||
let json = serde_json::to_value(CompletionItem::from(s)).unwrap();
|
||||
assert_eq!(json["truncated"], true);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn stamp_whole_line_range_covers_full_text() {
|
||||
let mut results = vec![
|
||||
ranked(10, SuggestionSource::History, true, "git commit"),
|
||||
ranked(9, SuggestionSource::History, false, "git checkout"),
|
||||
];
|
||||
stamp_whole_line_range(&mut results, 5);
|
||||
assert!(results.iter().all(|s| s.replace_range == Some((0, 5))));
|
||||
// Whole-line items double as their own span replacement.
|
||||
assert!(results.iter().all(|s| s.token_text.is_none()));
|
||||
}
|
||||
|
||||
/// Token-valued suggestions become the wire pair: token in `token_text`,
|
||||
/// `insert_text` rebuilt as the full line (safe for old pagers).
|
||||
#[test]
|
||||
fn splice_token_into_line_builds_compat_pair() {
|
||||
let mut results = vec![ranked(0, SuggestionSource::Path, false, "grep")];
|
||||
splice_token_into_line(&mut results, "ls | gr | wc -l", (5, 7));
|
||||
assert_eq!(results[0].insert_text, "ls | grep | wc -l");
|
||||
assert_eq!(results[0].token_text.as_deref(), Some("grep"));
|
||||
}
|
||||
|
||||
/// Equal-priority items must keep their provider-internal order: the
|
||||
/// file provider ships pre-ranked rows (fuzzy tier/score/dirs-first) at
|
||||
/// ONE shared priority and its ranking reaches the wire only through
|
||||
/// this sort's stability. Deliberately non-alphabetical, larger than
|
||||
/// the small-slice insertion-sort threshold, and interleaved with a
|
||||
/// second priority class so a `sort_unstable_by` swap turns this red.
|
||||
#[test]
|
||||
fn aggregate_preserves_provider_order_within_equal_priority() {
|
||||
let file: Vec<_> = (0..32)
|
||||
.map(|i| {
|
||||
ranked(
|
||||
2,
|
||||
SuggestionSource::File,
|
||||
false,
|
||||
&format!("ranked_{:02}", 31 - i),
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
let expected: Vec<String> = file.iter().map(|s| s.display.clone()).collect();
|
||||
let path: Vec<_> = (0..32)
|
||||
.map(|i| ranked(0, SuggestionSource::Path, false, &format!("exe_{i:02}")))
|
||||
.collect();
|
||||
|
||||
let (_, completions) = aggregate(vec![], path, file, vec![], "r", 100);
|
||||
let file_order: Vec<String> = completions
|
||||
.iter()
|
||||
.filter(|c| c.source == "file")
|
||||
.map(|c| c.display.clone())
|
||||
.collect();
|
||||
assert_eq!(file_order, expected);
|
||||
// The boosted file rows all sort ahead of the priority-0 path rows.
|
||||
assert_eq!(
|
||||
completions[..32]
|
||||
.iter()
|
||||
.filter(|c| c.source == "file")
|
||||
.count(),
|
||||
32
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn completion_item_preserves_source_string() {
|
||||
let item = CompletionItem::from(ranked(10, SuggestionSource::History, true, "cmd"));
|
||||
assert_eq!(item.source, "history");
|
||||
assert_eq!(item.priority, 10);
|
||||
|
||||
let item = CompletionItem::from(ranked(0, SuggestionSource::Path, false, "cmd"));
|
||||
assert_eq!(item.source, "path");
|
||||
|
||||
let item = CompletionItem::from(ranked(5, SuggestionSource::File, false, "cmd"));
|
||||
assert_eq!(item.source, "file");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn aggregate_includes_file_results() {
|
||||
let history = vec![ranked(10, SuggestionSource::History, true, "cat ~/.bashrc")];
|
||||
let file = vec![ranked(5, SuggestionSource::File, false, ".bashrc")];
|
||||
let (_, completions) = aggregate(history, vec![], file, vec![], "cat", 10);
|
||||
assert_eq!(completions.len(), 2);
|
||||
assert_eq!(completions[0].priority, 10);
|
||||
assert_eq!(completions[1].priority, 5);
|
||||
assert_eq!(completions[1].source, "file");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,407 @@
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::sync::{Arc, OnceLock};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use arc_swap::ArcSwap;
|
||||
|
||||
use super::shell_token::{CurrentToken, build_insert_token, parse_current_token};
|
||||
use super::{RankedSuggestion, SuggestContext, SuggestionSource, splice_token_into_line};
|
||||
|
||||
const CACHE_TTL: Duration = Duration::from_secs(60);
|
||||
const MAX_RESULTS: usize = 10;
|
||||
|
||||
pub(crate) struct PathProvider;
|
||||
|
||||
impl PathProvider {
|
||||
pub async fn suggest(&self, ctx: &SuggestContext) -> Vec<RankedSuggestion> {
|
||||
// shell_token quoting is POSIX-only: cmd/pwsh would misparse the
|
||||
// escaped line, so Windows serves no deterministic completions.
|
||||
if cfg!(windows) {
|
||||
return Vec::new();
|
||||
}
|
||||
let tok = match extract_command_token(ctx.prefix()) {
|
||||
Some(t) => t,
|
||||
None => return Vec::new(),
|
||||
};
|
||||
let token_range = (tok.start, ctx.prefix().len());
|
||||
|
||||
let cache = get_or_refresh_path_cache().await;
|
||||
let mut results = filter_executables(&tok, token_range, &cache.executables);
|
||||
splice_token_into_line(&mut results, &ctx.text, token_range);
|
||||
results
|
||||
}
|
||||
}
|
||||
|
||||
/// The command token being typed, via the canonical tokenizer: quotes hide
|
||||
/// separators (`echo "a | gr` is quoted data, not a command position), the
|
||||
/// cursor must sit in the segment's first word, and — mirroring the file
|
||||
/// provider — flag-looking tokens never complete.
|
||||
fn extract_command_token(prefix: &str) -> Option<CurrentToken> {
|
||||
let tok = parse_current_token(prefix);
|
||||
if tok.tokens_before != 0 || tok.after_redirect || tok.value.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if tok.value.starts_with('-') {
|
||||
return None;
|
||||
}
|
||||
Some(tok)
|
||||
}
|
||||
|
||||
fn filter_executables(
|
||||
tok: &CurrentToken,
|
||||
token_range: (usize, usize),
|
||||
executables: &[String],
|
||||
) -> Vec<RankedSuggestion> {
|
||||
let prefix = tok.value.as_str();
|
||||
// Binary search to the first entry >= prefix, then take while starts_with.
|
||||
let start = executables.partition_point(|e| e.as_str() < prefix);
|
||||
|
||||
let mut results: Vec<RankedSuggestion> = Vec::new();
|
||||
let mut truncated = false;
|
||||
for exe in executables[start..]
|
||||
.iter()
|
||||
.take_while(|exe| exe.starts_with(prefix))
|
||||
{
|
||||
if results.len() == MAX_RESULTS {
|
||||
// An uncapped match remains: the set is not exhaustive.
|
||||
truncated = true;
|
||||
break;
|
||||
}
|
||||
results.push(RankedSuggestion {
|
||||
display: exe.clone(),
|
||||
// Re-quoted like filenames: an executable named `zz;echo PWNED`
|
||||
// must insert as ONE word, never a second command.
|
||||
insert_text: build_insert_token(tok, "", exe, false),
|
||||
description: String::new(),
|
||||
source: SuggestionSource::Path,
|
||||
priority: 0,
|
||||
is_ghost_candidate: false,
|
||||
replace_range: Some(token_range),
|
||||
token_text: None,
|
||||
truncated: false,
|
||||
});
|
||||
}
|
||||
if truncated {
|
||||
results.iter_mut().for_each(|s| s.truncated = true);
|
||||
}
|
||||
results
|
||||
}
|
||||
|
||||
// --- PATH cache ---
|
||||
|
||||
struct PathCacheInner {
|
||||
executables: Vec<String>,
|
||||
updated_at: Instant,
|
||||
path_env: String,
|
||||
}
|
||||
|
||||
static PATH_CACHE: OnceLock<ArcSwap<PathCacheInner>> = OnceLock::new();
|
||||
static PATH_REFRESHING: AtomicBool = AtomicBool::new(false);
|
||||
|
||||
async fn get_or_refresh_path_cache() -> Arc<PathCacheInner> {
|
||||
let swap = PATH_CACHE.get_or_init(|| {
|
||||
ArcSwap::from_pointee(PathCacheInner {
|
||||
executables: Vec::new(),
|
||||
updated_at: Instant::now() - CACHE_TTL - Duration::from_secs(1),
|
||||
path_env: String::new(),
|
||||
})
|
||||
});
|
||||
|
||||
let current = swap.load_full();
|
||||
let current_path = std::env::var("PATH").unwrap_or_default();
|
||||
|
||||
if current.updated_at.elapsed() < CACHE_TTL && current.path_env == current_path {
|
||||
return current;
|
||||
}
|
||||
|
||||
if PATH_REFRESHING
|
||||
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Relaxed)
|
||||
.is_err()
|
||||
{
|
||||
return current;
|
||||
}
|
||||
|
||||
let result = match tokio::task::spawn_blocking(scan_path_dirs).await {
|
||||
Ok(executables) => {
|
||||
let new = Arc::new(PathCacheInner {
|
||||
executables,
|
||||
updated_at: Instant::now(),
|
||||
path_env: current_path,
|
||||
});
|
||||
swap.store(Arc::clone(&new));
|
||||
new
|
||||
}
|
||||
Err(_) => current,
|
||||
};
|
||||
|
||||
PATH_REFRESHING.store(false, Ordering::Release);
|
||||
result
|
||||
}
|
||||
|
||||
fn scan_path_dirs() -> Vec<String> {
|
||||
let path_var = std::env::var("PATH").unwrap_or_default();
|
||||
scan_path_from(&path_var)
|
||||
}
|
||||
|
||||
fn scan_path_from(path_var: &str) -> Vec<String> {
|
||||
let mut executables = Vec::new();
|
||||
|
||||
for dir in std::env::split_paths(path_var) {
|
||||
let entries = match std::fs::read_dir(&dir) {
|
||||
Ok(e) => e,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
for entry in entries.filter_map(Result::ok) {
|
||||
let meta = match entry.metadata() {
|
||||
Ok(m) => m,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
if !meta.is_file() {
|
||||
continue;
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
if meta.permissions().mode() & 0o111 == 0 {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
if let Ok(name) = entry.file_name().into_string() {
|
||||
executables.push(name);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
executables.sort_unstable();
|
||||
executables.dedup();
|
||||
executables
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// --- extract_command_token ---
|
||||
|
||||
fn cmd(prefix: &str) -> Option<(usize, String)> {
|
||||
extract_command_token(prefix).map(|t| (t.start, t.value))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_at_start_of_line() {
|
||||
assert_eq!(cmd("git"), Some((0, "git".into())));
|
||||
assert_eq!(cmd("gi"), Some((0, "gi".into())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn no_prefix_at_argument_position() {
|
||||
assert_eq!(cmd("git com"), None);
|
||||
assert_eq!(cmd("ls | grep foo"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn prefix_after_separators() {
|
||||
assert_eq!(cmd("ls | gr"), Some((5, "gr".into())));
|
||||
assert_eq!(cmd("make && gi"), Some((8, "gi".into())));
|
||||
assert_eq!(cmd("cd /tmp; ls"), Some((9, "ls".into())));
|
||||
assert_eq!(cmd("false || tr"), Some((9, "tr".into())));
|
||||
assert_eq!(cmd("sleep 10 & ls"), Some((11, "ls".into())));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_when_empty_after_separator() {
|
||||
assert_eq!(cmd("ls | "), None);
|
||||
assert_eq!(cmd("sleep 10 & "), None);
|
||||
assert_eq!(cmd("&&"), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn none_for_empty_or_whitespace_input() {
|
||||
assert_eq!(cmd(""), None);
|
||||
assert_eq!(cmd(" "), None);
|
||||
}
|
||||
|
||||
/// A separator inside quotes is data, not a command position — the old
|
||||
/// naive segment scan offered executables inside quoted strings.
|
||||
#[test]
|
||||
fn none_inside_quoted_data() {
|
||||
assert_eq!(cmd("echo \"x | gr"), None);
|
||||
}
|
||||
|
||||
/// Flag-looking and redirect-target tokens never command-complete.
|
||||
#[test]
|
||||
fn none_for_flags_and_redirect_targets() {
|
||||
assert_eq!(cmd("-gr"), None);
|
||||
assert_eq!(cmd("> lo"), None);
|
||||
}
|
||||
|
||||
// --- filter_executables ---
|
||||
|
||||
fn tok(prefix: &str) -> CurrentToken {
|
||||
parse_current_token(prefix)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_matches_prefix() {
|
||||
let exes = vec![
|
||||
"gcc".into(),
|
||||
"git".into(),
|
||||
"grep".into(),
|
||||
"less".into(),
|
||||
"ls".into(),
|
||||
];
|
||||
let results = filter_executables(&tok("gr"), (0, 2), &exes);
|
||||
assert_eq!(results.len(), 1);
|
||||
assert_eq!(results[0].insert_text, "grep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_multiple_matches() {
|
||||
let exes = vec!["git".into(), "grep".into(), "gzip".into()];
|
||||
let results = filter_executables(&tok("g"), (0, 1), &exes);
|
||||
assert_eq!(results.len(), 3);
|
||||
assert!(results.iter().all(|s| !s.truncated));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_no_match() {
|
||||
let exes = vec!["git".into(), "grep".into()];
|
||||
assert!(filter_executables(&tok("docker"), (0, 6), &exes).is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filter_path_suggestions_are_not_ghost() {
|
||||
let exes = vec!["git".into()];
|
||||
let results = filter_executables(&tok("g"), (0, 1), &exes);
|
||||
assert!(!results[0].is_ghost_candidate);
|
||||
assert_eq!(results[0].source, SuggestionSource::Path);
|
||||
assert_eq!(results[0].priority, 0);
|
||||
}
|
||||
|
||||
/// Capped sets mark every row truncated so the pager keeps
|
||||
/// dropdown-only semantics (an unshown match could disprove an LCP).
|
||||
#[test]
|
||||
fn filter_caps_at_max_and_marks_truncated() {
|
||||
let exes: Vec<String> = (0..20).map(|i| format!("test_{i:03}")).collect();
|
||||
let results = filter_executables(&tok("test"), (0, 4), &exes);
|
||||
assert_eq!(results.len(), MAX_RESULTS);
|
||||
assert!(results.iter().all(|s| s.truncated));
|
||||
|
||||
let exes: Vec<String> = (0..MAX_RESULTS).map(|i| format!("test_{i:03}")).collect();
|
||||
let results = filter_executables(&tok("test"), (0, 4), &exes);
|
||||
assert_eq!(results.len(), MAX_RESULTS);
|
||||
assert!(results.iter().all(|s| !s.truncated));
|
||||
}
|
||||
|
||||
/// The segment-after-pipe token range: accepting `grep` for `ls | gr`
|
||||
/// must target only the `gr` token, never the whole line.
|
||||
#[test]
|
||||
fn filter_stamps_segment_token_range() {
|
||||
let t = extract_command_token("ls | gr").unwrap();
|
||||
let exes = vec!["grep".into()];
|
||||
let results = filter_executables(&t, (t.start, 7), &exes);
|
||||
assert_eq!(results[0].replace_range, Some((5, 7)));
|
||||
assert_eq!(results[0].insert_text, "grep");
|
||||
}
|
||||
|
||||
/// Metacharacter executable names insert as ONE word — accepting
|
||||
/// `zz;echo PWNED` must never put a second command on the line.
|
||||
#[test]
|
||||
fn filter_escapes_metacharacter_executable_names() {
|
||||
let exes = vec!["zz;echo PWNED".into()];
|
||||
let results = filter_executables(&tok("zz"), (0, 2), &exes);
|
||||
assert_eq!(results[0].display, "zz;echo PWNED");
|
||||
assert_eq!(results[0].insert_text, "zz\\;echo\\ PWNED");
|
||||
}
|
||||
|
||||
/// A quoted command prefix completes inside its quote style.
|
||||
#[test]
|
||||
fn filter_requotes_quoted_command_prefix() {
|
||||
let exes = vec!["grep".into()];
|
||||
let results = filter_executables(&tok("\"gr"), (0, 3), &exes);
|
||||
assert_eq!(results[0].insert_text, "\"grep\"");
|
||||
}
|
||||
|
||||
// --- scan_path_from ---
|
||||
|
||||
#[test]
|
||||
fn scan_nonexistent_dir() {
|
||||
assert!(scan_path_from("/nonexistent/path/that/doesnt/exist").is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scan_creates_sorted_deduped_list() {
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let bin = dir.path().join("bin");
|
||||
fs::create_dir(&bin).unwrap();
|
||||
|
||||
for name in &["zzz_cmd", "aaa_cmd", "mmm_cmd"] {
|
||||
let path = bin.join(name);
|
||||
fs::write(&path, "").unwrap();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
fs::set_permissions(&path, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
let result = scan_path_from(bin.to_str().unwrap());
|
||||
assert_eq!(result, vec!["aaa_cmd", "mmm_cmd", "zzz_cmd"]);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[test]
|
||||
fn scan_skips_nonexecutable_files() {
|
||||
use std::fs;
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
use tempfile::TempDir;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let bin = dir.path().join("bin");
|
||||
fs::create_dir(&bin).unwrap();
|
||||
|
||||
let exec_path = bin.join("my_exec");
|
||||
fs::write(&exec_path, "").unwrap();
|
||||
fs::set_permissions(&exec_path, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
|
||||
let data_path = bin.join("my_data");
|
||||
fs::write(&data_path, "").unwrap();
|
||||
fs::set_permissions(&data_path, fs::Permissions::from_mode(0o644)).unwrap();
|
||||
|
||||
let result = scan_path_from(bin.to_str().unwrap());
|
||||
assert_eq!(result, vec!["my_exec"]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn scan_deduplicates_across_dirs() {
|
||||
use std::fs;
|
||||
use tempfile::TempDir;
|
||||
|
||||
let dir = TempDir::new().unwrap();
|
||||
let bin1 = dir.path().join("bin1");
|
||||
let bin2 = dir.path().join("bin2");
|
||||
fs::create_dir(&bin1).unwrap();
|
||||
fs::create_dir(&bin2).unwrap();
|
||||
|
||||
for bin in [&bin1, &bin2] {
|
||||
let path = bin.join("shared_cmd");
|
||||
fs::write(&path, "").unwrap();
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
fs::set_permissions(&path, fs::Permissions::from_mode(0o755)).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
let path_var = format!("{}:{}", bin1.to_str().unwrap(), bin2.to_str().unwrap());
|
||||
let result = scan_path_from(&path_var);
|
||||
assert_eq!(result, vec!["shared_cmd"]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,640 @@
|
||||
//! Minimal shell-token syntax for completion: find the token under the
|
||||
//! cursor and re-quote completed components to match how it was typed.
|
||||
//! Consumed by the file provider; the natural home for the $PATH provider's
|
||||
//! segmentation too, once it learns quoting.
|
||||
//!
|
||||
//! ## Scope (deliberately minimal)
|
||||
//!
|
||||
//! [`parse_current_token`] understands just enough POSIX-shell syntax to
|
||||
//! find the token under the cursor: double/single quotes, backslash
|
||||
//! escapes, whitespace, the segment separators `|`/`;`/`&`, and the
|
||||
//! redirection operators `<`/`>`. It does NOT model the full grammar:
|
||||
//! no subshells or `$(…)`, no here-docs, no brace/glob expansion, no
|
||||
//! `~user` home lookup.
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub(super) enum QuoteStyle {
|
||||
None,
|
||||
Double,
|
||||
Single,
|
||||
}
|
||||
|
||||
/// The shell token under the cursor plus the segment facts completion needs.
|
||||
#[derive(Debug)]
|
||||
pub(super) struct CurrentToken {
|
||||
/// Byte offset in the request text where the token starts (an opening
|
||||
/// quote is part of the token) — the start of the replace range.
|
||||
pub(super) start: usize,
|
||||
/// Unquoted/unescaped value typed so far.
|
||||
pub(super) value: String,
|
||||
/// `value` length right after its last `/` (`None`: no slash).
|
||||
pub(super) dir_value_len: Option<usize>,
|
||||
/// Byte offset just after the last raw `/` (== `start` without one).
|
||||
pub(super) dir_raw_end: usize,
|
||||
/// Quote state at the cursor.
|
||||
pub(super) quote: QuoteStyle,
|
||||
/// Byte offset of the still-open quote (meaningful when `quote != None`).
|
||||
pub(super) open_quote_idx: usize,
|
||||
/// Quote structure of the REPLACED component when no quote is open at
|
||||
/// the cursor: the style of a quote whose closer sits inside the
|
||||
/// component (at/after `dir_raw_end`), plus whether its opener does too
|
||||
/// (and so needs re-emitting). `None`: quote-free component.
|
||||
pub(super) closed_quote: Option<(QuoteStyle, bool)>,
|
||||
/// Byte-aligned with `value`: `true` where the char was consumed
|
||||
/// unquoted and unescaped — the only spellings the shell expands
|
||||
/// `~`/`$` in (`'$HOME'`, `\$HOME`, and `"~/…` are literal to it).
|
||||
pub(super) plain_mask: Vec<bool>,
|
||||
/// Completed tokens before this one in the current segment.
|
||||
pub(super) tokens_before: usize,
|
||||
/// The segment's command word (first non-redirect-target token value).
|
||||
pub(super) command: Option<String>,
|
||||
/// Token directly follows `<`/`>` — always a file argument.
|
||||
pub(super) after_redirect: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct TokenBuild {
|
||||
start: usize,
|
||||
value: String,
|
||||
dir_value_len: Option<usize>,
|
||||
dir_raw_end: usize,
|
||||
after_redirect: bool,
|
||||
/// Last quote closed within this token: `(open_idx, close_idx, style)`.
|
||||
last_close: Option<(usize, usize, QuoteStyle)>,
|
||||
plain_mask: Vec<bool>,
|
||||
}
|
||||
|
||||
impl TokenBuild {
|
||||
/// `plain`: the char reached `value` unquoted and unescaped — the only
|
||||
/// provenance the shell expands `~`/`$` in.
|
||||
fn push(&mut self, i: usize, c: char, plain: bool) {
|
||||
self.value.push(c);
|
||||
self.plain_mask
|
||||
.extend(std::iter::repeat_n(plain, c.len_utf8()));
|
||||
if c == '/' {
|
||||
self.dir_value_len = Some(self.value.len());
|
||||
self.dir_raw_end = i + 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_token<'a>(
|
||||
cur: &'a mut Option<TokenBuild>,
|
||||
i: usize,
|
||||
pending_redirect: &mut bool,
|
||||
) -> &'a mut TokenBuild {
|
||||
cur.get_or_insert_with(|| TokenBuild {
|
||||
start: i,
|
||||
dir_raw_end: i,
|
||||
after_redirect: std::mem::take(pending_redirect),
|
||||
..TokenBuild::default()
|
||||
})
|
||||
}
|
||||
|
||||
fn finish_token(
|
||||
cur: &mut Option<TokenBuild>,
|
||||
tokens_before: &mut usize,
|
||||
command: &mut Option<String>,
|
||||
) {
|
||||
if let Some(t) = cur.take() {
|
||||
*tokens_before += 1;
|
||||
if command.is_none() && !t.after_redirect {
|
||||
*command = Some(t.value);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Scan the cursor prefix and return the token being typed. Quotes hide
|
||||
/// separators (`echo "a | b` is one segment), backslashes escape the next
|
||||
/// char, and `|`/`;`/`&` reset the segment. A dangling backslash at the
|
||||
/// cursor contributes no value char but stays inside the token extent.
|
||||
pub(super) fn parse_current_token(prefix: &str) -> CurrentToken {
|
||||
let mut cur: Option<TokenBuild> = None;
|
||||
let mut quote = QuoteStyle::None;
|
||||
let mut open_quote_idx = 0usize;
|
||||
let mut escape = false;
|
||||
let mut tokens_before = 0usize;
|
||||
let mut command: Option<String> = None;
|
||||
let mut pending_redirect = false;
|
||||
|
||||
// Escape/quote state implies a token exists (`ensure_token` ran when the
|
||||
// state was entered), so the `ensure_token` calls below are no-op
|
||||
// lookups on valid input — but this parses arbitrary wire text, and a
|
||||
// mis-tokenized line must degrade, never panic the agent.
|
||||
for (i, c) in prefix.char_indices() {
|
||||
if escape {
|
||||
escape = false;
|
||||
let t = ensure_token(&mut cur, i, &mut pending_redirect);
|
||||
match quote {
|
||||
// `\X` outside quotes: literal X.
|
||||
QuoteStyle::None => t.push(i, c, false),
|
||||
// Inside double quotes `\` only escapes `"` `\` `$` `` ` ``.
|
||||
QuoteStyle::Double => {
|
||||
if !matches!(c, '"' | '\\' | '$' | '`') {
|
||||
t.push(i, '\\', false);
|
||||
}
|
||||
t.push(i, c, false);
|
||||
}
|
||||
// No escapes exist inside single quotes; keep the char.
|
||||
QuoteStyle::Single => t.push(i, c, false),
|
||||
}
|
||||
continue;
|
||||
}
|
||||
match quote {
|
||||
QuoteStyle::Single => match c {
|
||||
'\'' => {
|
||||
quote = QuoteStyle::None;
|
||||
if let Some(t) = cur.as_mut() {
|
||||
t.last_close = Some((open_quote_idx, i, QuoteStyle::Single));
|
||||
}
|
||||
}
|
||||
_ => ensure_token(&mut cur, i, &mut pending_redirect).push(i, c, false),
|
||||
},
|
||||
QuoteStyle::Double => match c {
|
||||
'"' => {
|
||||
quote = QuoteStyle::None;
|
||||
if let Some(t) = cur.as_mut() {
|
||||
t.last_close = Some((open_quote_idx, i, QuoteStyle::Double));
|
||||
}
|
||||
}
|
||||
'\\' => escape = true,
|
||||
_ => ensure_token(&mut cur, i, &mut pending_redirect).push(i, c, false),
|
||||
},
|
||||
QuoteStyle::None => match c {
|
||||
'\\' => {
|
||||
ensure_token(&mut cur, i, &mut pending_redirect);
|
||||
escape = true;
|
||||
}
|
||||
'\'' => {
|
||||
ensure_token(&mut cur, i, &mut pending_redirect);
|
||||
quote = QuoteStyle::Single;
|
||||
open_quote_idx = i;
|
||||
}
|
||||
'"' => {
|
||||
ensure_token(&mut cur, i, &mut pending_redirect);
|
||||
quote = QuoteStyle::Double;
|
||||
open_quote_idx = i;
|
||||
}
|
||||
c if c.is_whitespace() => {
|
||||
finish_token(&mut cur, &mut tokens_before, &mut command);
|
||||
}
|
||||
'|' | ';' | '&' => {
|
||||
finish_token(&mut cur, &mut tokens_before, &mut command);
|
||||
tokens_before = 0;
|
||||
command = None;
|
||||
pending_redirect = false;
|
||||
}
|
||||
'<' | '>' => {
|
||||
finish_token(&mut cur, &mut tokens_before, &mut command);
|
||||
pending_redirect = true;
|
||||
}
|
||||
_ => ensure_token(&mut cur, i, &mut pending_redirect).push(i, c, true),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
let (start, value, dir_value_len, dir_raw_end, after_redirect, last_close, plain_mask) =
|
||||
match cur {
|
||||
Some(t) => (
|
||||
t.start,
|
||||
t.value,
|
||||
t.dir_value_len,
|
||||
t.dir_raw_end,
|
||||
t.after_redirect,
|
||||
t.last_close,
|
||||
t.plain_mask,
|
||||
),
|
||||
// Cursor sits after a separator: a fresh empty token starts here.
|
||||
None => (
|
||||
prefix.len(),
|
||||
String::new(),
|
||||
None,
|
||||
prefix.len(),
|
||||
pending_redirect,
|
||||
None,
|
||||
Vec::new(),
|
||||
),
|
||||
};
|
||||
// A closure before the component boundary is raw-dir-internal
|
||||
// (balanced, kept verbatim) — only closers the component consumed
|
||||
// constrain how it re-renders.
|
||||
let closed_quote = last_close.and_then(|(open, close, style)| {
|
||||
(close >= dir_raw_end).then_some((style, open >= dir_raw_end))
|
||||
});
|
||||
CurrentToken {
|
||||
start,
|
||||
value,
|
||||
dir_value_len,
|
||||
dir_raw_end,
|
||||
quote,
|
||||
open_quote_idx,
|
||||
closed_quote,
|
||||
plain_mask,
|
||||
tokens_before,
|
||||
command,
|
||||
after_redirect,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Insert-token construction (quoting) ─────────────────────────────────
|
||||
|
||||
/// Build the replacement for the whole token: the user's verbatim directory
|
||||
/// prefix plus the completed component escaped for the quote context at the
|
||||
/// cursor. Files close an open quote; directories keep it open (and get the
|
||||
/// trailing `/`) so the next Tab drills down, bash-style.
|
||||
pub(super) fn build_insert_token(
|
||||
tok: &CurrentToken,
|
||||
raw_dir: &str,
|
||||
name: &str,
|
||||
is_dir: bool,
|
||||
) -> String {
|
||||
let mut out = String::with_capacity(raw_dir.len() + name.len() + 4);
|
||||
out.push_str(raw_dir);
|
||||
// A completed component starting with `-` would otherwise insert a
|
||||
// flag-looking argument (`rm ` + Tab → `rm -rf`, invisible when the
|
||||
// single-candidate insta-accept skips the dropdown); quoting wouldn't
|
||||
// help (`rm "-rf"` is still a flag to rm). Anchor bare names as
|
||||
// explicit paths — deliberately stricter than bash.
|
||||
if raw_dir.is_empty() && name.starts_with('-') {
|
||||
out.push_str("./");
|
||||
}
|
||||
// The quote context the component renders in: the quote still open at
|
||||
// the cursor, or one the component CLOSED (`cat "My Dir/fi"` — raw_dir
|
||||
// keeps the dangling opener, so dropping the closer would emit an
|
||||
// unbalanced line). A quote opened INSIDE the component (after the
|
||||
// last `/`) is not part of `raw_dir` — re-emit it.
|
||||
let (style, reopen) = match tok.quote {
|
||||
QuoteStyle::None => tok.closed_quote.unwrap_or((QuoteStyle::None, false)),
|
||||
open => (open, tok.open_quote_idx >= tok.dir_raw_end),
|
||||
};
|
||||
match style {
|
||||
QuoteStyle::None => out.push_str(&escape_unquoted(name)),
|
||||
QuoteStyle::Double => {
|
||||
if reopen {
|
||||
out.push('"');
|
||||
}
|
||||
out.push_str(&escape_double_quoted(name));
|
||||
if !is_dir {
|
||||
out.push('"');
|
||||
}
|
||||
}
|
||||
QuoteStyle::Single => {
|
||||
if reopen {
|
||||
out.push('\'');
|
||||
}
|
||||
out.push_str(&escape_single_quoted(name));
|
||||
if !is_dir {
|
||||
out.push('\'');
|
||||
}
|
||||
}
|
||||
}
|
||||
if is_dir {
|
||||
out.push('/');
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Bash-ish set of characters that need a backslash outside quotes.
|
||||
/// Deliberately generous: over-escaping is harmless to the shell,
|
||||
/// under-escaping breaks the command.
|
||||
fn needs_backslash(c: char) -> bool {
|
||||
matches!(
|
||||
c,
|
||||
' ' | '\t'
|
||||
| '"'
|
||||
| '\''
|
||||
| '\\'
|
||||
| '$'
|
||||
| '`'
|
||||
| '&'
|
||||
| '|'
|
||||
| ';'
|
||||
| '('
|
||||
| ')'
|
||||
| '<'
|
||||
| '>'
|
||||
| '*'
|
||||
| '?'
|
||||
| '['
|
||||
| ']'
|
||||
| '#'
|
||||
| '!'
|
||||
| '{'
|
||||
| '}'
|
||||
| '~'
|
||||
)
|
||||
}
|
||||
|
||||
fn escape_unquoted(name: &str) -> String {
|
||||
// Control chars (newlines…) can't be backslash-escaped portably (`\` +
|
||||
// newline is a line continuation) — single-quote the whole component.
|
||||
if name.chars().any(char::is_control) {
|
||||
return format!("'{}'", escape_single_quoted(name));
|
||||
}
|
||||
let mut out = String::with_capacity(name.len());
|
||||
for c in name.chars() {
|
||||
if needs_backslash(c) {
|
||||
out.push('\\');
|
||||
}
|
||||
out.push(c);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn escape_double_quoted(name: &str) -> String {
|
||||
let mut out = String::with_capacity(name.len());
|
||||
for c in name.chars() {
|
||||
if matches!(c, '"' | '\\' | '$' | '`') {
|
||||
out.push('\\');
|
||||
}
|
||||
out.push(c);
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
fn escape_single_quoted(name: &str) -> String {
|
||||
// `'` cannot appear inside single quotes: close, escape, reopen.
|
||||
name.replace('\'', "'\\''")
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
// --- parse_current_token ---
|
||||
|
||||
#[test]
|
||||
fn parse_after_pipe_and_semicolon() {
|
||||
let tok = parse_current_token("echo hi | cat foo");
|
||||
assert_eq!(tok.value, "foo");
|
||||
assert_eq!(tok.start, 14);
|
||||
assert_eq!(tok.command.as_deref(), Some("cat"));
|
||||
|
||||
let tok = parse_current_token("cd /tmp; ls ");
|
||||
assert_eq!(tok.value, "");
|
||||
assert_eq!(tok.start, 12);
|
||||
assert_eq!(tok.command.as_deref(), Some("ls"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_after_double_ampersand() {
|
||||
let tok = parse_current_token("make && cat foo");
|
||||
assert_eq!(tok.value, "foo");
|
||||
assert_eq!(tok.start, 12);
|
||||
}
|
||||
|
||||
/// Quotes hide segment separators: the pipe is data, not a new command.
|
||||
#[test]
|
||||
fn parse_quoted_pipe_is_one_token() {
|
||||
let tok = parse_current_token("echo \"a | b");
|
||||
assert_eq!(tok.value, "a | b");
|
||||
assert_eq!(tok.start, 5);
|
||||
assert_eq!(tok.quote, QuoteStyle::Double);
|
||||
assert_eq!(tok.command.as_deref(), Some("echo"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_open_double_quote_token() {
|
||||
let tok = parse_current_token("cat \"My Fi");
|
||||
assert_eq!(tok.value, "My Fi");
|
||||
assert_eq!(tok.start, 4);
|
||||
assert_eq!(tok.quote, QuoteStyle::Double);
|
||||
assert_eq!(tok.open_quote_idx, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_backslash_escaped_space_token() {
|
||||
let tok = parse_current_token("cat My\\ Fi");
|
||||
assert_eq!(tok.value, "My Fi");
|
||||
assert_eq!(tok.start, 4);
|
||||
assert_eq!(tok.quote, QuoteStyle::None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_open_single_quote_token() {
|
||||
let tok = parse_current_token("cat 'sing le");
|
||||
assert_eq!(tok.value, "sing le");
|
||||
assert_eq!(tok.quote, QuoteStyle::Single);
|
||||
}
|
||||
|
||||
/// Closed-quote token: quote state returns to None at the cursor and the
|
||||
/// raw dir keeps the user's quoting verbatim.
|
||||
#[test]
|
||||
fn parse_closed_quote_dir_prefix() {
|
||||
let tok = parse_current_token("cat \"My Dir\"/fi");
|
||||
assert_eq!(tok.value, "My Dir/fi");
|
||||
assert_eq!(tok.quote, QuoteStyle::None);
|
||||
assert_eq!(tok.dir_raw_end, 13);
|
||||
assert_eq!(tok.dir_value_len, Some(7));
|
||||
}
|
||||
|
||||
/// `<`/`>` end the preceding token and flag the next as a redirect
|
||||
/// target without resetting the segment (the command survives).
|
||||
#[test]
|
||||
fn parse_redirect_sets_flag_and_boundary() {
|
||||
let tok = parse_current_token("echo hi > lo");
|
||||
assert_eq!(tok.value, "lo");
|
||||
assert!(tok.after_redirect);
|
||||
assert_eq!(tok.command.as_deref(), Some("echo"));
|
||||
|
||||
let tok = parse_current_token("> lo");
|
||||
assert_eq!(tok.value, "lo");
|
||||
assert!(tok.after_redirect);
|
||||
assert_eq!(tok.tokens_before, 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_multibyte_whitespace() {
|
||||
let tok = parse_current_token("cat\u{3000}foo");
|
||||
assert_eq!(tok.value, "foo");
|
||||
}
|
||||
|
||||
/// The mask records per-byte quote/escape provenance: quoted and escaped
|
||||
/// chars are not `plain` (the shell would not expand `~`/`$` there).
|
||||
#[test]
|
||||
fn parse_plain_mask_tracks_quote_and_escape_provenance() {
|
||||
let tok = parse_current_token("cat '$A'/b\\$c");
|
||||
assert_eq!(tok.value, "$A/b$c");
|
||||
assert_eq!(tok.plain_mask, [false, false, true, true, false, true]);
|
||||
|
||||
let tok = parse_current_token("cat \"~/do");
|
||||
assert_eq!(tok.value, "~/do");
|
||||
assert!(tok.plain_mask.iter().all(|p| !p));
|
||||
|
||||
let tok = parse_current_token("cat ~/do");
|
||||
assert!(tok.plain_mask.iter().all(|p| *p));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dangling_backslash_keeps_token() {
|
||||
let tok = parse_current_token("cat Notes\\");
|
||||
assert_eq!(tok.value, "Notes");
|
||||
assert_eq!(tok.start, 4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_multiple_args_takes_last() {
|
||||
let tok = parse_current_token("cp src/a.txt dst/b");
|
||||
assert_eq!(tok.value, "dst/b");
|
||||
assert_eq!(tok.start, 13);
|
||||
}
|
||||
|
||||
// --- escaping / insert-token construction ---
|
||||
|
||||
#[test]
|
||||
fn escape_unquoted_space_and_specials() {
|
||||
assert_eq!(escape_unquoted("My File.txt"), "My\\ File.txt");
|
||||
assert_eq!(escape_unquoted("a\"b'c"), "a\\\"b\\'c");
|
||||
assert_eq!(escape_unquoted("a$b"), "a\\$b");
|
||||
assert_eq!(escape_unquoted("plain.txt"), "plain.txt");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_unquoted_control_chars_single_quote_fallback() {
|
||||
assert_eq!(escape_unquoted("a\nb"), "'a\nb'");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_double_quoted_minimal_set() {
|
||||
assert_eq!(escape_double_quoted("My File.txt"), "My File.txt");
|
||||
assert_eq!(escape_double_quoted("a\"b$c"), "a\\\"b\\$c");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn escape_single_quoted_embedded_quote() {
|
||||
assert_eq!(escape_single_quoted("it's"), "it'\\''s");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_token_backslash_style_dir_stays_open() {
|
||||
let tok = parse_current_token("cat No");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "Notes Archive", true),
|
||||
"Notes\\ Archive/"
|
||||
);
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "My File.txt", false),
|
||||
"My\\ File.txt"
|
||||
);
|
||||
}
|
||||
|
||||
/// Open double quote: files close it, directories keep it open for
|
||||
/// drill-down (bash behavior). Slashless tokens have an empty raw dir —
|
||||
/// the still-open quote sits inside the replaced component (reopen path).
|
||||
#[test]
|
||||
fn insert_token_preserves_open_double_quote() {
|
||||
let tok = parse_current_token("cat \"My Fi");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "My File.txt", false),
|
||||
"\"My File.txt\""
|
||||
);
|
||||
let tok = parse_current_token("cat \"No");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "Notes Archive", true),
|
||||
"\"Notes Archive/"
|
||||
);
|
||||
}
|
||||
|
||||
/// A quote opened after the last `/` sits inside the replaced component
|
||||
/// and must be re-emitted.
|
||||
#[test]
|
||||
fn insert_token_reopens_quote_after_slash() {
|
||||
let tok = parse_current_token("cat dir/\"fi");
|
||||
// raw_dir covers `dir/`; the quote reopened inside the component.
|
||||
assert_eq!(tok.dir_raw_end, 8);
|
||||
assert_eq!(tok.open_quote_idx, 8);
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "dir/", "file name.txt", false),
|
||||
"dir/\"file name.txt\""
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn insert_token_single_quote_style() {
|
||||
let tok = parse_current_token("cat 'My Fi");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "My File.txt", false),
|
||||
"'My File.txt'"
|
||||
);
|
||||
}
|
||||
|
||||
/// THE closed-at-cursor case: the closer sits inside the replaced
|
||||
/// component while `raw_dir` keeps the opener — the rebuilt insert must
|
||||
/// still close it (files) or keep drilling (dirs), never emit
|
||||
/// `"My Dir/file.txt` with a dangling opener.
|
||||
#[test]
|
||||
fn insert_token_quote_closed_at_cursor_keeps_closer() {
|
||||
let tok = parse_current_token("cat \"My Dir/fi\"");
|
||||
assert_eq!(tok.quote, QuoteStyle::None);
|
||||
assert_eq!(tok.closed_quote, Some((QuoteStyle::Double, false)));
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "\"My Dir/", "file.txt", false),
|
||||
"\"My Dir/file.txt\""
|
||||
);
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "\"My Dir/", "subdir", true),
|
||||
"\"My Dir/subdir/"
|
||||
);
|
||||
|
||||
let tok = parse_current_token("cat 'My Dir/fi'");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "'My Dir/", "file.txt", false),
|
||||
"'My Dir/file.txt'"
|
||||
);
|
||||
}
|
||||
|
||||
/// Cursor still INSIDE the quotes (closer not part of the prefix): the
|
||||
/// open-quote path is unchanged by the closed-quote tracking.
|
||||
#[test]
|
||||
fn insert_token_cursor_inside_quotes_unchanged() {
|
||||
let tok = parse_current_token("cat \"My Dir/fi");
|
||||
assert_eq!(tok.quote, QuoteStyle::Double);
|
||||
assert_eq!(tok.closed_quote, None);
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "\"My Dir/", "file.txt", false),
|
||||
"\"My Dir/file.txt\""
|
||||
);
|
||||
}
|
||||
|
||||
/// Balanced quotes entirely inside a slashless component are replaced
|
||||
/// wholesale: the insert re-opens AND closes them.
|
||||
#[test]
|
||||
fn insert_token_balanced_slashless_quotes_reopen_and_close() {
|
||||
let tok = parse_current_token("cat \"fi\"");
|
||||
assert_eq!(tok.closed_quote, Some((QuoteStyle::Double, true)));
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "file.txt", false),
|
||||
"\"file.txt\""
|
||||
);
|
||||
}
|
||||
|
||||
/// A quote closed BEFORE the last `/` is raw-dir-internal (kept
|
||||
/// verbatim) and must not force quote rendering on the component.
|
||||
#[test]
|
||||
fn insert_token_quote_closed_in_raw_dir_stays_plain() {
|
||||
let tok = parse_current_token("cat \"My Dir\"/fi");
|
||||
assert_eq!(tok.closed_quote, None);
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "\"My Dir\"/", "file.txt", false),
|
||||
"\"My Dir\"/file.txt"
|
||||
);
|
||||
}
|
||||
|
||||
/// Dash-leading names anchor as `./`-relative paths so a completed bare
|
||||
/// component can never parse as a flag (`rm ` + Tab must not become
|
||||
/// `rm -rf`); quoting alone would not help. Directory-prefixed
|
||||
/// components are already anchored.
|
||||
#[test]
|
||||
fn insert_token_anchors_dash_leading_names() {
|
||||
let tok = parse_current_token("rm ");
|
||||
assert_eq!(build_insert_token(&tok, "", "-rf", false), "./-rf");
|
||||
assert_eq!(
|
||||
build_insert_token(&tok, "", "-flag dir", true),
|
||||
"./-flag\\ dir/"
|
||||
);
|
||||
let tok = parse_current_token("rm \"");
|
||||
assert_eq!(build_insert_token(&tok, "", "-rf", false), "./\"-rf\"");
|
||||
let tok = parse_current_token("rm sub/");
|
||||
assert_eq!(build_insert_token(&tok, "sub/", "-rf", false), "sub/-rf");
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,366 @@
|
||||
use crate::agent::mvp_agent::MvpAgent;
|
||||
use crate::extensions::routing::RequestMeta;
|
||||
use crate::session::ExtMethodResult;
|
||||
use crate::terminal::{self, KillOutcome};
|
||||
use agent_client_protocol as acp;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct EnvVar {
|
||||
pub name: String,
|
||||
pub value: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CreateTerminalRequest {
|
||||
pub session_id: String,
|
||||
pub command: String,
|
||||
#[serde(default)]
|
||||
pub args: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub env: Vec<EnvVar>,
|
||||
pub cwd: Option<String>,
|
||||
pub output_byte_limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TerminalIdRequest {
|
||||
pub session_id: String,
|
||||
pub terminal_id: String,
|
||||
}
|
||||
|
||||
/// Response for any terminal creation — piped or PTY. Both return just a `terminalId`.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct CreateTerminalResponse {
|
||||
pub terminal_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PtyCreateRequest {
|
||||
pub shell: Option<String>,
|
||||
pub cwd: Option<String>,
|
||||
#[serde(default)]
|
||||
pub session_id: Option<acp::SessionId>,
|
||||
#[serde(default)]
|
||||
pub env: Vec<EnvVar>,
|
||||
pub rows: Option<u16>,
|
||||
pub cols: Option<u16>,
|
||||
pub name: Option<String>,
|
||||
#[serde(default, rename = "_meta")]
|
||||
pub meta: Option<RequestMeta>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PtyLoadRequest {
|
||||
pub terminal_id: String,
|
||||
#[serde(default, rename = "_meta")]
|
||||
pub meta: Option<RequestMeta>,
|
||||
}
|
||||
|
||||
/// Terminal kill request — `session_id` is required for piped terminals,
|
||||
/// ignored for PTY terminals (looked up by `terminal_id` alone).
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct KillTerminalRequest {
|
||||
pub terminal_id: String,
|
||||
pub session_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PtyResizeRequest {
|
||||
pub terminal_id: String,
|
||||
pub rows: u16,
|
||||
pub cols: u16,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct PtyInputNotification {
|
||||
pub terminal_id: String,
|
||||
pub data: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TerminalListResponse {
|
||||
pub terminals: Vec<terminal::TerminalInfo>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ExitStatusResponse {
|
||||
pub exit_code: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub signal: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TerminalOutputResponse {
|
||||
pub output: String,
|
||||
pub truncated: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub exit_status: Option<ExitStatusResponse>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, Serialize)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum KillOutcomeResponse {
|
||||
Killed,
|
||||
AlreadyExited,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct KillTerminalResponse {
|
||||
pub outcome: KillOutcomeResponse,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct ReleaseTerminalResponse {}
|
||||
|
||||
fn parse<T: serde::de::DeserializeOwned>(args: &acp::ExtRequest) -> Result<T, acp::Error> {
|
||||
serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::invalid_params().data(format!("invalid params: {e}")))
|
||||
}
|
||||
|
||||
fn respond<T: Serialize>(result: Result<T, impl std::fmt::Display>) -> ExtResult {
|
||||
ExtMethodResult::from_result(result)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
/// Like `respond`, but converts `TerminalExtError` into a structured
|
||||
/// `{ code, message, data }` error instead of stringifying it.
|
||||
fn respond_pty<T: Serialize>(result: Result<T, terminal::TerminalExtError>) -> ExtResult {
|
||||
let ext_result: ExtMethodResult<T> = match result {
|
||||
Ok(value) => ExtMethodResult::success(value),
|
||||
Err(err) => err.into(),
|
||||
};
|
||||
ext_result
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
const ERR_TERMINAL_NOT_FOUND: &str = "terminal not found";
|
||||
|
||||
impl From<terminal::ExitStatus> for ExitStatusResponse {
|
||||
fn from(s: terminal::ExitStatus) -> Self {
|
||||
Self {
|
||||
exit_code: s.exit_code,
|
||||
signal: s.signal,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<terminal::OutputSnapshot> for TerminalOutputResponse {
|
||||
fn from(s: terminal::OutputSnapshot) -> Self {
|
||||
Self {
|
||||
output: s.output,
|
||||
truncated: s.truncated,
|
||||
exit_status: s.exit_status.map(Into::into),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<KillOutcome> for KillOutcomeResponse {
|
||||
fn from(o: KillOutcome) -> Self {
|
||||
match o {
|
||||
KillOutcome::Killed => Self::Killed,
|
||||
KillOutcome::AlreadyExited => Self::AlreadyExited,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn handle(agent: &MvpAgent, args: &acp::ExtRequest) -> ExtResult {
|
||||
match args.method.as_ref() {
|
||||
"x.ai/terminal/create" => {
|
||||
let req: CreateTerminalRequest = parse(args)?;
|
||||
let env: HashMap<String, String> = req
|
||||
.env
|
||||
.iter()
|
||||
.map(|e| (e.name.clone(), e.value.clone()))
|
||||
.collect();
|
||||
|
||||
let result = terminal::create_terminal(
|
||||
&req.session_id,
|
||||
&req.command,
|
||||
&req.args,
|
||||
env,
|
||||
req.cwd.as_deref(),
|
||||
req.output_byte_limit,
|
||||
)
|
||||
.await
|
||||
.map(|terminal_id| CreateTerminalResponse { terminal_id });
|
||||
|
||||
respond(result)
|
||||
}
|
||||
|
||||
"x.ai/terminal/kill" => {
|
||||
// Try PTY registry first, then piped terminal registry.
|
||||
let req: KillTerminalRequest = parse(args)?;
|
||||
|
||||
// PTY registry (agent-scoped, keyed by terminal_id alone)
|
||||
if terminal::pty_session::get_pty(&req.terminal_id)
|
||||
.await
|
||||
.is_some()
|
||||
{
|
||||
let already_exited = terminal::pty_session::is_exited(&req.terminal_id).await;
|
||||
terminal::pty_session::close_pty(&req.terminal_id)
|
||||
.await
|
||||
.ok();
|
||||
return respond(Ok::<_, String>(KillTerminalResponse {
|
||||
outcome: if already_exited {
|
||||
KillOutcomeResponse::AlreadyExited
|
||||
} else {
|
||||
KillOutcomeResponse::Killed
|
||||
},
|
||||
}));
|
||||
}
|
||||
|
||||
// Piped terminal registry (session-scoped)
|
||||
let session_id = match &req.session_id {
|
||||
Some(session_id) => Some(session_id.clone()),
|
||||
None => terminal::find_terminal_session_id(&req.terminal_id).await,
|
||||
};
|
||||
|
||||
if let Some(session_id) = session_id {
|
||||
let result = terminal::kill_terminal(&session_id, &req.terminal_id)
|
||||
.await
|
||||
.and_then(|opt| {
|
||||
opt.map(|outcome| KillTerminalResponse {
|
||||
outcome: outcome.into(),
|
||||
})
|
||||
.ok_or_else(|| ERR_TERMINAL_NOT_FOUND.to_string())
|
||||
});
|
||||
respond(result)
|
||||
} else {
|
||||
respond(Err::<KillTerminalResponse, _>(ERR_TERMINAL_NOT_FOUND))
|
||||
}
|
||||
}
|
||||
|
||||
"x.ai/terminal/output" => {
|
||||
let req: TerminalIdRequest = parse(args)?;
|
||||
let result = terminal::get_terminal_output(&req.session_id, &req.terminal_id)
|
||||
.await
|
||||
.map(TerminalOutputResponse::from)
|
||||
.ok_or(ERR_TERMINAL_NOT_FOUND);
|
||||
respond(result)
|
||||
}
|
||||
|
||||
"x.ai/terminal/wait_for_exit" => {
|
||||
let req: TerminalIdRequest = parse(args)?;
|
||||
let result = terminal::wait_for_terminal_exit(&req.session_id, &req.terminal_id)
|
||||
.await
|
||||
.map(ExitStatusResponse::from)
|
||||
.ok_or(ERR_TERMINAL_NOT_FOUND);
|
||||
respond(result)
|
||||
}
|
||||
|
||||
"x.ai/terminal/release" => {
|
||||
let req: TerminalIdRequest = parse(args)?;
|
||||
terminal::release_terminal(&req.session_id, &req.terminal_id).await;
|
||||
ExtMethodResult::success(ReleaseTerminalResponse {})
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
"x.ai/terminal/background" => {
|
||||
// Mark a terminal as backgrounded - the process keeps running but
|
||||
// waiting callers are notified so the agent can continue.
|
||||
//
|
||||
// Route through the session's tool bridge so the LocalTerminalBackend
|
||||
// actor unblocks the foreground waiter (BashTool::run). Also try the
|
||||
// StreamingLocalTerminalRunner registry for AcpTerminalAdapter-based sessions.
|
||||
let req: TerminalIdRequest = parse(args)?;
|
||||
agent
|
||||
.background_foreground_command(&req.session_id, &req.terminal_id)
|
||||
.await;
|
||||
terminal::background_terminal(&req.session_id, &req.terminal_id).await;
|
||||
ExtMethodResult::success(ReleaseTerminalResponse {})
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
"x.ai/terminal/pty/create" => {
|
||||
let req: PtyCreateRequest = parse(args)?;
|
||||
let env: HashMap<String, String> = req
|
||||
.env
|
||||
.iter()
|
||||
.map(|e| (e.name.clone(), e.value.clone()))
|
||||
.collect();
|
||||
let target_client_id = req.meta.map(|m| m.client_id).unwrap_or_default();
|
||||
let cwd = req.cwd.or_else(|| {
|
||||
req.session_id
|
||||
.as_ref()
|
||||
.and_then(|sid| agent.get_session_cwd(sid))
|
||||
.map(|p| p.to_string_lossy().into_owned())
|
||||
});
|
||||
|
||||
let result = terminal::pty_session::create_pty(
|
||||
req.shell.as_deref(),
|
||||
cwd.as_deref(),
|
||||
env,
|
||||
req.rows.unwrap_or(24),
|
||||
req.cols.unwrap_or(80),
|
||||
req.name.as_deref(),
|
||||
agent.gateway.clone(),
|
||||
target_client_id,
|
||||
)
|
||||
.await
|
||||
.map(|id| CreateTerminalResponse { terminal_id: id });
|
||||
|
||||
respond_pty(result)
|
||||
}
|
||||
|
||||
"x.ai/terminal/pty/load" => {
|
||||
let req: PtyLoadRequest = parse(args)?;
|
||||
let target_client_id = req.meta.map(|m| m.client_id).unwrap_or_default();
|
||||
let result =
|
||||
terminal::pty_session::load(&req.terminal_id, &agent.gateway, target_client_id)
|
||||
.await;
|
||||
respond_pty(result)
|
||||
}
|
||||
|
||||
"x.ai/terminal/pty/resize" => {
|
||||
let req: PtyResizeRequest = parse(args)?;
|
||||
respond_pty(
|
||||
terminal::pty_session::resize_pty(&req.terminal_id, req.rows, req.cols).await,
|
||||
)
|
||||
}
|
||||
|
||||
"x.ai/terminal/list" => {
|
||||
let terminals = terminal::list_terminals().await;
|
||||
respond(Ok::<_, String>(TerminalListResponse { terminals }))
|
||||
}
|
||||
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn handle_pty_input(params: &serde_json::Value) {
|
||||
use base64::Engine as _;
|
||||
|
||||
let Ok(input) = serde_json::from_value::<PtyInputNotification>(params.clone()) else {
|
||||
tracing::warn!("failed to parse pty input notification");
|
||||
return;
|
||||
};
|
||||
let Ok(bytes) = base64::engine::general_purpose::STANDARD.decode(&input.data) else {
|
||||
tracing::warn!("failed to decode pty input base64");
|
||||
return;
|
||||
};
|
||||
|
||||
if let Err(e) = terminal::pty_session::write_pty_input(&input.terminal_id, &bytes).await {
|
||||
tracing::warn!("pty input write failed: {e}");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,613 @@
|
||||
//! Handler for x.ai/git/worktree/* extension methods.
|
||||
|
||||
use agent_client_protocol as acp;
|
||||
use kigi_acp_lib::AcpAgentGatewaySender as GatewaySender;
|
||||
|
||||
use crate::agent::mvp_agent::MvpAgent;
|
||||
use crate::session::ExtMethodResult;
|
||||
use crate::session::persistence::LocalSessionResolutionKind;
|
||||
use crate::session::worktree::{
|
||||
ApplyWorktreeRequest, CreateWorktreeFromWorktreeRequest, CreateWorktreeRequest,
|
||||
CreateWorktreeResponse, RehydrateSessionRequest, RemoveWorktreeRequest,
|
||||
ResumeSessionInWorktreeRequest, WorktreeNotificationSender, WorktreeStatus, WorktreeType,
|
||||
create_jj_workspace, create_worktree_async, create_worktree_from_worktree_async,
|
||||
rehydrate_session_in_worktree, resolve_session_repo_wide, resume_session_in_worktree,
|
||||
};
|
||||
|
||||
type ExtResult = Result<acp::ExtResponse, acp::Error>;
|
||||
|
||||
const WORKTREE_EXT_LOG: &str = "xai_worktree";
|
||||
|
||||
/// Wrapper to send worktree progress notifications via gateway.
|
||||
#[derive(Clone)]
|
||||
struct GatewayWorktreeNotifier {
|
||||
gateway: GatewaySender,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl WorktreeNotificationSender for GatewayWorktreeNotifier {
|
||||
async fn send_worktree_status(&self, progress: WorktreeStatus) {
|
||||
let params = match serde_json::value::to_raw_value(&progress) {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
tracing::warn!("Failed to serialize worktree progress: {}", e);
|
||||
return;
|
||||
}
|
||||
};
|
||||
let notification = acp::ExtNotification::new("x.ai/git/worktree/status", params.into());
|
||||
if let Err(e) = self.gateway.send(notification).await {
|
||||
tracing::warn!("Failed to send worktree progress notification: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn to_response<T: serde::Serialize>(result: anyhow::Result<T>) -> ExtResult {
|
||||
ExtMethodResult::from_result(result)
|
||||
.to_ext_response()
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))
|
||||
}
|
||||
|
||||
/// Extract the worktree path from a `Creating` response for pinning.
|
||||
fn extract_creating_path(resp: &anyhow::Result<CreateWorktreeResponse>) -> Option<String> {
|
||||
if let Ok(CreateWorktreeResponse::Creating { worktree_path, .. }) = resp {
|
||||
Some(worktree_path.clone())
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
// ── ACP request types for worktree management ──────────────────────────────────────────────────
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ListWorktreeRequest {
|
||||
#[serde(default)]
|
||||
pub repo: Option<String>,
|
||||
#[serde(default)]
|
||||
pub r#type: Vec<String>,
|
||||
#[serde(default)]
|
||||
pub include_all: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ShowWorktreeRequest {
|
||||
pub id_or_path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct GcWorktreeRequest {
|
||||
#[serde(default)]
|
||||
pub dry_run: bool,
|
||||
/// Duration string like "7d", "24h", "30m", "60s".
|
||||
#[serde(default)]
|
||||
pub max_age: Option<String>,
|
||||
#[serde(default)]
|
||||
pub force: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct WorktreeDbPathResponse {
|
||||
pub path: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ResolveLocalForWorktreeResumeRequest {
|
||||
pub session_id: String,
|
||||
pub cwd: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ResolveLocalForWorktreeResumeResponse {
|
||||
pub found: bool,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub resolved_session_id: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub resolved_cwd: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub resolution_kind: Option<LocalSessionResolutionKind>,
|
||||
}
|
||||
fn parse_duration(s: &str) -> Result<i64, acp::Error> {
|
||||
let s = s.trim();
|
||||
let (num, mult) = if let Some(n) = s.strip_suffix('d') {
|
||||
(n, 86400i64)
|
||||
} else if let Some(n) = s.strip_suffix('h') {
|
||||
(n, 3600)
|
||||
} else if let Some(n) = s.strip_suffix('m') {
|
||||
(n, 60)
|
||||
} else if let Some(n) = s.strip_suffix('s') {
|
||||
(n, 1)
|
||||
} else {
|
||||
return Err(acp::Error::invalid_params().data(format!(
|
||||
"invalid duration: {s} (expected e.g. 7d, 24h, 30m, 60s)"
|
||||
)));
|
||||
};
|
||||
num.parse::<i64>()
|
||||
.map(|v| v * mult)
|
||||
.map_err(|_| acp::Error::invalid_params().data(format!("invalid number in duration: {s}")))
|
||||
}
|
||||
|
||||
fn log_effective_worktree_type(
|
||||
method: &str,
|
||||
request_worktree_type: Option<WorktreeType>,
|
||||
agent_default_worktree_type: crate::util::config::WorktreeType,
|
||||
effective_worktree_type: WorktreeType,
|
||||
) {
|
||||
tracing::info!(
|
||||
target: WORKTREE_EXT_LOG,
|
||||
method,
|
||||
request_worktree_type = ?request_worktree_type,
|
||||
agent_default_worktree_type = ?agent_default_worktree_type,
|
||||
effective_worktree_type = ?effective_worktree_type,
|
||||
"WORKTREE_REQUEST_SHELL: resolved effective worktree type"
|
||||
);
|
||||
}
|
||||
pub async fn handle(
|
||||
agent: &MvpAgent,
|
||||
ops: &kigi_workspace::WorkspaceOps,
|
||||
args: &acp::ExtRequest,
|
||||
) -> ExtResult {
|
||||
let worktree_type_default = agent.worktree_type;
|
||||
let restore_code_default = agent.restore_code;
|
||||
|
||||
match args.method.as_ref() {
|
||||
"x.ai/git/worktree/create" => {
|
||||
let mut req = serde_json::from_str::<CreateWorktreeRequest>(args.params.get())?;
|
||||
// Pre-dispatch: apply worktree_type default
|
||||
let request_worktree_type = req.worktree_type;
|
||||
if req.worktree_type.is_none() {
|
||||
req.worktree_type = Some(worktree_type_default.into());
|
||||
}
|
||||
log_effective_worktree_type(
|
||||
"x.ai/git/worktree/create",
|
||||
request_worktree_type,
|
||||
worktree_type_default,
|
||||
req.worktree_type.unwrap_or(worktree_type_default.into()),
|
||||
);
|
||||
let result = ops
|
||||
.dispatch(&req, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
// Post-dispatch: spawn async task for Creating variant.
|
||||
if let Ok(resp) = serde_json::from_value::<CreateWorktreeResponse>(result.clone())
|
||||
&& matches!(&resp, CreateWorktreeResponse::Creating { .. })
|
||||
{
|
||||
req.worktree_path = extract_creating_path(&Ok(resp));
|
||||
let notifier = GatewayWorktreeNotifier {
|
||||
gateway: agent.gateway.clone(),
|
||||
};
|
||||
let copy_context = agent.background_copy_context();
|
||||
tokio::task::spawn_local(async move {
|
||||
create_worktree_async(req, notifier, copy_context).await;
|
||||
});
|
||||
}
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/remove" => {
|
||||
let req = serde_json::from_str::<RemoveWorktreeRequest>(args.params.get())?;
|
||||
let result = ops
|
||||
.dispatch(&req, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/apply" => {
|
||||
let req = serde_json::from_str::<ApplyWorktreeRequest>(args.params.get())?;
|
||||
let result = ops
|
||||
.dispatch(&req, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
// Create a worktree from an existing worktree (used during session fork)
|
||||
"x.ai/git/worktree/create_from_worktree" => {
|
||||
let mut req =
|
||||
serde_json::from_str::<CreateWorktreeFromWorktreeRequest>(args.params.get())?;
|
||||
let request_worktree_type = req.worktree_type;
|
||||
// Apply default if not explicitly set in request
|
||||
if req.worktree_type.is_none() {
|
||||
req.worktree_type = Some(worktree_type_default.into());
|
||||
}
|
||||
log_effective_worktree_type(
|
||||
"x.ai/git/worktree/create_from_worktree",
|
||||
request_worktree_type,
|
||||
worktree_type_default,
|
||||
req.worktree_type.unwrap_or(worktree_type_default.into()),
|
||||
);
|
||||
// Dispatch prepare through workspace
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::PrepareWorktreeFromWorktreeReq {
|
||||
inner: req.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
|
||||
// Convert the serialized response back
|
||||
if let Some(err) = result.error {
|
||||
return Err(acp::Error::internal_error().data(err));
|
||||
}
|
||||
let response_value = result.response.unwrap_or(serde_json::Value::Null);
|
||||
let response: CreateWorktreeResponse =
|
||||
serde_json::from_value(response_value).map_err(|e| {
|
||||
acp::Error::internal_error()
|
||||
.data(format!("failed to deserialize response: {e}"))
|
||||
})?;
|
||||
|
||||
if result.spawn_task {
|
||||
// Pin the resolved path so the async task reuses it instead of
|
||||
// generating a new UUID via auto_label().
|
||||
req.resolved_dest_path = extract_creating_path(&Ok(response.clone()));
|
||||
let notifier = GatewayWorktreeNotifier {
|
||||
gateway: agent.gateway.clone(),
|
||||
};
|
||||
tokio::task::spawn_local(async move {
|
||||
create_worktree_from_worktree_async(req, notifier).await;
|
||||
});
|
||||
}
|
||||
|
||||
to_response(Ok(response))
|
||||
}
|
||||
// Synchronous variant - waits for worktree creation to complete
|
||||
"x.ai/git/worktree/create_from_worktree_sync" => {
|
||||
let mut req =
|
||||
serde_json::from_str::<CreateWorktreeFromWorktreeRequest>(args.params.get())?;
|
||||
|
||||
// For jj repos, use jj workspace add instead of git worktree
|
||||
let source_path = std::path::Path::new(&req.source_worktree_path);
|
||||
let resolved_root = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::GitResolveRootReq {
|
||||
cwd: source_path.to_path_buf(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.ok()
|
||||
.flatten();
|
||||
if let Some(git_root) = resolved_root {
|
||||
let vcs_kind = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::DetectVcsKindReq {
|
||||
path: git_root.clone(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.unwrap_or(kigi_workspace::session::git::VcsKind::Git);
|
||||
if vcs_kind.is_jj() {
|
||||
tracing::info!("using jj workspace for subagent isolation");
|
||||
return to_response(create_jj_workspace(&req).await);
|
||||
}
|
||||
}
|
||||
|
||||
let request_worktree_type = req.worktree_type;
|
||||
// Apply default if not explicitly set in request
|
||||
if req.worktree_type.is_none() {
|
||||
req.worktree_type = Some(worktree_type_default.into());
|
||||
}
|
||||
log_effective_worktree_type(
|
||||
"x.ai/git/worktree/create_from_worktree_sync",
|
||||
request_worktree_type,
|
||||
worktree_type_default,
|
||||
req.worktree_type.unwrap_or(worktree_type_default.into()),
|
||||
);
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::CreateWorktreeFromWorktreeSyncReq {
|
||||
inner: req.into_wire(),
|
||||
},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
// Resume a session in a fresh worktree.
|
||||
"x.ai/git/worktree/resume_session" => {
|
||||
let req = serde_json::from_str::<ResumeSessionInWorktreeRequest>(args.params.get())?;
|
||||
log_effective_worktree_type(
|
||||
"x.ai/git/worktree/resume_session",
|
||||
req.worktree_type,
|
||||
worktree_type_default,
|
||||
req.worktree_type.unwrap_or(worktree_type_default.into()),
|
||||
);
|
||||
let registry_client = agent.session_registry_client();
|
||||
let agent_id = crate::util::agent_id::agent_id();
|
||||
|
||||
to_response(
|
||||
resume_session_in_worktree(
|
||||
&req,
|
||||
ops,
|
||||
worktree_type_default,
|
||||
restore_code_default,
|
||||
registry_client.as_ref(),
|
||||
Some(agent.auth_manager.clone()),
|
||||
&agent_id,
|
||||
)
|
||||
.await,
|
||||
)
|
||||
}
|
||||
// ── Repo-wide session resolution ─────────────────────────────────
|
||||
"x.ai/session/resolve_local_for_worktree_resume" => {
|
||||
let req =
|
||||
serde_json::from_str::<ResolveLocalForWorktreeResumeRequest>(args.params.get())?;
|
||||
let result = resolve_session_repo_wide(&req.session_id, std::path::Path::new(&req.cwd));
|
||||
match result {
|
||||
Ok(Some(resolved)) => to_response(Ok(ResolveLocalForWorktreeResumeResponse {
|
||||
found: true,
|
||||
resolved_session_id: Some(resolved.session_id),
|
||||
resolved_cwd: Some(resolved.cwd),
|
||||
resolution_kind: Some(resolved.resolution_kind),
|
||||
})),
|
||||
Ok(None) => to_response(Ok(ResolveLocalForWorktreeResumeResponse {
|
||||
found: false,
|
||||
resolved_session_id: None,
|
||||
resolved_cwd: None,
|
||||
resolution_kind: None,
|
||||
})),
|
||||
Err(e) => {
|
||||
Err(acp::Error::internal_error()
|
||||
.data(format!("repo-wide resolution failed: {e}")))
|
||||
}
|
||||
}
|
||||
}
|
||||
// ── Session rehydration (devbox recovery) ─────────────────────────
|
||||
"x.ai/session/rehydrate" => {
|
||||
let req = serde_json::from_str::<RehydrateSessionRequest>(args.params.get())?;
|
||||
let registry_client = agent.session_registry_client();
|
||||
|
||||
to_response(rehydrate_session_in_worktree(&req, ops, registry_client.as_ref()).await)
|
||||
}
|
||||
// ── Worktree management methods ──────────────────────────────────
|
||||
"x.ai/git/worktree/list" => {
|
||||
let req: kigi_workspace::workspace_ops::WorktreeListReq =
|
||||
serde_json::from_str(args.params.get())
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
let result = ops
|
||||
.dispatch(&req, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/show" => {
|
||||
let req = serde_json::from_str::<ShowWorktreeRequest>(args.params.get())?;
|
||||
let op = kigi_workspace::workspace_ops::WorktreeShowReq {
|
||||
id_or_path: req.id_or_path,
|
||||
};
|
||||
let result = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/gc" => {
|
||||
let req = serde_json::from_str::<GcWorktreeRequest>(args.params.get())?;
|
||||
let max_age_secs = req.max_age.as_deref().map(parse_duration).transpose()?;
|
||||
let op = kigi_workspace::workspace_ops::WorktreeGcReq {
|
||||
dry_run: req.dry_run,
|
||||
max_age_secs,
|
||||
force: req.force,
|
||||
};
|
||||
let result = ops
|
||||
.dispatch(&op, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/db/stats" => {
|
||||
let result = ops
|
||||
.dispatch(&kigi_workspace::workspace_ops::WorktreeDbStatsReq {}, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/db/rebuild" => {
|
||||
let result = ops
|
||||
.dispatch(
|
||||
&kigi_workspace::workspace_ops::WorktreeDbRebuildReq {},
|
||||
None,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
"x.ai/git/worktree/db/path" => {
|
||||
let result = ops
|
||||
.dispatch(&kigi_workspace::workspace_ops::WorktreeDbPathReq {}, None)
|
||||
.await
|
||||
.map_err(|e| acp::Error::internal_error().data(e.to_string()))?;
|
||||
to_response(Ok(result))
|
||||
}
|
||||
_ => Err(acp::Error::method_not_found()),
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn list_request_all_defaults() {
|
||||
let req: ListWorktreeRequest = serde_json::from_str("{}").unwrap();
|
||||
assert!(req.repo.is_none());
|
||||
assert!(req.r#type.is_empty());
|
||||
assert!(!req.include_all);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_request_with_filters() {
|
||||
let json = r#"{"repo": "xai", "type": ["session", "fork"], "includeAll": true}"#;
|
||||
let req: ListWorktreeRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.repo.as_deref(), Some("xai"));
|
||||
assert_eq!(req.r#type, vec!["session", "fork"]);
|
||||
assert!(req.include_all);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn show_request_deserializes() {
|
||||
let json = r#"{"idOrPath": "wt-abc"}"#;
|
||||
let req: ShowWorktreeRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.id_or_path, "wt-abc");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gc_request_all_defaults() {
|
||||
let req: GcWorktreeRequest = serde_json::from_str("{}").unwrap();
|
||||
assert!(!req.dry_run);
|
||||
assert!(req.max_age.is_none());
|
||||
assert!(!req.force);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn gc_request_with_all_fields() {
|
||||
let json = r#"{"dryRun": true, "maxAge": "7d", "force": true}"#;
|
||||
let req: GcWorktreeRequest = serde_json::from_str(json).unwrap();
|
||||
assert!(req.dry_run);
|
||||
assert_eq!(req.max_age.as_deref(), Some("7d"));
|
||||
assert!(req.force);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_duration_valid_values() {
|
||||
assert_eq!(parse_duration("7d").unwrap(), 7 * 86400);
|
||||
assert_eq!(parse_duration("24h").unwrap(), 24 * 3600);
|
||||
assert_eq!(parse_duration("30m").unwrap(), 30 * 60);
|
||||
assert_eq!(parse_duration("60s").unwrap(), 60);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_duration_rejects_invalid() {
|
||||
assert!(parse_duration("bad").is_err());
|
||||
assert!(parse_duration("").is_err());
|
||||
assert!(parse_duration("7x").is_err());
|
||||
assert!(parse_duration("abcd").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn db_path_response_serializes() {
|
||||
let resp = WorktreeDbPathResponse {
|
||||
path: "/home/user/.kigi/worktrees.db".into(),
|
||||
};
|
||||
let json = serde_json::to_string(&resp).unwrap();
|
||||
assert!(json.contains("\"path\":\"/home/user/.kigi/worktrees.db\""));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_request_rejects_both_fields_set() {
|
||||
use crate::session::worktree::{
|
||||
BackgroundCopyContext, RemoveWorktreeRequest, remove_worktree,
|
||||
};
|
||||
|
||||
let req = RemoveWorktreeRequest {
|
||||
worktree_path: Some("/a".into()),
|
||||
id_or_path: Some("b".into()),
|
||||
force: false,
|
||||
dry_run: false,
|
||||
};
|
||||
let ctx = BackgroundCopyContext::new();
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let result = rt.block_on(remove_worktree(&req, &ctx));
|
||||
assert!(result.is_err());
|
||||
let msg = result.unwrap_err().to_string();
|
||||
assert!(msg.contains("exactly one"), "unexpected error: {msg}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_request_rejects_neither_field_set() {
|
||||
use crate::session::worktree::{
|
||||
BackgroundCopyContext, RemoveWorktreeRequest, remove_worktree,
|
||||
};
|
||||
|
||||
let req = RemoveWorktreeRequest {
|
||||
worktree_path: None,
|
||||
id_or_path: None,
|
||||
force: false,
|
||||
dry_run: false,
|
||||
};
|
||||
let ctx = BackgroundCopyContext::new();
|
||||
let rt = tokio::runtime::Builder::new_current_thread()
|
||||
.enable_all()
|
||||
.build()
|
||||
.unwrap();
|
||||
let result = rt.block_on(remove_worktree(&req, &ctx));
|
||||
assert!(result.is_err());
|
||||
let msg = result.unwrap_err().to_string();
|
||||
assert!(
|
||||
msg.contains("either worktreePath or idOrPath"),
|
||||
"unexpected error: {msg}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn db_rebuild_response_carries_report_not_null() {
|
||||
// Regression: forwarding `()` instead of the report yields `result: null`,
|
||||
// which the CLI rejects with "ACP response missing result field".
|
||||
let report = serde_json::json!({
|
||||
"discovered": 5,
|
||||
"registered": 3,
|
||||
"already_tracked": 2,
|
||||
});
|
||||
let resp = to_response(Ok(report)).expect("rebuild response should serialize");
|
||||
let wire: serde_json::Value = serde_json::from_str(resp.0.get()).unwrap();
|
||||
|
||||
let result = wire
|
||||
.get("result")
|
||||
.expect("envelope must carry a result field");
|
||||
assert!(!result.is_null(), "rebuild result must not be null");
|
||||
assert_eq!(result["discovered"], 5);
|
||||
assert_eq!(result["registered"], 3);
|
||||
assert_eq!(result["already_tracked"], 2);
|
||||
assert!(wire.get("error").is_none() || wire["error"].is_null());
|
||||
}
|
||||
|
||||
// === Tests for repo-wide session resolution ACP types ===
|
||||
|
||||
#[test]
|
||||
fn resolve_local_request_deserializes() {
|
||||
let json = r#"{"sessionId": "sess-abc", "cwd": "/repo/main"}"#;
|
||||
let req: ResolveLocalForWorktreeResumeRequest = serde_json::from_str(json).unwrap();
|
||||
assert_eq!(req.session_id, "sess-abc");
|
||||
assert_eq!(req.cwd, "/repo/main");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_local_response_found_serializes() {
|
||||
use crate::session::persistence::LocalSessionResolutionKind;
|
||||
let resp = ResolveLocalForWorktreeResumeResponse {
|
||||
found: true,
|
||||
resolved_session_id: Some("sess-123".into()),
|
||||
resolved_cwd: Some("/repo/wt-1".into()),
|
||||
resolution_kind: Some(LocalSessionResolutionKind::SameRepoDifferentCwd),
|
||||
};
|
||||
let json = serde_json::to_string(&resp).unwrap();
|
||||
assert!(json.contains("\"found\":true"));
|
||||
assert!(json.contains("\"resolvedSessionId\":\"sess-123\""));
|
||||
assert!(json.contains("\"resolvedCwd\":\"/repo/wt-1\""));
|
||||
assert!(json.contains("\"resolutionKind\":"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn resolve_local_response_not_found_omits_optional_fields() {
|
||||
let resp = ResolveLocalForWorktreeResumeResponse {
|
||||
found: false,
|
||||
resolved_session_id: None,
|
||||
resolved_cwd: None,
|
||||
resolution_kind: None,
|
||||
};
|
||||
let json = serde_json::to_string(&resp).unwrap();
|
||||
assert!(json.contains("\"found\":false"));
|
||||
assert!(!json.contains("resolvedSessionId"));
|
||||
assert!(!json.contains("resolvedCwd"));
|
||||
assert!(!json.contains("resolutionKind"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user