Show reasoning effort in the model's own vocabulary (K3: max, not xhigh)
The welcome and prompt model labels (and /effort's 'current' hint) rendered the canonical internal level name, so K3 at its default effort showed 'K3 (xhigh)' even though the server's vocabulary for that level is 'max' (live /models think_efforts: low/high/max). New ModelState::reasoning_effort_display() resolves the current effort through the model's own effort menu (option id whose value matches), falling back to the canonical name only when the model has no entry for the level. All three display sites route through it; test pins the K3 mapping (Xhigh → 'max', Low → 'low', no-menu-entry → canonical).
This commit is contained in:
-97
@@ -1,97 +0,0 @@
|
||||
use super::*;
|
||||
use crate::session::events::ToolOutcome;
|
||||
use kigi_tool_protocol::session_event::ToolCallOutcome;
|
||||
use kigi_tool_protocol::turn_hook::TurnHookOutcome;
|
||||
#[test]
|
||||
fn map_tool_outcome_success() {
|
||||
assert_eq!(
|
||||
map_tool_outcome(ToolOutcome::Success),
|
||||
ToolCallOutcome::Success
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn map_tool_outcome_errors() {
|
||||
assert_eq!(map_tool_outcome(ToolOutcome::Error), ToolCallOutcome::Error);
|
||||
assert_eq!(
|
||||
map_tool_outcome(ToolOutcome::InvalidTool),
|
||||
ToolCallOutcome::Error
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn map_tool_outcome_cancellations() {
|
||||
for variant in [
|
||||
ToolOutcome::PermissionRejected,
|
||||
ToolOutcome::PermissionCancelled,
|
||||
ToolOutcome::Followup,
|
||||
ToolOutcome::HookDenied,
|
||||
ToolOutcome::Cancelled,
|
||||
] {
|
||||
assert_eq!(
|
||||
map_tool_outcome(variant),
|
||||
ToolCallOutcome::Cancelled,
|
||||
"expected Cancelled for {variant:?}",
|
||||
);
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn turn_result_completed() {
|
||||
let result: Result<TurnOutcome, acp::Error> = Ok(TurnOutcome::Completed {
|
||||
snapshot: Box::new(None),
|
||||
tools_called: vec![],
|
||||
structured_output: None,
|
||||
refusal: false,
|
||||
});
|
||||
assert_eq!(
|
||||
turn_result_to_hook_outcome(&result),
|
||||
TurnHookOutcome::Completed
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn turn_result_cancelled() {
|
||||
let result: Result<TurnOutcome, acp::Error> = Ok(TurnOutcome::Cancelled {
|
||||
category: None,
|
||||
context: None,
|
||||
});
|
||||
assert_eq!(
|
||||
turn_result_to_hook_outcome(&result),
|
||||
TurnHookOutcome::Cancelled
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn turn_result_error() {
|
||||
let result: Result<TurnOutcome, acp::Error> = Err(acp::Error::internal_error());
|
||||
assert_eq!(turn_result_to_hook_outcome(&result), TurnHookOutcome::Error);
|
||||
}
|
||||
#[test]
|
||||
fn is_remote_image_url_classifies_schemes() {
|
||||
assert!(is_remote_image_url("https://example.com/x.png"));
|
||||
assert!(is_remote_image_url("http://example.com/x.png"));
|
||||
assert!(!is_remote_image_url("file:///Users/me/x.png"));
|
||||
assert!(!is_remote_image_url("data:image/png;base64,AAAA"));
|
||||
assert!(!is_remote_image_url(""));
|
||||
assert!(!is_remote_image_url("FILE:///Users/me/x.png"));
|
||||
}
|
||||
#[test]
|
||||
fn pick_image_url_prefers_base64_over_file_uri() {
|
||||
let img = agent_client_protocol::ImageContent::new("AAAA", "image/png")
|
||||
.uri(Some("file:///Users/me/Downloads/screenshot.png".into()));
|
||||
assert_eq!(pick_user_image_url(&img), "data:image/png;base64,AAAA");
|
||||
}
|
||||
#[test]
|
||||
fn pick_image_url_prefers_base64_when_https_uri_also_present() {
|
||||
let img = agent_client_protocol::ImageContent::new("BBBB", "image/jpeg")
|
||||
.uri(Some("https://example.com/x.jpg".into()));
|
||||
assert_eq!(pick_user_image_url(&img), "data:image/jpeg;base64,BBBB");
|
||||
}
|
||||
#[test]
|
||||
fn pick_image_url_falls_back_to_https_uri_when_data_empty() {
|
||||
let img = agent_client_protocol::ImageContent::new(String::new(), "image/png")
|
||||
.uri(Some("https://example.com/x.png".into()));
|
||||
assert_eq!(pick_user_image_url(&img), "https://example.com/x.png");
|
||||
}
|
||||
#[test]
|
||||
fn pick_image_url_ignores_file_uri_when_data_empty() {
|
||||
let img = agent_client_protocol::ImageContent::new(String::new(), "image/png")
|
||||
.uri(Some("file:///Users/me/missing.png".into()));
|
||||
assert_eq!(pick_user_image_url(&img), "data:image/png;base64,");
|
||||
}
|
||||
@@ -204,6 +204,22 @@ impl ModelState {
|
||||
parse_reasoning_efforts_meta(info.meta.as_ref()).unwrap_or_else(legacy_effort_options)
|
||||
}
|
||||
|
||||
/// Display token for the current reasoning effort, in the MODEL's own
|
||||
/// vocabulary: the menu option id whose value matches (e.g. K3 spells
|
||||
/// `Xhigh` as `max`), falling back to the canonical name only when the
|
||||
/// model has no menu entry for it. Displays must never leak the internal
|
||||
/// canonical spelling for a level the server names differently.
|
||||
pub fn reasoning_effort_display(&self) -> Option<String> {
|
||||
let effort = self.reasoning_effort?;
|
||||
let display = self
|
||||
.reasoning_effort_options()
|
||||
.into_iter()
|
||||
.find(|opt| opt.value == effort)
|
||||
.map(|opt| opt.id)
|
||||
.unwrap_or_else(|| effort.as_str().to_owned());
|
||||
Some(display)
|
||||
}
|
||||
|
||||
/// Map a typed/selected effort token to its canonical value for the current
|
||||
/// model. Accepts a menu option id (case-insensitive) or a canonical level
|
||||
/// that appears as a **value** in that model's menu. Levels the model does
|
||||
@@ -479,6 +495,30 @@ mod tests {
|
||||
assert_eq!(opts[1].description.as_deref(), Some("Max"));
|
||||
}
|
||||
|
||||
/// K3 spells `Xhigh` as `max` (live /models `think_efforts`): every
|
||||
/// display surface must show the model's own vocabulary, never the
|
||||
/// internal canonical name.
|
||||
#[test]
|
||||
fn reasoning_effort_display_uses_model_vocabulary() {
|
||||
let mut state = state_with_meta(Some(serde_json::json!({
|
||||
"supportsReasoningEffort": true,
|
||||
"reasoningEfforts": [
|
||||
{ "id": "low", "value": "low", "label": "Low" },
|
||||
{ "id": "high", "value": "high", "label": "High" },
|
||||
{ "id": "max", "value": "xhigh", "label": "Max" },
|
||||
],
|
||||
})));
|
||||
state.reasoning_effort = Some(ReasoningEffort::Xhigh);
|
||||
assert_eq!(state.reasoning_effort_display().as_deref(), Some("max"));
|
||||
state.reasoning_effort = Some(ReasoningEffort::Low);
|
||||
assert_eq!(state.reasoning_effort_display().as_deref(), Some("low"));
|
||||
// No menu entry for the level → canonical fallback (never hide it).
|
||||
state.reasoning_effort = Some(ReasoningEffort::Medium);
|
||||
assert_eq!(state.reasoning_effort_display().as_deref(), Some("medium"));
|
||||
state.reasoning_effort = None;
|
||||
assert_eq!(state.reasoning_effort_display(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn reasoning_effort_options_gate_first_empty_when_unsupported() {
|
||||
// No current model → empty.
|
||||
|
||||
@@ -2038,7 +2038,7 @@ impl AgentView {
|
||||
}
|
||||
let mode_flags: &[PromptFlag] = &mode_flags_vec;
|
||||
let multiline = self.multiline_mode;
|
||||
let model_label = match self.session.models.reasoning_effort {
|
||||
let model_label = match self.session.models.reasoning_effort_display() {
|
||||
Some(eff) => format!("{model_id} ({eff})"),
|
||||
None => model_id,
|
||||
};
|
||||
|
||||
@@ -3060,7 +3060,7 @@ impl AppView {
|
||||
self.tip.as_deref()
|
||||
};
|
||||
let model_name_base = self.models.current_model_name().unwrap_or_default();
|
||||
let model_name = match self.models.reasoning_effort {
|
||||
let model_name = match self.models.reasoning_effort_display() {
|
||||
Some(eff) => format!("{model_name_base} ({eff})"),
|
||||
None => model_name_base,
|
||||
};
|
||||
|
||||
@@ -69,7 +69,7 @@ impl SlashCommand for EffortCommand {
|
||||
.collect();
|
||||
let current = ctx
|
||||
.models
|
||||
.reasoning_effort
|
||||
.reasoning_effort_display()
|
||||
.map(|e| format!(" (current: {e})"))
|
||||
.unwrap_or_default();
|
||||
let levels = if offered.is_empty() {
|
||||
|
||||
@@ -1,28 +0,0 @@
|
||||
[package]
|
||||
license = "Apache-2.0"
|
||||
name = "kigi-workspace-client"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
description = "Lightweight typed client for hub-proxied workspace.* RPCs (shared by kigi-shell proxy mode and other consumers)"
|
||||
|
||||
[dependencies]
|
||||
serde = { workspace = true, features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
tokio = { workspace = true, features = ["time"] }
|
||||
tracing = { workspace = true }
|
||||
kigi-computer-hub-sdk = { workspace = true }
|
||||
kigi-workspace-types = { workspace = true }
|
||||
kigi-tool-protocol = { workspace = true }
|
||||
kigi-tool-runtime = { workspace = true }
|
||||
|
||||
[features]
|
||||
default = []
|
||||
default-bazel = []
|
||||
[dev-dependencies]
|
||||
schemars = { workspace = true }
|
||||
tokio = { workspace = true, features = ["macros", "rt", "test-util"] }
|
||||
kigi-tool-types = { workspace = true }
|
||||
|
||||
[lints]
|
||||
workspace = true
|
||||
@@ -1,810 +0,0 @@
|
||||
#![allow(
|
||||
unused_imports,
|
||||
unused_variables,
|
||||
unused_mut,
|
||||
unreachable_code,
|
||||
dead_code
|
||||
)]
|
||||
//! Typed client for hub-proxied `workspace.*` RPC methods — the single
|
||||
//! transport for the `workspace_rpc` channel, shared by `WorkspaceOps`
|
||||
//! proxy mode and by consumers that cannot depend on
|
||||
//! `kigi-workspace`. Wire types live in
|
||||
//! `kigi_workspace_types::rpc`; this crate adds the connected-state
|
||||
//! latch, the generic [`WorkspaceClient::rpc`] core, and error mapping.
|
||||
//!
|
||||
//! No deadline is imposed by default ([`WorkspaceClient::with_deadline`]
|
||||
//! opts in), preserving `WorkspaceOps::rpc_raw` semantics where callers
|
||||
//! own their timeouts.
|
||||
use kigi_computer_hub_sdk::harness::ToolHarness;
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolStreamItem, TypedToolOutput};
|
||||
use kigi_workspace_types::rpc::agents_md::{AgentConfigFile, DiscoverAgentsMdReq};
|
||||
use kigi_workspace_types::rpc::code_nav::{
|
||||
CodeFindDefinitionsReq, CodeFindReferencesReq, CodeGotoDefinitionReq, CodeGotoReferencesReq,
|
||||
CodeIndexStatusReq, CodeIndexStatusResponse, CodeNavResponse,
|
||||
};
|
||||
use kigi_workspace_types::rpc::fs::{
|
||||
FsDeleteFileReq, FsExistsData, FsExistsReq, FsListData, FsListReq, FsReadFileData,
|
||||
FsReadFileReq, FsWriteFileReq, GetFilesReq, GetFilesRes, PutFilesReq, PutFilesRes,
|
||||
};
|
||||
use kigi_workspace_types::rpc::git::{
|
||||
CheckoutCommitResponse, CommitResult, DetectVcsKindReq, GitBranchInfoReq, GitBranchListData,
|
||||
GitBranchesReq, GitCheckoutCommitReq, GitCheckoutReq, GitCollectChangesReq,
|
||||
GitCollectChangesResponse, GitCommitReq, GitCurrentCommitReq, GitDiffReq, GitDiffsData,
|
||||
GitDiscardReq, GitFilesReq, GitInfoData, GitInfoReq, GitMetadataReq, GitReadFilesData,
|
||||
GitResolveRootReq, GitStageContentReq, GitStageReq, GitStashReq, GitStatusExtReq,
|
||||
GitStatusExtResponse, GitStatusReq, GitUnstageReq, StageData, VcsKind,
|
||||
};
|
||||
use kigi_workspace_types::rpc::hunks::{
|
||||
BulkHunkActionResponse, FileSummary, HunkActionResponse, HunkAllActionReq, HunkFileActionReq,
|
||||
HunkGetFileSummariesReq, HunkGetStagedFilesReq, HunkSingleActionReq, HunkTurnActionReq,
|
||||
};
|
||||
use kigi_workspace_types::rpc::search::{
|
||||
ContentSearchData, ContentSearchRequest, FuzzyChangeReq, FuzzyCloseReq, FuzzyOpenReq,
|
||||
FuzzyStatusReq,
|
||||
};
|
||||
use kigi_workspace_types::rpc::session::{
|
||||
BeginPromptReq, EndPromptReq, FileRewindResponse, RewindToReq,
|
||||
};
|
||||
use kigi_workspace_types::rpc::skills::{DiscoverPluginsReq, DiscoverSkillsReq, SkillInfo};
|
||||
use kigi_workspace_types::rpc::workspace::{
|
||||
ConfigureMcpReq, DropSessionReq, InstallPluginReq, LoadEnvrcReq, LoadPermissionsReq,
|
||||
LoadProjectConfigReq, RefreshPluginsReq, ResolveFileReferencesReq, ToolDefinitionsReq,
|
||||
UpdateToolConfigReq, WorkspaceInfo, WorkspaceInfoReq,
|
||||
};
|
||||
use kigi_workspace_types::rpc::worktree::{
|
||||
ApplyWorktreeRequest, CreateWorktreeRequest, RemoveWorktreeRequest, WorktreeCreateSyncReq,
|
||||
WorktreeDbPathReq, WorktreeDbPathResponse, WorktreeDbRebuildReq, WorktreeDbStatsReq,
|
||||
WorktreeGcReq, WorktreeListReq, WorktreeShowReq,
|
||||
};
|
||||
use kigi_workspace_types::rpc::{RpcEnvelope, RpcError, WORKSPACE_RPC_TOOL_ID, WorkspaceRpc};
|
||||
use serde_json::Value;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
use std::time::Duration;
|
||||
#[derive(Debug, thiserror::Error)]
|
||||
pub enum WorkspaceClientError {
|
||||
/// A previous call observed a fatal transport error and no
|
||||
/// reconnect has been signalled since.
|
||||
#[error("hub connection lost (previously disconnected)")]
|
||||
NotConnected,
|
||||
#[error("rpc failed: {0}")]
|
||||
Transport(String),
|
||||
#[error("{method} timed out after {after:?}")]
|
||||
Timeout { method: String, after: Duration },
|
||||
#[error("{method}: response decode: {source}")]
|
||||
Decode {
|
||||
method: String,
|
||||
#[source]
|
||||
source: serde_json::Error,
|
||||
},
|
||||
/// The server returned an error envelope.
|
||||
#[error("workspace rpc error: {0}")]
|
||||
Rpc(RpcError),
|
||||
}
|
||||
/// Consume a `ToolStream<TypedToolOutput>` to its terminal item,
|
||||
/// discarding progress frames.
|
||||
///
|
||||
/// Returns the terminal result, or a `ToolError::NetworkError` if the
|
||||
/// stream ended without producing a terminal item.
|
||||
pub async fn consume_stream_terminal(
|
||||
stream: &mut kigi_tool_runtime::ToolStream<TypedToolOutput>,
|
||||
) -> Result<TypedToolOutput, kigi_tool_runtime::ToolError> {
|
||||
loop {
|
||||
let item = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await;
|
||||
match item {
|
||||
Some(ToolStreamItem::Progress(_)) => {}
|
||||
Some(ToolStreamItem::Terminal(result)) => return result,
|
||||
None => {
|
||||
return Err(kigi_tool_runtime::ToolError::network_error(
|
||||
"stream ended without terminal item",
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Check whether a [`ToolError`](kigi_tool_runtime::ToolError) indicates
|
||||
/// a fatal transport failure that should mark the hub as disconnected.
|
||||
///
|
||||
/// Returns `true` for:
|
||||
/// - `NetworkError` — direct transport failure (socket dropped, stream
|
||||
/// ended without terminal item, etc.)
|
||||
/// - `Custom` with `details.code == "protocol_error"` — half-closed
|
||||
/// WebSocket producing malformed frames
|
||||
pub fn is_transport_fatal(err: &kigi_tool_runtime::ToolError) -> bool {
|
||||
match err.kind {
|
||||
kigi_tool_runtime::ToolErrorKind::NetworkError => true,
|
||||
kigi_tool_runtime::ToolErrorKind::Custom => err
|
||||
.details
|
||||
.as_ref()
|
||||
.and_then(|d| d.get("code"))
|
||||
.and_then(|c| c.as_str())
|
||||
.is_some_and(|c| c == "protocol_error"),
|
||||
_ => false,
|
||||
}
|
||||
}
|
||||
/// Typed client over a bound [`ToolHarness`] for `workspace.*` RPCs.
|
||||
///
|
||||
/// Clones share the harness and the connected latch, which fast-fails
|
||||
/// calls after a fatal transport error until
|
||||
/// [`mark_connected`](Self::mark_connected) resets it (e.g. from an SDK
|
||||
/// `on_reconnect` callback sharing the flag via
|
||||
/// [`with_connected_flag`](Self::with_connected_flag)).
|
||||
#[derive(Clone)]
|
||||
pub struct WorkspaceClient {
|
||||
harness: ToolHarness,
|
||||
connected: Arc<AtomicBool>,
|
||||
deadline: Option<Duration>,
|
||||
}
|
||||
impl std::fmt::Debug for WorkspaceClient {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
f.debug_struct("WorkspaceClient")
|
||||
.field("connected", &self.is_connected())
|
||||
.field("deadline", &self.deadline)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
impl WorkspaceClient {
|
||||
pub fn new(harness: ToolHarness) -> Self {
|
||||
Self {
|
||||
harness,
|
||||
connected: Arc::new(AtomicBool::new(true)),
|
||||
deadline: None,
|
||||
}
|
||||
}
|
||||
/// Shares a pre-created connected flag, so an SDK `on_reconnect`
|
||||
/// callback holding the same `Arc` can reset it.
|
||||
pub fn with_connected_flag(harness: ToolHarness, connected: Arc<AtomicBool>) -> Self {
|
||||
Self {
|
||||
harness,
|
||||
connected,
|
||||
deadline: None,
|
||||
}
|
||||
}
|
||||
/// Set a per-call deadline covering dispatch and stream consumption.
|
||||
pub fn with_deadline(mut self, deadline: Duration) -> Self {
|
||||
self.deadline = Some(deadline);
|
||||
self
|
||||
}
|
||||
pub fn harness(&self) -> &ToolHarness {
|
||||
&self.harness
|
||||
}
|
||||
/// Whether the hub connection is believed to be alive.
|
||||
pub fn is_connected(&self) -> bool {
|
||||
self.connected.load(Ordering::Relaxed)
|
||||
}
|
||||
pub fn mark_disconnected(&self) {
|
||||
self.connected.store(false, Ordering::Relaxed);
|
||||
}
|
||||
/// Reset after an SDK reconnect.
|
||||
pub fn mark_connected(&self) {
|
||||
self.connected.store(true, Ordering::Relaxed);
|
||||
}
|
||||
/// Untyped RPC call: `{"method": .., "params": ..}` through the
|
||||
/// `workspace_rpc` hub tool, returning the raw envelope value.
|
||||
pub async fn rpc_raw(
|
||||
&self,
|
||||
method: &str,
|
||||
params: Value,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
if !self.is_connected() {
|
||||
return Err(WorkspaceClientError::NotConnected);
|
||||
}
|
||||
let tool_id = kigi_tool_protocol::ToolId::new(WORKSPACE_RPC_TOOL_ID)
|
||||
.expect("constant tool id is valid");
|
||||
let args = serde_json::json!({ "method" : method, "params" : params });
|
||||
tracing::debug!(method, "WorkspaceClient::rpc");
|
||||
let fut = async {
|
||||
let mut stream = self
|
||||
.harness
|
||||
.call(tool_id, args, ToolCallContext::default())
|
||||
.await;
|
||||
consume_stream_terminal(&mut stream).await
|
||||
};
|
||||
let result = match self.deadline {
|
||||
Some(deadline) => match tokio::time::timeout(deadline, fut).await {
|
||||
Ok(r) => r,
|
||||
Err(_) => {
|
||||
return Err(WorkspaceClientError::Timeout {
|
||||
method: method.to_owned(),
|
||||
after: deadline,
|
||||
});
|
||||
}
|
||||
},
|
||||
None => fut.await,
|
||||
};
|
||||
let typed = result.map_err(|e| {
|
||||
if is_transport_fatal(&e) {
|
||||
self.mark_disconnected();
|
||||
}
|
||||
WorkspaceClientError::Transport(e.to_string())
|
||||
})?;
|
||||
Ok(typed.value)
|
||||
}
|
||||
/// Typed RPC call: derives the method and response type from the
|
||||
/// request type's [`WorkspaceRpc`] impl and decodes the envelope.
|
||||
pub async fn rpc<R: WorkspaceRpc>(&self, req: &R) -> Result<R::Response, WorkspaceClientError> {
|
||||
let params = serde_json::to_value(req).map_err(|e| WorkspaceClientError::Decode {
|
||||
method: R::METHOD.to_owned(),
|
||||
source: e,
|
||||
})?;
|
||||
let raw = self.rpc_raw(R::METHOD, params).await?;
|
||||
let envelope: RpcEnvelope<R::Response> =
|
||||
serde_json::from_value(raw).map_err(|e| WorkspaceClientError::Decode {
|
||||
method: R::METHOD.to_owned(),
|
||||
source: e,
|
||||
})?;
|
||||
envelope.into_result().map_err(WorkspaceClientError::Rpc)
|
||||
}
|
||||
/// `workspace.info`, decoded into the typed shape
|
||||
/// (`WorkspaceInfoReq::Response` is the raw `Value` for
|
||||
/// `WorkspaceOps` compat).
|
||||
pub async fn info(&self) -> Result<WorkspaceInfo, WorkspaceClientError> {
|
||||
let raw = self.rpc(&WorkspaceInfoReq {}).await?;
|
||||
serde_json::from_value(raw).map_err(|e| WorkspaceClientError::Decode {
|
||||
method: WorkspaceInfoReq::METHOD.to_owned(),
|
||||
source: e,
|
||||
})
|
||||
}
|
||||
/// `workspace.git_status` (JSON string value, ~1 KB server-side cap).
|
||||
pub async fn git_status(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&GitStatusReq {}).await
|
||||
}
|
||||
pub async fn discover_skills(&self) -> Result<Vec<SkillInfo>, WorkspaceClientError> {
|
||||
self.rpc(&DiscoverSkillsReq {}).await
|
||||
}
|
||||
pub async fn discover_agents_md(&self) -> Result<Vec<AgentConfigFile>, WorkspaceClientError> {
|
||||
self.rpc(&DiscoverAgentsMdReq {}).await
|
||||
}
|
||||
pub async fn git_status_ext(
|
||||
&self,
|
||||
req: &GitStatusExtReq,
|
||||
) -> Result<GitStatusExtResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_files(
|
||||
&self,
|
||||
req: &GitFilesReq,
|
||||
) -> Result<GitReadFilesData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_diff(&self, req: &GitDiffReq) -> Result<GitDiffsData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_stage(&self, req: &GitStageReq) -> Result<StageData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_stage_content(
|
||||
&self,
|
||||
req: &GitStageContentReq,
|
||||
) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_unstage(&self, req: &GitUnstageReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_discard(&self, req: &GitDiscardReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_commit(
|
||||
&self,
|
||||
req: &GitCommitReq,
|
||||
) -> Result<CommitResult, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_checkout(&self, req: &GitCheckoutReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_stash(&self, req: &GitStashReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_info(&self, req: &GitInfoReq) -> Result<GitInfoData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_branches(
|
||||
&self,
|
||||
req: &GitBranchesReq,
|
||||
) -> Result<GitBranchListData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_resolve_root(
|
||||
&self,
|
||||
req: &GitResolveRootReq,
|
||||
) -> Result<Option<std::path::PathBuf>, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_current_commit(
|
||||
&self,
|
||||
req: &GitCurrentCommitReq,
|
||||
) -> Result<Option<String>, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn detect_vcs_kind(
|
||||
&self,
|
||||
req: &DetectVcsKindReq,
|
||||
) -> Result<VcsKind, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_checkout_commit(
|
||||
&self,
|
||||
req: &GitCheckoutCommitReq,
|
||||
) -> Result<CheckoutCommitResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn git_branch_info(&self) -> Result<Option<GitInfoData>, WorkspaceClientError> {
|
||||
self.rpc(&GitBranchInfoReq {}).await
|
||||
}
|
||||
pub async fn git_metadata(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&GitMetadataReq {}).await
|
||||
}
|
||||
/// `workspace.git_collect_changes` — collect repository changes for serialization.
|
||||
pub async fn git_collect_changes(
|
||||
&self,
|
||||
req: &GitCollectChangesReq,
|
||||
) -> Result<GitCollectChangesResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn put_files(&self, req: &PutFilesReq) -> Result<PutFilesRes, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn get_files(&self, req: &GetFilesReq) -> Result<GetFilesRes, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fs_list(&self, req: &FsListReq) -> Result<FsListData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fs_exists(&self, req: &FsExistsReq) -> Result<FsExistsData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fs_read_file(
|
||||
&self,
|
||||
req: &FsReadFileReq,
|
||||
) -> Result<FsReadFileData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fs_write_file(&self, req: &FsWriteFileReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fs_delete_file(&self, req: &FsDeleteFileReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn hunk_action(
|
||||
&self,
|
||||
req: &HunkSingleActionReq,
|
||||
) -> Result<HunkActionResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn hunk_file_action(
|
||||
&self,
|
||||
req: &HunkFileActionReq,
|
||||
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn hunk_turn_action(
|
||||
&self,
|
||||
req: &HunkTurnActionReq,
|
||||
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn hunk_all_action(
|
||||
&self,
|
||||
req: &HunkAllActionReq,
|
||||
) -> Result<BulkHunkActionResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn hunk_get_staged_files(&self) -> Result<Vec<String>, WorkspaceClientError> {
|
||||
self.rpc(&HunkGetStagedFilesReq {}).await
|
||||
}
|
||||
pub async fn hunk_get_file_summaries(&self) -> Result<Vec<FileSummary>, WorkspaceClientError> {
|
||||
self.rpc(&HunkGetFileSummariesReq {}).await
|
||||
}
|
||||
pub async fn code_goto_definition(
|
||||
&self,
|
||||
req: &CodeGotoDefinitionReq,
|
||||
) -> Result<CodeNavResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn code_goto_references(
|
||||
&self,
|
||||
req: &CodeGotoReferencesReq,
|
||||
) -> Result<CodeNavResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn code_find_definitions(
|
||||
&self,
|
||||
req: &CodeFindDefinitionsReq,
|
||||
) -> Result<CodeNavResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn code_find_references(
|
||||
&self,
|
||||
req: &CodeFindReferencesReq,
|
||||
) -> Result<CodeNavResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn code_index_status(
|
||||
&self,
|
||||
req: &CodeIndexStatusReq,
|
||||
) -> Result<CodeIndexStatusResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn ripgrep(
|
||||
&self,
|
||||
req: &ContentSearchRequest,
|
||||
) -> Result<ContentSearchData, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fuzzy_open(&self, req: &FuzzyOpenReq) -> Result<String, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fuzzy_change(&self, req: &FuzzyChangeReq) -> Result<bool, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fuzzy_close(&self, req: &FuzzyCloseReq) -> Result<bool, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn fuzzy_status(&self, req: &FuzzyStatusReq) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn create_worktree(
|
||||
&self,
|
||||
req: &CreateWorktreeRequest,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn worktree_create_sync(
|
||||
&self,
|
||||
req: &WorktreeCreateSyncReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn remove_worktree(
|
||||
&self,
|
||||
req: &RemoveWorktreeRequest,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn apply_worktree(
|
||||
&self,
|
||||
req: &ApplyWorktreeRequest,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn worktree_list(
|
||||
&self,
|
||||
req: &WorktreeListReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn worktree_show(
|
||||
&self,
|
||||
req: &WorktreeShowReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn worktree_gc(&self, req: &WorktreeGcReq) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn worktree_db_rebuild(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&WorktreeDbRebuildReq {}).await
|
||||
}
|
||||
pub async fn worktree_db_path(&self) -> Result<WorktreeDbPathResponse, WorkspaceClientError> {
|
||||
self.rpc(&WorktreeDbPathReq {}).await
|
||||
}
|
||||
pub async fn worktree_db_stats(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&WorktreeDbStatsReq {}).await
|
||||
}
|
||||
pub async fn begin_prompt(&self, req: &BeginPromptReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn end_prompt(&self, req: &EndPromptReq) -> Result<(), WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn rewind_to(
|
||||
&self,
|
||||
req: &RewindToReq,
|
||||
) -> Result<FileRewindResponse, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn load_project_config(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&LoadProjectConfigReq {}).await
|
||||
}
|
||||
pub async fn load_permissions(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&LoadPermissionsReq {}).await
|
||||
}
|
||||
pub async fn load_envrc(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&LoadEnvrcReq {}).await
|
||||
}
|
||||
pub async fn tool_definitions(
|
||||
&self,
|
||||
req: &ToolDefinitionsReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn resolve_file_references(
|
||||
&self,
|
||||
req: &ResolveFileReferencesReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn update_tool_config(
|
||||
&self,
|
||||
req: &UpdateToolConfigReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn drop_session(&self, req: &DropSessionReq) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn configure_mcp(
|
||||
&self,
|
||||
req: &ConfigureMcpReq,
|
||||
) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(req).await
|
||||
}
|
||||
pub async fn install_plugin(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&InstallPluginReq {}).await
|
||||
}
|
||||
pub async fn refresh_plugins(&self) -> Result<Value, WorkspaceClientError> {
|
||||
self.rpc(&RefreshPluginsReq {}).await
|
||||
}
|
||||
pub async fn discover_plugins(&self) -> Result<Vec<Value>, WorkspaceClientError> {
|
||||
self.rpc(&DiscoverPluginsReq {}).await
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_computer_hub_sdk::harness::LocalRegistry;
|
||||
use kigi_tool_protocol::{SessionId, ToolId};
|
||||
use kigi_tool_runtime::{Tool, ToolError};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use kigi_workspace_types::rpc::skills::SkillScope;
|
||||
use schemars::JsonSchema;
|
||||
use serde::Deserialize;
|
||||
#[derive(Debug, Deserialize, JsonSchema)]
|
||||
struct RpcArgs {
|
||||
method: String,
|
||||
params: serde_json::Value,
|
||||
}
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
#[serde(transparent)]
|
||||
struct RawOut(serde_json::Value);
|
||||
impl kigi_tool_runtime::ToolOutput for RawOut {}
|
||||
#[derive(Debug)]
|
||||
struct FakeWorkspaceRpc;
|
||||
impl Tool for FakeWorkspaceRpc {
|
||||
type Args = RpcArgs;
|
||||
type Output = RawOut;
|
||||
fn id(&self) -> ToolId {
|
||||
ToolId::new(WORKSPACE_RPC_TOOL_ID).unwrap()
|
||||
}
|
||||
fn description(&self, _ctx: &::kigi_tool_runtime::ListToolsContext) -> ToolDescription {
|
||||
ToolDescription::new(WORKSPACE_RPC_TOOL_ID, "fake workspace rpc")
|
||||
}
|
||||
async fn run(&self, _ctx: ToolCallContext, args: Self::Args) -> Result<RawOut, ToolError> {
|
||||
let ok = |v: serde_json::Value| Ok(RawOut(serde_json::json!({ "ok" : v })));
|
||||
match args.method.as_str() {
|
||||
"workspace.info" => ok(serde_json::json!(
|
||||
{ "os" : "linux", "shell" : "bash", "cwd" : "/workspace", }
|
||||
)),
|
||||
"workspace.git_status" => ok(serde_json::json!("On branch main")),
|
||||
"workspace.discover_skills" => ok(serde_json::json!(
|
||||
[{ "name" : "my-skill", "description" : "A test skill",
|
||||
"path" : "/workspace/.kigi/skills/my-skill/SKILL.md", "scope"
|
||||
: "local", }]
|
||||
)),
|
||||
"workspace.discover_agents_md" => ok(serde_json::json!(
|
||||
[{ "file_name" : "AGENTS.md", "file_path" :
|
||||
"/workspace/AGENTS.md", "content" : "# Project instructions",
|
||||
}]
|
||||
)),
|
||||
"workspace.echo_params" => ok(args.params),
|
||||
"workspace.err" => Ok(RawOut(serde_json::json!(
|
||||
{ "err" : { "code" : "session_not_found", "message" :
|
||||
"ghost" }, }
|
||||
))),
|
||||
"workspace.malformed" => Ok(RawOut(serde_json::json!({ "neither" : true }))),
|
||||
"workspace.slow" => {
|
||||
tokio::time::sleep(Duration::from_secs(60)).await;
|
||||
ok(serde_json::Value::Null)
|
||||
}
|
||||
"workspace.netfail" => Err(ToolError::network_error("socket dropped")),
|
||||
"workspace.toolfail" => Err(ToolError::custom("some_code", "boom")),
|
||||
other => panic!("unexpected method {other}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
fn client() -> WorkspaceClient {
|
||||
let registry = LocalRegistry::new();
|
||||
registry.register(FakeWorkspaceRpc);
|
||||
let harness = ToolHarness::local_only_with(
|
||||
registry,
|
||||
SessionId::new("test").unwrap(),
|
||||
Default::default(),
|
||||
);
|
||||
WorkspaceClient::new(harness)
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn info_decodes_typed_response() {
|
||||
let info = client().info().await.unwrap();
|
||||
assert_eq!(
|
||||
info,
|
||||
WorkspaceInfo {
|
||||
os: "linux".into(),
|
||||
shell: "bash".into(),
|
||||
cwd: "/workspace".into(),
|
||||
}
|
||||
);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn git_status_returns_raw_value() {
|
||||
let v = client().git_status().await.unwrap();
|
||||
assert_eq!(v, serde_json::json!("On branch main"));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn discover_skills_decodes_typed_list() {
|
||||
let skills = client().discover_skills().await.unwrap();
|
||||
assert_eq!(skills.len(), 1);
|
||||
assert_eq!(skills[0].name, "my-skill");
|
||||
assert_eq!(skills[0].scope, SkillScope::Local);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn discover_agents_md_decodes_typed_list() {
|
||||
let files = client().discover_agents_md().await.unwrap();
|
||||
assert_eq!(files.len(), 1);
|
||||
assert_eq!(files[0].file_name, "AGENTS.md");
|
||||
assert_eq!(files[0].file_path, "/workspace/AGENTS.md");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn typed_request_params_round_trip() {
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
struct EchoReq {
|
||||
flag: bool,
|
||||
n: u32,
|
||||
}
|
||||
impl WorkspaceRpc for EchoReq {
|
||||
const METHOD: &'static str = "workspace.echo_params";
|
||||
type Response = Value;
|
||||
}
|
||||
let echoed = client().rpc(&EchoReq { flag: true, n: 7 }).await.unwrap();
|
||||
assert_eq!(echoed, serde_json::json!({ "flag" : true, "n" : 7 }));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn err_envelope_maps_to_rpc_error() {
|
||||
let c = client();
|
||||
let raw = c.rpc_raw("workspace.err", Value::Null).await.unwrap();
|
||||
assert!(raw.get("err").is_some());
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
struct ErrReq;
|
||||
impl WorkspaceRpc for ErrReq {
|
||||
const METHOD: &'static str = "workspace.err";
|
||||
type Response = Value;
|
||||
}
|
||||
let err = c.rpc(&ErrReq).await.unwrap_err();
|
||||
match err {
|
||||
WorkspaceClientError::Rpc(e) => {
|
||||
assert_eq!(e.code, "session_not_found");
|
||||
assert_eq!(e.message, "ghost");
|
||||
}
|
||||
other => panic!("expected Rpc, got {other:?}"),
|
||||
}
|
||||
assert!(c.is_connected(), "envelope errors must not trip the latch");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn malformed_envelope_maps_to_decode() {
|
||||
let c = client();
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
struct MalformedReq;
|
||||
impl WorkspaceRpc for MalformedReq {
|
||||
const METHOD: &'static str = "workspace.malformed";
|
||||
type Response = Value;
|
||||
}
|
||||
let err = c.rpc(&MalformedReq).await.unwrap_err();
|
||||
assert!(
|
||||
matches!(err, WorkspaceClientError::Decode { .. }),
|
||||
"{err:?}"
|
||||
);
|
||||
assert!(c.is_connected());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn network_error_trips_latch_and_fast_fails() {
|
||||
let c = client();
|
||||
let err = c
|
||||
.rpc_raw("workspace.netfail", Value::Null)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, WorkspaceClientError::Transport(_)), "{err:?}");
|
||||
assert!(!c.is_connected(), "network error must trip the latch");
|
||||
let err = c.rpc_raw("workspace.info", Value::Null).await.unwrap_err();
|
||||
assert!(matches!(err, WorkspaceClientError::NotConnected));
|
||||
c.mark_connected();
|
||||
assert!(c.rpc_raw("workspace.info", Value::Null).await.is_ok());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn non_fatal_tool_error_leaves_latch_up() {
|
||||
let c = client();
|
||||
let err = c
|
||||
.rpc_raw("workspace.toolfail", Value::Null)
|
||||
.await
|
||||
.unwrap_err();
|
||||
assert!(matches!(err, WorkspaceClientError::Transport(_)), "{err:?}");
|
||||
assert!(
|
||||
c.is_connected(),
|
||||
"non-fatal tool errors must not trip the latch"
|
||||
);
|
||||
}
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn deadline_times_out_slow_calls() {
|
||||
let c = client().with_deadline(Duration::from_secs(3));
|
||||
let err = c.rpc_raw("workspace.slow", Value::Null).await.unwrap_err();
|
||||
match err {
|
||||
WorkspaceClientError::Timeout { method, after } => {
|
||||
assert_eq!(method, "workspace.slow");
|
||||
assert_eq!(after, Duration::from_secs(3));
|
||||
}
|
||||
other => panic!("expected Timeout, got {other:?}"),
|
||||
}
|
||||
assert!(c.is_connected(), "timeouts must not trip the latch");
|
||||
}
|
||||
#[tokio::test(start_paused = true)]
|
||||
async fn no_deadline_by_default_waits_out_slow_calls() {
|
||||
let v = client()
|
||||
.rpc_raw("workspace.slow", Value::Null)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(v.get("ok").is_some());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn shared_connected_flag_is_externally_controllable() {
|
||||
let registry = LocalRegistry::new();
|
||||
registry.register(FakeWorkspaceRpc);
|
||||
let harness = ToolHarness::local_only_with(
|
||||
registry,
|
||||
SessionId::new("test").unwrap(),
|
||||
Default::default(),
|
||||
);
|
||||
let flag = Arc::new(AtomicBool::new(true));
|
||||
let c = WorkspaceClient::with_connected_flag(harness, flag.clone());
|
||||
flag.store(false, Ordering::Relaxed);
|
||||
let err = c.rpc_raw("workspace.info", Value::Null).await.unwrap_err();
|
||||
assert!(matches!(err, WorkspaceClientError::NotConnected));
|
||||
flag.store(true, Ordering::Relaxed);
|
||||
assert!(c.rpc_raw("workspace.info", Value::Null).await.is_ok());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn clones_share_the_latch() {
|
||||
let a = client();
|
||||
let b = a.clone();
|
||||
a.mark_disconnected();
|
||||
assert!(!b.is_connected());
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn consume_stream_terminal_returns_ok() {
|
||||
let value = serde_json::json!({ "result" : "hello" });
|
||||
let typed = TypedToolOutput::from_value(ToolId::new("t").unwrap(), value.clone());
|
||||
let mut stream = kigi_tool_runtime::terminal_only(Ok(typed));
|
||||
assert_eq!(
|
||||
consume_stream_terminal(&mut stream).await.unwrap().value,
|
||||
value
|
||||
);
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn consume_stream_terminal_returns_err() {
|
||||
let mut stream: kigi_tool_runtime::ToolStream<TypedToolOutput> =
|
||||
kigi_tool_runtime::terminal_only(Err(ToolError::network_error("oops")));
|
||||
let err = consume_stream_terminal(&mut stream).await.unwrap_err();
|
||||
assert!(err.to_string().contains("oops"));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn consume_stream_terminal_exhausted_stream_is_network_error() {
|
||||
let typed = TypedToolOutput::from_value(ToolId::new("t").unwrap(), Value::Null);
|
||||
let mut stream = kigi_tool_runtime::terminal_only::<TypedToolOutput>(Ok(typed));
|
||||
let _ = consume_stream_terminal(&mut stream).await;
|
||||
let err = consume_stream_terminal(&mut stream).await.unwrap_err();
|
||||
assert!(
|
||||
err.to_string()
|
||||
.contains("stream ended without terminal item")
|
||||
);
|
||||
assert!(is_transport_fatal(&err));
|
||||
}
|
||||
}
|
||||
@@ -1,583 +0,0 @@
|
||||
//! Standalone workspace ToolServer for remote sandboxes.
|
||||
//!
|
||||
//! Reads OIDC credentials from `~/.kigi/auth.json`, connects to a
|
||||
//! server, exposes workspace tools, and refreshes tokens
|
||||
//! automatically.
|
||||
use clap::Parser;
|
||||
use kigi_workspace::config::WorkspaceServerMetadata;
|
||||
use kigi_workspace::daemonize;
|
||||
use kigi_workspace::diag_server;
|
||||
use kigi_workspace::preview_supervisor::{self, PreviewArgs, PreviewVisibility};
|
||||
use std::path::PathBuf;
|
||||
use url::Url;
|
||||
/// OTLP `service.name` for this binary's exported traces/logs/metrics and
|
||||
/// direct-OTLP fastrace export. Single source so the call sites can't drift.
|
||||
const SERVICE_NAME: &str = "prod_grok_workspace";
|
||||
const EXIT_SERVER_ID_INVALID: i32 = 3;
|
||||
const INVALID_SERVER_ID_MARKER: &str = "workspace-server: invalid --server-id";
|
||||
fn server_id_startup_error(id: &str) -> Option<String> {
|
||||
id.parse::<kigi_tool_protocol::ServerId>()
|
||||
.err()
|
||||
.map(|e| format!("{INVALID_SERVER_ID_MARKER} {id:?}: {e}"))
|
||||
}
|
||||
#[derive(Parser)]
|
||||
#[command(name = "kigi-workspace-server")]
|
||||
#[command(about = "Standalone workspace ToolServer for the server connection")]
|
||||
struct Args {
|
||||
/// Print the capability manifest as JSON to stdout and exit 0. Legacy
|
||||
/// binaries reject the unknown flag via clap (non-zero exit), giving the
|
||||
/// launcher a definitive feature probe.
|
||||
#[arg(long)]
|
||||
capabilities: bool,
|
||||
#[arg(long, default_value = "wss://computer-hub.kigi.com/v1/tools")]
|
||||
hub_url: String,
|
||||
#[arg(long)]
|
||||
auth_config: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
cwd: Option<PathBuf>,
|
||||
/// Stable server identity for hub registration. Used as the
|
||||
/// `server_id` in `servers.list` and `server.bind` so clients
|
||||
/// can address this specific workspace server.
|
||||
/// When omitted, the SDK default ("workspace-server") is used.
|
||||
#[arg(long)]
|
||||
server_id: Option<String>,
|
||||
/// JSON metadata attached to the tool server registration.
|
||||
/// Propagated to `ServerInfo.metadata` in `servers.list` responses.
|
||||
#[arg(long)]
|
||||
metadata: Option<String>,
|
||||
/// Path to write a PID file once the server connection is established.
|
||||
/// The sandbox service polls this file to determine readiness.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_READY_PATH)]
|
||||
ready_file: PathBuf,
|
||||
/// Unix-socket path for the in-guest diagnostics HTTP server
|
||||
/// (`/ready`, `/statusz`).
|
||||
#[cfg(unix)]
|
||||
#[arg(long, default_value = diag_server::DEFAULT_DIAG_SOCKET_PATH)]
|
||||
diag_socket: PathBuf,
|
||||
/// Loopback TCP port for the diagnostics HTTP server (Windows guests,
|
||||
/// which lack a reliable Unix-socket HTTP client).
|
||||
#[cfg(windows)]
|
||||
#[arg(long, default_value_t = diag_server::DEFAULT_DIAG_PORT)]
|
||||
diag_port: u16,
|
||||
/// Permit a plaintext `ws://` hub on a non-loopback host. Only for a
|
||||
/// mesh-secured transport; the bearer crosses the network otherwise.
|
||||
#[arg(long)]
|
||||
allow_insecure_ws: bool,
|
||||
/// Fail `session.bind`s without an explicit toolset closed (RPC-only)
|
||||
/// instead of widening to the built-in default catalog. Passed by the
|
||||
/// sandbox service; doubles as a version tripwire (a stale revived binary
|
||||
/// rejects the argv and never reports ready).
|
||||
#[arg(long)]
|
||||
require_explicit_toolset: bool,
|
||||
/// Confine `x.ai/fs/*` resolution to the workspace root (reject `..`,
|
||||
/// absolute-outside-root, symlink escapes). On by default: the standalone
|
||||
/// server always backs a remote-sandbox workspace, a real tenant boundary.
|
||||
/// Override with `KIGI_WORKSPACE_CONFINE_FS_TO_ROOT=false` (e.g. local dev).
|
||||
#[arg(
|
||||
long,
|
||||
env = "KIGI_WORKSPACE_CONFINE_FS_TO_ROOT",
|
||||
default_value_t = true,
|
||||
action = clap::ArgAction::Set,
|
||||
)]
|
||||
confine_fs_to_workspace_root: bool,
|
||||
/// Self-daemonize at startup: double-fork + `setsid()` into a new session
|
||||
/// and process group (escaping the launcher's process-group reap),
|
||||
/// redirect stdio to a log file, and hold a single-instance pidfile lock.
|
||||
///
|
||||
/// Off by default — opt-in, passed by the launcher in the supervised
|
||||
/// deployment mode. With the flag absent, startup is unchanged.
|
||||
#[arg(long)]
|
||||
daemonize: bool,
|
||||
/// Where `--daemonize` redirects stdout+stderr. Ignored without
|
||||
/// `--daemonize`.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_LOG_PATH)]
|
||||
log_file: PathBuf,
|
||||
/// Single-instance pidfile lock path used with `--daemonize`. Ignored
|
||||
/// without `--daemonize`.
|
||||
#[arg(long, default_value = daemonize::DEFAULT_PIDFILE_PATH)]
|
||||
pid_file: PathBuf,
|
||||
#[command(flatten)]
|
||||
preview: PreviewCliArgs,
|
||||
}
|
||||
/// Preview-proxy supervision flags. Forwarded 1:1 to the
|
||||
/// `/usr/local/bin/xai-grok-preview-proxy` child (see `cli.rs` for the proxy's
|
||||
/// flag names). Off by default — when `--preview-enabled` is absent the
|
||||
/// supervisor is never started and startup is byte-for-byte the non-preview
|
||||
/// path.
|
||||
#[derive(clap::Args, Debug)]
|
||||
struct PreviewCliArgs {
|
||||
/// Spawn and supervise the in-sandbox preview-proxy. The launcher passes
|
||||
/// this only when the proxy binary was mounted into this container.
|
||||
#[arg(long)]
|
||||
preview_enabled: bool,
|
||||
/// Proxy `--preview-port` (externally exposed listener). Absent ⇒ proxy default.
|
||||
#[arg(long)]
|
||||
preview_port: Option<u16>,
|
||||
/// Proxy `--control-port` (loopback control). Absent ⇒ proxy default.
|
||||
#[arg(long)]
|
||||
preview_control_port: Option<u16>,
|
||||
/// Proxy `--visibility` (`owner` | `public`). Absent ⇒ proxy default.
|
||||
/// Validated here so a bad value fails fast instead of crash-looping the proxy.
|
||||
#[arg(long, value_enum)]
|
||||
preview_visibility: Option<PreviewVisibility>,
|
||||
/// Proxy `--instance-suffix` for inbound Host validation.
|
||||
#[arg(long)]
|
||||
preview_instance_suffix: Option<String>,
|
||||
/// Proxy `--auth-redirect`: URL the unauthenticated handshake redirects to.
|
||||
/// Required for the default `owner` gate to redirect rather than deny.
|
||||
#[arg(long)]
|
||||
preview_auth_redirect: Option<String>,
|
||||
/// Proxy `--allow-public` org public-policy gate (forwarded only when set).
|
||||
#[arg(long)]
|
||||
preview_allow_public: bool,
|
||||
/// Proxy `--workspace-server-port` (added to the discovery denylist).
|
||||
#[arg(long)]
|
||||
preview_workspace_server_port: Option<u16>,
|
||||
}
|
||||
impl PreviewCliArgs {
|
||||
fn into_preview_args(self, workspace_dir: PathBuf) -> PreviewArgs {
|
||||
PreviewArgs {
|
||||
enabled: self.preview_enabled,
|
||||
port: self.preview_port,
|
||||
control_port: self.preview_control_port,
|
||||
visibility: self.preview_visibility,
|
||||
instance_suffix: self.preview_instance_suffix,
|
||||
auth_redirect: self.preview_auth_redirect,
|
||||
allow_public: self.preview_allow_public,
|
||||
workspace_server_port: self.preview_workspace_server_port,
|
||||
workspace_dir,
|
||||
}
|
||||
}
|
||||
}
|
||||
/// Capability manifest printed by `--capabilities`, consumed by the sandbox
|
||||
/// launcher to pick a launch protocol. Additions are backward-compatible.
|
||||
#[derive(Debug, serde::Serialize)]
|
||||
struct Capabilities {
|
||||
/// The in-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`).
|
||||
diag: bool,
|
||||
}
|
||||
const CAPABILITIES: Capabilities = Capabilities { diag: true };
|
||||
fn main() -> anyhow::Result<()> {
|
||||
let mut args = Args::parse();
|
||||
if args.capabilities {
|
||||
println!("{}", serde_json::to_string(&CAPABILITIES)?);
|
||||
return Ok(());
|
||||
}
|
||||
if let Some(msg) = args.server_id.as_deref().and_then(server_id_startup_error) {
|
||||
eprintln!("{msg}");
|
||||
std::process::exit(EXIT_SERVER_ID_INVALID);
|
||||
}
|
||||
let cwd = match args.cwd {
|
||||
Some(ref p) => dunce::canonicalize(p)?,
|
||||
None => std::env::current_dir()?,
|
||||
};
|
||||
let _pidfile_guard = if args.daemonize {
|
||||
let anchor = |p: PathBuf| if p.is_absolute() { p } else { cwd.join(p) };
|
||||
args.log_file = anchor(std::mem::take(&mut args.log_file));
|
||||
args.pid_file = anchor(std::mem::take(&mut args.pid_file));
|
||||
args.ready_file = anchor(std::mem::take(&mut args.ready_file));
|
||||
#[cfg(unix)]
|
||||
{
|
||||
args.diag_socket = anchor(std::mem::take(&mut args.diag_socket));
|
||||
}
|
||||
args.auth_config = args.auth_config.take().map(anchor);
|
||||
daemonize::daemonize(&args.log_file)?;
|
||||
match daemonize::PidFile::acquire_or_take_over(&args.pid_file, daemonize::TAKEOVER_GRACE)? {
|
||||
Some(guard) => Some(guard),
|
||||
None => return Ok(()),
|
||||
}
|
||||
} else {
|
||||
None
|
||||
};
|
||||
tokio::runtime::Builder::new_multi_thread()
|
||||
.enable_all()
|
||||
.build()?
|
||||
.block_on(run(args, cwd))
|
||||
}
|
||||
async fn run(args: Args, cwd: PathBuf) -> anyhow::Result<()> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
use tracing_subscriber::layer::SubscriberExt as _;
|
||||
use tracing_subscriber::util::SubscriberInitExt as _;
|
||||
let env_filter = tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"));
|
||||
let donating = kigi_computer_hub_sdk::DonatingLogLayer::new_inert();
|
||||
tracing_subscriber::registry()
|
||||
.with(env_filter)
|
||||
.with(tracing_subscriber::fmt::layer())
|
||||
.with(donating.clone())
|
||||
.init();
|
||||
let direct_otlp = match std::env::var("KIGI_WORKSPACE_OTLP_ENDPOINT") {
|
||||
Ok(endpoint) if !endpoint.is_empty() => {
|
||||
match kigi_tracing::init_fastrace(endpoint.clone(), SERVICE_NAME.to_owned(), None) {
|
||||
Ok(()) => {
|
||||
tracing::info!(% endpoint, "trace export enabled (direct OTLP)");
|
||||
true
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = % e, "direct OTLP trace export init failed");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => false,
|
||||
};
|
||||
let url = Url::parse(&args.hub_url).map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?;
|
||||
{
|
||||
use kigi_sandbox::{ProfileName, SandboxManager};
|
||||
let profile = match std::env::var("KIGI_SANDBOX_PROFILE").ok() {
|
||||
Some(val) => {
|
||||
let parsed = val
|
||||
.parse::<ProfileName>()
|
||||
.expect("ProfileName::from_str is infallible");
|
||||
if matches!(parsed, ProfileName::Custom(_)) {
|
||||
tracing::warn!(
|
||||
value = % val,
|
||||
"Unrecognized KIGI_SANDBOX_PROFILE, defaulting to workspace"
|
||||
);
|
||||
ProfileName::Workspace
|
||||
} else {
|
||||
parsed
|
||||
}
|
||||
}
|
||||
None if kigi_sandbox::trust_bwrap_marker_for_devbox() => ProfileName::Devbox,
|
||||
None => ProfileName::Workspace,
|
||||
};
|
||||
let profile_name = profile.to_string();
|
||||
if profile == ProfileName::Off {
|
||||
tracing::info!(
|
||||
profile = % profile_name,
|
||||
"Sandbox explicitly disabled via KIGI_SANDBOX_PROFILE=off"
|
||||
);
|
||||
} else {
|
||||
let mut sandbox = SandboxManager::new(profile, &cwd);
|
||||
if let Err(e) = sandbox.apply(&cwd) {
|
||||
tracing::warn!(
|
||||
error = % e, "Sandbox apply returned error, continuing unsandboxed"
|
||||
);
|
||||
} else if !sandbox.is_applied() {
|
||||
tracing::warn!("Sandbox could not be applied (unsupported platform)");
|
||||
}
|
||||
sandbox.install();
|
||||
let active = kigi_sandbox::is_active();
|
||||
let status_msg = if active {
|
||||
"Workspace server sandbox active"
|
||||
} else {
|
||||
"Workspace server sandbox NOT active"
|
||||
};
|
||||
tracing::info!(
|
||||
profile = % profile_name, active, restrict_network =
|
||||
kigi_sandbox::should_restrict_child_network(), "{status_msg}"
|
||||
);
|
||||
}
|
||||
}
|
||||
let auth_provider = kigi_workspace::hub_auth::provider(&url, args.auth_config.as_deref())?;
|
||||
tracing::info!(hub_url = % url, cwd = % cwd.display(), "Starting workspace server");
|
||||
let cwd_display = cwd.display().to_string();
|
||||
let session_id = std::env::var("KIGI_SESSION_ID").ok();
|
||||
let parsed_metadata = match args.metadata {
|
||||
Some(json_str) => Some(
|
||||
serde_json::from_str(&json_str)
|
||||
.map_err(|e| anyhow::anyhow!("invalid --metadata JSON: {e}"))?,
|
||||
),
|
||||
None => None,
|
||||
};
|
||||
let metadata = WorkspaceServerMetadata::merge_session_metadata(parsed_metadata, session_id);
|
||||
let launch_id = metadata
|
||||
.as_ref()
|
||||
.and_then(|v| v.get("launch_id"))
|
||||
.and_then(serde_json::Value::as_str)
|
||||
.map(str::to_owned);
|
||||
let diag_handle = diag_server::DiagHandle::new(launch_id);
|
||||
#[cfg(unix)]
|
||||
let diag_listener = diag_server::DiagListener::Unix(args.diag_socket);
|
||||
#[cfg(windows)]
|
||||
let diag_listener = diag_server::DiagListener::Tcp(args.diag_port);
|
||||
let diag_log_file = args.daemonize.then_some(args.log_file);
|
||||
let _diag_server =
|
||||
match diag_server::serve(diag_listener, diag_handle.clone(), diag_log_file).await {
|
||||
Ok(bound) => {
|
||||
tracing::info!(addr = % bound.addr, "diagnostics server listening");
|
||||
Some(bound)
|
||||
}
|
||||
Err(e) => {
|
||||
if args.daemonize {
|
||||
tracing::error!(error = % e, "{}", diag_server::DIAG_BIND_FAILED_MARKER);
|
||||
std::process::exit(diag_server::EXIT_DIAG_BIND_FAILED);
|
||||
}
|
||||
tracing::warn!(
|
||||
error = % e, "{} (continuing without)",
|
||||
diag_server::DIAG_BIND_FAILED_MARKER
|
||||
);
|
||||
None
|
||||
}
|
||||
};
|
||||
tracing::info!(
|
||||
cwd = % cwd_display,
|
||||
"Workspace server starting — sessions created dynamically via server bind"
|
||||
);
|
||||
let server_id = args.server_id.clone();
|
||||
let status_config = kigi_workspace::StatusConfig::from_env();
|
||||
let preview_shutdown = if args.preview.preview_enabled {
|
||||
let control_port = args.preview.preview_control_port;
|
||||
let cfg = args.preview.into_preview_args(cwd.clone());
|
||||
let (tx, rx) = tokio::sync::watch::channel(false);
|
||||
tokio::spawn(preview_supervisor::supervise_preview(cfg, rx));
|
||||
Some((tx, control_port))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let project_lsp_trusted = true;
|
||||
let preview_scrape_interval = status_config.preview_activity_scrape_interval;
|
||||
kigi_workspace::init_metrics();
|
||||
let ws_handle = kigi_workspace::handle::connect_local_workspace(
|
||||
cwd,
|
||||
url,
|
||||
auth_provider,
|
||||
metadata,
|
||||
server_id.clone(),
|
||||
None,
|
||||
args.allow_insecure_ws,
|
||||
status_config,
|
||||
project_lsp_trusted,
|
||||
Some(args.ready_file.clone()),
|
||||
Some(diag_handle.clone()),
|
||||
args.require_explicit_toolset,
|
||||
args.confine_fs_to_workspace_root,
|
||||
)
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("failed to connect workspace to hub: {e}"))?;
|
||||
if let Some((tx, control_port)) = &preview_shutdown {
|
||||
tokio::spawn(preview_supervisor::supervise_preview_activity(
|
||||
*control_port,
|
||||
ws_handle.activity_tracker().clone(),
|
||||
preview_scrape_interval,
|
||||
tx.subscribe(),
|
||||
));
|
||||
}
|
||||
let mut donation_pump = None;
|
||||
if !direct_otlp {
|
||||
match ws_handle.trace_donation_reporter(SERVICE_NAME).await {
|
||||
Some((reporter, pump)) => {
|
||||
fastrace::set_reporter(reporter, fastrace::collector::Config::default());
|
||||
donation_pump = Some(pump);
|
||||
tracing::info!("trace export enabled");
|
||||
}
|
||||
None => tracing::info!("trace export disabled (not connected)"),
|
||||
}
|
||||
}
|
||||
let mut log_donation_pump = None;
|
||||
match ws_handle.log_donation_layer(SERVICE_NAME).await {
|
||||
Some((sender, pump)) => {
|
||||
donating.activate(sender);
|
||||
log_donation_pump = Some(pump);
|
||||
tracing::info!("log export enabled");
|
||||
}
|
||||
None => tracing::info!("log export disabled (not connected)"),
|
||||
}
|
||||
let mut metric_donation_pump = None;
|
||||
match ws_handle.metric_donation_reporter(SERVICE_NAME).await {
|
||||
Some(pump) => {
|
||||
metric_donation_pump = Some(pump);
|
||||
tracing::info!("metric export enabled");
|
||||
}
|
||||
None => tracing::info!("metric export disabled (not connected)"),
|
||||
}
|
||||
tracing::info!(
|
||||
server_id = ? server_id, "Workspace server connected to hub. Serving tools."
|
||||
);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use tokio::signal::unix::{SignalKind, signal};
|
||||
let mut sigterm = signal(SignalKind::terminate())?;
|
||||
tokio::select! {
|
||||
_ = tokio::signal::ctrl_c() => {} _ = sigterm.recv() => {}
|
||||
}
|
||||
}
|
||||
#[cfg(not(unix))]
|
||||
tokio::signal::ctrl_c().await?;
|
||||
if let Some((tx, _)) = &preview_shutdown {
|
||||
let _ = tx.send(true);
|
||||
}
|
||||
let _ = std::fs::remove_file(&args.ready_file);
|
||||
diag_handle.set_shutting_down();
|
||||
tracing::info!("Received shutdown signal, draining...");
|
||||
let tracker = ws_handle.activity_tracker().clone();
|
||||
let grace_budget = kigi_workspace::handle::termination_grace_from_env();
|
||||
ws_handle
|
||||
.two_phase_drain(grace_budget, kigi_workspace::handle::DrainReason::Sigterm)
|
||||
.await;
|
||||
tracker.set_shutting_down();
|
||||
tracing::info!("Shutting down...");
|
||||
fastrace::flush();
|
||||
if let Some(pump) = &donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
kigi_computer_hub_sdk::flush_log_layer();
|
||||
if let Some(pump) = &log_donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
if let Some(pump) = &metric_donation_pump {
|
||||
pump.drain().await;
|
||||
}
|
||||
ws_handle.shutdown_hub().await;
|
||||
kigi_sandbox::flush();
|
||||
Ok(())
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
#[test]
|
||||
fn capabilities_flag_parses_and_defaults_off() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.capabilities);
|
||||
let args = Args::try_parse_from(["kigi-workspace-server", "--capabilities"]).unwrap();
|
||||
assert!(args.capabilities);
|
||||
}
|
||||
#[test]
|
||||
fn capabilities_manifest_shape() {
|
||||
let value = serde_json::to_value(CAPABILITIES).unwrap();
|
||||
assert_eq!(value, serde_json::json!({ "diag" : true }));
|
||||
}
|
||||
#[test]
|
||||
fn capabilities_probe_of_legacy_binary_exits_nonzero() {
|
||||
/// A stand-in for a pre-`--capabilities` Args surface.
|
||||
#[derive(Debug, Parser)]
|
||||
struct LegacyArgs {
|
||||
#[arg(long)]
|
||||
daemonize: bool,
|
||||
}
|
||||
let err = LegacyArgs::try_parse_from(["kigi-workspace-server", "--capabilities"])
|
||||
.expect_err("a legacy binary must reject the flag");
|
||||
assert_eq!(err.kind(), clap::error::ErrorKind::UnknownArgument);
|
||||
assert_ne!(
|
||||
err.exit_code(),
|
||||
0,
|
||||
"the probe relies on a non-zero exit distinguishing legacy binaries"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn daemonize_defaults_are_inert() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.daemonize);
|
||||
assert_eq!(args.log_file, PathBuf::from(daemonize::DEFAULT_LOG_PATH));
|
||||
assert_eq!(
|
||||
args.pid_file,
|
||||
PathBuf::from(daemonize::DEFAULT_PIDFILE_PATH)
|
||||
);
|
||||
assert_eq!(
|
||||
args.ready_file,
|
||||
PathBuf::from(daemonize::DEFAULT_READY_PATH)
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn invalid_server_id_produces_the_marker_line() {
|
||||
for bad in ["auto:tool:x", ""] {
|
||||
let msg = server_id_startup_error(bad)
|
||||
.unwrap_or_else(|| panic!("server id {bad:?} must be rejected"));
|
||||
assert!(
|
||||
msg.starts_with(INVALID_SERVER_ID_MARKER),
|
||||
"startup error must START with the greppable marker prefix: {msg}"
|
||||
);
|
||||
}
|
||||
assert_eq!(server_id_startup_error("session-abc-123"), None);
|
||||
}
|
||||
#[test]
|
||||
fn argv_rejection_exit_code_is_distinct_from_server_id_exit_code() {
|
||||
let err = Args::try_parse_from(["kigi-workspace-server", "--flag-from-the-future"])
|
||||
.err()
|
||||
.expect("unknown argv must be rejected");
|
||||
assert_eq!(err.exit_code(), 2, "clap argv rejection exits 2");
|
||||
assert_ne!(err.exit_code(), EXIT_SERVER_ID_INVALID);
|
||||
assert_ne!(EXIT_SERVER_ID_INVALID, 0);
|
||||
assert_ne!(EXIT_SERVER_ID_INVALID, 1);
|
||||
}
|
||||
#[test]
|
||||
fn preview_defaults_are_inert() {
|
||||
let args = Args::try_parse_from(["kigi-workspace-server"]).unwrap();
|
||||
assert!(!args.preview.preview_enabled);
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert!(!cfg.enabled);
|
||||
assert!(
|
||||
cfg.to_argv().is_empty(),
|
||||
"an inert config forwards no proxy args"
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn preview_flags_parse_and_lower_to_supervisor_config() {
|
||||
let args = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-port",
|
||||
"6014",
|
||||
"--preview-control-port",
|
||||
"6015",
|
||||
"--preview-visibility",
|
||||
"public",
|
||||
"--preview-instance-suffix",
|
||||
".inst.example",
|
||||
"--preview-auth-redirect",
|
||||
"https://grok.com/preview-auth",
|
||||
"--preview-allow-public",
|
||||
"--preview-workspace-server-port",
|
||||
"8470",
|
||||
])
|
||||
.unwrap();
|
||||
assert!(args.preview.preview_enabled);
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert!(cfg.enabled);
|
||||
assert_eq!(cfg.port, Some(6014));
|
||||
assert_eq!(cfg.control_port, Some(6015));
|
||||
assert_eq!(cfg.visibility, Some(PreviewVisibility::Public));
|
||||
assert_eq!(cfg.instance_suffix.as_deref(), Some(".inst.example"));
|
||||
assert_eq!(
|
||||
cfg.auth_redirect.as_deref(),
|
||||
Some("https://grok.com/preview-auth")
|
||||
);
|
||||
assert!(cfg.allow_public);
|
||||
assert_eq!(cfg.workspace_server_port, Some(8470));
|
||||
assert_eq!(cfg.workspace_dir, PathBuf::from("/workspace"));
|
||||
assert_eq!(
|
||||
cfg.to_argv(),
|
||||
vec![
|
||||
"--preview-port",
|
||||
"6014",
|
||||
"--control-port",
|
||||
"6015",
|
||||
"--visibility",
|
||||
"public",
|
||||
"--instance-suffix",
|
||||
".inst.example",
|
||||
"--auth-redirect",
|
||||
"https://grok.com/preview-auth",
|
||||
"--allow-public",
|
||||
"--workspace-server-port",
|
||||
"8470",
|
||||
],
|
||||
);
|
||||
}
|
||||
#[test]
|
||||
fn preview_visibility_rejects_invalid_value() {
|
||||
let err = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-visibility",
|
||||
"nobody",
|
||||
])
|
||||
.err()
|
||||
.expect("an invalid --preview-visibility must be rejected");
|
||||
assert_eq!(err.kind(), clap::error::ErrorKind::InvalidValue);
|
||||
}
|
||||
#[test]
|
||||
fn preview_visibility_owner_parses_and_lowers() {
|
||||
let args = Args::try_parse_from([
|
||||
"kigi-workspace-server",
|
||||
"--preview-enabled",
|
||||
"--preview-visibility",
|
||||
"owner",
|
||||
])
|
||||
.unwrap();
|
||||
let cfg = args.preview.into_preview_args(PathBuf::from("/workspace"));
|
||||
assert_eq!(cfg.visibility, Some(PreviewVisibility::Owner));
|
||||
assert_eq!(cfg.to_argv(), vec!["--visibility", "owner"]);
|
||||
}
|
||||
}
|
||||
@@ -1,214 +0,0 @@
|
||||
//! Probe that verifies a running workspace-server actually serves tools
|
||||
//! over the server connection (not just that it reached READY).
|
||||
//!
|
||||
//! Connects to the server as a client harness, binds the workspace-server's
|
||||
//! session (its `server_id` equals the sandbox `session_id`), then
|
||||
//! invokes real tools and asserts the results:
|
||||
//! - `run_terminal_command` echoes a nonce that must come back,
|
||||
//! - `read_file` reads a file the first command wrote.
|
||||
//!
|
||||
//! Exits 0 and prints `PROBE_OK` on success; non-zero with a diagnostic
|
||||
//! on failure. Intended for sandbox end-to-end tests.
|
||||
//!
|
||||
//! The client connects to the *local* server (the same one the
|
||||
//! workspace-server reaches back to, e.g. `ws://localhost:10030/v1/tools`)
|
||||
//! using a bearer token. `servers.list` is scoped per-user on the server, so
|
||||
//! the bearer must resolve to the same user that owns the session — the
|
||||
//! access token from `~/.kigi/auth.json` does (same identity).
|
||||
|
||||
use base64::Engine;
|
||||
use clap::Parser;
|
||||
use kigi_computer_hub_sdk::pool::HubConnectionPool;
|
||||
use kigi_computer_hub_sdk::{AuthCredential, ToolHarnessBuilder};
|
||||
use kigi_tool_protocol::{SessionId, ToolId};
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolStreamItem, TypedToolOutput};
|
||||
use serde_json::{Value, json};
|
||||
use url::Url;
|
||||
use uuid::Uuid;
|
||||
|
||||
#[derive(Parser)]
|
||||
#[command(name = "workspace-server-probe")]
|
||||
#[command(about = "Invoke tools on a running workspace-server via the server connection")]
|
||||
struct Args {
|
||||
/// Server WebSocket URL the client connects to (the local server).
|
||||
#[arg(long, default_value = "ws://localhost:10030/v1/tools")]
|
||||
hub_url: String,
|
||||
|
||||
/// The workspace-server's server_id (equals the sandbox session_id).
|
||||
#[arg(long)]
|
||||
session_id: String,
|
||||
|
||||
/// Hub user to authenticate as. Must match the user the
|
||||
/// workspace-server registered under (`local-dev` in local-auth-dev).
|
||||
#[arg(long, default_value = "local-dev")]
|
||||
user_id: String,
|
||||
|
||||
/// Explicit bearer token (overrides --user-id). For real tokens.
|
||||
#[arg(long)]
|
||||
bearer: Option<String>,
|
||||
|
||||
/// Workspace directory to bind on the server.
|
||||
#[arg(long, default_value = "/workspace")]
|
||||
cwd: String,
|
||||
}
|
||||
|
||||
/// Mint an unsigned JWT carrying `{"sub": user_id}`. The local-auth-dev
|
||||
/// hub derives the principal's user from the bearer's JWT `sub` and does
|
||||
/// NOT verify the signature, so this is sufficient to authenticate as a
|
||||
/// specific local-dev user. Not usable against a real (verifying) hub.
|
||||
fn dev_bearer(user_id: &str) -> String {
|
||||
let b64 = base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
let header = b64.encode(br#"{"alg":"none","typ":"JWT"}"#);
|
||||
let payload = b64.encode(json!({ "sub": user_id }).to_string());
|
||||
format!("{header}.{payload}.sig")
|
||||
}
|
||||
|
||||
fn bearer(args: &Args) -> String {
|
||||
args.bearer
|
||||
.clone()
|
||||
.unwrap_or_else(|| dev_bearer(&args.user_id))
|
||||
}
|
||||
|
||||
/// Drive a tool call to its terminal item, discarding progress.
|
||||
async fn call_tool(
|
||||
harness: &kigi_computer_hub_sdk::ToolHarness,
|
||||
name: &str,
|
||||
args: Value,
|
||||
) -> anyhow::Result<Value> {
|
||||
let tool_id = ToolId::new(name).map_err(|e| anyhow::anyhow!("invalid tool id {name}: {e}"))?;
|
||||
let mut stream = harness
|
||||
.call(tool_id, args, ToolCallContext::default())
|
||||
.await;
|
||||
loop {
|
||||
let item = std::future::poll_fn(|cx| stream.as_mut().poll_next(cx)).await;
|
||||
match item {
|
||||
Some(ToolStreamItem::Progress(_)) => {}
|
||||
Some(ToolStreamItem::Terminal(Ok(typed))) => {
|
||||
let typed: TypedToolOutput = typed;
|
||||
return Ok(typed.value);
|
||||
}
|
||||
Some(ToolStreamItem::Terminal(Err(e))) => {
|
||||
anyhow::bail!("tool `{name}` returned error: {e}")
|
||||
}
|
||||
None => anyhow::bail!("tool `{name}` stream ended without a terminal item"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(
|
||||
tracing_subscriber::EnvFilter::try_from_default_env()
|
||||
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("warn")),
|
||||
)
|
||||
.init();
|
||||
let args = Args::parse();
|
||||
|
||||
let credential = AuthCredential::bearer(bearer(&args));
|
||||
|
||||
// The server connection can drop once right after connect (the SDK
|
||||
// reconnects with backoff). build()'s eager session_open doesn't
|
||||
// survive that first blip, so retry the connect + bind a few times.
|
||||
let mut last_err = None;
|
||||
for attempt in 1..=8u32 {
|
||||
match connect_and_bind(&args, &credential).await {
|
||||
Ok(harness) => return run_checks(&harness, &args).await,
|
||||
Err(e) => {
|
||||
eprintln!("[probe] attempt {attempt} failed: {e}");
|
||||
last_err = Some(e);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1500)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(last_err.unwrap_or_else(|| anyhow::anyhow!("probe failed")))
|
||||
}
|
||||
|
||||
/// Connect to the server, open a session, and bind the workspace-server.
|
||||
async fn connect_and_bind(
|
||||
args: &Args,
|
||||
credential: &AuthCredential,
|
||||
) -> anyhow::Result<kigi_computer_hub_sdk::ToolHarness> {
|
||||
let harness_session = SessionId::new(format!("probe-{}", Uuid::new_v4()))
|
||||
.map_err(|e| anyhow::anyhow!("invalid harness session id: {e}"))?;
|
||||
let url = Url::parse(&format!("{}?role=harness", args.hub_url))
|
||||
.map_err(|e| anyhow::anyhow!("invalid --hub-url: {e}"))?;
|
||||
|
||||
let harness = ToolHarnessBuilder::default()
|
||||
.pool(HubConnectionPool::new())
|
||||
.url(url)
|
||||
.auth(credential.clone())
|
||||
.session(harness_session)
|
||||
.build()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("failed to build server harness: {e}"))?;
|
||||
|
||||
let servers = harness
|
||||
.list_servers()
|
||||
.await
|
||||
.map_err(|e| anyhow::anyhow!("servers.list failed: {e}"))?;
|
||||
let server_ids: Vec<String> = servers
|
||||
.iter()
|
||||
.map(|s| s.server_id.as_str().to_owned())
|
||||
.collect();
|
||||
eprintln!("[probe] servers visible to this user: {server_ids:?}");
|
||||
|
||||
let server_id = args.session_id.as_str();
|
||||
// Strict servers (`--require-explicit-toolset`) fail metadata-less binds
|
||||
// closed: bind with exactly the tools the checks below invoke.
|
||||
let metadata = json!({
|
||||
"tools": [
|
||||
{"id": "GrokBuild:run_terminal_cmd", "name_override": "run_terminal_command"},
|
||||
{"id": "GrokBuild:read_file"},
|
||||
],
|
||||
});
|
||||
let tools = harness
|
||||
.session_bind(server_id, Some(&args.cwd), Some(metadata))
|
||||
.await
|
||||
.map_err(|e| {
|
||||
anyhow::anyhow!(
|
||||
"session_bind({server_id}) failed: {e} (is the workspace-server registered \
|
||||
for this user? visible servers: {server_ids:?})"
|
||||
)
|
||||
})?;
|
||||
eprintln!(
|
||||
"[probe] bound workspace-server {server_id}: {} tools available",
|
||||
tools.len()
|
||||
);
|
||||
Ok(harness)
|
||||
}
|
||||
|
||||
/// Invoke real tools on the bound workspace-server and assert results.
|
||||
async fn run_checks(
|
||||
harness: &kigi_computer_hub_sdk::ToolHarness,
|
||||
_args: &Args,
|
||||
) -> anyhow::Result<()> {
|
||||
// 1) run_terminal_command must echo our nonce back.
|
||||
let nonce = format!("probe-nonce-{}", Uuid::new_v4());
|
||||
let marker = format!("/tmp/{nonce}.txt");
|
||||
let out = call_tool(
|
||||
harness,
|
||||
"run_terminal_command",
|
||||
json!({ "command": format!("echo {nonce} | tee {marker}") }),
|
||||
)
|
||||
.await?;
|
||||
let out_text = serde_json::to_string(&out).unwrap_or_default();
|
||||
anyhow::ensure!(
|
||||
out_text.contains(&nonce),
|
||||
"run_terminal_command output did not contain the nonce; got: {out_text}"
|
||||
);
|
||||
eprintln!("[probe] run_terminal_command OK (nonce echoed)");
|
||||
|
||||
// 2) read_file must read the file the command wrote.
|
||||
let read_out = call_tool(harness, "read_file", json!({ "target_file": marker })).await?;
|
||||
let read_text = serde_json::to_string(&read_out).unwrap_or_default();
|
||||
anyhow::ensure!(
|
||||
read_text.contains(&nonce),
|
||||
"read_file did not return the written nonce; got: {read_text}"
|
||||
);
|
||||
eprintln!("[probe] read_file OK (read back written nonce)");
|
||||
|
||||
println!("PROBE_OK");
|
||||
Ok(())
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,561 +0,0 @@
|
||||
//! In-guest diagnostics HTTP server (`/ready`, `/statusz`, `/logs`) for the
|
||||
//! standalone workspace-server.
|
||||
//!
|
||||
//! The surface is reachable by any process inside the user's own sandbox
|
||||
//! (loopback-only TCP, or a 0600 Unix socket) and is never exposed through
|
||||
//! the sandbox port mapping. `/logs` returns the raw daemon log: treat its
|
||||
//! output as sensitive and keep the log stream free of secrets.
|
||||
|
||||
use std::io::{self, Read as _, Seek as _, SeekFrom};
|
||||
use std::net::Ipv4Addr;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::{env, fs, process};
|
||||
|
||||
use anyhow::anyhow;
|
||||
use axum::Router;
|
||||
use axum::extract::{Query, State};
|
||||
use axum::http::{StatusCode, header};
|
||||
use axum::response::{IntoResponse, Response};
|
||||
use axum::routing::get;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::net::TcpListener;
|
||||
#[cfg(unix)]
|
||||
use tokio::net::UnixListener;
|
||||
use tokio::task::JoinHandle;
|
||||
|
||||
/// Default Unix socket path (resolves host-visibly under the production
|
||||
/// launcher's `/tmp` bind mount, next to the ready/pid files).
|
||||
#[cfg(unix)]
|
||||
pub const DEFAULT_DIAG_SOCKET_PATH: &str = "/tmp/workspace-server.sock";
|
||||
|
||||
/// Default loopback TCP port for Windows guests.
|
||||
pub const DEFAULT_DIAG_PORT: u16 = 6016;
|
||||
|
||||
/// Grep-able daemon-log marker for a diagnostics bind failure.
|
||||
pub const DIAG_BIND_FAILED_MARKER: &str = "diagnostics server bind failed";
|
||||
|
||||
/// Process exit code for a fatal diagnostics bind failure in `--daemonize` mode.
|
||||
pub const EXIT_DIAG_BIND_FAILED: i32 = 5;
|
||||
|
||||
/// Default `/logs` tail size when `tail_bytes` is not given.
|
||||
pub const DEFAULT_LOG_TAIL_BYTES: u64 = 64 * 1024;
|
||||
|
||||
/// Hard cap on a `/logs` response; larger `tail_bytes` values are clamped.
|
||||
pub const MAX_LOG_TAIL_BYTES: u64 = 256 * 1024;
|
||||
|
||||
/// Hub connection state as reported on `/ready`.
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DiagState {
|
||||
Starting,
|
||||
Connected,
|
||||
Disconnected,
|
||||
}
|
||||
|
||||
/// Response body for `/ready`. The field set is a frozen contract with the
|
||||
/// sandbox readiness gate: never rename or remove fields; additions are
|
||||
/// backward-compatible.
|
||||
#[derive(Debug, Serialize)]
|
||||
struct ReadyBody {
|
||||
/// Serialized as an explicit `null` (never omitted) for nonce-less
|
||||
/// launches.
|
||||
launch_id: Option<String>,
|
||||
state: DiagState,
|
||||
pid: u32,
|
||||
connected_at: Option<u64>,
|
||||
state_changed_at: u64,
|
||||
version: &'static str,
|
||||
}
|
||||
|
||||
/// Response body for `/statusz`: the `/ready` fields plus debug extras.
|
||||
#[derive(Debug, Serialize)]
|
||||
struct StatuszBody {
|
||||
#[serde(flatten)]
|
||||
ready: ReadyBody,
|
||||
os: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct Inner {
|
||||
state: DiagState,
|
||||
connected_at: Option<u64>,
|
||||
state_changed_at: u64,
|
||||
shutting_down: bool,
|
||||
}
|
||||
|
||||
/// Cloneable handle publishing hub lifecycle transitions to the server.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct DiagHandle {
|
||||
launch_id: Option<String>,
|
||||
inner: Arc<Mutex<Inner>>,
|
||||
}
|
||||
|
||||
impl DiagHandle {
|
||||
/// `launch_id` is the caller-minted per-spawn nonce, echoed verbatim on
|
||||
/// `/ready` (`null` for nonce-less local launches).
|
||||
pub fn new(launch_id: Option<String>) -> Self {
|
||||
Self {
|
||||
launch_id,
|
||||
inner: Arc::new(Mutex::new(Inner {
|
||||
state: DiagState::Starting,
|
||||
connected_at: None,
|
||||
state_changed_at: now_ms(),
|
||||
shutting_down: false,
|
||||
})),
|
||||
}
|
||||
}
|
||||
|
||||
/// Initial hello completed, or a reconnect's serve replay settled.
|
||||
/// Ignored after [`Self::set_shutting_down`]: a reconnect that settles
|
||||
/// during the shutdown drain must not republish `connected`.
|
||||
pub fn set_connected(&self) {
|
||||
let mut inner = self.lock();
|
||||
if inner.shutting_down {
|
||||
return;
|
||||
}
|
||||
inner.state = DiagState::Connected;
|
||||
let now = now_ms();
|
||||
inner.connected_at.get_or_insert(now);
|
||||
inner.state_changed_at = now;
|
||||
}
|
||||
|
||||
/// Server socket dropped.
|
||||
pub fn set_disconnected(&self) {
|
||||
let mut inner = self.lock();
|
||||
inner.state = DiagState::Disconnected;
|
||||
inner.state_changed_at = now_ms();
|
||||
}
|
||||
|
||||
/// Latch `disconnected` for process shutdown: reported as `disconnected`
|
||||
/// on `/ready`, and later `set_connected` calls become no-ops.
|
||||
pub fn set_shutting_down(&self) {
|
||||
let mut inner = self.lock();
|
||||
inner.shutting_down = true;
|
||||
inner.state = DiagState::Disconnected;
|
||||
inner.state_changed_at = now_ms();
|
||||
}
|
||||
|
||||
/// True after [`Self::set_shutting_down`].
|
||||
pub fn is_shutting_down(&self) -> bool {
|
||||
self.lock().shutting_down
|
||||
}
|
||||
|
||||
fn lock(&self) -> MutexGuard<'_, Inner> {
|
||||
self.inner.lock().unwrap_or_else(PoisonError::into_inner)
|
||||
}
|
||||
|
||||
fn ready_body(&self) -> ReadyBody {
|
||||
let inner = self.lock();
|
||||
ReadyBody {
|
||||
launch_id: self.launch_id.clone(),
|
||||
state: inner.state,
|
||||
pid: process::id(),
|
||||
connected_at: inner.connected_at,
|
||||
state_changed_at: inner.state_changed_at,
|
||||
version: kigi_version::VERSION,
|
||||
}
|
||||
}
|
||||
|
||||
fn statusz_body(&self) -> StatuszBody {
|
||||
StatuszBody {
|
||||
ready: self.ready_body(),
|
||||
os: env::consts::OS,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn now_ms() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|d| d.as_millis() as u64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
/// Where the diagnostics server listens: a Unix socket on Linux, loopback TCP
|
||||
/// on Windows. Both variants compile everywhere so the TCP path is testable
|
||||
/// on Linux.
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum DiagListener {
|
||||
#[cfg(unix)]
|
||||
Unix(PathBuf),
|
||||
Tcp(u16),
|
||||
}
|
||||
|
||||
/// Shared request state: the lifecycle handle plus the daemon log path
|
||||
/// (`None` when logs go to a terminal instead of a file — `/logs` is 404).
|
||||
#[derive(Debug, Clone)]
|
||||
struct DiagContext {
|
||||
handle: DiagHandle,
|
||||
log_file: Option<Arc<PathBuf>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct LogsQuery {
|
||||
tail_bytes: Option<u64>,
|
||||
}
|
||||
|
||||
async fn logs(State(ctx): State<DiagContext>, Query(query): Query<LogsQuery>) -> Response {
|
||||
let Some(path) = ctx.log_file else {
|
||||
return StatusCode::NOT_FOUND.into_response();
|
||||
};
|
||||
let tail = query
|
||||
.tail_bytes
|
||||
.unwrap_or(DEFAULT_LOG_TAIL_BYTES)
|
||||
.min(MAX_LOG_TAIL_BYTES);
|
||||
match tokio::task::spawn_blocking(move || tail_file(&path, tail)).await {
|
||||
Ok(Ok(bytes)) => (
|
||||
[(header::CONTENT_TYPE, "text/plain; charset=utf-8")],
|
||||
String::from_utf8_lossy(&bytes).into_owned(),
|
||||
)
|
||||
.into_response(),
|
||||
Ok(Err(e)) if e.kind() == io::ErrorKind::NotFound => StatusCode::NOT_FOUND.into_response(),
|
||||
// Generic body: an io::Error would echo the log-file path to clients.
|
||||
Ok(Err(e)) => {
|
||||
tracing::warn!(error = %e, "failed to read log tail");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response()
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, "log tail task failed");
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "failed to read log").into_response()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Read at most the last `max` bytes of `path`.
|
||||
fn tail_file(path: &Path, max: u64) -> io::Result<Vec<u8>> {
|
||||
let mut file = fs::File::open(path)?;
|
||||
let len = file.metadata()?.len();
|
||||
file.seek(SeekFrom::Start(len.saturating_sub(max)))?;
|
||||
let mut buf = Vec::new();
|
||||
// `take` bounds the read even if the file grows underneath us.
|
||||
file.take(max).read_to_end(&mut buf)?;
|
||||
Ok(buf)
|
||||
}
|
||||
|
||||
fn router(ctx: DiagContext) -> Router {
|
||||
Router::new()
|
||||
.route(
|
||||
"/ready",
|
||||
get(|State(ctx): State<DiagContext>| async move {
|
||||
let body = ctx.handle.ready_body();
|
||||
// Non-2xx for "not ready" so naive HTTP probes agree with
|
||||
// consumers that parse `state`. The body is served either way.
|
||||
let status = if body.state == DiagState::Connected {
|
||||
StatusCode::OK
|
||||
} else {
|
||||
StatusCode::SERVICE_UNAVAILABLE
|
||||
};
|
||||
(status, axum::Json(body))
|
||||
}),
|
||||
)
|
||||
.route(
|
||||
"/statusz",
|
||||
get(|State(ctx): State<DiagContext>| async move { axum::Json(ctx.handle.statusz_body()) }),
|
||||
)
|
||||
.route("/logs", get(logs))
|
||||
.with_state(ctx)
|
||||
}
|
||||
|
||||
/// Bind the listener and spawn the server task. Binding happens before this
|
||||
/// returns, so a bind failure surfaces synchronously. `log_file` is the
|
||||
/// daemon log served by `/logs` (`None` ⇒ `/logs` is 404).
|
||||
pub async fn serve(
|
||||
listener: DiagListener,
|
||||
handle: DiagHandle,
|
||||
log_file: Option<PathBuf>,
|
||||
) -> anyhow::Result<BoundDiag> {
|
||||
let ctx = DiagContext {
|
||||
handle,
|
||||
log_file: log_file.map(Arc::new),
|
||||
};
|
||||
match listener {
|
||||
#[cfg(unix)]
|
||||
DiagListener::Unix(path) => {
|
||||
let _ = fs::remove_file(&path);
|
||||
let listener =
|
||||
UnixListener::bind(&path).map_err(|e| anyhow!("bind {}: {e}", path.display()))?;
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
if let Err(e) = fs::set_permissions(&path, fs::Permissions::from_mode(0o600)) {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"failed to restrict diagnostics socket permissions"
|
||||
);
|
||||
}
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(e) = axum::serve(listener, router(ctx)).await {
|
||||
tracing::warn!(error = %e, "diagnostics server exited");
|
||||
}
|
||||
});
|
||||
Ok(BoundDiag {
|
||||
addr: format!("unix:{}", path.display()),
|
||||
port: None,
|
||||
task,
|
||||
})
|
||||
}
|
||||
DiagListener::Tcp(port) => {
|
||||
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, port))
|
||||
.await
|
||||
.map_err(|e| anyhow!("bind 127.0.0.1:{port}: {e}"))?;
|
||||
let local = listener.local_addr()?;
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(e) = axum::serve(listener, router(ctx)).await {
|
||||
tracing::warn!(error = %e, "diagnostics server exited");
|
||||
}
|
||||
});
|
||||
Ok(BoundDiag {
|
||||
addr: format!("http://{local}"),
|
||||
port: Some(local.port()),
|
||||
task,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// A successfully bound diagnostics server.
|
||||
#[derive(Debug)]
|
||||
pub struct BoundDiag {
|
||||
/// Human-readable bound address for the startup log line.
|
||||
pub addr: String,
|
||||
/// Bound TCP port (`None` for Unix sockets).
|
||||
pub port: Option<u16>,
|
||||
/// The serve task; held by the production launcher for the process lifetime.
|
||||
pub task: JoinHandle<()>,
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use serde_json::Value;
|
||||
|
||||
use super::*;
|
||||
|
||||
async fn get_json(port: u16, path: &str) -> (u16, Value) {
|
||||
let response = reqwest::get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.await
|
||||
.expect("request");
|
||||
let status = response.status().as_u16();
|
||||
let body = response.text().await.expect("body");
|
||||
(status, serde_json::from_str(&body).expect("json body"))
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn ready_response_contract_is_frozen() {
|
||||
let handle = DiagHandle::new(Some("nonce-1".to_owned()));
|
||||
let bound = serve(DiagListener::Tcp(0), handle, None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let (status, body) = get_json(bound.port.expect("tcp port"), "/ready").await;
|
||||
|
||||
assert_eq!(status, 503, "not yet connected must not probe as ready");
|
||||
let obj = body.as_object().expect("object");
|
||||
for key in [
|
||||
"launch_id",
|
||||
"state",
|
||||
"pid",
|
||||
"connected_at",
|
||||
"state_changed_at",
|
||||
"version",
|
||||
] {
|
||||
assert!(obj.contains_key(key), "missing frozen key {key}");
|
||||
}
|
||||
assert_eq!(body["launch_id"], "nonce-1");
|
||||
assert_eq!(body["state"], "starting");
|
||||
assert_eq!(body["connected_at"], Value::Null);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn state_follows_hub_lifecycle_and_freezes_connected_at() {
|
||||
let handle = DiagHandle::new(None);
|
||||
let bound = serve(DiagListener::Tcp(0), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let port = bound.port.expect("tcp port");
|
||||
|
||||
handle.set_connected();
|
||||
let (status, connected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(connected["state"], "connected");
|
||||
assert!(connected["connected_at"].is_u64());
|
||||
assert_eq!(connected["launch_id"], Value::Null);
|
||||
|
||||
handle.set_disconnected();
|
||||
let (status, disconnected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 503);
|
||||
assert_eq!(disconnected["state"], "disconnected");
|
||||
assert_eq!(
|
||||
disconnected["connected_at"], connected["connected_at"],
|
||||
"connected_at is frozen at first connect and echoed on disconnect"
|
||||
);
|
||||
|
||||
handle.set_connected();
|
||||
let (status, reconnected) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(reconnected["state"], "connected");
|
||||
assert_eq!(
|
||||
reconnected["connected_at"], connected["connected_at"],
|
||||
"reconnect must not re-mint connected_at"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn shutting_down_latches_disconnected_across_reconnects() {
|
||||
let handle = DiagHandle::new(None);
|
||||
let bound = serve(DiagListener::Tcp(0), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind");
|
||||
let port = bound.port.expect("tcp port");
|
||||
|
||||
handle.set_connected();
|
||||
handle.set_shutting_down();
|
||||
// A reconnect settling during the shutdown drain must not republish
|
||||
// `connected`.
|
||||
handle.set_connected();
|
||||
|
||||
let (status, body) = get_json(port, "/ready").await;
|
||||
assert_eq!(status, 503);
|
||||
assert_eq!(body["state"], "disconnected");
|
||||
assert!(handle.is_shutting_down());
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn unix_socket_serves_ready_and_rebinds_over_stale_socket() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let sock = dir.path().join("ws.sock");
|
||||
|
||||
drop(std::os::unix::net::UnixListener::bind(&sock).expect("stale bind"));
|
||||
|
||||
let handle = DiagHandle::new(Some("nonce-uds".to_owned()));
|
||||
let _bound = serve(DiagListener::Unix(sock.clone()), handle.clone(), None)
|
||||
.await
|
||||
.expect("bind over stale socket");
|
||||
handle.set_connected();
|
||||
|
||||
use std::os::unix::fs::PermissionsExt as _;
|
||||
let mode = fs::metadata(&sock)
|
||||
.expect("socket meta")
|
||||
.permissions()
|
||||
.mode();
|
||||
assert_eq!(mode & 0o777, 0o600, "socket must be owner-only");
|
||||
|
||||
use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _};
|
||||
let mut stream = tokio::net::UnixStream::connect(&sock)
|
||||
.await
|
||||
.expect("connect");
|
||||
stream
|
||||
.write_all(b"GET /ready HTTP/1.1\r\nHost: ws\r\nConnection: close\r\n\r\n")
|
||||
.await
|
||||
.expect("write request");
|
||||
let mut response = Vec::new();
|
||||
stream
|
||||
.read_to_end(&mut response)
|
||||
.await
|
||||
.expect("read response");
|
||||
let response = String::from_utf8_lossy(&response);
|
||||
assert!(response.starts_with("HTTP/1.1 200"), "got: {response}");
|
||||
let body = response.split("\r\n\r\n").nth(1).expect("body");
|
||||
let json_start = body.find('{').expect("json start");
|
||||
let json_end = body.rfind('}').expect("json end");
|
||||
let parsed: Value = serde_json::from_str(&body[json_start..=json_end]).expect("json");
|
||||
assert_eq!(parsed["launch_id"], "nonce-uds");
|
||||
assert_eq!(parsed["state"], "connected");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tcp_bind_conflict_surfaces_as_error() {
|
||||
let first = serve(DiagListener::Tcp(0), DiagHandle::new(None), None)
|
||||
.await
|
||||
.expect("first bind");
|
||||
let port = first.port.expect("tcp port");
|
||||
let err = serve(DiagListener::Tcp(port), DiagHandle::new(None), None).await;
|
||||
assert!(err.is_err(), "second bind on the same port must fail");
|
||||
}
|
||||
|
||||
async fn get_text(port: u16, path: &str) -> (u16, Option<String>, String) {
|
||||
let response = reqwest::get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.await
|
||||
.expect("request");
|
||||
let status = response.status().as_u16();
|
||||
let content_type = response
|
||||
.headers()
|
||||
.get(reqwest::header::CONTENT_TYPE)
|
||||
.map(|v| v.to_str().expect("content-type").to_owned());
|
||||
(status, content_type, response.text().await.expect("body"))
|
||||
}
|
||||
|
||||
async fn serve_with_log(log_file: Option<PathBuf>) -> u16 {
|
||||
let bound = serve(DiagListener::Tcp(0), DiagHandle::new(None), log_file)
|
||||
.await
|
||||
.expect("bind");
|
||||
bound.port.expect("tcp port")
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_tails_requested_bytes_as_plain_text() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
fs::write(&log, "0123456789").expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, content_type, body) = get_text(port, "/logs?tail_bytes=4").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(content_type.as_deref(), Some("text/plain; charset=utf-8"));
|
||||
assert_eq!(body, "6789");
|
||||
|
||||
// A tail larger than the file returns the whole file.
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=1000").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body, "0123456789");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_default_and_hard_cap_bound_the_response() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
// Larger than the hard cap; ends with a marker to prove we got the tail.
|
||||
let mut content = vec![b'a'; (MAX_LOG_TAIL_BYTES + 4096) as usize];
|
||||
content.extend_from_slice(b"END-MARKER");
|
||||
fs::write(&log, &content).expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body.len() as u64, DEFAULT_LOG_TAIL_BYTES);
|
||||
assert!(body.ends_with("END-MARKER"));
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=999999999").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body.len() as u64, MAX_LOG_TAIL_BYTES, "hard cap applies");
|
||||
assert!(body.ends_with("END-MARKER"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_tail_is_lossy_utf8() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let log = dir.path().join("ws.log");
|
||||
// 'é' is 0xC3 0xA9; a 1-byte tail cuts the sequence mid-char.
|
||||
fs::write(&log, "aé").expect("write log");
|
||||
let port = serve_with_log(Some(log)).await;
|
||||
|
||||
let (status, _, body) = get_text(port, "/logs?tail_bytes=1").await;
|
||||
assert_eq!(status, 200);
|
||||
assert_eq!(body, "\u{FFFD}", "a torn UTF-8 boundary must be lossy");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_without_log_file_is_404() {
|
||||
let port = serve_with_log(None).await;
|
||||
let (status, _, _) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 404);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn logs_missing_log_file_is_404() {
|
||||
let dir = tempfile::tempdir().expect("tempdir");
|
||||
let port = serve_with_log(Some(dir.path().join("never-created.log"))).await;
|
||||
let (status, _, _) = get_text(port, "/logs").await;
|
||||
assert_eq!(status, 404);
|
||||
}
|
||||
}
|
||||
@@ -1,775 +0,0 @@
|
||||
//! Read-only filesystem helpers backing the client-facing
|
||||
//! `workspace.client_fs_*` RPCs (the grok.com conversation-files UI,
|
||||
//! tunneled through the server).
|
||||
//!
|
||||
//! Deliberately separate from the shell-facing ext ops in
|
||||
//! [`ext_fs`](super::ext_fs): every path here is workspace-root-relative
|
||||
//! and resolves through the root-confinement helper
|
||||
//! (`WorkspaceHandle::resolve_service_path`), the list walk excludes
|
||||
//! symlinks that resolve outside the root (and never descends into
|
||||
//! them), listings paginate with stable post-sort slices, and reads are
|
||||
//! binary-safe (base64 chunks).
|
||||
//!
|
||||
//! Wire types live in `kigi_workspace_types::rpc::fs` (the
|
||||
//! `ClientFs*` types), shared with the backend caller.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use kigi_workspace_types::rpc::fs::{
|
||||
ClientFsListNode as FsListNode, ClientFsListReq as FsListReq, ClientFsListRes as FsListRes,
|
||||
ClientFsReadFileReq as FsReadFileReq, ClientFsReadFileRes as FsReadFileRes,
|
||||
ClientFsStatReq as FsStatReq, ClientFsStatRes as FsStatRes, FsContentType, FsNodeType,
|
||||
};
|
||||
|
||||
use crate::error::{WorkspaceError, WorkspaceResult};
|
||||
use crate::handle::WorkspaceHandle;
|
||||
|
||||
/// Hard cap on entries collected per list call before sorting (shared
|
||||
/// across all fs surfaces; see [`super::walk::MAX_LIST_COLLECT`]).
|
||||
const MAX_LIST_COLLECT: usize = super::walk::MAX_LIST_COLLECT;
|
||||
|
||||
/// Server-side cap on `FsListReq::limit`.
|
||||
const MAX_LIST_LIMIT: u32 = 1000;
|
||||
|
||||
/// Server-side cap on a single read's effective byte budget (shared
|
||||
/// across all fs surfaces; see [`super::walk::MAX_READ_BYTES`]). Only
|
||||
/// referenced by tests now that the clamp lives in `walk::clamp_read_length`.
|
||||
#[cfg(test)]
|
||||
const MAX_READ_BYTES: u64 = super::walk::MAX_READ_BYTES;
|
||||
|
||||
/// Bound on memoized hashes; the memo is cleared (not LRU-evicted) when
|
||||
/// full — entries simply re-hash on next use.
|
||||
const HASH_MEMO_CAPACITY: usize = 4096;
|
||||
|
||||
// =========================================================================
|
||||
// (path, size, mtime_ms) → hash memo
|
||||
// =========================================================================
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MemoEntry {
|
||||
size: u64,
|
||||
mtime_ms: i64,
|
||||
hash: String,
|
||||
}
|
||||
|
||||
/// Memo of full-content SHA-256 digests keyed by absolute path and
|
||||
/// validated against `(size, mtime_ms)`, so unchanged files hash once
|
||||
/// instead of on every `client_fs_stat`. The memo only avoids redundant
|
||||
/// hashing — it never substitutes mtime for content addressing: a
|
||||
/// `(size, mtime_ms)` mismatch is a miss and the caller re-hashes.
|
||||
#[derive(Debug, Default)]
|
||||
pub(crate) struct FileHashMemo {
|
||||
entries: parking_lot::Mutex<HashMap<PathBuf, MemoEntry>>,
|
||||
}
|
||||
|
||||
impl FileHashMemo {
|
||||
/// Return the memoized hash when `(size, mtime_ms)` still match.
|
||||
pub(crate) fn lookup(&self, path: &Path, size: u64, mtime_ms: i64) -> Option<String> {
|
||||
let entries = self.entries.lock();
|
||||
let entry = entries.get(path)?;
|
||||
(entry.size == size && entry.mtime_ms == mtime_ms).then(|| entry.hash.clone())
|
||||
}
|
||||
|
||||
/// Record a freshly computed hash, replacing any stale entry for the
|
||||
/// same path. Clears the whole memo when inserting a new path would
|
||||
/// exceed [`HASH_MEMO_CAPACITY`].
|
||||
pub(crate) fn store(&self, path: &Path, size: u64, mtime_ms: i64, hash: String) {
|
||||
let mut entries = self.entries.lock();
|
||||
if !entries.contains_key(path) && entries.len() >= HASH_MEMO_CAPACITY {
|
||||
entries.clear();
|
||||
}
|
||||
entries.insert(
|
||||
path.to_path_buf(),
|
||||
MemoEntry {
|
||||
size,
|
||||
mtime_ms,
|
||||
hash,
|
||||
},
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// Path resolution
|
||||
// =========================================================================
|
||||
|
||||
/// Resolve a root-relative request path through the workspace's
|
||||
/// root-confinement helper, returning the resolved path together with the
|
||||
/// canonical root it was checked against. `""` and `"."` mean the
|
||||
/// workspace root; absolute paths, `..` escapes, and symlink escapes are
|
||||
/// rejected there.
|
||||
async fn resolve_with_root(
|
||||
ws: &WorkspaceHandle,
|
||||
path: &str,
|
||||
) -> WorkspaceResult<(PathBuf, PathBuf)> {
|
||||
let rel = if path.is_empty() { "." } else { path };
|
||||
let canonical_root = ws.canonical_root().await?;
|
||||
let abs = ws.resolve_service_path(rel, &canonical_root).await?;
|
||||
Ok((abs, canonical_root))
|
||||
}
|
||||
|
||||
/// [`resolve_with_root`] for callers that don't need the canonical root.
|
||||
async fn resolve(ws: &WorkspaceHandle, path: &str) -> WorkspaceResult<PathBuf> {
|
||||
resolve_with_root(ws, path).await.map(|(abs, _)| abs)
|
||||
}
|
||||
|
||||
fn system_time_ms(st: std::time::SystemTime) -> i64 {
|
||||
match st.duration_since(std::time::UNIX_EPOCH) {
|
||||
Ok(d) => i64::try_from(d.as_millis()).unwrap_or(i64::MAX),
|
||||
Err(e) => -i64::try_from(e.duration().as_millis()).unwrap_or(i64::MAX),
|
||||
}
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// list
|
||||
// =========================================================================
|
||||
|
||||
/// List `req.path` (workspace-root-relative) with stable pagination:
|
||||
/// collect the full walk (bounded by [`MAX_LIST_COLLECT`]), sort
|
||||
/// directories-first / case-insensitive by name, then slice
|
||||
/// `[offset, offset + limit)`. Symlinks resolving outside the workspace
|
||||
/// root are excluded from the walk (and never descended into).
|
||||
pub(crate) async fn list(ws: &WorkspaceHandle, req: &FsListReq) -> WorkspaceResult<FsListRes> {
|
||||
let (abs, canonical_root) = resolve_with_root(ws, &req.path).await?;
|
||||
let root = ws.root_cwd()?;
|
||||
let req = req.clone();
|
||||
// The walk does synchronous traversal + metadata syscalls; run it off
|
||||
// the async executor (matching the ext_fs ops).
|
||||
tokio::task::spawn_blocking(move || {
|
||||
list_blocking(&abs, &root, &canonical_root, &req, MAX_LIST_COLLECT)
|
||||
})
|
||||
.await
|
||||
.map_err(|e| WorkspaceError::JoinError(e.to_string()))?
|
||||
}
|
||||
|
||||
fn list_blocking(
|
||||
abs_dir: &Path,
|
||||
root: &Path,
|
||||
canonical_root: &Path,
|
||||
req: &FsListReq,
|
||||
max_collect: usize,
|
||||
) -> WorkspaceResult<FsListRes> {
|
||||
// Root confinement also holds mid-walk: a symlink inside the root
|
||||
// pointing outside must not enumerate outside metadata.
|
||||
let page = super::walk::list_directory_paged(
|
||||
abs_dir,
|
||||
super::walk::ListOptions {
|
||||
depth: req.depth as usize,
|
||||
follow_symlinks: req.follow_symlinks,
|
||||
respect_git_ignore: req.respect_git_ignore,
|
||||
include_hidden: req.include_hidden,
|
||||
include_globs: &req.include_globs,
|
||||
exclude_globs: &req.exclude_globs,
|
||||
offset: req.offset,
|
||||
limit: req.limit.min(MAX_LIST_LIMIT) as usize,
|
||||
confine_to_canonical_root: Some(canonical_root.to_path_buf()),
|
||||
},
|
||||
max_collect,
|
||||
);
|
||||
|
||||
let nodes: Vec<FsListNode> = page
|
||||
.entries
|
||||
.into_iter()
|
||||
.map(|e| FsListNode {
|
||||
node_type: if e.is_dir {
|
||||
FsNodeType::Directory
|
||||
} else {
|
||||
FsNodeType::File
|
||||
},
|
||||
size: e.size,
|
||||
mtime_ms: e.modified.map(system_time_ms),
|
||||
is_symlink: e.is_symlink.then_some(true),
|
||||
// Root-relative path (divergent from the shell's absolute path).
|
||||
// A walk under a symlinked root yields canonical-root-spelled
|
||||
// entries, so strip either spelling.
|
||||
path: e
|
||||
.abs_path
|
||||
.strip_prefix(root)
|
||||
.or_else(|_| e.abs_path.strip_prefix(canonical_root))
|
||||
.unwrap_or(&e.abs_path)
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
name: e.name,
|
||||
})
|
||||
.collect();
|
||||
|
||||
Ok(FsListRes {
|
||||
nodes,
|
||||
truncated: page.truncated,
|
||||
})
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// stat
|
||||
// =========================================================================
|
||||
|
||||
/// Stat `req.path`: existence, kind, size, mtime, and — for files — a
|
||||
/// full-content SHA-256 served through the workspace hash memo.
|
||||
pub(crate) async fn stat(ws: &WorkspaceHandle, req: &FsStatReq) -> WorkspaceResult<FsStatRes> {
|
||||
let abs = resolve(ws, &req.path).await?;
|
||||
let md = match tokio::fs::metadata(&abs).await {
|
||||
Ok(md) => md,
|
||||
// NotADirectory: a *file* sits mid-path (e.g. `a.txt/sub`) — for an
|
||||
// existence probe that is a miss, not an RPC error.
|
||||
Err(e)
|
||||
if e.kind() == std::io::ErrorKind::NotFound
|
||||
|| e.kind() == std::io::ErrorKind::NotADirectory =>
|
||||
{
|
||||
return Ok(FsStatRes {
|
||||
exists: false,
|
||||
node_type: None,
|
||||
size: None,
|
||||
mtime_ms: None,
|
||||
hash: None,
|
||||
});
|
||||
}
|
||||
Err(e) => {
|
||||
return Err(WorkspaceError::HubError(format!(
|
||||
"stat failed for {}: {e}",
|
||||
req.path
|
||||
)));
|
||||
}
|
||||
};
|
||||
let mtime_ms = md.modified().ok().map(system_time_ms);
|
||||
if md.is_dir() {
|
||||
return Ok(FsStatRes {
|
||||
exists: true,
|
||||
node_type: Some(FsNodeType::Directory),
|
||||
size: None,
|
||||
mtime_ms,
|
||||
hash: None,
|
||||
});
|
||||
}
|
||||
let size = md.len();
|
||||
let memo = &ws.shared.client_fs_hash_memo;
|
||||
let hash = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) {
|
||||
Some(hash) => hash,
|
||||
None => {
|
||||
let (hash, _, _) = crate::handle::stream_hash_and_range(&abs, 0, 0)
|
||||
.await
|
||||
.map_err(|e| {
|
||||
WorkspaceError::HubError(format!("hash failed for {}: {e}", req.path))
|
||||
})?;
|
||||
if let Some(m) = mtime_ms {
|
||||
memo.store(&abs, size, m, hash.clone());
|
||||
}
|
||||
hash
|
||||
}
|
||||
};
|
||||
Ok(FsStatRes {
|
||||
exists: true,
|
||||
node_type: Some(FsNodeType::File),
|
||||
size: Some(size),
|
||||
mtime_ms,
|
||||
hash: Some(hash),
|
||||
})
|
||||
}
|
||||
|
||||
// =========================================================================
|
||||
// read_file
|
||||
// =========================================================================
|
||||
|
||||
/// Read a byte range of `req.path` (binary-safe, capped at
|
||||
/// `min(req.max_bytes, MAX_READ_BYTES)`) together with the full-file
|
||||
/// SHA-256. When the hash is memoized for the current `(size, mtime)`
|
||||
/// only the requested range is read; otherwise the whole file streams
|
||||
/// once (via the shared [`crate::handle::stream_hash_and_range`]) to
|
||||
/// hash it.
|
||||
pub(crate) async fn read_file(
|
||||
ws: &WorkspaceHandle,
|
||||
req: &FsReadFileReq,
|
||||
) -> WorkspaceResult<FsReadFileRes> {
|
||||
let abs = resolve(ws, &req.path).await?;
|
||||
let read_err =
|
||||
|e: std::io::Error| WorkspaceError::HubError(format!("read failed for {}: {e}", req.path));
|
||||
let md = tokio::fs::metadata(&abs).await.map_err(read_err)?;
|
||||
if md.is_dir() {
|
||||
return Err(WorkspaceError::HubError(format!(
|
||||
"not a file: {}",
|
||||
req.path
|
||||
)));
|
||||
}
|
||||
let size = md.len();
|
||||
let mtime_ms = md.modified().ok().map(system_time_ms);
|
||||
let offset = req.offset.unwrap_or(0);
|
||||
// Server-side clamp: a hostile/buggy caller cannot lift the per-chunk
|
||||
// budget past MAX_READ_BYTES regardless of `maxBytes`.
|
||||
let length = super::walk::clamp_read_length(req.length, req.max_bytes);
|
||||
|
||||
let memo = &ws.shared.client_fs_hash_memo;
|
||||
let (hash, chunk, size) = match mtime_ms.and_then(|m| memo.lookup(&abs, size, m)) {
|
||||
Some(hash) => {
|
||||
let chunk = super::walk::read_range(&abs, offset, length)
|
||||
.await
|
||||
.map_err(read_err)?;
|
||||
(hash, chunk, size)
|
||||
}
|
||||
None => {
|
||||
let (hash, chunk, streamed) =
|
||||
crate::handle::stream_hash_and_range(&abs, offset, length)
|
||||
.await
|
||||
.map_err(read_err)?;
|
||||
if let Some(m) = mtime_ms {
|
||||
memo.store(&abs, streamed, m, hash.clone());
|
||||
}
|
||||
(hash, chunk, streamed)
|
||||
}
|
||||
};
|
||||
|
||||
// Shared encoder keeps the paired wire fields coherent with one UTF-8
|
||||
// validation pass; `type` is text iff the bytes were valid UTF-8.
|
||||
let (payload, is_text) = super::walk::encode_chunk(chunk, req.encoding);
|
||||
let (content, content_base64) = match payload {
|
||||
super::walk::ChunkPayload::Text(t) => (Some(t), None),
|
||||
super::walk::ChunkPayload::Base64(b) => (None, Some(b)),
|
||||
};
|
||||
let content_type = if is_text {
|
||||
FsContentType::Text
|
||||
} else {
|
||||
FsContentType::Binary
|
||||
};
|
||||
Ok(FsReadFileRes {
|
||||
content,
|
||||
content_base64,
|
||||
size,
|
||||
hash,
|
||||
content_type,
|
||||
})
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use base64::Engine;
|
||||
use kigi_workspace_types::rpc::fs::FsReadEncoding;
|
||||
|
||||
use super::*;
|
||||
use crate::handle::tests::make_handle;
|
||||
|
||||
fn list_req(path: &str) -> FsListReq {
|
||||
FsListReq {
|
||||
path: path.to_owned(),
|
||||
depth: 1,
|
||||
include_hidden: true,
|
||||
limit: 1000,
|
||||
offset: 0,
|
||||
follow_symlinks: true,
|
||||
respect_git_ignore: false,
|
||||
include_globs: vec![],
|
||||
exclude_globs: vec![],
|
||||
}
|
||||
}
|
||||
|
||||
/// Fixture: root with files `b.txt`, `A.txt`, `c.txt` and dirs
|
||||
/// `Zeta`, `alpha`. Expected order: dirs first case-insensitive
|
||||
/// (`alpha`, `Zeta`), then files (`A.txt`, `b.txt`, `c.txt`).
|
||||
fn populate(root: &Path) {
|
||||
std::fs::write(root.join("b.txt"), b"bb").unwrap();
|
||||
std::fs::write(root.join("A.txt"), b"a").unwrap();
|
||||
std::fs::write(root.join("c.txt"), b"ccc").unwrap();
|
||||
std::fs::create_dir(root.join("Zeta")).unwrap();
|
||||
std::fs::create_dir(root.join("alpha")).unwrap();
|
||||
}
|
||||
|
||||
/// `list_blocking` against `dir` as both walk root and workspace root
|
||||
/// (canonicalized for the confinement check, like production).
|
||||
fn list_dir(dir: &Path, req: &FsListReq, max_collect: usize) -> FsListRes {
|
||||
let canonical = dunce::canonicalize(dir).unwrap();
|
||||
list_blocking(dir, dir, &canonical, req, max_collect).unwrap()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_sorts_dirs_first_case_insensitive() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let res = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT);
|
||||
let names: Vec<&str> = res.nodes.iter().map(|n| n.name.as_str()).collect();
|
||||
assert_eq!(names, ["alpha", "Zeta", "A.txt", "b.txt", "c.txt"]);
|
||||
assert!(!res.truncated);
|
||||
assert_eq!(res.nodes[0].node_type, FsNodeType::Directory);
|
||||
assert_eq!(res.nodes[2].node_type, FsNodeType::File);
|
||||
assert_eq!(res.nodes[2].size, Some(1));
|
||||
assert!(res.nodes[2].mtime_ms.is_some());
|
||||
// Paths are workspace-root-relative.
|
||||
assert_eq!(res.nodes[2].path, "A.txt");
|
||||
}
|
||||
|
||||
/// Pagination slices the *sorted* listing, so consecutive pages have
|
||||
/// stable boundaries and concatenate to the full listing.
|
||||
#[test]
|
||||
fn list_paginates_post_sort_with_stable_boundaries() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let full = list_dir(dir.path(), &list_req(""), MAX_LIST_COLLECT);
|
||||
|
||||
let mut paged = Vec::new();
|
||||
for page_start in [0u64, 2, 4] {
|
||||
let req = FsListReq {
|
||||
limit: 2,
|
||||
offset: page_start,
|
||||
..list_req("")
|
||||
};
|
||||
let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
// truncated while more entries remain past this page.
|
||||
assert_eq!(page.truncated, page_start + 2 < full.nodes.len() as u64);
|
||||
paged.extend(page.nodes);
|
||||
}
|
||||
assert_eq!(paged, full.nodes);
|
||||
|
||||
// Offset past the end yields an empty, non-truncated page.
|
||||
let req = FsListReq {
|
||||
offset: 100,
|
||||
..list_req("")
|
||||
};
|
||||
let page = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
assert!(page.nodes.is_empty());
|
||||
assert!(!page.truncated);
|
||||
}
|
||||
|
||||
/// The collection cap marks the result truncated even when the page
|
||||
/// itself is not full.
|
||||
#[test]
|
||||
fn list_collection_cap_truncates() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let res = list_dir(dir.path(), &list_req(""), 2);
|
||||
assert_eq!(res.nodes.len(), 2);
|
||||
assert!(res.truncated);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn list_caps_limit_at_server_max() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
populate(dir.path());
|
||||
let req = FsListReq {
|
||||
limit: u32::MAX,
|
||||
..list_req("")
|
||||
};
|
||||
// Must not panic / overflow; the page is everything (< 1000).
|
||||
let res = list_dir(dir.path(), &req, MAX_LIST_COLLECT);
|
||||
assert_eq!(res.nodes.len(), 5);
|
||||
}
|
||||
|
||||
/// Regression: the list walk must not traverse — or
|
||||
/// even surface — in-root symlinks that resolve outside the workspace
|
||||
/// root, while symlinks staying inside the root keep working.
|
||||
#[test]
|
||||
#[cfg(unix)]
|
||||
fn list_excludes_symlink_escapes_mid_walk() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let root = dir.path();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap();
|
||||
// Escaping symlink: root/escape_link -> <outside>.
|
||||
std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
// In-root symlink: root/good_link -> root/real_dir.
|
||||
std::fs::create_dir(root.join("real_dir")).unwrap();
|
||||
std::fs::write(root.join("real_dir/inner.txt"), b"inner").unwrap();
|
||||
std::os::unix::fs::symlink(root.join("real_dir"), root.join("good_link")).unwrap();
|
||||
|
||||
let req = FsListReq {
|
||||
depth: 2,
|
||||
follow_symlinks: true,
|
||||
..list_req("")
|
||||
};
|
||||
let res = list_dir(root, &req, MAX_LIST_COLLECT);
|
||||
let paths: Vec<&str> = res.nodes.iter().map(|n| n.path.as_str()).collect();
|
||||
assert!(
|
||||
!paths.iter().any(|p| p.contains("escape_link")),
|
||||
"escaping symlink (and its subtree) must be excluded: {paths:?}"
|
||||
);
|
||||
assert!(
|
||||
!paths.iter().any(|p| p.contains("secret.txt")),
|
||||
"outside entries must not be enumerated: {paths:?}"
|
||||
);
|
||||
// Confinement must not over-filter: in-root symlinks survive,
|
||||
// including descent through them.
|
||||
assert!(paths.contains(&"good_link"), "{paths:?}");
|
||||
assert!(paths.contains(&"good_link/inner.txt"), "{paths:?}");
|
||||
assert!(paths.contains(&"real_dir/inner.txt"), "{paths:?}");
|
||||
let good = res.nodes.iter().find(|n| n.path == "good_link").unwrap();
|
||||
assert_eq!(good.is_symlink, Some(true));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn memo_lookup_hits_and_invalidates_on_mismatch() {
|
||||
let memo = FileHashMemo::default();
|
||||
let path = Path::new("/ws/a.txt");
|
||||
memo.store(path, 10, 1000, "h1".into());
|
||||
assert_eq!(memo.lookup(path, 10, 1000).as_deref(), Some("h1"));
|
||||
// Size change ⇒ miss.
|
||||
assert_eq!(memo.lookup(path, 11, 1000), None);
|
||||
// Mtime change ⇒ miss.
|
||||
assert_eq!(memo.lookup(path, 10, 2000), None);
|
||||
// Re-store replaces the stale entry.
|
||||
memo.store(path, 11, 2000, "h2".into());
|
||||
assert_eq!(memo.lookup(path, 11, 2000).as_deref(), Some("h2"));
|
||||
assert_eq!(memo.lookup(path, 10, 1000), None);
|
||||
}
|
||||
|
||||
/// `stat` consults the memo (no re-hash for an unchanged file) and
|
||||
/// recomputes when `(size, mtime)` no longer match.
|
||||
#[tokio::test]
|
||||
async fn stat_uses_memo_until_file_changes() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("data.txt"), b"hello world").unwrap();
|
||||
|
||||
let req = FsStatReq {
|
||||
path: "data.txt".into(),
|
||||
};
|
||||
let first = stat(&ws, &req).await.unwrap();
|
||||
assert!(first.exists);
|
||||
assert_eq!(first.node_type, Some(FsNodeType::File));
|
||||
assert_eq!(first.size, Some(11));
|
||||
let real_hash = first.hash.clone().expect("hash for files");
|
||||
|
||||
// Plant a sentinel hash for the file's current (size, mtime). A
|
||||
// second stat must return the sentinel — proof it did not re-hash.
|
||||
let abs = root.join("data.txt");
|
||||
let md = std::fs::metadata(&abs).unwrap();
|
||||
let mtime = system_time_ms(md.modified().unwrap());
|
||||
ws.shared
|
||||
.client_fs_hash_memo
|
||||
.store(&abs, md.len(), mtime, "sentinel".into());
|
||||
let memoized = stat(&ws, &req).await.unwrap();
|
||||
assert_eq!(memoized.hash.as_deref(), Some("sentinel"));
|
||||
|
||||
// A size change invalidates the memo entry and re-hashes.
|
||||
std::fs::write(&abs, b"hello brave new world").unwrap();
|
||||
let rehashed = stat(&ws, &req).await.unwrap();
|
||||
let new_hash = rehashed.hash.expect("hash for files");
|
||||
assert_ne!(new_hash, "sentinel");
|
||||
assert_ne!(new_hash, real_hash);
|
||||
}
|
||||
|
||||
/// A path with a *file* as an intermediate component (`ENOTDIR`) is an
|
||||
/// existence miss, not an RPC error.
|
||||
#[tokio::test]
|
||||
async fn stat_enotdir_intermediate_reports_not_exists() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("file.txt"), b"x").unwrap();
|
||||
let res = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "file.txt/nested".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!res.exists);
|
||||
assert_eq!(res.node_type, None);
|
||||
assert_eq!(res.hash, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stat_missing_path_reports_not_exists() {
|
||||
let ws = make_handle();
|
||||
let res = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "nope.txt".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(!res.exists);
|
||||
assert_eq!(res.node_type, None);
|
||||
assert_eq!(res.hash, None);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_file_chunks_are_binary_safe_and_capped() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
// Non-UTF-8 payload: every byte value once.
|
||||
let payload: Vec<u8> = (0u8..=255).collect();
|
||||
std::fs::write(root.join("blob.bin"), &payload).unwrap();
|
||||
|
||||
let req = FsReadFileReq {
|
||||
path: "blob.bin".into(),
|
||||
// Bytes 200..210 are bare continuation bytes — never valid UTF-8.
|
||||
offset: Some(200),
|
||||
length: Some(50),
|
||||
max_bytes: 10, // cap below the requested length
|
||||
encoding: FsReadEncoding::Base64,
|
||||
};
|
||||
let res = read_file(&ws, &req).await.unwrap();
|
||||
assert_eq!(res.size, 256);
|
||||
assert_eq!(res.content, None);
|
||||
assert_eq!(res.content_type, FsContentType::Binary);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(res.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(bytes, payload[200..210], "maxBytes caps the chunk");
|
||||
|
||||
// Full-file hash regardless of the requested range.
|
||||
use sha2::{Digest, Sha256};
|
||||
assert_eq!(res.hash, format!("{:x}", Sha256::digest(&payload)));
|
||||
|
||||
// Memoized second read (range-only fast path) returns the
|
||||
// identical chunk + hash.
|
||||
let again = read_file(&ws, &req).await.unwrap();
|
||||
assert_eq!(again.hash, res.hash);
|
||||
let again_bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(again.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(again_bytes, payload[200..210]);
|
||||
}
|
||||
|
||||
/// `maxBytes` is server-capped at [`MAX_READ_BYTES`]:
|
||||
/// a caller-supplied huge budget cannot make the workspace buffer the
|
||||
/// whole file.
|
||||
#[tokio::test]
|
||||
async fn read_file_server_caps_max_bytes() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let payload = vec![0u8; (MAX_READ_BYTES + 100) as usize];
|
||||
std::fs::write(root.join("big.bin"), &payload).unwrap();
|
||||
|
||||
let res = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "big.bin".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: u64::MAX,
|
||||
encoding: FsReadEncoding::Base64,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(res.size, payload.len() as u64);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(res.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(
|
||||
bytes.len() as u64,
|
||||
MAX_READ_BYTES,
|
||||
"clamped to the server cap"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn read_file_utf8_default_and_binary_fallback() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("text.txt"), "héllo").unwrap();
|
||||
std::fs::write(root.join("bin.dat"), [0xff, 0xfe, 0x00]).unwrap();
|
||||
|
||||
let text = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "text.txt".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(text.content.as_deref(), Some("héllo"));
|
||||
assert_eq!(text.content_base64, None);
|
||||
assert_eq!(text.content_type, FsContentType::Text);
|
||||
|
||||
// Invalid UTF-8 under the utf8 default degrades to base64.
|
||||
let bin = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "bin.dat".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Utf8,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(bin.content, None);
|
||||
assert_eq!(bin.content_type, FsContentType::Binary);
|
||||
let bytes = base64::engine::general_purpose::STANDARD
|
||||
.decode(bin.content_base64.unwrap())
|
||||
.unwrap();
|
||||
assert_eq!(bytes, [0xff, 0xfe, 0x00]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn resolve_rejects_escapes() {
|
||||
let ws = make_handle();
|
||||
for path in ["/etc/passwd", "../escape.txt"] {
|
||||
let err = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: path.to_owned(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("escape must be rejected");
|
||||
assert!(matches!(err, WorkspaceError::HubError(_)), "{err:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// An absolute path *inside* the workspace root is accepted and stats
|
||||
/// the same file as its root-relative form.
|
||||
#[tokio::test]
|
||||
async fn resolve_accepts_absolute_within_root() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("data.txt"), b"hello").unwrap();
|
||||
|
||||
let rel = stat(
|
||||
&ws,
|
||||
&FsStatReq {
|
||||
path: "data.txt".into(),
|
||||
},
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
assert!(rel.exists);
|
||||
|
||||
let abs_path = root.join("data.txt").to_string_lossy().into_owned();
|
||||
let abs = stat(&ws, &FsStatReq { path: abs_path }).await.unwrap();
|
||||
assert!(abs.exists);
|
||||
assert_eq!(abs.node_type, rel.node_type);
|
||||
assert_eq!(abs.size, rel.size);
|
||||
assert_eq!(abs.hash, rel.hash);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
#[cfg(unix)]
|
||||
async fn resolve_rejects_symlink_escape() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
let outside = tempfile::tempdir().unwrap();
|
||||
std::fs::write(outside.path().join("secret.txt"), b"secret").unwrap();
|
||||
std::os::unix::fs::symlink(outside.path(), root.join("escape_link")).unwrap();
|
||||
|
||||
let err = read_file(
|
||||
&ws,
|
||||
&FsReadFileReq {
|
||||
path: "escape_link/secret.txt".into(),
|
||||
offset: None,
|
||||
length: None,
|
||||
max_bytes: 1_048_576,
|
||||
encoding: FsReadEncoding::Base64,
|
||||
},
|
||||
)
|
||||
.await
|
||||
.expect_err("symlink escape must be rejected");
|
||||
assert!(
|
||||
err.to_string().contains("symlink escape"),
|
||||
"unexpected error: {err}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn list_empty_path_lists_root() {
|
||||
let ws = make_handle();
|
||||
let root = ws.root_cwd().unwrap();
|
||||
std::fs::write(root.join("rooted.txt"), b"x").unwrap();
|
||||
let res = list(&ws, &list_req("")).await.unwrap();
|
||||
assert!(res.nodes.iter().any(|n| n.name == "rooted.txt"));
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,490 +0,0 @@
|
||||
//! Hub [`AuthProvider`] from `~/.kigi/auth.json` for the standalone
|
||||
//! `workspace_server` binary: loopback `ws://` uses a plain bearer, otherwise
|
||||
//! an auto-refreshing OIDC provider that persists rotated tokens to disk.
|
||||
//!
|
||||
//! The in-leader `grok workspace` exposure does NOT use this path — it sources
|
||||
//! an in-memory provider from the leader's `AuthManager` (see
|
||||
//! `LeaderAuthProvider`) to avoid racing the leader's own auth.json writer.
|
||||
|
||||
use std::collections::BTreeMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
use kigi_computer_hub_sdk::{
|
||||
AuthCredential, AuthIdentity, AuthProvider, OidcAuthProviderBuilder, RefreshEvent,
|
||||
};
|
||||
use url::Url;
|
||||
|
||||
/// Plain bearer provider that also carries the owner identity parsed from the
|
||||
/// same auth.json entry. Used for the loopback / local-dev path (no OIDC
|
||||
/// refresh) so consumers can still read the owner identity from the auth
|
||||
/// provider — without a second auth.json read.
|
||||
struct BearerWithIdentity {
|
||||
token: String,
|
||||
identity: AuthIdentity,
|
||||
}
|
||||
|
||||
impl std::fmt::Debug for BearerWithIdentity {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
// Never log the bearer token; surface only the (non-secret) identity.
|
||||
f.debug_struct("BearerWithIdentity")
|
||||
.field("identity", &self.identity)
|
||||
.finish_non_exhaustive()
|
||||
}
|
||||
}
|
||||
|
||||
impl AuthProvider for BearerWithIdentity {
|
||||
fn current(&self) -> AuthCredential {
|
||||
AuthCredential::bearer(self.token.clone())
|
||||
}
|
||||
|
||||
fn identity(&self) -> Option<AuthIdentity> {
|
||||
Some(self.identity.clone())
|
||||
}
|
||||
}
|
||||
|
||||
/// Owner identity parsed from an auth.json entry, for the [`AuthProvider`]s
|
||||
/// built here to surface via [`AuthProvider::identity`].
|
||||
fn identity_from_entry(entry: &AuthEntry) -> AuthIdentity {
|
||||
AuthIdentity {
|
||||
user_id: entry.user_id.clone(),
|
||||
principal_type: entry.principal_type.clone(),
|
||||
principal_id: entry.principal_id.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, serde::Deserialize)]
|
||||
struct AuthEntry {
|
||||
key: String,
|
||||
#[serde(default)]
|
||||
user_id: String,
|
||||
#[serde(default)]
|
||||
refresh_token: Option<String>,
|
||||
#[serde(default)]
|
||||
oidc_issuer: Option<String>,
|
||||
#[serde(default)]
|
||||
oidc_client_id: Option<String>,
|
||||
#[serde(default)]
|
||||
principal_type: Option<String>,
|
||||
#[serde(default)]
|
||||
principal_id: Option<String>,
|
||||
#[serde(default)]
|
||||
expires_at: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
fn default_auth_path() -> anyhow::Result<PathBuf> {
|
||||
let grok = kigi_config::user_kigi_home()
|
||||
.ok_or_else(|| anyhow::anyhow!("no user grok home (set $KIGI_SHARE_DIR or $HOME)"))?;
|
||||
Ok(grok.join("auth.json"))
|
||||
}
|
||||
|
||||
/// Read the active OIDC entry and its scope key. The key is threaded to the
|
||||
/// refresh write so rotation updates exactly the entry that was read.
|
||||
fn read_auth_entry(path: &Path) -> anyhow::Result<(String, AuthEntry)> {
|
||||
if !path.exists() {
|
||||
anyhow::bail!(
|
||||
"No auth credentials found at {}. Run `kigi login` first.",
|
||||
path.display()
|
||||
);
|
||||
}
|
||||
|
||||
let content = std::fs::read_to_string(path)
|
||||
.map_err(|e| anyhow::anyhow!("failed to read {}: {e}", path.display()))?;
|
||||
let entries: BTreeMap<String, AuthEntry> = serde_json::from_str(&content)
|
||||
.map_err(|e| anyhow::anyhow!("failed to parse {}: {e}", path.display()))?;
|
||||
|
||||
entries
|
||||
.into_iter()
|
||||
.find(|(_, e)| e.refresh_token.is_some() && e.oidc_issuer.is_some())
|
||||
.ok_or_else(|| {
|
||||
anyhow::anyhow!(
|
||||
"no OIDC auth entry found in {}. Run `kigi login` first.",
|
||||
path.display()
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
fn build_oidc_provider(
|
||||
scope_key: String,
|
||||
entry: &AuthEntry,
|
||||
auth_path: PathBuf,
|
||||
) -> anyhow::Result<Arc<dyn AuthProvider>> {
|
||||
let refresh_token = entry.refresh_token.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no refresh_token — cannot refresh expired tokens")
|
||||
})?;
|
||||
let issuer = entry.oidc_issuer.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no oidc_issuer — cannot refresh expired tokens")
|
||||
})?;
|
||||
let client_id = entry.oidc_client_id.as_ref().ok_or_else(|| {
|
||||
anyhow::anyhow!("auth entry has no oidc_client_id — cannot refresh expired tokens")
|
||||
})?;
|
||||
|
||||
let mut builder = OidcAuthProviderBuilder::new(&entry.key, refresh_token, issuer, client_id);
|
||||
|
||||
// Owner identity is surfaced via `AuthProvider::identity()` so consumers
|
||||
// read it from this provider — no separate auth.json read.
|
||||
builder = builder.user_id(&entry.user_id);
|
||||
if let Some(ref pt) = entry.principal_type {
|
||||
builder = builder.principal_type(pt);
|
||||
}
|
||||
if let Some(ref pid) = entry.principal_id {
|
||||
builder = builder.principal_id(pid);
|
||||
}
|
||||
if let Some(exp) = entry.expires_at {
|
||||
builder = builder.expires_at(exp);
|
||||
}
|
||||
|
||||
builder = builder.on_refresh(Arc::new(move |event: &RefreshEvent| {
|
||||
if let Err(e) = write_refreshed_token(&auth_path, &scope_key, event) {
|
||||
tracing::warn!(error = %e, "failed to persist refreshed token to auth.json");
|
||||
}
|
||||
}));
|
||||
|
||||
Ok(Arc::new(builder.build()))
|
||||
}
|
||||
|
||||
fn write_refreshed_token(path: &Path, scope_key: &str, event: &RefreshEvent) -> anyhow::Result<()> {
|
||||
let content = std::fs::read_to_string(path)?;
|
||||
let mut raw: serde_json::Value = serde_json::from_str(&content)?;
|
||||
|
||||
let Some(obj) = raw.get_mut(scope_key).and_then(|e| e.as_object_mut()) else {
|
||||
anyhow::bail!("auth entry '{scope_key}' not found while persisting refreshed token");
|
||||
};
|
||||
obj.insert(
|
||||
"key".to_owned(),
|
||||
serde_json::Value::String(event.access_token.clone()),
|
||||
);
|
||||
if let Some(ref rt) = event.new_refresh_token {
|
||||
obj.insert(
|
||||
"refresh_token".to_owned(),
|
||||
serde_json::Value::String(rt.clone()),
|
||||
);
|
||||
}
|
||||
if let Some(exp) = event.expires_at {
|
||||
obj.insert(
|
||||
"expires_at".to_owned(),
|
||||
serde_json::Value::String(exp.to_rfc3339()),
|
||||
);
|
||||
}
|
||||
|
||||
write_json_atomic(path, &raw)?;
|
||||
tracing::info!(path = %path.display(), "persisted refreshed token to auth.json");
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Atomically replace `path`: temp file (0600 on Unix) + fsync + rename. Avoids
|
||||
/// the truncate-in-place corruption window when the long-lived binary rewrites
|
||||
/// auth.json.
|
||||
fn write_json_atomic(path: &Path, value: &serde_json::Value) -> anyhow::Result<()> {
|
||||
use std::io::Write;
|
||||
|
||||
let json = serde_json::to_string_pretty(value)?;
|
||||
let tmp = path.with_extension(format!("json.{}.tmp", std::process::id()));
|
||||
|
||||
let mut opts = std::fs::OpenOptions::new();
|
||||
opts.write(true).create(true).truncate(true);
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::OpenOptionsExt;
|
||||
opts.mode(0o600);
|
||||
}
|
||||
|
||||
let mut file = opts
|
||||
.open(&tmp)
|
||||
.map_err(|e| anyhow::anyhow!("failed to open {}: {e}", tmp.display()))?;
|
||||
file.write_all(json.as_bytes())?;
|
||||
file.sync_all()?;
|
||||
drop(file);
|
||||
|
||||
#[cfg(windows)]
|
||||
let _ = std::fs::remove_file(path);
|
||||
|
||||
if let Err(e) = std::fs::rename(&tmp, path) {
|
||||
let _ = std::fs::remove_file(&tmp);
|
||||
return Err(anyhow::anyhow!("failed to replace {}: {e}", path.display()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build a hub auth provider for `hub_url`. `auth_config` overrides
|
||||
/// the default credential path (`~/.kigi/auth.json`).
|
||||
pub fn provider(
|
||||
hub_url: &Url,
|
||||
auth_config: Option<&Path>,
|
||||
) -> anyhow::Result<Arc<dyn AuthProvider>> {
|
||||
let auth_path = match auth_config {
|
||||
Some(p) => p.to_path_buf(),
|
||||
None => default_auth_path()?,
|
||||
};
|
||||
let (scope_key, entry) = read_auth_entry(&auth_path)?;
|
||||
|
||||
let is_loopback = hub_url.scheme() == "ws"
|
||||
&& matches!(hub_url.host_str(), Some("localhost" | "127.0.0.1" | "::1"));
|
||||
|
||||
if is_loopback {
|
||||
tracing::info!("Using local-dev auth (loopback hub)");
|
||||
Ok(Arc::new(BearerWithIdentity {
|
||||
identity: identity_from_entry(&entry),
|
||||
token: entry.key.clone(),
|
||||
}))
|
||||
} else {
|
||||
build_oidc_provider(scope_key, &entry, auth_path)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Write;
|
||||
|
||||
fn write_auth_json(dir: &std::path::Path, json: &str) -> PathBuf {
|
||||
let path = dir.join("auth.json");
|
||||
let mut f = std::fs::File::create(&path).unwrap();
|
||||
f.write_all(json.as_bytes()).unwrap();
|
||||
path
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_picks_oidc_entry() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"legacy": { "key": "xai-plainkey", "user_id": "u1" },
|
||||
"oidc": {
|
||||
"key": "eyJhbGciOiJFUzI1NiJ9.test",
|
||||
"user_id": "u2",
|
||||
"refresh_token": "rt",
|
||||
"oidc_issuer": "https://auth.example.com",
|
||||
"oidc_client_id": "client1"
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (key, entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(key, "oidc");
|
||||
assert_eq!(entry.refresh_token.as_deref(), Some("rt"));
|
||||
assert_eq!(
|
||||
entry.oidc_issuer.as_deref(),
|
||||
Some("https://auth.example.com")
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_rejects_non_oidc() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"api_key": { "key": "xai-plainkey", "user_id": "u1" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let err = read_auth_entry(&path).unwrap_err();
|
||||
assert!(err.to_string().contains("no OIDC auth entry"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_missing_file() {
|
||||
let path = PathBuf::from("/nonexistent/auth.json");
|
||||
let err = read_auth_entry(&path).unwrap_err();
|
||||
assert!(err.to_string().contains("No auth credentials"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn read_auth_entry_tolerates_extra_fields() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"scope": {
|
||||
"key": "eyJhbGciOiJFUzI1NiJ9.tok",
|
||||
"user_id": "u1",
|
||||
"auth_mode": "oauth",
|
||||
"create_time": "2026-01-01T00:00:00Z",
|
||||
"email": "test@x.ai",
|
||||
"first_name": "Test",
|
||||
"refresh_token": "rt1",
|
||||
"oidc_issuer": "https://auth.x.ai",
|
||||
"oidc_client_id": "c1",
|
||||
"some_future_field": true
|
||||
}
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (_key, entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(entry.refresh_token.as_deref(), Some("rt1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_refresh_token() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: None,
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("refresh_token"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_issuer() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: None,
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("oidc_issuer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_requires_client_id() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: None,
|
||||
principal_type: None,
|
||||
principal_id: None,
|
||||
expires_at: None,
|
||||
};
|
||||
let err = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap_err();
|
||||
assert!(err.to_string().contains("oidc_client_id"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn build_oidc_provider_succeeds_with_all_fields() {
|
||||
let entry = AuthEntry {
|
||||
key: "eyJ.tok".into(),
|
||||
user_id: "u1".into(),
|
||||
refresh_token: Some("rt".into()),
|
||||
oidc_issuer: Some("https://auth.x.ai".into()),
|
||||
oidc_client_id: Some("c1".into()),
|
||||
principal_type: Some("Team".into()),
|
||||
principal_id: Some("t1".into()),
|
||||
expires_at: Some(chrono::Utc::now() + chrono::Duration::hours(1)),
|
||||
};
|
||||
let provider = build_oidc_provider("oidc".into(), &entry, PathBuf::from("/tmp/x")).unwrap();
|
||||
let cred = provider.current();
|
||||
match cred {
|
||||
kigi_computer_hub_sdk::AuthCredential::Bearer { token } => {
|
||||
assert_eq!(token, "eyJ.tok");
|
||||
}
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
// Identity is surfaced from the parsed entry (no second auth.json read).
|
||||
let id = provider.identity().expect("identity present");
|
||||
assert_eq!(id.user_id, "u1");
|
||||
assert_eq!(id.principal_type.as_deref(), Some("Team"));
|
||||
assert_eq!(id.principal_id.as_deref(), Some("t1"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_updates_jwt_entry() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"legacy": { "key": "xai-old", "user_id": "u1" },
|
||||
"oidc": { "key": "eyJ.old", "user_id": "u2", "refresh_token": "rt-old", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.new".into(),
|
||||
new_refresh_token: Some("rt-new".into()),
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, "oidc", &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["oidc"]["key"], "eyJ.new");
|
||||
assert_eq!(updated["oidc"]["refresh_token"], "rt-new");
|
||||
assert_eq!(updated["legacy"]["key"], "xai-old"); // untouched
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_targets_exact_scope_key() {
|
||||
// Non-sorted order: refresh must update the read-selected key ("aaa"),
|
||||
// not the first in file order ("zzz").
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"zzz": { "key": "eyJ.z", "refresh_token": "rt-z", "oidc_issuer": "https://auth.x.ai" },
|
||||
"aaa": { "key": "eyJ.a", "refresh_token": "rt-a", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let (key, _entry) = read_auth_entry(&path).unwrap();
|
||||
assert_eq!(key, "aaa");
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.a-new".into(),
|
||||
new_refresh_token: Some("rt-a-new".into()),
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, &key, &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["aaa"]["key"], "eyJ.a-new");
|
||||
assert_eq!(updated["aaa"]["refresh_token"], "rt-a-new");
|
||||
assert_eq!(updated["zzz"]["key"], "eyJ.z");
|
||||
assert_eq!(updated["zzz"]["refresh_token"], "rt-z");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn write_refreshed_token_preserves_existing_rt_when_not_rotated() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{
|
||||
"oidc": { "key": "eyJ.old", "user_id": "u1", "refresh_token": "rt-keep", "oidc_issuer": "https://auth.x.ai" }
|
||||
}"#,
|
||||
);
|
||||
|
||||
let event = RefreshEvent {
|
||||
access_token: "eyJ.new".into(),
|
||||
new_refresh_token: None,
|
||||
expires_at: None,
|
||||
};
|
||||
write_refreshed_token(&path, "oidc", &event).unwrap();
|
||||
|
||||
let updated: serde_json::Value =
|
||||
serde_json::from_str(&std::fs::read_to_string(&path).unwrap()).unwrap();
|
||||
assert_eq!(updated["oidc"]["key"], "eyJ.new");
|
||||
assert_eq!(updated["oidc"]["refresh_token"], "rt-keep");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn provider_loopback_uses_bearer() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = write_auth_json(
|
||||
dir.path(),
|
||||
r#"{ "oidc": { "key": "eyJ.tok", "user_id": "u1", "refresh_token": "rt", "oidc_issuer": "https://auth.x.ai", "oidc_client_id": "c1" } }"#,
|
||||
);
|
||||
let url = Url::parse("ws://localhost:9988/v1/tools").unwrap();
|
||||
let auth = provider(&url, Some(&path)).unwrap();
|
||||
match auth.current() {
|
||||
AuthCredential::Bearer { token } => assert_eq!(token, "eyJ.tok"),
|
||||
_ => panic!("expected Bearer"),
|
||||
}
|
||||
// Loopback still surfaces identity from the same entry.
|
||||
let id = auth.identity().expect("loopback identity present");
|
||||
assert_eq!(id.user_id, "u1");
|
||||
}
|
||||
}
|
||||
@@ -1,127 +0,0 @@
|
||||
//! Server-proxied workspace utilities.
|
||||
//!
|
||||
//! Helper functions for consuming server tool streams and extracting typed
|
||||
//! notifications from server notification frames. These are used by both
|
||||
//! the workspace crate (hub_server) and the shell crate (proxy-mode
|
||||
//! session actors) to interact with the server.
|
||||
//!
|
||||
//! The `HubWorkspaceChannel` struct that previously lived here has been
|
||||
//! removed. Sessions now call the server harness directly via `ToolContext`.
|
||||
|
||||
use kigi_tools::notification::types::ToolNotification;
|
||||
use kigi_workspace_types::WorkspaceEvent;
|
||||
|
||||
pub use crate::hub_ids::WORKSPACE_RPC_TOOL_ID;
|
||||
// Canonical in the client crate; re-exported for existing importers.
|
||||
pub use kigi_workspace_client::consume_stream_terminal;
|
||||
|
||||
/// Extract a `WorkspaceEvent` from a custom `ToolNotificationFrame`.
|
||||
pub fn extract_workspace_event(
|
||||
frame: &kigi_tool_protocol::ToolNotificationFrame,
|
||||
) -> Option<WorkspaceEvent> {
|
||||
use kigi_tool_protocol::WireToolNotification;
|
||||
match &frame.notification {
|
||||
WireToolNotification::Custom(c) => serde_json::from_value(c.payload.clone()).ok(),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Extract a `ToolNotification` from a custom `ToolNotificationFrame`.
|
||||
///
|
||||
/// Currently has no producer; kept for a future per-session `tool.notify` path.
|
||||
pub fn extract_tool_notification(
|
||||
frame: &kigi_tool_protocol::ToolNotificationFrame,
|
||||
) -> Option<ToolNotification> {
|
||||
use kigi_tool_protocol::WireToolNotification;
|
||||
match &frame.notification {
|
||||
WireToolNotification::Custom(c) => {
|
||||
serde_json::from_value::<ToolNotification>(c.payload.clone()).ok()
|
||||
}
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use kigi_tool_protocol::{ToolId, ToolNotificationFrame, WireToolNotification};
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_custom_valid() {
|
||||
let event = WorkspaceEvent::ToolsChanged {
|
||||
session_id: "main".into(),
|
||||
};
|
||||
let payload = serde_json::to_value(&event).unwrap();
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_events").unwrap(),
|
||||
"workspace_event",
|
||||
payload,
|
||||
);
|
||||
let result = extract_workspace_event(&frame);
|
||||
assert!(result.is_some(), "should parse valid WorkspaceEvent");
|
||||
match result.unwrap() {
|
||||
WorkspaceEvent::ToolsChanged { session_id } => {
|
||||
assert_eq!(session_id, "main");
|
||||
}
|
||||
other => panic!("expected ToolsChanged, got {other:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_invalid_payload() {
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_events").unwrap(),
|
||||
"workspace_event",
|
||||
serde_json::json!({"not_a_valid_event": true}),
|
||||
);
|
||||
assert!(
|
||||
extract_workspace_event(&frame).is_none(),
|
||||
"invalid payload should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_workspace_event_known_variant_returns_none() {
|
||||
let frame = ToolNotificationFrame {
|
||||
tool_call_id: None,
|
||||
tool_id: Some(ToolId::new("workspace_events").unwrap()),
|
||||
notification: WireToolNotification::Known(serde_json::json!({})),
|
||||
};
|
||||
assert!(
|
||||
extract_workspace_event(&frame).is_none(),
|
||||
"Known variant should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_tool_notification_invalid_payload() {
|
||||
let frame = ToolNotificationFrame::custom(
|
||||
ToolId::new("workspace_tool_notifications").unwrap(),
|
||||
"tool_notification",
|
||||
serde_json::json!({"not_a_notification": true}),
|
||||
);
|
||||
assert!(
|
||||
extract_tool_notification(&frame).is_none(),
|
||||
"invalid payload should return None"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn extract_tool_notification_known_variant_returns_none() {
|
||||
let frame = ToolNotificationFrame {
|
||||
tool_call_id: None,
|
||||
tool_id: Some(ToolId::new("workspace_tool_notifications").unwrap()),
|
||||
notification: WireToolNotification::Known(serde_json::json!({})),
|
||||
};
|
||||
assert!(
|
||||
extract_tool_notification(&frame).is_none(),
|
||||
"Known variant should return None"
|
||||
);
|
||||
}
|
||||
|
||||
// consume_stream_terminal tests live in kigi-workspace-client.
|
||||
}
|
||||
@@ -1,10 +0,0 @@
|
||||
//! Hub tool ID constants, canonical in `kigi_workspace_types::rpc`
|
||||
//! and re-exported here for existing importers.
|
||||
//!
|
||||
//! `WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID` is intentionally producer-less today;
|
||||
//! see [`crate::hub_channel::extract_tool_notification`].
|
||||
|
||||
pub use kigi_workspace_types::rpc::{
|
||||
WORKSPACE_CLIENT_EXT_NOTIFICATIONS_TOOL_ID, WORKSPACE_EVENTS_TOOL_ID, WORKSPACE_RPC_TOOL_ID,
|
||||
WORKSPACE_TOOL_NOTIFICATIONS_TOOL_ID,
|
||||
};
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,220 +0,0 @@
|
||||
//! MCP integration for the workspace server.
|
||||
//!
|
||||
//! Bridges [`McpClient`] to the server's [`McpTransport`] trait and wraps
|
||||
//! tool handlers with qualified `server__tool` names.
|
||||
|
||||
use std::sync::Arc;
|
||||
|
||||
use async_trait::async_trait;
|
||||
use kigi_computer_hub_mcp_adapter::{
|
||||
McpBridgeConfig, McpCallResult, McpContent, McpServerInfo, McpToolDefinition, McpToolHandler,
|
||||
McpTransport,
|
||||
};
|
||||
use kigi_computer_hub_sdk::ToolServerHandler;
|
||||
use kigi_mcp::rmcp;
|
||||
use kigi_mcp::servers::McpClient;
|
||||
use kigi_tool_protocol::ToolId;
|
||||
use kigi_tool_runtime::{ToolCallContext, ToolStream, TypedToolOutput};
|
||||
use kigi_tool_types::ToolDescription;
|
||||
use serde_json::Value;
|
||||
|
||||
/// Adapts [`McpClient`] to the [`McpTransport`] trait for [`McpBridge`].
|
||||
pub(crate) struct McpClientTransportAdapter {
|
||||
client: Arc<McpClient>,
|
||||
}
|
||||
|
||||
impl McpClientTransportAdapter {
|
||||
pub fn new(client: Arc<McpClient>) -> Self {
|
||||
Self { client }
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl McpTransport for McpClientTransportAdapter {
|
||||
async fn initialize(&self) -> Result<McpServerInfo, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
let info = service.peer_info().ok_or_else(|| {
|
||||
kigi_computer_hub_mcp_adapter::McpError::Transport("no peer info after init".into())
|
||||
})?;
|
||||
Ok(McpServerInfo {
|
||||
name: info.server_info.name.clone(),
|
||||
version: info.server_info.version.clone(),
|
||||
capabilities: serde_json::to_value(&info.capabilities).unwrap_or_default(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn list_tools(
|
||||
&self,
|
||||
) -> Result<Vec<McpToolDefinition>, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
let mut all_tools = Vec::new();
|
||||
let mut cursor: Option<String> = None;
|
||||
loop {
|
||||
let result = service
|
||||
.list_tools(Some(
|
||||
rmcp::model::PaginatedRequestParams::default().with_cursor(cursor.clone()),
|
||||
))
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
all_tools.extend(result.tools.into_iter().map(|t| McpToolDefinition {
|
||||
name: t.name.to_string(),
|
||||
description: t.description.map(|d| d.to_string()),
|
||||
input_schema: serde_json::to_value(&t.input_schema).ok(),
|
||||
}));
|
||||
|
||||
match result.next_cursor {
|
||||
Some(next) => cursor = Some(next),
|
||||
None => break,
|
||||
}
|
||||
}
|
||||
|
||||
Ok(all_tools)
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
name: &str,
|
||||
arguments: Value,
|
||||
) -> Result<McpCallResult, kigi_computer_hub_mcp_adapter::McpError> {
|
||||
let service = self
|
||||
.client
|
||||
.ensure_initialized()
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
// MCP spec requires arguments to be an object; coerce if needed.
|
||||
let args_object = match arguments {
|
||||
Value::Object(map) => Some(map),
|
||||
Value::Null => None,
|
||||
other => {
|
||||
let mut wrapper = serde_json::Map::new();
|
||||
wrapper.insert("value".to_string(), other);
|
||||
Some(wrapper)
|
||||
}
|
||||
};
|
||||
let result = service
|
||||
.call_tool({
|
||||
let mut params = rmcp::model::CallToolRequestParams::new(name.to_string());
|
||||
params.arguments = args_object;
|
||||
params
|
||||
})
|
||||
.await
|
||||
.map_err(|e| kigi_computer_hub_mcp_adapter::McpError::Transport(e.to_string()))?;
|
||||
|
||||
Ok(McpCallResult {
|
||||
content: result
|
||||
.content
|
||||
.into_iter()
|
||||
.map(|c| match c {
|
||||
rmcp::model::ContentBlock::Text(t) => McpContent::Text { text: t.text },
|
||||
rmcp::model::ContentBlock::Image(img) => McpContent::Image {
|
||||
mime_type: img.mime_type,
|
||||
data: img.data,
|
||||
},
|
||||
_ => McpContent::Text {
|
||||
text: "[unsupported content type]".to_string(),
|
||||
},
|
||||
})
|
||||
.collect(),
|
||||
is_error: result.is_error.unwrap_or(false),
|
||||
})
|
||||
}
|
||||
|
||||
async fn close(&self) -> Result<(), kigi_computer_hub_mcp_adapter::McpError> {
|
||||
// No-op: cleanup happens when McpClient is dropped.
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Wraps an [`McpToolHandler`] to qualify tool names as `server__tool`.
|
||||
pub(crate) struct QualifiedMcpToolHandler {
|
||||
qualified_id: ToolId,
|
||||
qualified_name: String,
|
||||
inner: Arc<McpToolHandler>,
|
||||
}
|
||||
|
||||
impl QualifiedMcpToolHandler {
|
||||
/// Returns `None` if the qualified name is not a valid `ToolId`.
|
||||
pub fn try_new(qualified_name: String, inner: Arc<McpToolHandler>) -> Option<Self> {
|
||||
let qualified_id = match ToolId::new(&qualified_name) {
|
||||
Ok(id) => id,
|
||||
Err(err) => {
|
||||
tracing::warn!(
|
||||
qualified_name = %qualified_name,
|
||||
error = %err,
|
||||
"skipping MCP tool: qualified name is not a valid ToolId"
|
||||
);
|
||||
return None;
|
||||
}
|
||||
};
|
||||
Some(Self {
|
||||
qualified_id,
|
||||
qualified_name,
|
||||
inner,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ToolServerHandler for QualifiedMcpToolHandler {
|
||||
fn tool_id(&self) -> ToolId {
|
||||
self.qualified_id.clone()
|
||||
}
|
||||
|
||||
fn description(&self) -> ToolDescription {
|
||||
let inner_desc = self.inner.description();
|
||||
ToolDescription::new(self.qualified_name.clone(), inner_desc.description)
|
||||
}
|
||||
|
||||
fn input_schema(&self) -> Option<Value> {
|
||||
self.inner.input_schema()
|
||||
}
|
||||
|
||||
async fn handle_call(&self, ctx: ToolCallContext, args: Value) -> ToolStream<TypedToolOutput> {
|
||||
self.inner.handle_call(ctx, args).await
|
||||
}
|
||||
}
|
||||
|
||||
/// Result of a `workspace.configure_mcp` RPC call.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct McpStartResult {
|
||||
/// Server names that started successfully.
|
||||
pub started: Vec<String>,
|
||||
/// Servers that failed to start.
|
||||
pub failed: Vec<McpStartFailure>,
|
||||
}
|
||||
|
||||
/// A single MCP server startup failure.
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct McpStartFailure {
|
||||
/// Server name.
|
||||
pub name: String,
|
||||
/// Human-readable error description.
|
||||
pub error: String,
|
||||
}
|
||||
|
||||
/// Extract a server name from an [`McpError`](kigi_mcp::servers::McpError),
|
||||
/// falling back to `"unknown"`.
|
||||
pub(crate) fn server_name_from_mcp_error(e: &kigi_mcp::servers::McpError) -> &str {
|
||||
e.server_name().unwrap_or("unknown")
|
||||
}
|
||||
|
||||
/// Bridge config factory for MCP bridge connections.
|
||||
pub(crate) fn make_bridge_config(
|
||||
session_id: kigi_tool_protocol::SessionId,
|
||||
server_name: &str,
|
||||
) -> McpBridgeConfig {
|
||||
McpBridgeConfig {
|
||||
session_id,
|
||||
namespace: Some(server_name.to_owned()),
|
||||
}
|
||||
}
|
||||
@@ -1,498 +0,0 @@
|
||||
//! Tool-permission emit: when the rules engine returns "ask" for a guarded
|
||||
//! tool, request the decision from chat over the server instead of prompting a
|
||||
//! local ACP client, then map chat's reply back onto a [`PromptOutcome`] so the
|
||||
//! manager's existing decision + `ALWAYS_*` persistence applies unchanged.
|
||||
use crate::permission::prompter::{PromptOutcome, tool_name_for_access};
|
||||
use crate::permission::types::AccessKind;
|
||||
use async_trait::async_trait;
|
||||
use kigi_computer_hub_sdk::harness::PERMISSION_REQUEST_KIND;
|
||||
use kigi_computer_hub_sdk::{ToolServer, WeakToolServer};
|
||||
use kigi_tool_protocol::SessionId;
|
||||
use prometheus::{HistogramVec, IntCounter, register_histogram_vec, register_int_counter};
|
||||
use serde_json::Value;
|
||||
use std::sync::LazyLock;
|
||||
/// Wall-clock time the workspace awaits chat's decision on a `permission_request`
|
||||
/// hook. `outcome` is `ok` (chat replied) or `error` (transport failure /
|
||||
/// backstop deadline).
|
||||
static PERMISSION_REPLY_DURATION: LazyLock<HistogramVec> = LazyLock::new(|| {
|
||||
register_histogram_vec!(
|
||||
"grok_workspace_permission_reply_seconds",
|
||||
"Wall-clock time awaiting chat's reply to a permission_request hook",
|
||||
&["outcome"],
|
||||
vec![0.5, 1.0, 2.0, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0, 600.0]
|
||||
)
|
||||
.expect("grok_workspace_permission_reply_seconds must register once")
|
||||
});
|
||||
/// Permission requests whose reply timed out (the server backstop deadline fired).
|
||||
/// A subset of the histogram's `error` outcome, promoted to its own counter so a
|
||||
/// stuck/lost reply is distinguishable from other transport failures.
|
||||
static PERMISSION_TIMEOUT_TOTAL: LazyLock<IntCounter> = LazyLock::new(|| {
|
||||
register_int_counter!(
|
||||
"grok_workspace_permission_timeout_total",
|
||||
"permission_request hooks whose reply timed out (backstop deadline fired)"
|
||||
)
|
||||
.expect("grok_workspace_permission_timeout_total must register once")
|
||||
});
|
||||
/// Zero-init this module's metric families. See [`crate::init_metrics`].
|
||||
pub(crate) fn init_metrics() {
|
||||
for outcome in ["ok", "error"] {
|
||||
let _ = PERMISSION_REPLY_DURATION.with_label_values(&[outcome]);
|
||||
}
|
||||
PERMISSION_TIMEOUT_TOTAL.inc_by(0);
|
||||
}
|
||||
/// Identifies the reply backstop-deadline timeout by its rendered message; the
|
||||
/// server SDK exposes no typed timeout variant to match on. If that message text
|
||||
/// changes, such a reply is recorded under the histogram's `error` outcome but
|
||||
/// not counted in `permission_timeout_total`.
|
||||
fn is_timeout_err(msg: &str) -> bool {
|
||||
msg.contains("timed out")
|
||||
}
|
||||
/// Env var that enables the HITL-live **tool-permission** emit (workspace →
|
||||
/// chat over the server) for local e2e and gradual rollout. Prefer server capability
|
||||
/// negotiation long-term; this is the interim gate so tool-permission can be
|
||||
/// exercised without waiting on that wire format.
|
||||
pub const HITL_PERMISSION_LIVE_ENV: &str = "KIGI_HITL_PERMISSION_LIVE";
|
||||
/// Whether the HITL-live permission path is enabled.
|
||||
///
|
||||
/// Intended long-term gate: the chat flag `grok_chat_enable_hitl_live_path`,
|
||||
/// propagated by the server at session-bind (capability negotiation). Until that
|
||||
/// lands, honor [`HITL_PERMISSION_LIVE_ENV`] (`1` / `true` / `yes`) so local
|
||||
/// stacks and e2e can turn the emit on explicitly. Default remains **off**
|
||||
/// (fail closed to the local ACP prompt).
|
||||
pub fn hitl_permission_live_enabled() -> bool {
|
||||
match std::env::var(HITL_PERMISSION_LIVE_ENV) {
|
||||
Ok(v) => {
|
||||
matches!(
|
||||
v.trim().to_ascii_lowercase().as_str(),
|
||||
"1" | "true" | "yes" | "on"
|
||||
)
|
||||
}
|
||||
Err(_) => false,
|
||||
}
|
||||
}
|
||||
/// Sends a `permission_request` hook to chat and awaits the decision reply.
|
||||
#[async_trait]
|
||||
pub trait PermissionHookTransport: Send + Sync {
|
||||
/// Emit the permission-request `payload` and return chat's decision reply.
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String>;
|
||||
}
|
||||
/// Hub-backed permission transport (weak server handle; upgrades per request).
|
||||
pub struct ToolServerPermissionTransport {
|
||||
server: WeakToolServer,
|
||||
session_id: SessionId,
|
||||
}
|
||||
impl ToolServerPermissionTransport {
|
||||
pub fn new(server: ToolServer, session_id: SessionId) -> Self {
|
||||
Self {
|
||||
server: server.downgrade(),
|
||||
session_id,
|
||||
}
|
||||
}
|
||||
/// Build from a session id held as a string; `None` if it is not a valid
|
||||
/// [`SessionId`].
|
||||
pub fn from_session_id(server: ToolServer, session_id: &str) -> Option<Self> {
|
||||
SessionId::new(session_id)
|
||||
.ok()
|
||||
.map(|sid| Self::new(server, sid))
|
||||
}
|
||||
}
|
||||
#[async_trait]
|
||||
impl PermissionHookTransport for ToolServerPermissionTransport {
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String> {
|
||||
let start = std::time::Instant::now();
|
||||
let Some(server) = self.server.upgrade() else {
|
||||
PERMISSION_REPLY_DURATION
|
||||
.with_label_values(&["error"])
|
||||
.observe(start.elapsed().as_secs_f64());
|
||||
return Err("tool server gone (weak upgrade failed)".to_owned());
|
||||
};
|
||||
let raw = server
|
||||
.request_hook(
|
||||
self.session_id.clone(),
|
||||
PERMISSION_REQUEST_KIND.to_owned(),
|
||||
payload,
|
||||
)
|
||||
.await;
|
||||
let outcome = match &raw {
|
||||
Ok(_) => "ok",
|
||||
Err(e) => {
|
||||
if is_timeout_err(&e.to_string()) {
|
||||
PERMISSION_TIMEOUT_TOTAL.inc();
|
||||
}
|
||||
"error"
|
||||
}
|
||||
};
|
||||
PERMISSION_REPLY_DURATION
|
||||
.with_label_values(&[outcome])
|
||||
.observe(start.elapsed().as_secs_f64());
|
||||
raw.map_err(|e| e.to_string())
|
||||
}
|
||||
}
|
||||
fn scope_for_access(access: &AccessKind) -> &'static str {
|
||||
match access {
|
||||
AccessKind::Bash(_) | AccessKind::Edit(_) | AccessKind::MCPTool { .. } => "write",
|
||||
AccessKind::Read(_)
|
||||
| AccessKind::Grep { .. }
|
||||
| AccessKind::WebFetch(_)
|
||||
| AccessKind::WebSearch(_) => "read",
|
||||
}
|
||||
}
|
||||
fn describe_access(access: &AccessKind) -> String {
|
||||
match access {
|
||||
AccessKind::Bash(_) => "Run a terminal command".to_owned(),
|
||||
AccessKind::Edit(path) => format!("Edit {path}"),
|
||||
AccessKind::MCPTool { name, .. } => format!("Run MCP tool {name}"),
|
||||
AccessKind::WebFetch(url) => format!("Fetch {url}"),
|
||||
AccessKind::WebSearch(query) => format!("Search the web for {query}"),
|
||||
AccessKind::Read(_) => "Read a file".to_owned(),
|
||||
AccessKind::Grep { .. } => "Search file contents".to_owned(),
|
||||
}
|
||||
}
|
||||
/// Build the server → chat `permission_request` payload. The field set matches
|
||||
/// chat's `PermissionRequestPayload` parser: `tool_call_id`, `tool_name`,
|
||||
/// `description`, `scope`, and the bash/edit context.
|
||||
pub(crate) fn build_permission_payload(access: &AccessKind, tool_call_id: &str) -> Value {
|
||||
let mut payload = serde_json::json!(
|
||||
{ "tool_call_id" : tool_call_id, "tool_name" : tool_name_for_access(access),
|
||||
"description" : describe_access(access), "scope" : scope_for_access(access), }
|
||||
);
|
||||
if let Some(map) = payload.as_object_mut() {
|
||||
match access {
|
||||
AccessKind::Bash(command) => {
|
||||
map.insert("bash_command".to_owned(), Value::from(command.clone()));
|
||||
}
|
||||
AccessKind::Edit(path) => {
|
||||
map.insert(
|
||||
"edit_file_paths".to_owned(),
|
||||
Value::from(vec![path.clone()]),
|
||||
);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
payload
|
||||
}
|
||||
/// Decode chat's decision reply onto a [`PromptOutcome`]. The reply is chat's
|
||||
/// `permission_answer_to_json` output: `{ "outcome", "scope"?, "followup_message"? }`.
|
||||
/// An unknown / `unspecified` outcome fails closed (reject).
|
||||
pub(crate) fn reply_to_outcome(reply: &Value) -> PromptOutcome {
|
||||
let outcome = match reply.get("outcome") {
|
||||
Some(Value::String(s)) => s.as_str(),
|
||||
Some(Value::Number(n)) => match n.as_i64() {
|
||||
Some(1) => "approve",
|
||||
Some(2) => "reject",
|
||||
Some(3) => "always_approve",
|
||||
Some(4) => "always_reject",
|
||||
_ => "",
|
||||
},
|
||||
_ => "",
|
||||
};
|
||||
let followup = reply
|
||||
.get("followup_message")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.is_empty());
|
||||
match outcome {
|
||||
"approve" => PromptOutcome::AllowOnce,
|
||||
"always_approve" => match scope_kind_value(reply) {
|
||||
Some(("bash_command", Some(value))) => PromptOutcome::AllowAlwaysBashCommand(value),
|
||||
Some(("server_prefix", Some(value))) => PromptOutcome::AllowAlwaysMcpServer(value),
|
||||
Some(("domain", Some(value))) => PromptOutcome::AllowAlwaysDomain(value),
|
||||
_ => PromptOutcome::AllowAlways,
|
||||
},
|
||||
"reject" => match followup {
|
||||
Some(message) => PromptOutcome::FollowupMessage(message.to_owned()),
|
||||
None => PromptOutcome::RejectOnce,
|
||||
},
|
||||
"always_reject" => match scope_kind_value(reply) {
|
||||
Some(("bash_command", Some(value))) => PromptOutcome::RejectAlwaysBashCommand(value),
|
||||
_ => PromptOutcome::RejectOnce,
|
||||
},
|
||||
"cancelled" => PromptOutcome::Cancelled,
|
||||
_ => PromptOutcome::RejectOnce,
|
||||
}
|
||||
}
|
||||
fn scope_kind_value(reply: &Value) -> Option<(&str, Option<String>)> {
|
||||
let scope = reply.get("scope")?;
|
||||
let kind = scope.get("kind").and_then(Value::as_str)?;
|
||||
let value = scope
|
||||
.get("value")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_owned);
|
||||
Some((kind, value))
|
||||
}
|
||||
/// Map a hub-served tool name + JSON args onto an [`AccessKind`] for the
|
||||
/// permission gate in [`crate::hub::SessionRoutedToolHandler`]. Returns `None`
|
||||
/// for tools that never need a user prompt (reads / todos / dynamic).
|
||||
pub fn access_kind_for_hub_tool(tool_name: &str, args: &Value) -> Option<AccessKind> {
|
||||
let name = tool_name.rsplit(':').next().unwrap_or(tool_name);
|
||||
let name = name.strip_prefix("GrokBuild:").unwrap_or(name);
|
||||
match name {
|
||||
"run_terminal_command" | "run_terminal_cmd" | "bash" | "shell" => {
|
||||
let cmd = args
|
||||
.get("command")
|
||||
.or_else(|| args.get("full_command"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
Some(AccessKind::Bash(cmd))
|
||||
}
|
||||
"search_replace" | "hashline_edit" => {
|
||||
let path = args
|
||||
.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("unknown")
|
||||
.to_owned();
|
||||
Some(AccessKind::Edit(path))
|
||||
}
|
||||
"write" | "write_file" => {
|
||||
let path = args
|
||||
.get("file_path")
|
||||
.or_else(|| args.get("path"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("unknown")
|
||||
.to_owned();
|
||||
Some(AccessKind::Edit(path))
|
||||
}
|
||||
"apply_patch" => Some(AccessKind::Edit("apply_patch".to_owned())),
|
||||
"web_fetch" => {
|
||||
let url = args
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_owned();
|
||||
Some(AccessKind::WebFetch(url))
|
||||
}
|
||||
n if n.contains("__") || n.starts_with("mcp") => Some(AccessKind::MCPTool {
|
||||
name: tool_name.to_owned(),
|
||||
input: args.clone(),
|
||||
}),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
/// Whether a [`PromptOutcome`] allows the tool call to proceed.
|
||||
pub fn prompt_outcome_allows(outcome: &PromptOutcome) -> bool {
|
||||
matches!(
|
||||
outcome,
|
||||
PromptOutcome::AllowOnce
|
||||
| PromptOutcome::AllowAlways
|
||||
| PromptOutcome::AllowEditsForSession
|
||||
| PromptOutcome::AllowAlwaysBashCommand(_)
|
||||
| PromptOutcome::AllowAlwaysDomain(_)
|
||||
| PromptOutcome::AllowAlwaysMcpTool(_)
|
||||
| PromptOutcome::AllowAlwaysMcpServer(_)
|
||||
)
|
||||
}
|
||||
/// Request a permission decision from chat over `transport` and map the reply
|
||||
/// to a [`PromptOutcome`]. A transport error fails closed (the manager turns an
|
||||
/// `Error` outcome into a reject) so a lost server connection never silently runs
|
||||
/// a guarded tool.
|
||||
pub async fn request_permission_via_hub(
|
||||
transport: &dyn PermissionHookTransport,
|
||||
access: &AccessKind,
|
||||
tool_call_id: &str,
|
||||
) -> PromptOutcome {
|
||||
let payload = build_permission_payload(access, tool_call_id);
|
||||
match transport.request_permission(payload).await {
|
||||
Ok(reply) => match reply_to_outcome(&reply) {
|
||||
PromptOutcome::AllowAlways if matches!(access, AccessKind::Edit(_)) => {
|
||||
PromptOutcome::AllowEditsForSession
|
||||
}
|
||||
other => other,
|
||||
},
|
||||
Err(e) => {
|
||||
tracing::error!(error = % e, "hub permission request failed; rejecting");
|
||||
PromptOutcome::Error(format!("hub permission request failed: {e}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::sync::Mutex;
|
||||
/// Pins the current SDK timeout wording the classifier matches on.
|
||||
#[test]
|
||||
fn is_timeout_err_matches_backstop_wording_only() {
|
||||
assert!(is_timeout_err("request timed out after 600s"));
|
||||
assert!(is_timeout_err("request timed out after 600.0s"));
|
||||
assert!(!is_timeout_err("connection lost"));
|
||||
assert!(!is_timeout_err("tool server gone (weak upgrade failed)"));
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_bash_carries_command_and_write_scope() {
|
||||
let payload = build_permission_payload(&AccessKind::Bash("rm -rf /tmp/x".into()), "tc-1");
|
||||
assert_eq!(payload["tool_call_id"], "tc-1");
|
||||
assert_eq!(payload["tool_name"], "run_terminal_command");
|
||||
assert_eq!(payload["description"], "Run a terminal command");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert_eq!(payload["bash_command"], "rm -rf /tmp/x");
|
||||
assert!(payload.get("edit_file_paths").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_edit_carries_file_paths() {
|
||||
let payload = build_permission_payload(&AccessKind::Edit("src/main.rs".into()), "tc-2");
|
||||
assert_eq!(payload["tool_name"], "search_replace");
|
||||
assert_eq!(payload["description"], "Edit src/main.rs");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert_eq!(
|
||||
payload["edit_file_paths"],
|
||||
serde_json::json!(["src/main.rs"])
|
||||
);
|
||||
assert!(payload.get("bash_command").is_none());
|
||||
assert!(payload.get("edit_kind").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn payload_for_mcp_has_no_tool_context() {
|
||||
let payload = build_permission_payload(
|
||||
&AccessKind::MCPTool {
|
||||
name: "linear__list".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
"tc-3",
|
||||
);
|
||||
assert_eq!(payload["tool_name"], "mcp:linear__list");
|
||||
assert_eq!(payload["description"], "Run MCP tool linear__list");
|
||||
assert_eq!(payload["scope"], "write");
|
||||
assert!(payload.get("bash_command").is_none());
|
||||
assert!(payload.get("edit_file_paths").is_none());
|
||||
}
|
||||
#[test]
|
||||
fn reply_outcomes_map_to_prompt_outcomes() {
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "approve" })),
|
||||
PromptOutcome::AllowOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "reject" })),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "cancelled" })),
|
||||
PromptOutcome::Cancelled
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "unspecified"
|
||||
})),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({})),
|
||||
PromptOutcome::RejectOnce
|
||||
));
|
||||
}
|
||||
#[test]
|
||||
fn reject_with_followup_routes_message_to_model() {
|
||||
let reply = serde_json::json!(
|
||||
{ "outcome" : "reject", "followup_message" : "use cargo instead" }
|
||||
);
|
||||
match reply_to_outcome(&reply) {
|
||||
PromptOutcome::FollowupMessage(m) => assert_eq!(m, "use cargo instead"),
|
||||
other => panic!("expected FollowupMessage, got {other:?}"),
|
||||
}
|
||||
}
|
||||
#[test]
|
||||
fn always_approve_maps_scope_to_persistent_outcome() {
|
||||
let bash = serde_json::json!(
|
||||
{ "outcome" : "always_approve", "scope" : { "kind" : "bash_command", "value"
|
||||
: "cargo build" }, }
|
||||
);
|
||||
match reply_to_outcome(&bash) {
|
||||
PromptOutcome::AllowAlwaysBashCommand(v) => assert_eq!(v, "cargo build"),
|
||||
other => panic!("expected AllowAlwaysBashCommand, got {other:?}"),
|
||||
}
|
||||
let server = serde_json::json!(
|
||||
{ "outcome" : "always_approve", "scope" : { "kind" : "server_prefix", "value"
|
||||
: "linear" }, }
|
||||
);
|
||||
match reply_to_outcome(&server) {
|
||||
PromptOutcome::AllowAlwaysMcpServer(v) => assert_eq!(v, "linear"),
|
||||
other => panic!("expected AllowAlwaysMcpServer, got {other:?}"),
|
||||
}
|
||||
assert!(matches!(
|
||||
reply_to_outcome(&serde_json::json!({ "outcome" : "always_approve"
|
||||
})),
|
||||
PromptOutcome::AllowAlways
|
||||
));
|
||||
}
|
||||
#[test]
|
||||
fn always_reject_with_bash_scope_persists_the_denied_prefix() {
|
||||
let reply = serde_json::json!(
|
||||
{ "outcome" : "always_reject", "scope" : { "kind" : "bash_command", "value" :
|
||||
"curl" }, }
|
||||
);
|
||||
match reply_to_outcome(&reply) {
|
||||
PromptOutcome::RejectAlwaysBashCommand(v) => assert_eq!(v, "curl"),
|
||||
other => panic!("expected RejectAlwaysBashCommand, got {other:?}"),
|
||||
}
|
||||
}
|
||||
struct StubTransport {
|
||||
reply: Result<Value, String>,
|
||||
seen: Mutex<Option<Value>>,
|
||||
}
|
||||
#[async_trait]
|
||||
impl PermissionHookTransport for StubTransport {
|
||||
async fn request_permission(&self, payload: Value) -> Result<Value, String> {
|
||||
*self.seen.lock().unwrap() = Some(payload);
|
||||
self.reply.clone()
|
||||
}
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn request_sends_payload_and_decodes_reply() {
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Bash("ls -la".into()), "tc-7")
|
||||
.await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowOnce));
|
||||
let seen = transport
|
||||
.seen
|
||||
.lock()
|
||||
.unwrap()
|
||||
.clone()
|
||||
.expect("payload sent");
|
||||
assert_eq!(seen["tool_call_id"], "tc-7");
|
||||
assert_eq!(seen["bash_command"], "ls -la");
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn transport_error_fails_closed() {
|
||||
let transport = StubTransport {
|
||||
reply: Err("connection lost".to_owned()),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-8").await;
|
||||
assert!(matches!(outcome, PromptOutcome::Error(_)));
|
||||
}
|
||||
#[tokio::test]
|
||||
async fn edit_always_approve_maps_to_session_scope() {
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "always_approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome =
|
||||
request_permission_via_hub(&transport, &AccessKind::Edit("a.rs".into()), "tc-9").await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowEditsForSession));
|
||||
let transport = StubTransport {
|
||||
reply: Ok(serde_json::json!({ "outcome" : "always_approve" })),
|
||||
seen: Mutex::new(None),
|
||||
};
|
||||
let outcome = request_permission_via_hub(
|
||||
&transport,
|
||||
&AccessKind::MCPTool {
|
||||
name: "x".into(),
|
||||
input: serde_json::Value::Null,
|
||||
},
|
||||
"tc-10",
|
||||
)
|
||||
.await;
|
||||
assert!(matches!(outcome, PromptOutcome::AllowAlways));
|
||||
}
|
||||
#[test]
|
||||
fn hitl_permission_live_defaults_off_without_env() {
|
||||
if std::env::var(HITL_PERMISSION_LIVE_ENV).is_err() {
|
||||
assert!(!hitl_permission_live_enabled());
|
||||
}
|
||||
}
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,158 +0,0 @@
|
||||
//! Startup janitor for workspace-owned on-disk session state.
|
||||
|
||||
use std::path::Path;
|
||||
use std::time::Duration;
|
||||
|
||||
/// Default maximum age for a per-session state directory before the
|
||||
/// [`cleanup_stale_sessions`] janitor reclaims it (7 days).
|
||||
pub const DEFAULT_SESSION_MAX_AGE: Duration = Duration::from_secs(7 * 24 * 60 * 60);
|
||||
|
||||
/// Remove per-session state directories under `<workspace_home>/sessions/`
|
||||
/// whose mtime is older than `max_age`, bounding the unbounded growth a
|
||||
/// long-lived workspace (or reused sandbox) would otherwise accumulate.
|
||||
///
|
||||
/// A directory's mtime advances on every atomic-rename persistence write (the
|
||||
/// rename mutates the directory entry), so it tracks last activity closely
|
||||
/// enough for a best-effort reclaim. All errors are swallowed — a startup
|
||||
/// janitor must never fail boot. Only directories with a resolvable, expired
|
||||
/// mtime are removed; stray files and future-mtime entries are left untouched.
|
||||
pub async fn cleanup_stale_sessions(workspace_home: &Path, max_age: Duration) {
|
||||
let sessions_dir = workspace_home.join("sessions");
|
||||
let Ok(mut entries) = tokio::fs::read_dir(&sessions_dir).await else {
|
||||
return; // No sessions dir yet (first boot).
|
||||
};
|
||||
|
||||
let mut removed = 0u32;
|
||||
while let Ok(Some(entry)) = entries.next_entry().await {
|
||||
let Ok(ft) = entry.file_type().await else {
|
||||
continue;
|
||||
};
|
||||
if !ft.is_dir() {
|
||||
continue;
|
||||
}
|
||||
let path = entry.path();
|
||||
if let Ok(metadata) = tokio::fs::metadata(&path).await
|
||||
&& let Ok(modified) = metadata.modified()
|
||||
&& let Ok(age) = modified.elapsed()
|
||||
&& age > max_age
|
||||
{
|
||||
match tokio::fs::remove_dir_all(&path).await {
|
||||
Ok(()) => {
|
||||
removed += 1;
|
||||
tracing::info!(
|
||||
path = %path.display(),
|
||||
age_secs = age.as_secs(),
|
||||
"cleanup_stale_sessions: removed stale session dir"
|
||||
);
|
||||
}
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
path = %path.display(),
|
||||
error = %e,
|
||||
"cleanup_stale_sessions: failed to remove stale session dir"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if removed > 0 {
|
||||
tracing::info!(
|
||||
removed,
|
||||
"cleanup_stale_sessions: stale session-dir sweep complete"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
/// A session dir older than `max_age` is removed.
|
||||
#[tokio::test]
|
||||
async fn cleanup_removes_stale_session_dir() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let stale = home.path().join("sessions").join("sess-old");
|
||||
std::fs::create_dir_all(&stale).unwrap();
|
||||
std::fs::write(stale.join("tool_state.json"), b"{}").unwrap();
|
||||
// Guarantee `age > Duration::ZERO` regardless of mtime resolution.
|
||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
|
||||
assert!(
|
||||
!stale.exists(),
|
||||
"a session dir older than max_age must be removed"
|
||||
);
|
||||
}
|
||||
|
||||
/// A session dir younger than `max_age` is kept.
|
||||
#[tokio::test]
|
||||
async fn cleanup_keeps_fresh_session_dir() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let fresh = home.path().join("sessions").join("sess-new");
|
||||
std::fs::create_dir_all(&fresh).unwrap();
|
||||
std::fs::write(fresh.join("tool_state.json"), b"{}").unwrap();
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::from_secs(3600)).await;
|
||||
|
||||
assert!(
|
||||
fresh.exists(),
|
||||
"a session dir younger than max_age must be kept"
|
||||
);
|
||||
assert!(
|
||||
fresh.join("tool_state.json").exists(),
|
||||
"a kept session dir must retain its contents"
|
||||
);
|
||||
}
|
||||
|
||||
/// No `sessions/` directory (first boot) is a silent no-op, not a panic.
|
||||
#[tokio::test]
|
||||
async fn cleanup_missing_sessions_dir_is_noop() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
assert!(home.path().exists(), "home is left untouched");
|
||||
}
|
||||
|
||||
/// Stray files under `sessions/` are never removed — directories only.
|
||||
#[tokio::test]
|
||||
async fn cleanup_ignores_non_dir_entries() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let sessions = home.path().join("sessions");
|
||||
std::fs::create_dir_all(&sessions).unwrap();
|
||||
let stray = sessions.join("stray.txt");
|
||||
std::fs::write(&stray, b"not a session dir").unwrap();
|
||||
tokio::time::sleep(Duration::from_millis(15)).await;
|
||||
|
||||
cleanup_stale_sessions(home.path(), Duration::ZERO).await;
|
||||
|
||||
assert!(
|
||||
stray.exists(),
|
||||
"stray files under sessions/ must never be removed"
|
||||
);
|
||||
}
|
||||
|
||||
/// Mixed sweep: the `max_age` comparison is per-entry, not all-or-nothing.
|
||||
#[tokio::test]
|
||||
async fn cleanup_removes_only_expired_dirs() {
|
||||
let home = tempfile::TempDir::new().unwrap();
|
||||
let sessions = home.path().join("sessions");
|
||||
let old = sessions.join("sess-old");
|
||||
std::fs::create_dir_all(&old).unwrap();
|
||||
// Wide gap (100ms) vs a 50ms threshold so neither side is sensitive to
|
||||
// scheduler jitter.
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
let max_age = Duration::from_millis(50);
|
||||
let fresh = sessions.join("sess-fresh");
|
||||
std::fs::create_dir_all(&fresh).unwrap();
|
||||
|
||||
cleanup_stale_sessions(home.path(), max_age).await;
|
||||
|
||||
assert!(!old.exists(), "the >max_age dir must be removed");
|
||||
assert!(
|
||||
fresh.exists(),
|
||||
"the <max_age dir must survive the same sweep"
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -1,227 +0,0 @@
|
||||
//! `WorkspaceError` <-> wire-code mapping for the workspace RPC envelope
|
||||
//! (the envelope types are canonical in `kigi_workspace_types::rpc`).
|
||||
//!
|
||||
//! `error_code` uses a non-wildcard match so the compiler enforces
|
||||
//! coverage of new `WorkspaceError` variants.
|
||||
|
||||
pub use kigi_workspace_types::rpc::{RpcEnvelope, RpcError};
|
||||
|
||||
use crate::error::WorkspaceError;
|
||||
|
||||
/// Build an error envelope from a `WorkspaceError`.
|
||||
pub fn envelope_err<T>(error: &WorkspaceError) -> RpcEnvelope<T> {
|
||||
RpcEnvelope::err_parts(error_code(error), error.to_string())
|
||||
}
|
||||
|
||||
/// Map a `WorkspaceError` to its wire code string.
|
||||
///
|
||||
/// Uses an exhaustive match with no wildcard -- the compiler will
|
||||
/// error when new variants are added to `WorkspaceError`, forcing
|
||||
/// the implementer to assign a wire code.
|
||||
pub fn error_code(err: &WorkspaceError) -> &'static str {
|
||||
match err {
|
||||
WorkspaceError::ParentSessionNotFound(_) => "parent_session_not_found",
|
||||
WorkspaceError::SessionNotFound(_) => "session_not_found",
|
||||
WorkspaceError::SessionAlreadyExists(_) => "session_already_exists",
|
||||
WorkspaceError::EmptyAgentId => "empty_agent_id",
|
||||
WorkspaceError::CannotDropMainSession => "cannot_drop_main",
|
||||
WorkspaceError::Finalize(_) => "finalize",
|
||||
WorkspaceError::CapabilityWidening { .. } => "capability_widening",
|
||||
WorkspaceError::Unauthorized { .. } => "unauthorized",
|
||||
WorkspaceError::TurnActive(_) => kigi_workspace_types::rpc::envelope::TURN_ACTIVE,
|
||||
WorkspaceError::MaxDepthExceeded { .. } => "max_depth_exceeded",
|
||||
WorkspaceError::JoinError(_) => "join_error",
|
||||
WorkspaceError::InvalidHunkAction(_) => "invalid_hunk_action",
|
||||
WorkspaceError::HunkActionFailed(_) => "hunk_action_failed",
|
||||
WorkspaceError::HubError(_) => "hub_error",
|
||||
WorkspaceError::DeployError { kind, .. } => kind.wire_code(),
|
||||
WorkspaceError::ShuttingDown => "shutting_down",
|
||||
WorkspaceError::ToolsetExternallyOwned(_) => "toolset_externally_owned",
|
||||
}
|
||||
}
|
||||
|
||||
/// Map a wire [`RpcError`] back to a [`WorkspaceError`].
|
||||
///
|
||||
/// Known codes are mapped to their specific variants. Unknown codes
|
||||
/// degrade gracefully to `WorkspaceError::HubError`, ensuring
|
||||
/// forward compatibility when a newer workspace sends codes an older
|
||||
/// shell does not recognise.
|
||||
///
|
||||
/// # Intentional degradation
|
||||
///
|
||||
/// The structured variants `CapabilityWidening`, `Unauthorized`, and
|
||||
/// `MaxDepthExceeded` carry multiple fields that are not preserved in
|
||||
/// the wire `message` string. These are mapped to `HubError` on the
|
||||
/// deserializing side because reconstructing the original struct fields
|
||||
/// from a flattened `Display` string would be fragile and error-prone.
|
||||
/// Callers that need to distinguish these errors can match on the
|
||||
/// `HubError` message which contains the original error code as a
|
||||
/// prefix (e.g. `"capability_widening: ..."`).
|
||||
pub fn rpc_error_to_workspace(err: RpcError) -> WorkspaceError {
|
||||
if let Some(kind) = kigi_workspace_types::rpc::deploy::DeployError::from_wire_code(&err.code) {
|
||||
return WorkspaceError::DeployError {
|
||||
kind,
|
||||
message: err.message,
|
||||
};
|
||||
}
|
||||
match err.code.as_str() {
|
||||
"parent_session_not_found" => WorkspaceError::ParentSessionNotFound(err.message),
|
||||
"session_not_found" => WorkspaceError::SessionNotFound(err.message),
|
||||
"session_already_exists" => WorkspaceError::SessionAlreadyExists(err.message),
|
||||
"empty_agent_id" => WorkspaceError::EmptyAgentId,
|
||||
"cannot_drop_main" => WorkspaceError::CannotDropMainSession,
|
||||
"finalize" => WorkspaceError::Finalize(err.message),
|
||||
"capability_widening" => {
|
||||
WorkspaceError::HubError(format!("capability_widening: {}", err.message))
|
||||
}
|
||||
"unauthorized" => WorkspaceError::HubError(format!("unauthorized: {}", err.message)),
|
||||
"turn_active" => WorkspaceError::TurnActive(err.message),
|
||||
"max_depth_exceeded" => {
|
||||
WorkspaceError::HubError(format!("max_depth_exceeded: {}", err.message))
|
||||
}
|
||||
"join_error" => WorkspaceError::JoinError(err.message),
|
||||
"invalid_hunk_action" => WorkspaceError::InvalidHunkAction(err.message),
|
||||
"hunk_action_failed" => WorkspaceError::HunkActionFailed(err.message),
|
||||
"hub_error" => WorkspaceError::HubError(err.message),
|
||||
"shutting_down" => WorkspaceError::ShuttingDown,
|
||||
"toolset_externally_owned" => WorkspaceError::ToolsetExternallyOwned(err.message),
|
||||
unknown => {
|
||||
WorkspaceError::HubError(format!("unknown error code: {unknown}: {}", err.message))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::capability::CapabilityMode;
|
||||
|
||||
/// Verify round-trip fidelity for every `WorkspaceError` variant.
|
||||
#[test]
|
||||
fn error_code_round_trip_all_variants() {
|
||||
let mut variants: Vec<WorkspaceError> = vec![
|
||||
WorkspaceError::ParentSessionNotFound("p".into()),
|
||||
WorkspaceError::SessionNotFound("s".into()),
|
||||
WorkspaceError::SessionAlreadyExists("s".into()),
|
||||
WorkspaceError::EmptyAgentId,
|
||||
WorkspaceError::CannotDropMainSession,
|
||||
WorkspaceError::Finalize("f".into()),
|
||||
WorkspaceError::CapabilityWidening {
|
||||
parent: CapabilityMode::ReadOnly,
|
||||
child: CapabilityMode::All,
|
||||
},
|
||||
WorkspaceError::Unauthorized {
|
||||
caller: "a".into(),
|
||||
target: "b".into(),
|
||||
},
|
||||
WorkspaceError::TurnActive("s".into()),
|
||||
WorkspaceError::MaxDepthExceeded { parent: "p".into() },
|
||||
WorkspaceError::JoinError("j".into()),
|
||||
WorkspaceError::InvalidHunkAction("h".into()),
|
||||
WorkspaceError::HunkActionFailed("h".into()),
|
||||
WorkspaceError::HubError("hub".into()),
|
||||
WorkspaceError::ShuttingDown,
|
||||
WorkspaceError::ToolsetExternallyOwned("s".into()),
|
||||
];
|
||||
variants.extend(
|
||||
kigi_workspace_types::rpc::deploy::DeployError::ALL
|
||||
.into_iter()
|
||||
.map(|kind| WorkspaceError::DeployError {
|
||||
kind,
|
||||
message: "deploy".into(),
|
||||
}),
|
||||
);
|
||||
|
||||
for err in &variants {
|
||||
let code = error_code(err);
|
||||
assert!(!code.is_empty(), "code must not be empty for {err:?}");
|
||||
|
||||
// Round-trip through RpcError
|
||||
let rpc_err = RpcError {
|
||||
code: code.to_owned(),
|
||||
message: err.to_string(),
|
||||
};
|
||||
let recovered = rpc_error_to_workspace(rpc_err);
|
||||
// The recovered error's code should match the original code
|
||||
let recovered_code = error_code(&recovered);
|
||||
|
||||
// Structured variants (CapabilityWidening, Unauthorized,
|
||||
// MaxDepthExceeded) lose their fields on the wire and
|
||||
// degrade to HubError, which is the expected behavior.
|
||||
// Their error messages are preserved in the HubError string.
|
||||
match err {
|
||||
WorkspaceError::CapabilityWidening { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("capability_widening"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
WorkspaceError::Unauthorized { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("unauthorized"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
WorkspaceError::MaxDepthExceeded { .. } => {
|
||||
assert_eq!(recovered_code, "hub_error");
|
||||
let msg = recovered.to_string();
|
||||
assert!(
|
||||
msg.contains("max_depth_exceeded"),
|
||||
"degraded error should contain original code: {msg}"
|
||||
);
|
||||
}
|
||||
_ => {
|
||||
assert_eq!(
|
||||
recovered_code, code,
|
||||
"round-trip mismatch for {err:?}: got code {recovered_code}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify unknown codes degrade to HubError.
|
||||
#[test]
|
||||
fn unknown_code_degrades_to_hub_error() {
|
||||
let rpc_err = RpcError {
|
||||
code: "future_new_variant".into(),
|
||||
message: "something new".into(),
|
||||
};
|
||||
let recovered = rpc_error_to_workspace(rpc_err);
|
||||
assert!(matches!(recovered, WorkspaceError::HubError(_)));
|
||||
let msg = recovered.to_string();
|
||||
assert!(msg.contains("future_new_variant"));
|
||||
}
|
||||
|
||||
/// Verify serde round-trip of RpcEnvelope.
|
||||
#[test]
|
||||
fn envelope_serde_round_trip_ok() {
|
||||
let env: RpcEnvelope<String> = RpcEnvelope::ok("hello".into());
|
||||
let json = serde_json::to_value(&env).unwrap();
|
||||
let recovered: RpcEnvelope<String> = serde_json::from_value(json).unwrap();
|
||||
match recovered.into_result() {
|
||||
Ok(v) => assert_eq!(v, "hello"),
|
||||
Err(e) => panic!("expected Ok, got {e:?}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Verify serde round-trip of RpcEnvelope error, through the
|
||||
/// `WorkspaceError` mapping in both directions.
|
||||
#[test]
|
||||
fn envelope_serde_round_trip_err() {
|
||||
let err = WorkspaceError::SessionNotFound("ghost".into());
|
||||
let env: RpcEnvelope<String> = envelope_err(&err);
|
||||
let json = serde_json::to_value(&env).unwrap();
|
||||
let recovered: RpcEnvelope<String> = serde_json::from_value(json).unwrap();
|
||||
match recovered.into_result().map_err(rpc_error_to_workspace) {
|
||||
Ok(_) => panic!("expected Err"),
|
||||
Err(e) => {
|
||||
assert_eq!(error_code(&e), "session_not_found");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user