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,327 @@
|
||||
//! Remote sandbox client for cli-chat-proxy.
|
||||
//!
|
||||
//! This module provides an HTTP client to interact with cli-chat-proxy
|
||||
//! for managing sandbox sessions and environments via REST API.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::auth::{AuthManager, GrokComConfig};
|
||||
use anyhow::{Context, Result, bail};
|
||||
use serde::de::DeserializeOwned;
|
||||
|
||||
// Re-export sandbox API types from cli-chat-proxy-types for convenience.
|
||||
// Sorted alphabetically; see sandbox_types.rs for logical grouping.
|
||||
pub use prod_mc_cli_chat_proxy_types::{
|
||||
SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse,
|
||||
SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest,
|
||||
SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse,
|
||||
SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse,
|
||||
SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse,
|
||||
SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse,
|
||||
SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse,
|
||||
SandboxTerminateRequest, SandboxUpdateEnvironmentRequest,
|
||||
};
|
||||
|
||||
// ============================================================================
|
||||
// Sandbox Client
|
||||
// ============================================================================
|
||||
|
||||
/// HTTP client for interacting with the sandbox API via cli-chat-proxy.
|
||||
///
|
||||
/// Path parameters (`session_id`, `environment_id`) are interpolated directly
|
||||
/// into URLs without percent-encoding. This is safe because these IDs are
|
||||
/// UUIDs in practice. If ID formats ever change to include URL-unsafe
|
||||
/// characters, the `format!()` calls should be updated to use percent-encoding.
|
||||
pub struct SandboxClient {
|
||||
client: reqwest::Client,
|
||||
base_url: String,
|
||||
auth_manager: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl SandboxClient {
|
||||
pub fn new(base_url: impl Into<String>, auth_manager: Arc<AuthManager>) -> Self {
|
||||
Self {
|
||||
client: crate::http::shared_client(),
|
||||
base_url: base_url.into(),
|
||||
auth_manager,
|
||||
}
|
||||
}
|
||||
|
||||
/// Returns the base URL.
|
||||
pub fn base_url(&self) -> &str {
|
||||
&self.base_url
|
||||
}
|
||||
|
||||
// Do not set Content-Type — callers use .json() and reqwest .header() appends.
|
||||
async fn auth_headers(
|
||||
&self,
|
||||
builder: reqwest::RequestBuilder,
|
||||
) -> Result<reqwest::RequestBuilder> {
|
||||
let auth = self
|
||||
.auth_manager
|
||||
.auth()
|
||||
.await
|
||||
.context("failed to resolve sandbox auth")?;
|
||||
let mut builder = builder
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header("X-XAI-Token-Auth", GrokComConfig::default().token_header)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION);
|
||||
|
||||
if let Some(email) = &auth.email {
|
||||
builder = builder.header("x-email", email);
|
||||
}
|
||||
|
||||
builder = builder
|
||||
.header(
|
||||
"x-grok-client-identifier",
|
||||
crate::http::process_client_identifier(),
|
||||
)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
);
|
||||
|
||||
Ok(kigi_file_utils::trace_context::inject_trace_context_into_request(builder))
|
||||
}
|
||||
|
||||
/// Check an HTTP response for errors, then deserialize the JSON body.
|
||||
async fn parse_response<T: DeserializeOwned>(
|
||||
response: reqwest::Response,
|
||||
operation: &str,
|
||||
) -> Result<T> {
|
||||
if !response.status().is_success() {
|
||||
let status = response.status().as_u16();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
bail!("{operation} failed: {status} - {body}");
|
||||
}
|
||||
response
|
||||
.json()
|
||||
.await
|
||||
.with_context(|| format!("failed to parse {operation} response"))
|
||||
}
|
||||
|
||||
/// Check an HTTP response for errors, discarding the body.
|
||||
async fn check_response(response: reqwest::Response, operation: &str) -> Result<()> {
|
||||
if !response.status().is_success() {
|
||||
let status = response.status().as_u16();
|
||||
let body = response.text().await.unwrap_or_default();
|
||||
bail!("{operation} failed: {status} - {body}");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fork an existing sandbox session.
|
||||
pub async fn fork_session(&self, request: &SandboxForkRequest) -> Result<SandboxForkResponse> {
|
||||
let url = format!("{}/sandbox/sessions/fork", self.base_url);
|
||||
let response = self
|
||||
.auth_headers(self.client.post(&url))
|
||||
.await?
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send fork session request")?;
|
||||
Self::parse_response(response, "fork session").await
|
||||
}
|
||||
|
||||
/// Terminate a sandbox session.
|
||||
pub async fn terminate_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
request: &SandboxTerminateRequest,
|
||||
) -> Result<()> {
|
||||
let mut url = format!("{}/sandbox/sessions/{}", self.base_url, session_id);
|
||||
if let Some(env_id) = &request.environment_id {
|
||||
url = format!("{}?environmentId={}", url, env_id);
|
||||
}
|
||||
|
||||
let response = self
|
||||
.auth_headers(self.client.delete(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send terminate session request")?;
|
||||
|
||||
if response.status().as_u16() == 404 {
|
||||
bail!("session not found: {session_id}");
|
||||
}
|
||||
Self::check_response(response, "terminate session").await
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Session Lifecycle
|
||||
// ========================================================================
|
||||
|
||||
/// Start a sandbox session (non-TUI).
|
||||
pub async fn start_session(
|
||||
&self,
|
||||
request: &SandboxStartRequest,
|
||||
) -> Result<SandboxStartResponse> {
|
||||
let url = format!("{}/sandbox/sessions/start", self.base_url);
|
||||
let response = self
|
||||
.auth_headers(self.client.post(&url))
|
||||
.await?
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send start session request")?;
|
||||
Self::parse_response(response, "start session").await
|
||||
}
|
||||
|
||||
/// Get sandbox session status.
|
||||
pub async fn get_session_status(&self, session_id: &str) -> Result<SandboxStatusResponse> {
|
||||
let url = format!("{}/sandbox/sessions/{}/status", self.base_url, session_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.get(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send get session status request")?;
|
||||
Self::parse_response(response, "get session status").await
|
||||
}
|
||||
|
||||
/// Get sandbox session logs.
|
||||
pub async fn get_session_logs(&self, session_id: &str) -> Result<SandboxLogsResponse> {
|
||||
let url = format!("{}/sandbox/sessions/{}/logs", self.base_url, session_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.get(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send get session logs request")?;
|
||||
Self::parse_response(response, "get session logs").await
|
||||
}
|
||||
|
||||
/// Hibernate a sandbox session (snapshot rootfs to GCS and terminate).
|
||||
pub async fn hibernate_session(&self, session_id: &str) -> Result<SandboxHibernateResponse> {
|
||||
let url = format!(
|
||||
"{}/sandbox/sessions/{}/hibernate",
|
||||
self.base_url, session_id
|
||||
);
|
||||
let response = self
|
||||
.auth_headers(self.client.post(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send hibernate session request")?;
|
||||
Self::parse_response(response, "hibernate session").await
|
||||
}
|
||||
|
||||
/// Restore a previously hibernated sandbox session from its snapshot.
|
||||
pub async fn restore_session(
|
||||
&self,
|
||||
session_id: &str,
|
||||
request: &SandboxRestoreRequest,
|
||||
) -> Result<SandboxRestoreResponse> {
|
||||
let url = format!("{}/sandbox/sessions/{}/restore", self.base_url, session_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.post(&url))
|
||||
.await?
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send restore session request")?;
|
||||
Self::parse_response(response, "restore session").await
|
||||
}
|
||||
|
||||
// ========================================================================
|
||||
// Environment CRUD
|
||||
// ========================================================================
|
||||
|
||||
/// List sandbox environments.
|
||||
pub async fn list_environments(
|
||||
&self,
|
||||
request: &SandboxListEnvironmentsRequest,
|
||||
) -> Result<SandboxListEnvironmentsResponse> {
|
||||
let url = format!("{}/sandbox/environments", self.base_url);
|
||||
let mut builder = self.auth_headers(self.client.get(&url)).await?;
|
||||
if let Some(page) = request.page {
|
||||
builder = builder.query(&[("page", page)]);
|
||||
}
|
||||
if let Some(page_size) = request.page_size {
|
||||
builder = builder.query(&[("pageSize", page_size)]);
|
||||
}
|
||||
let response = builder
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send list environments request")?;
|
||||
Self::parse_response(response, "list environments").await
|
||||
}
|
||||
|
||||
/// Create a new sandbox environment.
|
||||
pub async fn create_environment(
|
||||
&self,
|
||||
request: &SandboxCreateEnvironmentRequest,
|
||||
) -> Result<SandboxEnvironmentResponse> {
|
||||
let url = format!("{}/sandbox/environments", self.base_url);
|
||||
let response = self
|
||||
.auth_headers(self.client.post(&url))
|
||||
.await?
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send create environment request")?;
|
||||
Self::parse_response(response, "create environment").await
|
||||
}
|
||||
|
||||
/// Get a sandbox environment by ID.
|
||||
pub async fn get_environment(
|
||||
&self,
|
||||
environment_id: &str,
|
||||
) -> Result<SandboxEnvironmentResponse> {
|
||||
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.get(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send get environment request")?;
|
||||
Self::parse_response(response, "get environment").await
|
||||
}
|
||||
|
||||
/// Update a sandbox environment.
|
||||
pub async fn update_environment(
|
||||
&self,
|
||||
environment_id: &str,
|
||||
request: &SandboxUpdateEnvironmentRequest,
|
||||
) -> Result<SandboxEnvironmentResponse> {
|
||||
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.put(&url))
|
||||
.await?
|
||||
.json(request)
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send update environment request")?;
|
||||
Self::parse_response(response, "update environment").await
|
||||
}
|
||||
|
||||
/// Delete a sandbox environment.
|
||||
pub async fn delete_environment(&self, environment_id: &str) -> Result<()> {
|
||||
let url = format!("{}/sandbox/environments/{}", self.base_url, environment_id);
|
||||
let response = self
|
||||
.auth_headers(self.client.delete(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send delete environment request")?;
|
||||
Self::check_response(response, "delete environment").await
|
||||
}
|
||||
|
||||
/// List preinstalled packages available for sandbox environments.
|
||||
pub async fn list_preinstalled_packages(
|
||||
&self,
|
||||
) -> Result<SandboxListPreinstalledPackagesResponse> {
|
||||
let url = format!(
|
||||
"{}/sandbox/environments/preinstalled-packages",
|
||||
self.base_url
|
||||
);
|
||||
let response = self
|
||||
.auth_headers(self.client.get(&url))
|
||||
.await?
|
||||
.send()
|
||||
.await
|
||||
.context("failed to send list preinstalled packages request")?;
|
||||
Self::parse_response(response, "list preinstalled packages").await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
//! grok.com chat-product model catalog (`POST /rest/modes`) — the models
|
||||
//! grok-web's chat picker shows, distinct from the CLI `/v1/models` build
|
||||
//! catalog. Transport only; cache + ACP mapping live in
|
||||
//! [`crate::agent::chat_modes`].
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
const KIGI_WEB_URL: &str = "https://grok.com";
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Mode {
|
||||
#[serde(default)]
|
||||
pub id: String,
|
||||
#[serde(default)]
|
||||
pub title: String,
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
#[serde(default)]
|
||||
pub badge_text: Option<String>,
|
||||
#[serde(default)]
|
||||
pub availability: ModeAvailability,
|
||||
#[serde(default)]
|
||||
pub icon_hint: String,
|
||||
#[serde(default)]
|
||||
pub tags: Vec<String>,
|
||||
}
|
||||
|
||||
impl Mode {
|
||||
pub fn is_available(&self) -> bool {
|
||||
self.availability.available.is_some()
|
||||
}
|
||||
}
|
||||
|
||||
/// proto3-JSON oneof: exactly one field is present.
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ModeAvailability {
|
||||
#[serde(default)]
|
||||
pub available: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub unavailable: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub requires_upgrade: Option<serde_json::Value>,
|
||||
#[serde(default)]
|
||||
pub coming_soon: Option<serde_json::Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct ListModesResponse {
|
||||
#[serde(default)]
|
||||
pub modes: Vec<Mode>,
|
||||
#[serde(default)]
|
||||
pub default_mode_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ChatModelsError {
|
||||
#[error("no grok.com credentials")]
|
||||
NoAuth,
|
||||
#[error("request timed out")]
|
||||
Timeout,
|
||||
#[error("network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
#[error("request failed: {status}")]
|
||||
Http { status: u16 },
|
||||
#[error("parse error: {0}")]
|
||||
Parse(#[from] serde_json::Error),
|
||||
}
|
||||
|
||||
/// Stateless transport for `POST /rest/modes`; caching lives in
|
||||
/// [`crate::agent::chat_modes::ChatModesManager`].
|
||||
pub struct ChatModelsClient {
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
auth: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl ChatModelsClient {
|
||||
pub fn new(auth: Arc<AuthManager>) -> Self {
|
||||
let base_url = std::env::var("KIGI_MODES_BASE_URL")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.or_else(|| {
|
||||
std::env::var("KIGI_CONVERSATIONS_BASE_URL")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
})
|
||||
.or_else(|| {
|
||||
std::env::var("KIGI_CODE_WEB_URL")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
})
|
||||
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
|
||||
Self {
|
||||
http: crate::http::shared_client(),
|
||||
base_url,
|
||||
auth,
|
||||
}
|
||||
}
|
||||
|
||||
/// Gated only on a valid grok.com bearer — deliberately NOT `is_xai_auth()`
|
||||
/// (unlike workspaces/conversations), since `/rest/modes` is the public chat
|
||||
/// endpoint and that gate would exclude API-key / cached-token chat users.
|
||||
pub async fn list_modes(&self, locale: &str) -> Result<ListModesResponse, ChatModelsError> {
|
||||
let auth = self
|
||||
.auth
|
||||
.auth()
|
||||
.await
|
||||
.map_err(|_| ChatModelsError::NoAuth)?;
|
||||
|
||||
let url = format!("{}/rest/modes", self.base_url);
|
||||
let body = serde_json::json!({ "locale": locale });
|
||||
let mut builder = self
|
||||
.http
|
||||
.post(&url)
|
||||
.json(&body)
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header(
|
||||
"X-XAI-Token-Auth",
|
||||
self.auth.grok_com_config().token_header.clone(),
|
||||
)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
"x-grok-client-identifier",
|
||||
crate::http::process_client_identifier(),
|
||||
)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.header(reqwest::header::ACCEPT, "application/json");
|
||||
if let Some(email) = &auth.email {
|
||||
builder = builder.header("x-email", email);
|
||||
}
|
||||
let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder);
|
||||
|
||||
let response = builder.send().await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
return Err(ChatModelsError::Http {
|
||||
status: status.as_u16(),
|
||||
});
|
||||
}
|
||||
|
||||
let bytes = response.bytes().await?;
|
||||
let resp: ListModesResponse = serde_json::from_slice(&bytes)?;
|
||||
Ok(resp)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn modes_parse_camelcase_wire() {
|
||||
let json = serde_json::json!({
|
||||
"modes": [{
|
||||
"id": "auto",
|
||||
"title": "Auto",
|
||||
"description": "Picks the best model",
|
||||
"badgeText": "New",
|
||||
"availability": { "available": {} },
|
||||
"iconHint": "rocket",
|
||||
"tags": ["TAG_PRIMARY"]
|
||||
}, {
|
||||
"id": "heavy",
|
||||
"title": "Heavy",
|
||||
"availability": { "requiresUpgrade": { "message": "Upgrade" } }
|
||||
}],
|
||||
"defaultModeId": "auto"
|
||||
});
|
||||
let resp: ListModesResponse = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(resp.modes.len(), 2);
|
||||
assert_eq!(resp.default_mode_id, "auto");
|
||||
let auto = &resp.modes[0];
|
||||
assert_eq!(auto.id, "auto");
|
||||
assert_eq!(auto.title, "Auto");
|
||||
assert_eq!(auto.badge_text.as_deref(), Some("New"));
|
||||
assert_eq!(auto.icon_hint, "rocket");
|
||||
assert_eq!(auto.tags, vec!["TAG_PRIMARY".to_string()]);
|
||||
assert!(auto.is_available());
|
||||
assert!(!resp.modes[1].is_available());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_fields_default_gracefully() {
|
||||
let json = serde_json::json!({ "modes": [{ "id": "m1" }] });
|
||||
let resp: ListModesResponse = serde_json::from_value(json).unwrap();
|
||||
let m = &resp.modes[0];
|
||||
assert_eq!(m.id, "m1");
|
||||
assert!(m.title.is_empty());
|
||||
assert!(m.description.is_empty());
|
||||
assert!(m.badge_text.is_none());
|
||||
// No availability field on the wire → not selectable.
|
||||
assert!(!m.is_available());
|
||||
assert!(resp.default_mode_id.is_empty());
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,309 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use crate::auth::{AuthManager, GrokAuth};
|
||||
|
||||
const KIGI_WEB_URL: &str = "https://grok.com";
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Conversation {
|
||||
#[serde(default)]
|
||||
pub conversation_id: String,
|
||||
#[serde(default)]
|
||||
pub title: String,
|
||||
#[serde(default)]
|
||||
pub starred: bool,
|
||||
#[serde(default)]
|
||||
pub create_time: Option<String>,
|
||||
#[serde(default)]
|
||||
pub modify_time: Option<String>,
|
||||
#[serde(default)]
|
||||
pub workspaces: Vec<Workspace>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Workspace {
|
||||
#[serde(default)]
|
||||
pub workspace_id: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ConvQuery {
|
||||
pub page_size: i64,
|
||||
pub page_token: Option<String>,
|
||||
pub search_query: Option<String>,
|
||||
pub workspace_id: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ListConversationsPage {
|
||||
pub conversations: Vec<Conversation>,
|
||||
pub next_page_token: Option<String>,
|
||||
}
|
||||
|
||||
/// Body for `PUT /rest/app-chat/conversations/{id}` (grok-web `chatUpdateConversation`).
|
||||
#[derive(Debug, Clone, Default, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct UpdateConversationBody {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub title: Option<String>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub starred: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum ConvError {
|
||||
#[error("no OAuth credentials for conversations:read")]
|
||||
NoOauth,
|
||||
#[error("network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
#[error("request failed: {status}")]
|
||||
Http { status: u16 },
|
||||
#[error("parse error: {0}")]
|
||||
Parse(#[from] serde_json::Error),
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ListConversationsResponseWire {
|
||||
#[serde(default)]
|
||||
conversations: Vec<Conversation>,
|
||||
#[serde(default)]
|
||||
next_page_token: Option<String>,
|
||||
#[serde(default)]
|
||||
text_search_matches: Vec<ListConversationsMatchWire>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ListConversationsMatchWire {
|
||||
#[serde(default)]
|
||||
conversation: Option<Conversation>,
|
||||
}
|
||||
|
||||
pub struct ConversationsClient {
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
auth: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl ConversationsClient {
|
||||
pub fn new(auth: Arc<AuthManager>) -> Self {
|
||||
let base_url = std::env::var("KIGI_CONVERSATIONS_BASE_URL")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
.or_else(|| {
|
||||
std::env::var("KIGI_CODE_WEB_URL")
|
||||
.ok()
|
||||
.filter(|s| !s.is_empty())
|
||||
})
|
||||
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
|
||||
Self {
|
||||
http: crate::http::shared_client(),
|
||||
base_url,
|
||||
auth,
|
||||
}
|
||||
}
|
||||
|
||||
async fn require_xai_auth(&self) -> Result<GrokAuth, ConvError> {
|
||||
let auth = self.auth.auth().await.map_err(|_| ConvError::NoOauth)?;
|
||||
if !auth.is_xai_auth() {
|
||||
return Err(ConvError::NoOauth);
|
||||
}
|
||||
Ok(auth)
|
||||
}
|
||||
|
||||
fn apply_auth_headers(
|
||||
&self,
|
||||
builder: reqwest::RequestBuilder,
|
||||
auth: &GrokAuth,
|
||||
) -> reqwest::RequestBuilder {
|
||||
let mut builder = builder
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header(
|
||||
"X-XAI-Token-Auth",
|
||||
self.auth.grok_com_config().token_header.clone(),
|
||||
)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
"x-grok-client-identifier",
|
||||
crate::http::process_client_identifier(),
|
||||
)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.header(reqwest::header::ACCEPT, "application/json");
|
||||
if let Some(email) = &auth.email {
|
||||
builder = builder.header("x-email", email);
|
||||
}
|
||||
kigi_file_utils::trace_context::inject_trace_context_into_request(builder)
|
||||
}
|
||||
|
||||
pub async fn list_conversations(
|
||||
&self,
|
||||
q: &ConvQuery,
|
||||
) -> Result<ListConversationsPage, ConvError> {
|
||||
let auth = self.require_xai_auth().await?;
|
||||
|
||||
let url = format!("{}/rest/app-chat/conversations", self.base_url);
|
||||
let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())];
|
||||
if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("pageToken", token.to_owned()));
|
||||
}
|
||||
if let Some(search) = q.search_query.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("searchQuery", search.to_owned()));
|
||||
}
|
||||
if let Some(workspace) = q.workspace_id.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("workspaceId", workspace.to_owned()));
|
||||
}
|
||||
|
||||
let builder = self.apply_auth_headers(self.http.get(&url).query(&query), &auth);
|
||||
|
||||
let response = builder.send().await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
return Err(ConvError::Http {
|
||||
status: status.as_u16(),
|
||||
});
|
||||
}
|
||||
|
||||
let bytes = response.bytes().await?;
|
||||
let wire: ListConversationsResponseWire = serde_json::from_slice(&bytes)?;
|
||||
|
||||
let searching = q.search_query.as_deref().is_some_and(|s| !s.is_empty());
|
||||
// During an active search, results come exclusively from
|
||||
// `text_search_matches`. Never fall back to `wire.conversations` here:
|
||||
// an empty match set means "no hits", and the server may return
|
||||
// recent/unfiltered conversations in `conversations` that are NOT search
|
||||
// matches — surfacing those would be wrong.
|
||||
let conversations = if searching {
|
||||
wire.text_search_matches
|
||||
.into_iter()
|
||||
.filter_map(|m| m.conversation)
|
||||
.collect()
|
||||
} else {
|
||||
wire.conversations
|
||||
};
|
||||
|
||||
Ok(ListConversationsPage {
|
||||
conversations,
|
||||
next_page_token: wire.next_page_token.filter(|t| !t.is_empty()),
|
||||
})
|
||||
}
|
||||
|
||||
/// `PUT /rest/app-chat/conversations/{conversation_id}` — rename and/or star.
|
||||
pub async fn update_conversation(
|
||||
&self,
|
||||
conversation_id: &str,
|
||||
body: &UpdateConversationBody,
|
||||
) -> Result<(), ConvError> {
|
||||
let auth = self.require_xai_auth().await?;
|
||||
let url = format!(
|
||||
"{}/rest/app-chat/conversations/{}",
|
||||
self.base_url,
|
||||
urlencoding::encode(conversation_id)
|
||||
);
|
||||
let builder = self
|
||||
.apply_auth_headers(self.http.put(&url), &auth)
|
||||
.json(body);
|
||||
|
||||
let response = builder.send().await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
return Err(ConvError::Http {
|
||||
status: status.as_u16(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// `DELETE /rest/app-chat/conversations/soft/{conversation_id}` — soft-delete.
|
||||
pub async fn soft_delete_conversation(&self, conversation_id: &str) -> Result<(), ConvError> {
|
||||
let auth = self.require_xai_auth().await?;
|
||||
let url = format!(
|
||||
"{}/rest/app-chat/conversations/soft/{}",
|
||||
self.base_url,
|
||||
urlencoding::encode(conversation_id)
|
||||
);
|
||||
let builder = self.apply_auth_headers(self.http.delete(&url), &auth);
|
||||
|
||||
let response = builder.send().await?;
|
||||
let status = response.status();
|
||||
// 404 = already soft-deleted; keep deletion idempotent like the
|
||||
// build path's `classify_remote_delete`.
|
||||
if !status.is_success() && status.as_u16() != 404 {
|
||||
return Err(ConvError::Http {
|
||||
status: status.as_u16(),
|
||||
});
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn conversation_parses_camelcase_wire() {
|
||||
let json = serde_json::json!({
|
||||
"conversations": [{
|
||||
"conversationId": "conv_abc",
|
||||
"title": "Compare GPU vendors",
|
||||
"starred": true,
|
||||
"createTime": "2026-06-18T17:30:00Z",
|
||||
"modifyTime": "2026-06-18T18:02:00Z",
|
||||
"workspaces": [{ "workspaceId": "ws_9f3a" }]
|
||||
}],
|
||||
"nextPageToken": "tok2"
|
||||
});
|
||||
let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(wire.conversations.len(), 1);
|
||||
let c = &wire.conversations[0];
|
||||
assert_eq!(c.conversation_id, "conv_abc");
|
||||
assert_eq!(c.title, "Compare GPU vendors");
|
||||
assert!(c.starred);
|
||||
assert_eq!(c.modify_time.as_deref(), Some("2026-06-18T18:02:00Z"));
|
||||
assert_eq!(c.workspaces[0].workspace_id, "ws_9f3a");
|
||||
assert_eq!(wire.next_page_token.as_deref(), Some("tok2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_fields_default_gracefully() {
|
||||
let json = serde_json::json!({ "conversations": [{ "conversationId": "c1" }] });
|
||||
let wire: ListConversationsResponseWire = serde_json::from_value(json).unwrap();
|
||||
let c = &wire.conversations[0];
|
||||
assert_eq!(c.conversation_id, "c1");
|
||||
assert!(c.title.is_empty());
|
||||
assert!(c.modify_time.is_none());
|
||||
assert!(c.create_time.is_none());
|
||||
assert!(c.workspaces.is_empty());
|
||||
assert!(wire.next_page_token.is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn update_body_serializes_only_set_fields() {
|
||||
let title_only = UpdateConversationBody {
|
||||
title: Some("New title".into()),
|
||||
starred: None,
|
||||
};
|
||||
assert_eq!(
|
||||
serde_json::to_value(&title_only).unwrap(),
|
||||
serde_json::json!({ "title": "New title" })
|
||||
);
|
||||
|
||||
let both = UpdateConversationBody {
|
||||
title: Some("T".into()),
|
||||
starred: Some(true),
|
||||
};
|
||||
assert_eq!(
|
||||
serde_json::to_value(&both).unwrap(),
|
||||
serde_json::json!({ "title": "T", "starred": true })
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
//! Remote storage client for the backend.
|
||||
|
||||
pub mod agent;
|
||||
pub mod chat_models_client;
|
||||
pub mod client;
|
||||
pub mod conversations_client;
|
||||
pub mod pull;
|
||||
#[cfg(test)]
|
||||
mod pull_smoke_test;
|
||||
pub mod sync;
|
||||
pub mod workspaces_client;
|
||||
|
||||
pub use agent::{
|
||||
SandboxClient, SandboxCreateEnvironmentRequest, SandboxEnvironment, SandboxEnvironmentResponse,
|
||||
SandboxEnvironmentVariable, SandboxEnvironmentWithMetadata, SandboxForkRequest,
|
||||
SandboxForkResponse, SandboxForkedSession, SandboxHibernateResponse,
|
||||
SandboxListEnvironmentsRequest, SandboxListEnvironmentsResponse,
|
||||
SandboxListPreinstalledPackagesResponse, SandboxLogsExitCodes, SandboxLogsResponse,
|
||||
SandboxMode, SandboxPreinstalledPackage, SandboxRestoreRequest, SandboxRestoreResponse,
|
||||
SandboxSecretInput, SandboxStartRequest, SandboxStartResponse, SandboxStatusResponse,
|
||||
SandboxTerminateRequest, SandboxUpdateEnvironmentRequest,
|
||||
};
|
||||
pub use chat_models_client::{
|
||||
ChatModelsClient, ChatModelsError, ListModesResponse, Mode, ModeAvailability,
|
||||
};
|
||||
pub use client::{
|
||||
BackendClient, BackendError, FetchModelsResult, FetchedBundle, fetch_bundle,
|
||||
fetch_login_device_flow, fetch_settings_blocking, fetch_subagent_bundle, share_url,
|
||||
};
|
||||
pub(crate) use client::{DEFAULT_CONTEXT_WINDOW, fetch_models_blocking, models_list_url};
|
||||
pub use conversations_client::{
|
||||
ConvError, ConvQuery, Conversation, ConversationsClient, ListConversationsPage,
|
||||
UpdateConversationBody,
|
||||
};
|
||||
pub use pull::{PullResult, pull_session_to_local};
|
||||
pub use sync::RemoteSync;
|
||||
pub use workspaces_client::{ListWorkspacesPage, Workspace, WorkspacesClient, WsError, WsQuery};
|
||||
@@ -0,0 +1,772 @@
|
||||
//! Pull-on-miss: fetch a session from the backend and hydrate local JSONL storage.
|
||||
|
||||
use crate::remote::client::{BackendClient, BackendError};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum PullResult {
|
||||
/// Written to local storage. The [`Info`] cwd comes from the backend (may differ from caller's).
|
||||
Hydrated(crate::session::info::Info),
|
||||
/// Not found on the backend.
|
||||
NotFound,
|
||||
}
|
||||
|
||||
/// Fetch a session from the backend and hydrate local JSONL storage.
|
||||
pub async fn pull_session_to_local(
|
||||
session_id: &str,
|
||||
client: &BackendClient,
|
||||
) -> Result<PullResult, BackendError> {
|
||||
let loaded = match client.load_session_data(session_id).await {
|
||||
Ok(resp) => resp,
|
||||
Err(BackendError::SessionNotFound { .. }) => return Ok(PullResult::NotFound),
|
||||
Err(e) => return Err(e),
|
||||
};
|
||||
|
||||
let remote = match loaded.session.as_ref() {
|
||||
Some(s) => s,
|
||||
None => return Ok(PullResult::NotFound),
|
||||
};
|
||||
|
||||
// cwd required for local dir placement; null means pre-writeback session.
|
||||
let cwd = match remote.cwd.as_ref() {
|
||||
Some(cwd) => cwd,
|
||||
None => {
|
||||
tracing::warn!(session_id, "Cannot pull session: backend has cwd=null");
|
||||
return Ok(PullResult::NotFound);
|
||||
}
|
||||
};
|
||||
|
||||
let info = crate::session::info::Info {
|
||||
id: agent_client_protocol::SessionId::new(std::sync::Arc::from(session_id)),
|
||||
cwd: cwd.clone(),
|
||||
};
|
||||
let dir = crate::session::persistence::session_dir(&info);
|
||||
|
||||
let num_messages = hydrate::write_to_dir(&dir, &loaded)?;
|
||||
|
||||
tracing::info!(session_id, %cwd, num_messages, "Pulled session from backend");
|
||||
|
||||
Ok(PullResult::Hydrated(info))
|
||||
}
|
||||
|
||||
pub(crate) mod hydrate {
|
||||
use std::path::Path;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::remote::client::{BackendError, LoadDataResponse, LoadedMessage, SessionInfo};
|
||||
use crate::session::info::Info;
|
||||
use crate::session::persistence::{CHAT_FORMAT_VERSION, Summary, default_model_id};
|
||||
|
||||
fn io_err(path: &Path, source: std::io::Error) -> BackendError {
|
||||
BackendError::Hydration {
|
||||
path: path.to_path_buf(),
|
||||
source,
|
||||
}
|
||||
}
|
||||
|
||||
/// Write all session files to `dir`.
|
||||
pub(super) fn write_to_dir(
|
||||
dir: &Path,
|
||||
loaded: &LoadDataResponse,
|
||||
) -> Result<usize, BackendError> {
|
||||
let remote = loaded
|
||||
.session
|
||||
.as_ref()
|
||||
.expect("caller checked session.is_some()");
|
||||
|
||||
let info = Info {
|
||||
id: agent_client_protocol::SessionId::new(Arc::from(remote.session_id.as_str())),
|
||||
cwd: remote.cwd.clone().expect("caller verified cwd is Some"),
|
||||
};
|
||||
|
||||
std::fs::create_dir_all(dir).map_err(|e| io_err(dir, e))?;
|
||||
|
||||
let num_messages = loaded.messages.as_ref().map_or(0, |m| m.len());
|
||||
let mut num_chat_messages = 0;
|
||||
|
||||
if let Some(ref messages) = loaded.messages {
|
||||
write_updates(dir, messages)?;
|
||||
num_chat_messages = rebuild_chat_history(dir)?;
|
||||
}
|
||||
|
||||
write_summary(dir, &info, remote, num_messages, num_chat_messages)?;
|
||||
write_remote_origin_marker(dir);
|
||||
|
||||
Ok(num_messages)
|
||||
}
|
||||
|
||||
fn write_summary(
|
||||
dir: &Path,
|
||||
info: &Info,
|
||||
remote: &SessionInfo,
|
||||
num_messages: usize,
|
||||
num_chat_messages: usize,
|
||||
) -> Result<(), BackendError> {
|
||||
let meta = remote.metadata.as_ref();
|
||||
|
||||
let model_id = meta
|
||||
.and_then(|m| m.get("modelId"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(agent_client_protocol::ModelId::new)
|
||||
.unwrap_or_else(default_model_id);
|
||||
|
||||
let parent_session_id = meta
|
||||
.and_then(|m| m.get("parentSessionId"))
|
||||
.and_then(|v| v.as_str())
|
||||
.map(String::from);
|
||||
|
||||
let summary = Summary {
|
||||
info: info.clone(),
|
||||
session_summary: remote.title.clone().unwrap_or_default(),
|
||||
created_at: parse_rfc3339_or_now(remote.created_at.as_deref()),
|
||||
updated_at: parse_rfc3339_or_now(remote.updated_at.as_deref()),
|
||||
num_messages,
|
||||
num_chat_messages,
|
||||
current_model_id: model_id,
|
||||
parent_session_id,
|
||||
forked_at: None,
|
||||
collection_id: None,
|
||||
next_trace_turn: 0,
|
||||
chat_format_version: CHAT_FORMAT_VERSION,
|
||||
prompt_display_cwd: None,
|
||||
session_kind: None,
|
||||
fork_context_source: None,
|
||||
fork_parent_prompt_id: None,
|
||||
inherited_prefix_len: None,
|
||||
hidden: None,
|
||||
source_workspace_dir: None,
|
||||
git_root_dir: None,
|
||||
git_remotes: Vec::new(),
|
||||
head_commit: None,
|
||||
head_branch: None,
|
||||
request_id: None,
|
||||
// Record the *local* kigi_home (where this hydrated copy lives),
|
||||
// not the original remote session's, since reconstruction runs locally.
|
||||
kigi_home: crate::session::persistence::kigi_home_string(),
|
||||
last_active_at: None,
|
||||
generated_title: None,
|
||||
title_is_manual: false,
|
||||
worktree_label: None,
|
||||
agent_name: None,
|
||||
// Hydrated locally — record the profile this process runs under.
|
||||
sandbox_profile: kigi_sandbox::configured_profile_name().map(String::from),
|
||||
reasoning_effort: None,
|
||||
};
|
||||
|
||||
let json = serde_json::to_string_pretty(&summary)?;
|
||||
write_file(&dir.join("summary.json"), json.as_bytes())
|
||||
}
|
||||
|
||||
/// Convert backend JSON-RPC messages to local updates.jsonl (replayable methods only).
|
||||
pub(super) fn write_updates(
|
||||
dir: &Path,
|
||||
messages: &[LoadedMessage],
|
||||
) -> Result<(), BackendError> {
|
||||
use std::io::Write;
|
||||
|
||||
let path = dir.join("updates.jsonl");
|
||||
let file = std::fs::File::create(&path).map_err(|e| io_err(&path, e))?;
|
||||
let mut w = std::io::BufWriter::new(file);
|
||||
|
||||
for msg in messages {
|
||||
let parsed = match serde_json::from_str::<serde_json::Value>(&msg.content) {
|
||||
Ok(v) => v,
|
||||
Err(_) => continue,
|
||||
};
|
||||
if !is_session_update(&parsed) {
|
||||
continue;
|
||||
}
|
||||
if let Some(line) = to_envelope_line(&parsed) {
|
||||
let _ = w.write_all(line.as_bytes());
|
||||
let _ = w.write_all(b"\n");
|
||||
}
|
||||
}
|
||||
|
||||
w.flush().map_err(|e| io_err(&path, e))
|
||||
}
|
||||
|
||||
/// Rebuild `chat_history.jsonl` from `updates.jsonl` so pulled sessions are continuable.
|
||||
fn rebuild_chat_history(dir: &Path) -> Result<usize, BackendError> {
|
||||
use crate::session::storage::UpdatesIterator;
|
||||
use std::io::{Seek, Write};
|
||||
|
||||
let updates_path = dir.join("updates.jsonl");
|
||||
let Some(iter) =
|
||||
UpdatesIterator::open(&updates_path).map_err(|e| io_err(&updates_path, e))?
|
||||
else {
|
||||
return Ok(0);
|
||||
};
|
||||
|
||||
let chat_path = dir.join("chat_history.jsonl");
|
||||
let file = std::fs::File::create(&chat_path).map_err(|e| io_err(&chat_path, e))?;
|
||||
let mut writer = std::io::BufWriter::new(file);
|
||||
let mut reducer = ChatReducer::new();
|
||||
|
||||
for result in iter {
|
||||
let update = match result {
|
||||
Ok(u) => u,
|
||||
Err(_) => continue,
|
||||
};
|
||||
|
||||
for item in reducer.process(&update) {
|
||||
if let Ok(line) = serde_json::to_string(&item) {
|
||||
let _ = writer.write_all(line.as_bytes());
|
||||
let _ = writer.write_all(b"\n");
|
||||
}
|
||||
}
|
||||
|
||||
// CompactionCheckpoint: truncate file and reset
|
||||
if reducer.should_truncate() {
|
||||
reducer.clear_truncate_flag();
|
||||
let _ = writer.seek(std::io::SeekFrom::Start(0));
|
||||
let _ = writer.get_mut().set_len(0);
|
||||
}
|
||||
}
|
||||
|
||||
// Flush trailing state
|
||||
for item in reducer.flush() {
|
||||
if let Ok(line) = serde_json::to_string(&item) {
|
||||
let _ = writer.write_all(line.as_bytes());
|
||||
let _ = writer.write_all(b"\n");
|
||||
}
|
||||
}
|
||||
|
||||
writer.flush().map_err(|e| io_err(&chat_path, e))?;
|
||||
Ok(reducer.count())
|
||||
}
|
||||
|
||||
use crate::sampling::{AssistantItem, ContentPart, ConversationItem, ToolCall};
|
||||
use agent_client_protocol as acp;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
/// Reduces ACP session updates into conversation items.
|
||||
///
|
||||
/// Turn boundaries: User→Agent flushes user, Agent→User flushes agent,
|
||||
/// tool completion flushes agent before emitting result.
|
||||
struct ChatReducer {
|
||||
user_parts: Vec<ContentPart>,
|
||||
agent_text: String,
|
||||
agent_tool_calls: Vec<ToolCall>,
|
||||
|
||||
in_user_turn: bool,
|
||||
has_agent_content: bool,
|
||||
needs_truncate: bool,
|
||||
|
||||
tool_args: HashMap<String, String>,
|
||||
emitted_tool_results: HashSet<String>,
|
||||
item_count: usize,
|
||||
}
|
||||
|
||||
impl ChatReducer {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
user_parts: Vec::new(),
|
||||
agent_text: String::new(),
|
||||
agent_tool_calls: Vec::new(),
|
||||
in_user_turn: false,
|
||||
has_agent_content: false,
|
||||
needs_truncate: false,
|
||||
tool_args: HashMap::new(),
|
||||
emitted_tool_results: HashSet::new(),
|
||||
item_count: 0,
|
||||
}
|
||||
}
|
||||
|
||||
fn process(
|
||||
&mut self,
|
||||
update: &crate::session::storage::SessionUpdate,
|
||||
) -> Vec<ConversationItem> {
|
||||
use crate::session::storage::SessionUpdate;
|
||||
|
||||
match update {
|
||||
SessionUpdate::Acp(n) => self.handle_acp(&n.update),
|
||||
SessionUpdate::Xai(n) => self.handle_xai(&n.update),
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_acp(&mut self, update: &acp::SessionUpdate) -> Vec<ConversationItem> {
|
||||
match update {
|
||||
acp::SessionUpdate::UserMessageChunk(chunk) => self.on_user_chunk(chunk),
|
||||
acp::SessionUpdate::AgentMessageChunk(chunk) => self.on_agent_chunk(chunk),
|
||||
acp::SessionUpdate::ToolCall(tc) => self.on_tool_call(tc),
|
||||
acp::SessionUpdate::ToolCallUpdate(tc) => self.on_tool_call_update(tc),
|
||||
_ => Vec::new(), // AgentThoughtChunk, Retry, Plan not needed
|
||||
}
|
||||
}
|
||||
|
||||
fn handle_xai(
|
||||
&mut self,
|
||||
update: &crate::extensions::notification::SessionUpdate,
|
||||
) -> Vec<ConversationItem> {
|
||||
use crate::extensions::notification::SessionUpdate as XaiUpdate;
|
||||
|
||||
match update {
|
||||
XaiUpdate::CompactionCheckpoint(_) => {
|
||||
self.reset();
|
||||
self.needs_truncate = true;
|
||||
Vec::new()
|
||||
}
|
||||
_ => Vec::new(), // DiffReview, MemoryFlush, etc. not needed
|
||||
}
|
||||
}
|
||||
|
||||
fn on_user_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec<ConversationItem> {
|
||||
let mut out = Vec::new();
|
||||
|
||||
if !self.in_user_turn {
|
||||
out.extend(self.flush_agent());
|
||||
self.in_user_turn = true;
|
||||
}
|
||||
|
||||
match &chunk.content {
|
||||
acp::ContentBlock::Text(t) => {
|
||||
self.user_parts.push(ContentPart::Text {
|
||||
text: std::sync::Arc::<str>::from(t.text.clone()),
|
||||
});
|
||||
}
|
||||
acp::ContentBlock::Image(img) => {
|
||||
if let Some(uri) = &img.uri {
|
||||
self.user_parts.push(ContentPart::Image {
|
||||
url: std::sync::Arc::<str>::from(uri.clone()),
|
||||
});
|
||||
}
|
||||
}
|
||||
_ => {} // Audio, Resource, etc. not needed for chat replay
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn on_agent_chunk(&mut self, chunk: &acp::ContentChunk) -> Vec<ConversationItem> {
|
||||
let mut out = Vec::new();
|
||||
|
||||
if self.in_user_turn {
|
||||
out.extend(self.flush_user());
|
||||
self.in_user_turn = false;
|
||||
}
|
||||
|
||||
if let acp::ContentBlock::Text(t) = &chunk.content {
|
||||
self.agent_text.push_str(&t.text);
|
||||
self.has_agent_content = true;
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn on_tool_call(&mut self, tc: &acp::ToolCall) -> Vec<ConversationItem> {
|
||||
let id = tc.tool_call_id.0.to_string();
|
||||
let args = tc
|
||||
.raw_input
|
||||
.as_ref()
|
||||
.map(|v| v.to_string())
|
||||
.unwrap_or_default();
|
||||
|
||||
self.tool_args.insert(id.clone(), args.clone());
|
||||
self.agent_tool_calls.push(ToolCall {
|
||||
id: std::sync::Arc::<str>::from(id),
|
||||
name: tc.title.clone(),
|
||||
arguments: std::sync::Arc::<str>::from(args),
|
||||
});
|
||||
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
fn on_tool_call_update(&mut self, tc: &acp::ToolCallUpdate) -> Vec<ConversationItem> {
|
||||
let id = tc.tool_call_id.0.to_string();
|
||||
self.maybe_backfill_args(&id, &tc.fields);
|
||||
|
||||
if Self::is_completed(&tc.fields) && self.emitted_tool_results.insert(id.clone()) {
|
||||
return self.emit_tool_result(&id, &tc.fields);
|
||||
}
|
||||
Vec::new()
|
||||
}
|
||||
|
||||
/// Backfill tool arguments from ToolCallUpdate if ToolCall didn't have them.
|
||||
fn maybe_backfill_args(&mut self, id: &str, fields: &acp::ToolCallUpdateFields) {
|
||||
let Some(raw) = &fields.raw_input else { return };
|
||||
let needs_backfill = self.tool_args.get(id).is_none_or(String::is_empty);
|
||||
if !needs_backfill {
|
||||
return;
|
||||
}
|
||||
|
||||
let args = raw.to_string();
|
||||
self.tool_args.insert(id.to_string(), args.clone());
|
||||
|
||||
if let Some(call) = self
|
||||
.agent_tool_calls
|
||||
.iter_mut()
|
||||
.find(|c| c.id.as_ref() == id)
|
||||
{
|
||||
call.arguments = std::sync::Arc::<str>::from(args);
|
||||
}
|
||||
}
|
||||
|
||||
fn is_completed(fields: &acp::ToolCallUpdateFields) -> bool {
|
||||
matches!(
|
||||
fields.status,
|
||||
Some(acp::ToolCallStatus::Completed | acp::ToolCallStatus::Failed)
|
||||
)
|
||||
}
|
||||
|
||||
fn emit_tool_result(
|
||||
&mut self,
|
||||
id: &str,
|
||||
fields: &acp::ToolCallUpdateFields,
|
||||
) -> Vec<ConversationItem> {
|
||||
let mut out = Vec::new();
|
||||
out.extend(self.flush_agent());
|
||||
|
||||
let content = extract_tool_result_text(fields);
|
||||
let item = ConversationItem::tool_result(id.to_string(), content);
|
||||
self.item_count += 1;
|
||||
out.push(item);
|
||||
out
|
||||
}
|
||||
|
||||
fn flush_user(&mut self) -> Option<ConversationItem> {
|
||||
if self.user_parts.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let item = ConversationItem::user_with_parts(std::mem::take(&mut self.user_parts));
|
||||
self.item_count += 1;
|
||||
Some(item)
|
||||
}
|
||||
|
||||
fn flush_agent(&mut self) -> Option<ConversationItem> {
|
||||
if !self.has_agent_content && self.agent_tool_calls.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let item = ConversationItem::Assistant(AssistantItem {
|
||||
content: std::sync::Arc::<str>::from(std::mem::take(&mut self.agent_text)),
|
||||
tool_calls: std::mem::take(&mut self.agent_tool_calls),
|
||||
model_id: None,
|
||||
model_fingerprint: None,
|
||||
reasoning_effort: None,
|
||||
});
|
||||
self.has_agent_content = false;
|
||||
self.item_count += 1;
|
||||
Some(item)
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> Vec<ConversationItem> {
|
||||
let mut out = Vec::new();
|
||||
out.extend(self.flush_user());
|
||||
out.extend(self.flush_agent());
|
||||
out
|
||||
}
|
||||
|
||||
fn reset(&mut self) {
|
||||
self.user_parts.clear();
|
||||
self.agent_text.clear();
|
||||
self.agent_tool_calls.clear();
|
||||
self.tool_args.clear();
|
||||
self.emitted_tool_results.clear();
|
||||
self.in_user_turn = false;
|
||||
self.has_agent_content = false;
|
||||
self.item_count = 0;
|
||||
}
|
||||
|
||||
fn should_truncate(&self) -> bool {
|
||||
self.needs_truncate
|
||||
}
|
||||
|
||||
fn clear_truncate_flag(&mut self) {
|
||||
self.needs_truncate = false;
|
||||
}
|
||||
|
||||
fn count(&self) -> usize {
|
||||
self.item_count
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract displayable text from a completed ToolCallUpdate.
|
||||
fn extract_tool_result_text(fields: &agent_client_protocol::ToolCallUpdateFields) -> String {
|
||||
if let Some(content) = &fields.content {
|
||||
let text: String = content
|
||||
.iter()
|
||||
.filter_map(|c| match c {
|
||||
agent_client_protocol::ToolCallContent::Content(
|
||||
agent_client_protocol::Content {
|
||||
content: agent_client_protocol::ContentBlock::Text(t),
|
||||
..
|
||||
},
|
||||
) => Some(t.text.as_str()),
|
||||
_ => None,
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
.join("");
|
||||
if !text.is_empty() {
|
||||
return text;
|
||||
}
|
||||
}
|
||||
if let Some(raw) = &fields.raw_output {
|
||||
return raw.to_string();
|
||||
}
|
||||
String::new()
|
||||
}
|
||||
|
||||
fn write_remote_origin_marker(dir: &Path) {
|
||||
let _ = std::fs::write(
|
||||
dir.join(".remote_origin"),
|
||||
format!("pulled_at={}\n", chrono::Utc::now().to_rfc3339()),
|
||||
);
|
||||
}
|
||||
|
||||
/// Replayable JSON-RPC methods (excludes metadata like `prompt_complete`).
|
||||
const REPLAYABLE_METHODS: &[&str] = &["session/update", "_x.ai/session/update"];
|
||||
|
||||
fn is_session_update(json_rpc: &serde_json::Value) -> bool {
|
||||
json_rpc
|
||||
.get("method")
|
||||
.and_then(|v| v.as_str())
|
||||
.is_some_and(|m| REPLAYABLE_METHODS.contains(&m))
|
||||
}
|
||||
|
||||
fn to_envelope_line(json_rpc: &serde_json::Value) -> Option<String> {
|
||||
let method = json_rpc.get("method").and_then(|v| v.as_str())?;
|
||||
let params = json_rpc.get("params").cloned().unwrap_or_default();
|
||||
|
||||
serde_json::to_string(&serde_json::json!({
|
||||
"timestamp": 0u64,
|
||||
"method": method,
|
||||
"params": params,
|
||||
}))
|
||||
.ok()
|
||||
}
|
||||
|
||||
fn parse_rfc3339_or_now(s: Option<&str>) -> chrono::DateTime<chrono::Utc> {
|
||||
s.and_then(|s| chrono::DateTime::parse_from_rfc3339(s).ok())
|
||||
.map(|dt| dt.with_timezone(&chrono::Utc))
|
||||
.unwrap_or_else(chrono::Utc::now)
|
||||
}
|
||||
|
||||
fn write_file(path: &Path, data: &[u8]) -> Result<(), BackendError> {
|
||||
std::fs::write(path, data).map_err(|e| io_err(path, e))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::remote::client::LoadedMessage;
|
||||
|
||||
#[test]
|
||||
fn hydrate_writes_valid_updates_jsonl() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let messages = vec![
|
||||
LoadedMessage {
|
||||
id: "1".into(),
|
||||
content: r#"{"method":"session/update","params":{"update":"hello"}}"#.into(),
|
||||
timestamp: None,
|
||||
},
|
||||
LoadedMessage {
|
||||
id: "2".into(),
|
||||
content: r#"{"method":"session/update","params":{"update":"world"}}"#.into(),
|
||||
timestamp: None,
|
||||
},
|
||||
];
|
||||
|
||||
super::hydrate::write_updates(tmp.path(), &messages).unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap();
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
assert_eq!(lines.len(), 2);
|
||||
|
||||
for line in &lines {
|
||||
let v: serde_json::Value = serde_json::from_str(line).unwrap();
|
||||
assert_eq!(v["timestamp"], 0);
|
||||
assert_eq!(v["method"], "session/update");
|
||||
assert!(v["params"].is_object());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rebuild_chat_history_merges_chunks() {
|
||||
use crate::session::export::ExportedMessage;
|
||||
use agent_client_protocol::{ContentBlock, ContentChunk, SessionUpdate, TextContent};
|
||||
use std::sync::Arc;
|
||||
|
||||
// Build ACP notifications matching the RemoteSync path
|
||||
let sid = agent_client_protocol::SessionId::new(Arc::from("test"));
|
||||
let notifications = [
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("hello "),
|
||||
))),
|
||||
),
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("world"),
|
||||
))),
|
||||
),
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("hi back"),
|
||||
))),
|
||||
),
|
||||
];
|
||||
|
||||
// Serialize through ExportedMessage (writeback path)
|
||||
let messages: Vec<LoadedMessage> = notifications
|
||||
.iter()
|
||||
.map(|n| {
|
||||
let exported = ExportedMessage::from_notification(n);
|
||||
LoadedMessage {
|
||||
id: "x".into(),
|
||||
content: exported.content,
|
||||
timestamp: None,
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
|
||||
let data = crate::remote::client::LoadDataResponse {
|
||||
messages: Some(messages),
|
||||
session: Some(crate::remote::client::SessionInfo {
|
||||
session_id: "test".into(),
|
||||
title: None,
|
||||
cwd: Some("/tmp".into()),
|
||||
status: None,
|
||||
created_at: None,
|
||||
updated_at: None,
|
||||
metadata: None,
|
||||
}),
|
||||
};
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
super::hydrate::write_to_dir(tmp.path(), &data).unwrap();
|
||||
|
||||
let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap();
|
||||
let items: Vec<crate::sampling::ConversationItem> = chat
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.filter_map(|l| serde_json::from_str(l).ok())
|
||||
.collect();
|
||||
|
||||
assert_eq!(items.len(), 2, "should have 1 user + 1 agent item");
|
||||
assert!(matches!(
|
||||
&items[0],
|
||||
crate::sampling::ConversationItem::User(_)
|
||||
));
|
||||
assert!(matches!(
|
||||
&items[1],
|
||||
crate::sampling::ConversationItem::Assistant(_)
|
||||
));
|
||||
if let crate::sampling::ConversationItem::User(u) = &items[0] {
|
||||
let text: String = u
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|p| match p {
|
||||
crate::sampling::ContentPart::Text { text } => Some(text.as_ref()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(text, "hello world");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rebuild_chat_history_preserves_user_images() {
|
||||
use crate::session::export::ExportedMessage;
|
||||
use agent_client_protocol::{
|
||||
ContentBlock, ContentChunk, ImageContent, SessionUpdate, TextContent,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
let sid = agent_client_protocol::SessionId::new(Arc::from("test"));
|
||||
let notifications = [
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("look at this"),
|
||||
))),
|
||||
),
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Image(
|
||||
ImageContent::new(String::new(), String::new())
|
||||
.uri(Some("data:image/png;base64,abc".into())),
|
||||
))),
|
||||
),
|
||||
agent_client_protocol::SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("I see an image"),
|
||||
))),
|
||||
),
|
||||
];
|
||||
|
||||
let messages: Vec<LoadedMessage> = notifications
|
||||
.iter()
|
||||
.map(|n| LoadedMessage {
|
||||
id: "x".into(),
|
||||
content: ExportedMessage::from_notification(n).content,
|
||||
timestamp: None,
|
||||
})
|
||||
.collect();
|
||||
|
||||
let data = crate::remote::client::LoadDataResponse {
|
||||
messages: Some(messages),
|
||||
session: Some(crate::remote::client::SessionInfo {
|
||||
session_id: "test".into(),
|
||||
title: None,
|
||||
cwd: Some("/tmp".into()),
|
||||
status: None,
|
||||
created_at: None,
|
||||
updated_at: None,
|
||||
metadata: None,
|
||||
}),
|
||||
};
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
super::hydrate::write_to_dir(tmp.path(), &data).unwrap();
|
||||
|
||||
let chat = std::fs::read_to_string(tmp.path().join("chat_history.jsonl")).unwrap();
|
||||
let items: Vec<crate::sampling::ConversationItem> = chat
|
||||
.lines()
|
||||
.filter(|l| !l.is_empty())
|
||||
.filter_map(|l| serde_json::from_str(l).ok())
|
||||
.collect();
|
||||
|
||||
assert_eq!(items.len(), 2);
|
||||
if let crate::sampling::ConversationItem::User(u) = &items[0] {
|
||||
assert_eq!(u.content.len(), 2, "should have text + image parts");
|
||||
assert!(matches!(
|
||||
&u.content[0],
|
||||
crate::sampling::ContentPart::Text { .. }
|
||||
));
|
||||
assert!(matches!(
|
||||
&u.content[1],
|
||||
crate::sampling::ContentPart::Image { .. }
|
||||
));
|
||||
} else {
|
||||
panic!("expected User item");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn hydrate_skips_invalid_messages() {
|
||||
let tmp = tempfile::TempDir::new().unwrap();
|
||||
let messages = vec![
|
||||
LoadedMessage {
|
||||
id: "1".into(),
|
||||
content: r#"{"method":"session/update","params":{}}"#.into(),
|
||||
timestamp: None,
|
||||
},
|
||||
LoadedMessage {
|
||||
id: "bad".into(),
|
||||
content: "not valid json".into(),
|
||||
timestamp: None,
|
||||
},
|
||||
LoadedMessage {
|
||||
id: "3".into(),
|
||||
content: r#"{"method":"session/update","params":{"x":1}}"#.into(),
|
||||
timestamp: None,
|
||||
},
|
||||
];
|
||||
|
||||
super::hydrate::write_updates(tmp.path(), &messages).unwrap();
|
||||
|
||||
let content = std::fs::read_to_string(tmp.path().join("updates.jsonl")).unwrap();
|
||||
let lines: Vec<&str> = content.lines().collect();
|
||||
assert_eq!(lines.len(), 2, "invalid message should be skipped");
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
//! Push → pull round-trip smoke test against the live backend.
|
||||
//!
|
||||
//! Run with: `cargo test -p kigi-shell -- pull_smoke --ignored --nocapture`
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use crate::auth::GrokAuth;
|
||||
use crate::remote::client::BackendClient;
|
||||
use crate::session::storage::{JsonlStorageAdapter, StorageAdapter};
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
fn load_prod_auth() -> Option<GrokAuth> {
|
||||
let path = crate::util::kigi_home::kigi_home().join("auth.json");
|
||||
let contents = std::fs::read_to_string(&path).ok()?;
|
||||
let store: BTreeMap<String, GrokAuth> = serde_json::from_str(&contents).ok()?;
|
||||
let scope = crate::auth::GrokComConfig::default().auth_scope();
|
||||
crate::auth::lookup_auth(&store, &scope)
|
||||
}
|
||||
|
||||
/// Full round-trip using the real RemoteSync production code path:
|
||||
/// create RemoteSync → queue ACP notifications → flush → verify on
|
||||
/// backend → pull back → verify local hydration + storage adapter load.
|
||||
#[tokio::test]
|
||||
#[ignore]
|
||||
async fn smoke_push_pull_round_trip() {
|
||||
use crate::remote::sync::RemoteSync;
|
||||
use crate::session::export::ExportedMetadata;
|
||||
use agent_client_protocol::{
|
||||
ContentBlock, ContentChunk, SessionNotification, SessionUpdate, TextContent,
|
||||
};
|
||||
|
||||
let auth = load_prod_auth().expect("No auth.json — run `grok login`");
|
||||
let am = Arc::new(crate::auth::AuthManager::new(
|
||||
&crate::util::kigi_home::kigi_home(),
|
||||
crate::auth::GrokComConfig::default(),
|
||||
));
|
||||
am.hot_swap(auth);
|
||||
let client = BackendClient::new().with_auth_manager(am.clone());
|
||||
|
||||
let session_id = format!("test-rt-{}", uuid::Uuid::new_v4());
|
||||
let test_cwd = "/tmp/smoke-test".to_string();
|
||||
let test_title = "Push-Pull Round Trip Test";
|
||||
|
||||
// PUSH via RemoteSync (real production path)
|
||||
let metadata = ExportedMetadata {
|
||||
title: Some(test_title.into()),
|
||||
cwd: test_cwd.clone(),
|
||||
model_id: Some("grok-3".into()),
|
||||
created_at: Some(chrono::Utc::now().to_rfc3339()),
|
||||
updated_at: Some(chrono::Utc::now().to_rfc3339()),
|
||||
total_messages: None,
|
||||
parent_session_id: None,
|
||||
session_kind: None,
|
||||
subagent_type: None,
|
||||
subagent_persona: None,
|
||||
subagent_role: None,
|
||||
fork_context_source: None,
|
||||
subagent_depth: None,
|
||||
};
|
||||
|
||||
let sync = RemoteSync::new(
|
||||
session_id.clone(),
|
||||
metadata,
|
||||
BackendClient::new().with_auth_manager(am.clone()),
|
||||
);
|
||||
|
||||
let sid = agent_client_protocol::SessionId::new(Arc::from(session_id.as_str()));
|
||||
sync.queue(SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::UserMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("Hello from smoke test — user".to_string()),
|
||||
))),
|
||||
));
|
||||
sync.queue(SessionNotification::new(
|
||||
sid.clone(),
|
||||
SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text(
|
||||
TextContent::new("Hello from smoke test — agent".to_string()),
|
||||
))),
|
||||
));
|
||||
sync.flush();
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(3)).await;
|
||||
|
||||
// Verify backend has cwd, title, messages
|
||||
let loaded = client
|
||||
.load_session_data(&session_id)
|
||||
.await
|
||||
.expect("load after push failed");
|
||||
let remote = loaded.session.as_ref().expect("no session row");
|
||||
assert_eq!(remote.cwd.as_deref(), Some(test_cwd.as_str()));
|
||||
assert_eq!(remote.title.as_deref(), Some(test_title));
|
||||
assert!(loaded.messages.as_ref().map_or(0, |m| m.len()) >= 2);
|
||||
|
||||
// PULL back to local
|
||||
let result = crate::remote::pull_session_to_local(&session_id, &client)
|
||||
.await
|
||||
.expect("pull failed");
|
||||
let pulled = match result {
|
||||
crate::remote::PullResult::Hydrated(info) => info,
|
||||
crate::remote::PullResult::NotFound => panic!("pull returned NotFound"),
|
||||
};
|
||||
assert_eq!(pulled.cwd, test_cwd);
|
||||
|
||||
// Verify local storage loads
|
||||
let local_dir = crate::session::persistence::session_dir(&pulled);
|
||||
assert!(local_dir.join("summary.json").exists());
|
||||
assert!(local_dir.join("updates.jsonl").exists());
|
||||
|
||||
let storage = JsonlStorageAdapter::default();
|
||||
let data = storage
|
||||
.load_session_without_updates(&pulled)
|
||||
.await
|
||||
.expect("storage load failed");
|
||||
assert_eq!(data.summary.session_summary, test_title);
|
||||
|
||||
// Verify chat_history has both turns
|
||||
let chat =
|
||||
std::fs::read_to_string(local_dir.join("chat_history.jsonl")).unwrap_or_default();
|
||||
assert!(chat.contains("user"), "chat_history missing user turn");
|
||||
assert!(chat.contains("agent"), "chat_history missing agent turn");
|
||||
|
||||
// Cleanup
|
||||
drop(sync);
|
||||
let _ = client.delete_session_data(&session_id).await;
|
||||
let _ = std::fs::remove_dir_all(&local_dir);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
//! Writeback push: async queue that flushes session updates to the backend.
|
||||
//!
|
||||
//! `RemoteSync` runs a background tokio task that buffers ACP notifications
|
||||
//! and flushes them to the backend via [`BackendClient::save_session_data()`].
|
||||
//!
|
||||
//! ## Backpressure
|
||||
//!
|
||||
//! When the buffer exceeds [`MAX_PENDING`], the task attempts an emergency
|
||||
//! flush. If that also fails (network down), the oldest messages are dropped
|
||||
//! to prevent unbounded memory growth.
|
||||
//!
|
||||
//! ## Drop behavior
|
||||
//!
|
||||
//! When `RemoteSync` is dropped, the sender half of the channel closes and
|
||||
//! the background task exits. **Pending buffered messages are lost.** This
|
||||
//! is acceptable because the local JSONL files are the source of truth —
|
||||
//! writeback is best-effort.
|
||||
|
||||
use crate::remote::BackendClient;
|
||||
use crate::session::export::{ExportedMessage, ExportedMetadata};
|
||||
use agent_client_protocol as acp;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
/// Max buffered notifications before triggering an emergency flush.
|
||||
/// Sized to keep memory under ~50MB even with large notifications.
|
||||
const MAX_PENDING: usize = 512;
|
||||
|
||||
/// How many oldest messages to drop when an emergency flush fails.
|
||||
/// Dropping a batch (not one-by-one) avoids repeated failed flushes.
|
||||
const DROP_BATCH_SIZE: usize = 64;
|
||||
|
||||
enum SyncMsg {
|
||||
Queue(Box<acp::SessionNotification>),
|
||||
Flush,
|
||||
SetTitle(String),
|
||||
SetModelId(String),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct RemoteSync {
|
||||
tx: mpsc::UnboundedSender<SyncMsg>,
|
||||
}
|
||||
|
||||
impl RemoteSync {
|
||||
/// Metadata is included on every flush to keep the backend session row current.
|
||||
pub(crate) fn new(
|
||||
session_id: String,
|
||||
metadata: ExportedMetadata,
|
||||
client: BackendClient,
|
||||
) -> Self {
|
||||
let (tx, rx) = mpsc::unbounded_channel();
|
||||
tokio::spawn(sync_task(session_id, metadata, client, rx));
|
||||
Self { tx }
|
||||
}
|
||||
|
||||
pub fn queue(&self, notification: acp::SessionNotification) {
|
||||
let _ = self.tx.send(SyncMsg::Queue(Box::new(notification)));
|
||||
}
|
||||
|
||||
pub fn flush(&self) {
|
||||
let _ = self.tx.send(SyncMsg::Flush);
|
||||
}
|
||||
|
||||
pub fn set_title(&self, title: String) {
|
||||
let _ = self.tx.send(SyncMsg::SetTitle(title));
|
||||
}
|
||||
|
||||
pub fn set_model_id(&self, model_id: String) {
|
||||
let _ = self.tx.send(SyncMsg::SetModelId(model_id));
|
||||
}
|
||||
}
|
||||
|
||||
async fn do_flush(
|
||||
client: &BackendClient,
|
||||
session_id: &str,
|
||||
metadata: &ExportedMetadata,
|
||||
pending: &mut Vec<acp::SessionNotification>,
|
||||
) -> bool {
|
||||
if pending.is_empty() {
|
||||
return true;
|
||||
}
|
||||
|
||||
let messages: Vec<ExportedMessage> = pending
|
||||
.iter()
|
||||
.map(ExportedMessage::from_notification)
|
||||
.collect();
|
||||
|
||||
match client
|
||||
.save_session_data(session_id, &messages, Some(metadata))
|
||||
.await
|
||||
{
|
||||
Ok(()) => {
|
||||
tracing::debug!(count = pending.len(), "Writeback: synced");
|
||||
pending.clear();
|
||||
|
||||
// Link session to agent so the relay can route requests to it.
|
||||
if let Err(e) = client
|
||||
.upsert_session(session_id, metadata, &crate::util::agent_id::agent_id())
|
||||
.await
|
||||
{
|
||||
tracing::warn!(error = %e, "Writeback: failed to upsert session");
|
||||
}
|
||||
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, pending = pending.len(), "Writeback: flush failed");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn sync_task(
|
||||
session_id: String,
|
||||
mut metadata: ExportedMetadata,
|
||||
client: BackendClient,
|
||||
mut rx: mpsc::UnboundedReceiver<SyncMsg>,
|
||||
) {
|
||||
let mut pending: Vec<acp::SessionNotification> = Vec::new();
|
||||
|
||||
while let Some(msg) = rx.recv().await {
|
||||
match msg {
|
||||
SyncMsg::Queue(n) => {
|
||||
if pending.len() >= MAX_PENDING {
|
||||
tracing::warn!(
|
||||
pending = pending.len(),
|
||||
"Writeback: buffer full, attempting emergency flush"
|
||||
);
|
||||
|
||||
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
|
||||
if !do_flush(&client, &session_id, &metadata, &mut pending).await {
|
||||
let dropped = pending.drain(0..DROP_BATCH_SIZE.min(pending.len())).count();
|
||||
tracing::error!(
|
||||
dropped = dropped,
|
||||
"Writeback: emergency flush failed, dropping oldest messages"
|
||||
);
|
||||
}
|
||||
}
|
||||
pending.push(*n);
|
||||
}
|
||||
SyncMsg::Flush => {
|
||||
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
|
||||
do_flush(&client, &session_id, &metadata, &mut pending).await;
|
||||
}
|
||||
SyncMsg::SetTitle(title) => {
|
||||
metadata.title = Some(title);
|
||||
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
|
||||
if let Err(e) = client
|
||||
.save_session_data(&session_id, &[], Some(&metadata))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(?e, "Writeback: failed to sync title to backend");
|
||||
}
|
||||
}
|
||||
SyncMsg::SetModelId(id) => {
|
||||
metadata.model_id = Some(id);
|
||||
metadata.updated_at = Some(chrono::Utc::now().to_rfc3339());
|
||||
if let Err(e) = client
|
||||
.save_session_data(&session_id, &[], Some(&metadata))
|
||||
.await
|
||||
{
|
||||
tracing::warn!(?e, "Writeback: failed to sync model_id to backend");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,180 @@
|
||||
use std::sync::Arc;
|
||||
|
||||
use serde::Deserialize;
|
||||
|
||||
use crate::auth::AuthManager;
|
||||
|
||||
const KIGI_WEB_URL: &str = "https://grok.com";
|
||||
|
||||
#[derive(Debug, Clone, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Workspace {
|
||||
#[serde(default)]
|
||||
pub workspace_id: String,
|
||||
#[serde(default)]
|
||||
pub name: String,
|
||||
#[serde(default)]
|
||||
pub create_time: Option<String>,
|
||||
#[serde(default)]
|
||||
pub kind: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct WsQuery {
|
||||
pub page_size: i64,
|
||||
pub page_token: Option<String>,
|
||||
pub query: Option<String>,
|
||||
pub kind: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Default)]
|
||||
pub struct ListWorkspacesPage {
|
||||
pub workspaces: Vec<Workspace>,
|
||||
pub next_page_token: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum WsError {
|
||||
#[error("no OAuth credentials for workspaces:read")]
|
||||
NoOauth,
|
||||
#[error("network error: {0}")]
|
||||
Network(#[from] reqwest::Error),
|
||||
#[error("request failed: {status}")]
|
||||
Http { status: u16 },
|
||||
#[error("parse error: {0}")]
|
||||
Parse(#[from] serde_json::Error),
|
||||
}
|
||||
|
||||
#[derive(Debug, Default, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct ListWorkspacesResponseWire {
|
||||
#[serde(default)]
|
||||
workspaces: Vec<Workspace>,
|
||||
#[serde(default)]
|
||||
next_page_token: Option<String>,
|
||||
}
|
||||
|
||||
pub struct WorkspacesClient {
|
||||
http: reqwest::Client,
|
||||
base_url: String,
|
||||
auth: Arc<AuthManager>,
|
||||
}
|
||||
|
||||
impl WorkspacesClient {
|
||||
pub fn new(auth: Arc<AuthManager>) -> Self {
|
||||
let base_url = first_nonempty_env(&[
|
||||
"KIGI_WORKSPACES_BASE_URL",
|
||||
"KIGI_CONVERSATIONS_BASE_URL",
|
||||
"KIGI_CODE_WEB_URL",
|
||||
])
|
||||
.unwrap_or_else(|| KIGI_WEB_URL.to_string());
|
||||
Self {
|
||||
http: crate::http::shared_client(),
|
||||
base_url,
|
||||
auth,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_workspaces(&self, q: &WsQuery) -> Result<ListWorkspacesPage, WsError> {
|
||||
let auth = self.auth.auth().await.map_err(|_| WsError::NoOauth)?;
|
||||
if !auth.is_xai_auth() {
|
||||
return Err(WsError::NoOauth);
|
||||
}
|
||||
|
||||
let url = format!("{}/rest/workspaces", self.base_url);
|
||||
let mut query: Vec<(&str, String)> = vec![("pageSize", q.page_size.to_string())];
|
||||
if let Some(token) = q.page_token.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("pageToken", token.to_owned()));
|
||||
}
|
||||
if let Some(search) = q.query.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("query", search.to_owned()));
|
||||
}
|
||||
if let Some(kind) = q.kind.as_deref().filter(|s| !s.is_empty()) {
|
||||
query.push(("kind", kind.to_owned()));
|
||||
}
|
||||
|
||||
let mut builder = self
|
||||
.http
|
||||
.get(&url)
|
||||
.query(&query)
|
||||
.header("Authorization", format!("Bearer {}", auth.key))
|
||||
.header(
|
||||
"X-XAI-Token-Auth",
|
||||
self.auth.grok_com_config().token_header.clone(),
|
||||
)
|
||||
.header("x-userid", &auth.user_id)
|
||||
.header("x-grok-client-version", kigi_version::VERSION)
|
||||
.header(
|
||||
"x-grok-client-identifier",
|
||||
crate::http::process_client_identifier(),
|
||||
)
|
||||
.header(
|
||||
crate::http::CLIENT_MODE_HEADER,
|
||||
crate::http::process_client_mode(),
|
||||
)
|
||||
.header(reqwest::header::ACCEPT, "application/json");
|
||||
if let Some(email) = &auth.email {
|
||||
builder = builder.header("x-email", email);
|
||||
}
|
||||
let builder = kigi_file_utils::trace_context::inject_trace_context_into_request(builder);
|
||||
|
||||
let response = builder.send().await?;
|
||||
let status = response.status();
|
||||
if !status.is_success() {
|
||||
return Err(WsError::Http {
|
||||
status: status.as_u16(),
|
||||
});
|
||||
}
|
||||
|
||||
let bytes = response.bytes().await?;
|
||||
let wire: ListWorkspacesResponseWire = serde_json::from_slice(&bytes)?;
|
||||
|
||||
Ok(ListWorkspacesPage {
|
||||
workspaces: wire.workspaces,
|
||||
next_page_token: wire.next_page_token.filter(|t| !t.is_empty()),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn first_nonempty_env(keys: &[&str]) -> Option<String> {
|
||||
keys.iter()
|
||||
.find_map(|k| std::env::var(k).ok().filter(|s| !s.is_empty()))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn workspace_parses_camelcase_wire() {
|
||||
let json = serde_json::json!({
|
||||
"workspaces": [{
|
||||
"workspaceId": "ws_9f3a",
|
||||
"name": "GPU vendor research",
|
||||
"createTime": "2026-06-18T17:30:00Z",
|
||||
"kind": "WORKSPACE_KIND_IMAGINE"
|
||||
}],
|
||||
"nextPageToken": "tok2"
|
||||
});
|
||||
let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap();
|
||||
assert_eq!(wire.workspaces.len(), 1);
|
||||
let w = &wire.workspaces[0];
|
||||
assert_eq!(w.workspace_id, "ws_9f3a");
|
||||
assert_eq!(w.name, "GPU vendor research");
|
||||
assert_eq!(w.create_time.as_deref(), Some("2026-06-18T17:30:00Z"));
|
||||
assert_eq!(w.kind.as_deref(), Some("WORKSPACE_KIND_IMAGINE"));
|
||||
assert_eq!(wire.next_page_token.as_deref(), Some("tok2"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn missing_fields_default_gracefully() {
|
||||
let json = serde_json::json!({ "workspaces": [{ "workspaceId": "w1" }] });
|
||||
let wire: ListWorkspacesResponseWire = serde_json::from_value(json).unwrap();
|
||||
let w = &wire.workspaces[0];
|
||||
assert_eq!(w.workspace_id, "w1");
|
||||
assert!(w.name.is_empty());
|
||||
assert!(w.create_time.is_none());
|
||||
assert!(w.kind.is_none());
|
||||
assert!(wire.next_page_token.is_none());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user