ZacharyZhang-NY
ad3840f9ec
kigi allowed loopback unconditionally and missed several non-public ranges, and the SSRF check ran only on the initial URL. Policy (ssrf.rs): - loopback is blocked unless `[toolset.web_fetch] allow_local` (or KIGI_WEB_FETCH_ALLOW_LOCAL) is on, AND the URL names it explicitly, so a public name resolving to loopback stays blocked (DNS rebinding) - add 0.0.0.0/8, 100.64/10, 192.0.0.0/24, TEST-NET-1/2/3, 198.18/15, 240/4, IPv6 site-local and documentation prefixes - inherit the IPv4 verdict through mapped, compatible, NAT64 and 6to4 wrappers; network-specific NAT64 prefixes remain uncovered (see doc) Plumbing (client.rs), where the exploitable half lived: - re-check every redirect hop, not just the first - compare hosts exactly; a `www` sibling has its own A records, so it is a cross-host redirect rather than an auto-followed hop - run the check before the fetch service, so a blocked URL is never posted to an endpoint that egresses elsewhere - exempt explicit local hosts from the https upgrade and from the single-label filter, and re-upgrade each followed hop Wiring: allow_local reaches WebFetchParams from both construction paths; documented in the config guide and the README env table.
840 lines
33 KiB
Rust
840 lines
33 KiB
Rust
//! Tool config resolution pipeline.
|
|
//!
|
|
//! Four-step resolution:
|
|
//! 1. `effective_tool_config = config.tool_config.unwrap_or_else(|| parent.effective_tool_config.clone())`
|
|
//! 2. `merged = merge_mcp_tools(effective_tool_config, shared.mcp_servers.snapshot())`
|
|
//! 3. `filtered = config.capability_mode.filter(merged)`
|
|
//! 4. `toolset = build_finalized_toolset(filtered, &session.cwd, &session.session_env, ...)`
|
|
use crate::capability::{CapabilityMode, kind_allowed};
|
|
use crate::config::SessionContextFactory;
|
|
use crate::error::{WorkspaceError, WorkspaceResult};
|
|
use kigi_tools::registry::types::{
|
|
FinalizedToolset, ToolConfig, ToolRegistryBuilder, ToolServerConfig,
|
|
};
|
|
use kigi_tools::types::tool::ToolKind;
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
/// Create-shaped entry of the resolution pipeline: run
|
|
/// [`resolve_session_toolset_rebuild`] around a FRESH factory-built
|
|
/// session-lifetime terminal backend, and return that backend so the caller
|
|
/// can store it on the session it is creating.
|
|
pub(crate) fn resolve_session_toolset(
|
|
effective_tool_config: ToolServerConfig,
|
|
capability_mode: CapabilityMode,
|
|
mcp_snapshot: &[ToolConfig],
|
|
cwd: PathBuf,
|
|
session_env: Arc<HashMap<String, String>>,
|
|
session_id: &str,
|
|
factory: &dyn SessionContextFactory,
|
|
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
|
|
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
|
|
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
|
|
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
|
|
) -> WorkspaceResult<(
|
|
ToolServerConfig,
|
|
Arc<FinalizedToolset>,
|
|
crate::config::SessionTerminalBackend,
|
|
)> {
|
|
let terminal_backend = factory.build_terminal_backend();
|
|
let (effective, toolset) = resolve_session_toolset_rebuild(
|
|
effective_tool_config,
|
|
capability_mode,
|
|
mcp_snapshot,
|
|
cwd,
|
|
session_env,
|
|
session_id,
|
|
factory,
|
|
local_registry,
|
|
lsp,
|
|
viewer_ctx,
|
|
notification_handle,
|
|
terminal_backend.backend().clone(),
|
|
)?;
|
|
Ok((effective, toolset, terminal_backend))
|
|
}
|
|
/// Rebuild-shaped entry: run steps 2-5 of the resolution pipeline around an
|
|
/// EXISTING session-owned terminal backend. The parameter is non-optional on
|
|
/// purpose: every toolset-swap call site must state which backend it rebuilds
|
|
/// around, so background tasks and shell state can never be orphaned by a
|
|
/// resolve that silently built a fresh backend.
|
|
///
|
|
/// Returns the *unmodified* `effective_tool_config` (step-1 baseline) so
|
|
/// the caller can store it on the session. The FinalizedToolset reflects
|
|
/// MCP merging and capability filtering on top of that baseline.
|
|
///
|
|
/// **MCP-origin `kind: None` tools are dropped under every non-`All`
|
|
/// mode.** Baseline `kind: None` tools are always kept — but before
|
|
/// filtering, kind-less baseline entries whose id the binary's registry
|
|
/// knows get their [`ToolKind`] backfilled (see [`backfill_tool_kinds`]),
|
|
/// so the capability filter applies to pinned toolsets whose wire entries
|
|
/// cannot carry a kind.
|
|
pub(crate) fn resolve_session_toolset_rebuild(
|
|
effective_tool_config: ToolServerConfig,
|
|
capability_mode: CapabilityMode,
|
|
mcp_snapshot: &[ToolConfig],
|
|
cwd: PathBuf,
|
|
session_env: Arc<HashMap<String, String>>,
|
|
session_id: &str,
|
|
factory: &dyn SessionContextFactory,
|
|
local_registry: Option<kigi_tool_runtime::LocalRegistry>,
|
|
lsp: Option<std::sync::Arc<dyn kigi_tools::implementations::lsp::LspBackend>>,
|
|
viewer_ctx: Option<kigi_tool_runtime::WorkspaceViewerContext>,
|
|
notification_handle: Option<kigi_tools::notification::types::ToolNotificationHandle>,
|
|
terminal_backend: Arc<dyn kigi_tools::computer::types::TerminalBackend>,
|
|
) -> WorkspaceResult<(ToolServerConfig, Arc<FinalizedToolset>)> {
|
|
let mut builder = factory.registry_builder();
|
|
if let Some(lr) = local_registry {
|
|
builder = builder.with_local_registry(lr);
|
|
}
|
|
let baseline = backfill_tool_kinds(&effective_tool_config, &builder.known_tool_kinds());
|
|
let finalize_config = merge_and_filter(&baseline, mcp_snapshot, capability_mode, session_id);
|
|
let mut ctx = factory.build_session_context(session_id, cwd, session_env, terminal_backend);
|
|
if let Some(lsp_handle) = lsp {
|
|
ctx.lsp = Some(lsp_handle);
|
|
}
|
|
if let Some(handle) = notification_handle {
|
|
ctx.notification_handle = handle;
|
|
}
|
|
let toolset = builder
|
|
.finalize_with_trunc_config(
|
|
finalize_config,
|
|
ctx,
|
|
kigi_tools::types::context::TruncationConfig::default(),
|
|
viewer_ctx,
|
|
)
|
|
.map_err(|errs| {
|
|
let summary: Vec<String> = errs.iter().map(|e| e.summary()).collect();
|
|
WorkspaceError::Finalize(summary.join("; "))
|
|
})?;
|
|
Ok((effective_tool_config, Arc::new(toolset)))
|
|
}
|
|
/// Backfill `kind: None` baseline entries from the binary's own registry
|
|
/// (fully-qualified id -> declared [`ToolKind`]).
|
|
///
|
|
/// Ids unknown to the registry stay `None` and keep the always-kept
|
|
/// baseline behavior (ad-hoc `ToolConfig::simple` tools). Entries that
|
|
/// already carry a kind are left untouched.
|
|
fn backfill_tool_kinds(
|
|
config: &ToolServerConfig,
|
|
kinds: &HashMap<String, ToolKind>,
|
|
) -> ToolServerConfig {
|
|
ToolServerConfig {
|
|
tools: config
|
|
.tools
|
|
.iter()
|
|
.map(|tool| {
|
|
let mut tool = tool.clone();
|
|
if tool.kind.is_none() {
|
|
tool.kind = kinds.get(&tool.id).copied();
|
|
}
|
|
tool
|
|
})
|
|
.collect(),
|
|
behavior_preset: config.behavior_preset.clone(),
|
|
}
|
|
}
|
|
/// Steps 2-3 of the resolution pipeline, without the `finalize` step:
|
|
///
|
|
/// - **Step 2** -- MCP merge: append MCP-origin tools, skipping ID/name collisions with baseline.
|
|
/// - **Step 3** -- Capability filter: drop tools whose `kind` is not allowed by the mode.
|
|
/// External (MCP) `kind: None` tools are only kept under `CapabilityMode::All`.
|
|
///
|
|
/// Priority on ID/name collision: baseline wins over MCP.
|
|
pub(crate) fn merge_and_filter(
|
|
baseline: &ToolServerConfig,
|
|
mcp_snapshot: &[ToolConfig],
|
|
mode: CapabilityMode,
|
|
session_id: &str,
|
|
) -> ToolServerConfig {
|
|
if mcp_snapshot.is_empty() {
|
|
return mode.filter(baseline);
|
|
}
|
|
let baseline_ids: std::collections::HashSet<&str> =
|
|
baseline.tools.iter().map(|t| t.id.as_str()).collect();
|
|
let mut taken_names: std::collections::HashSet<String> = baseline
|
|
.tools
|
|
.iter()
|
|
.map(|t| {
|
|
let unqualified = t.id.rsplit_once(':').map_or(t.id.as_str(), |(_, n)| n);
|
|
t.resolve_client_name(unqualified)
|
|
})
|
|
.collect();
|
|
let mut tagged: Vec<(ToolConfig, bool)> =
|
|
baseline.tools.iter().cloned().map(|t| (t, false)).collect();
|
|
for mcp_tool in mcp_snapshot {
|
|
if baseline_ids.contains(mcp_tool.id.as_str()) {
|
|
tracing::warn!(
|
|
mcp_id = % mcp_tool.id, session = % session_id,
|
|
"skipping MCP tool: id collides with baseline"
|
|
);
|
|
continue;
|
|
}
|
|
let client_name = mcp_tool.resolve_client_name(&mcp_tool.id);
|
|
if !taken_names.insert(client_name.clone()) {
|
|
tracing::warn!(
|
|
mcp_id = % mcp_tool.id, client_name = % client_name, session = %
|
|
session_id,
|
|
"skipping MCP tool: resolved client name collides with another tool"
|
|
);
|
|
continue;
|
|
}
|
|
tagged.push((mcp_tool.clone(), true));
|
|
}
|
|
let kept: Vec<ToolConfig> = tagged
|
|
.into_iter()
|
|
.filter(|(tool, is_external)| match tool.kind {
|
|
Some(k) => kind_allowed(mode, k),
|
|
None if !*is_external => true,
|
|
None => matches!(mode, CapabilityMode::All),
|
|
})
|
|
.map(|(t, _)| t)
|
|
.collect();
|
|
ToolServerConfig {
|
|
tools: kept,
|
|
behavior_preset: baseline.behavior_preset.clone(),
|
|
}
|
|
}
|
|
/// Alias for backward compatibility.
|
|
pub type NoopSessionContextFactory = WorkspaceSessionContextFactory;
|
|
/// Sanitize a `session_id` into a single safe filesystem path segment: chars
|
|
/// outside `[A-Za-z0-9_-]` become `_`, empty becomes `anon`. When any
|
|
/// replacement happened, an 8-hex digest of the ORIGINAL id is appended so the
|
|
/// mapping stays injective — plain substitution would collide distinct ids
|
|
/// (`sess/1` and `sess_1`) into one directory, cross-contaminating
|
|
/// persistence and rehydration. Already-safe ids (the common UUID case)
|
|
/// map to themselves.
|
|
fn sanitize_session_id(session_id: &str) -> String {
|
|
let mut safe = String::with_capacity(session_id.len());
|
|
let mut modified = false;
|
|
for c in session_id.chars() {
|
|
if c.is_ascii_alphanumeric() || c == '_' || c == '-' {
|
|
safe.push(c);
|
|
} else {
|
|
safe.push('_');
|
|
modified = true;
|
|
}
|
|
}
|
|
if safe.is_empty() {
|
|
safe.push_str("anon");
|
|
modified = true;
|
|
}
|
|
if modified {
|
|
let digest = kigi_file_utils::sha256_hex(session_id.as_bytes());
|
|
safe.push('-');
|
|
safe.push_str(&digest[..8]);
|
|
}
|
|
safe
|
|
}
|
|
/// `<root>/sessions/<sanitized_id>`, creating the directory when possible.
|
|
fn ensure_session_dir(root: &std::path::Path, session_id: &str) -> (PathBuf, std::io::Result<()>) {
|
|
let dir = root.join("sessions").join(sanitize_session_id(session_id));
|
|
let created = std::fs::create_dir_all(&dir);
|
|
(dir, created)
|
|
}
|
|
/// Serializes tests (across modules) that mutate the process-global
|
|
/// `KIGI_WORKSPACE_TOOL_STATE_ENABLED`. Aliased to the crate-wide
|
|
/// [`crate::ENV_TEST_LOCK`] so ALL env-mutating tests share ONE lock (the
|
|
/// hazard is the global `environ` array, not the variable's value).
|
|
#[cfg(test)]
|
|
pub(crate) use crate::ENV_TEST_LOCK as TOOL_STATE_ENV_LOCK;
|
|
/// [`SessionContextFactory`] for workspace sessions.
|
|
///
|
|
/// [`SessionContext::session_folder`] is `/tmp/sessions/<sanitized_id>/`
|
|
/// (terminal logs and other tool artifacts — not the project `cwd`).
|
|
///
|
|
/// Terminal backends are persistent-shell [`LocalTerminalBackend`]s, built
|
|
/// once per session by [`build_terminal_backend`] and passed into every
|
|
/// [`build_session_context`] call.
|
|
///
|
|
/// [`build_terminal_backend`]: crate::config::SessionContextFactory::build_terminal_backend
|
|
/// [`build_session_context`]: crate::config::SessionContextFactory::build_session_context
|
|
/// [`LocalTerminalBackend`]: kigi_tools::computer::local::LocalTerminalBackend
|
|
#[derive(Default)]
|
|
pub struct WorkspaceSessionContextFactory;
|
|
impl WorkspaceSessionContextFactory {
|
|
pub fn new() -> Self {
|
|
Self
|
|
}
|
|
/// `/tmp/sessions/<sanitized_id>/` for terminal logs and other tool artifacts.
|
|
fn resolve_session_folder(session_id: &str) -> PathBuf {
|
|
let (dir, created) = ensure_session_dir(std::path::Path::new("/tmp"), session_id);
|
|
if let Err(e) = created {
|
|
tracing::warn!(
|
|
session = % session_id, dir = % dir.display(), error = % e,
|
|
"session_folder: failed to create dir; tools may create it on write"
|
|
);
|
|
}
|
|
dir
|
|
}
|
|
}
|
|
impl SessionContextFactory for WorkspaceSessionContextFactory {
|
|
fn build_session_context(
|
|
&self,
|
|
session_id: &str,
|
|
cwd: PathBuf,
|
|
session_env: Arc<HashMap<String, String>>,
|
|
backend: Arc<dyn kigi_tools::computer::types::TerminalBackend>,
|
|
) -> kigi_tools::registry::types::SessionContext {
|
|
let fs = Arc::new(kigi_tools::computer::local::LocalFs)
|
|
as Arc<dyn kigi_tools::computer::types::AsyncFileSystem>;
|
|
let notification_handle = kigi_tools::notification::ToolNotificationHandle::noop();
|
|
kigi_tools::registry::types::SessionContext {
|
|
backend,
|
|
fs,
|
|
cwd,
|
|
session_folder: Self::resolve_session_folder(session_id),
|
|
session_env,
|
|
notification_handle,
|
|
owner_session_id: None,
|
|
parent_scheduler_handle: None,
|
|
skills: vec![],
|
|
state_path: PathBuf::new(),
|
|
memory_backend: None,
|
|
web_search_config: kigi_tools::implementations::web_search::WebSearchConfig::default(),
|
|
web_fetch_config: build_web_fetch_config(),
|
|
lsp: None,
|
|
app_builder_deployer_config:
|
|
kigi_tools::implementations::kigi::deploy_app::AppBuilderDeployerConfig::default(),
|
|
api_key_provider: None,
|
|
attribution_callback: None,
|
|
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
|
|
}
|
|
}
|
|
fn build_terminal_backend(&self) -> crate::config::SessionTerminalBackend {
|
|
crate::config::SessionTerminalBackend::local(
|
|
kigi_tools::computer::local::LocalTerminalBackend::with_persistent_shell(),
|
|
)
|
|
}
|
|
fn registry_builder(&self) -> ToolRegistryBuilder {
|
|
ToolRegistryBuilder::new()
|
|
}
|
|
fn known_tool_ids(&self) -> Arc<std::collections::HashSet<String>> {
|
|
static IDS: std::sync::LazyLock<Arc<std::collections::HashSet<String>>> =
|
|
std::sync::LazyLock::new(|| Arc::new(ToolRegistryBuilder::new().known_tool_ids()));
|
|
IDS.clone()
|
|
}
|
|
}
|
|
/// Enabled with default params unless `KIGI_DISABLE_WEB_FETCH=1` is set.
|
|
fn build_web_fetch_config() -> kigi_tools::implementations::kigi::web_fetch::WebFetchConfig {
|
|
use kigi_tools::implementations::kigi::web_fetch::{WebFetchConfig, WebFetchParams};
|
|
if std::env::var("KIGI_DISABLE_WEB_FETCH").is_ok_and(|v| v == "1" || v == "true") {
|
|
return WebFetchConfig::Disabled;
|
|
}
|
|
let mut params = WebFetchParams::default();
|
|
if let Ok(proxy) = std::env::var("KIGI_WEB_FETCH_PROXY") {
|
|
params.proxy_endpoint = Some(proxy);
|
|
}
|
|
if kigi_config::env_bool("KIGI_WEB_FETCH_ALLOW_LOCAL") == Some(true) {
|
|
params.allow_local = Some(true);
|
|
}
|
|
WebFetchConfig::Enabled { params }
|
|
}
|
|
#[cfg(any(test, feature = "test-support"))]
|
|
pub mod test_support {
|
|
use crate::config::SessionContextFactory;
|
|
use kigi_tools::computer::local::{LocalFs, LocalTerminalBackend};
|
|
use kigi_tools::notification::ToolNotificationHandle;
|
|
use kigi_tools::registry::types::{
|
|
SessionContext, ToolConfig, ToolRegistryBuilder, ToolServerConfig,
|
|
};
|
|
use kigi_tools::types::tool::ToolKind;
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
use tempfile::TempDir;
|
|
/// Test factory: builds a `SessionContext` rooted at a per-test temp dir.
|
|
pub struct TestSessionContextFactory {
|
|
pub temp: TempDir,
|
|
}
|
|
impl Default for TestSessionContextFactory {
|
|
fn default() -> Self {
|
|
Self::new()
|
|
}
|
|
}
|
|
impl TestSessionContextFactory {
|
|
pub fn new() -> Self {
|
|
Self {
|
|
temp: TempDir::new().expect("create temp dir"),
|
|
}
|
|
}
|
|
}
|
|
impl SessionContextFactory for TestSessionContextFactory {
|
|
fn build_session_context(
|
|
&self,
|
|
session_id: &str,
|
|
cwd: PathBuf,
|
|
session_env: Arc<HashMap<String, String>>,
|
|
backend: Arc<dyn kigi_tools::computer::types::TerminalBackend>,
|
|
) -> SessionContext {
|
|
let session_root = self
|
|
.temp
|
|
.path()
|
|
.join(super::sanitize_session_id(session_id));
|
|
std::fs::create_dir_all(&session_root).expect("create session root");
|
|
SessionContext {
|
|
backend,
|
|
fs: Arc::new(LocalFs),
|
|
cwd,
|
|
session_folder: session_root.clone(),
|
|
session_env,
|
|
notification_handle: ToolNotificationHandle::noop(),
|
|
owner_session_id: None,
|
|
parent_scheduler_handle: None,
|
|
skills: vec![],
|
|
state_path: session_root.join("tool_state.json"),
|
|
memory_backend: None,
|
|
web_search_config: Default::default(),
|
|
web_fetch_config: Default::default(),
|
|
lsp: None,
|
|
app_builder_deployer_config: Default::default(),
|
|
api_key_provider: None,
|
|
attribution_callback: None,
|
|
system_reminder_tag: kigi_tools::reminders::DEFAULT_REMINDER_TAG,
|
|
}
|
|
}
|
|
fn build_terminal_backend(&self) -> crate::config::SessionTerminalBackend {
|
|
crate::config::SessionTerminalBackend::local(LocalTerminalBackend::new())
|
|
}
|
|
fn registry_builder(&self) -> ToolRegistryBuilder {
|
|
ToolRegistryBuilder::new()
|
|
}
|
|
}
|
|
/// `ToolConfig` builder helper.
|
|
pub fn tc(id: &str, kind: Option<ToolKind>) -> ToolConfig {
|
|
ToolConfig {
|
|
id: id.to_owned(),
|
|
params: None,
|
|
name_override: None,
|
|
params_name_overrides: None,
|
|
description_override: None,
|
|
behavior_version: None,
|
|
kind,
|
|
}
|
|
}
|
|
/// Minimal valid `ToolServerConfig` for finalize-time tests.
|
|
pub fn baseline_config() -> ToolServerConfig {
|
|
ToolServerConfig {
|
|
tools: vec![
|
|
tc("Kigi:read_file", Some(ToolKind::Read)),
|
|
tc("Kigi:search_replace", Some(ToolKind::Edit)),
|
|
tc("Kigi:grep", Some(ToolKind::Search)),
|
|
tc("Kigi:list_dir", Some(ToolKind::ListDir)),
|
|
],
|
|
behavior_preset: None,
|
|
}
|
|
}
|
|
}
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::config::SessionContextFactory;
|
|
use kigi_tools::types::tool::ToolKind;
|
|
use std::collections::HashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::Arc;
|
|
fn factory_for_test() -> Arc<dyn SessionContextFactory> {
|
|
Arc::new(test_support::TestSessionContextFactory::new())
|
|
}
|
|
fn empty_env() -> Arc<HashMap<String, String>> {
|
|
Arc::new(HashMap::new())
|
|
}
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_empty_mcp_snapshot_is_noop_for_baseline() {
|
|
let factory = factory_for_test();
|
|
let cwd = PathBuf::from("/tmp");
|
|
let baseline = test_support::baseline_config();
|
|
let baseline_ids: Vec<String> = baseline.tools.iter().map(|t| t.id.clone()).collect();
|
|
let (eff, ts, _backend) = resolve_session_toolset(
|
|
baseline,
|
|
CapabilityMode::ReadWrite,
|
|
&[],
|
|
cwd,
|
|
empty_env(),
|
|
"main",
|
|
factory.as_ref(),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("resolve");
|
|
assert_eq!(
|
|
eff.tools
|
|
.iter()
|
|
.map(|t| t.id.clone())
|
|
.collect::<Vec<String>>(),
|
|
baseline_ids
|
|
);
|
|
assert!(!ts.tool_definitions().is_empty());
|
|
}
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_mcp_merge_dedup_by_id_baseline_wins() {
|
|
let factory = factory_for_test();
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![test_support::tc("Kigi:read_file", Some(ToolKind::Read))],
|
|
behavior_preset: None,
|
|
};
|
|
let mut mcp_dup = test_support::tc("Kigi:read_file", Some(ToolKind::Read));
|
|
mcp_dup.name_override = Some("mcp_read".into());
|
|
let snapshot = vec![mcp_dup];
|
|
let (_eff, ts, _backend) = resolve_session_toolset(
|
|
baseline,
|
|
CapabilityMode::ReadWrite,
|
|
&snapshot,
|
|
PathBuf::from("/tmp"),
|
|
empty_env(),
|
|
"main",
|
|
factory.as_ref(),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("resolve");
|
|
let defs = ts.tool_definitions();
|
|
let names: Vec<&str> = defs.iter().map(|d| d.function.name.as_str()).collect();
|
|
assert!(
|
|
names.contains(&"read_file"),
|
|
"baseline read_file must survive: {names:?}"
|
|
);
|
|
assert!(
|
|
!names.contains(&"mcp_read"),
|
|
"MCP duplicate must be skipped: {names:?}"
|
|
);
|
|
}
|
|
#[test]
|
|
fn backfill_tool_kinds_fills_known_kindless_ids_only() {
|
|
let kinds = HashMap::from([
|
|
("Kigi:search_replace".to_owned(), ToolKind::Edit),
|
|
("Kigi:read_file".to_owned(), ToolKind::Read),
|
|
]);
|
|
let config = ToolServerConfig {
|
|
tools: vec![
|
|
test_support::tc("Kigi:search_replace", None),
|
|
test_support::tc("adhoc.opaque", None),
|
|
test_support::tc("Kigi:read_file", Some(ToolKind::Search)),
|
|
],
|
|
behavior_preset: Some("current".to_owned()),
|
|
};
|
|
let backfilled = backfill_tool_kinds(&config, &kinds);
|
|
let kind_of = |id: &str| {
|
|
backfilled
|
|
.tools
|
|
.iter()
|
|
.find(|t| t.id == id)
|
|
.expect("tool present")
|
|
.kind
|
|
};
|
|
assert_eq!(kind_of("Kigi:search_replace"), Some(ToolKind::Edit));
|
|
assert_eq!(
|
|
kind_of("adhoc.opaque"),
|
|
None,
|
|
"ids unknown to the registry stay kind-less"
|
|
);
|
|
assert_eq!(
|
|
kind_of("Kigi:read_file"),
|
|
Some(ToolKind::Search),
|
|
"an explicit kind wins over the registry's"
|
|
);
|
|
assert_eq!(backfilled.behavior_preset.as_deref(), Some("current"));
|
|
}
|
|
/// Regression: pinned server-bind toolsets arrive kind-less (the gRPC
|
|
/// `ToolConfigEntry` has no kind field), which used to make every
|
|
/// baseline entry bypass the capability filter — a `read_only`
|
|
/// sub-agent kept edit + execute tools. The registry backfill must
|
|
/// restore the filter.
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_readonly_filters_kindless_pinned_tools() {
|
|
let factory = factory_for_test();
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![
|
|
test_support::tc("Kigi:read_file", None),
|
|
test_support::tc("Kigi:grep", None),
|
|
test_support::tc("Kigi:list_dir", None),
|
|
test_support::tc("Kigi:search_replace", None),
|
|
test_support::tc("Kigi:run_terminal_cmd", None),
|
|
],
|
|
behavior_preset: None,
|
|
};
|
|
let (eff, ts, _backend) = resolve_session_toolset(
|
|
baseline,
|
|
CapabilityMode::ReadOnly,
|
|
&[],
|
|
PathBuf::from("/tmp"),
|
|
empty_env(),
|
|
"main",
|
|
factory.as_ref(),
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("resolve");
|
|
assert!(eff.tools.iter().all(|t| t.kind.is_none()));
|
|
let names: Vec<String> = ts
|
|
.tool_definitions()
|
|
.into_iter()
|
|
.map(|d| d.function.name)
|
|
.collect();
|
|
for kept in ["read_file", "grep", "list_dir"] {
|
|
assert!(
|
|
names.iter().any(|n| n == kept),
|
|
"{kept} must survive ReadOnly: {names:?}"
|
|
);
|
|
}
|
|
for dropped in ["search_replace", "run_terminal_cmd"] {
|
|
assert!(
|
|
!names.iter().any(|n| n == dropped),
|
|
"{dropped} must be dropped under ReadOnly: {names:?}"
|
|
);
|
|
}
|
|
}
|
|
#[test]
|
|
fn resolve_session_toolset_mcp_edit_dropped_under_readonly() {
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![test_support::tc("Kigi:read_file", Some(ToolKind::Read))],
|
|
behavior_preset: None,
|
|
};
|
|
let mcp_edit = test_support::tc("mcp.editor", Some(ToolKind::Edit));
|
|
let filtered = merge_and_filter(&baseline, &[mcp_edit], CapabilityMode::ReadOnly, "test");
|
|
assert!(!filtered.tools.iter().any(|t| t.id == "mcp.editor"));
|
|
}
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_mcp_kind_none_dropped_under_readonly() {
|
|
let factory = factory_for_test();
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![
|
|
test_support::tc("Kigi:read_file", Some(ToolKind::Read)),
|
|
test_support::tc("baseline.opaque", None),
|
|
],
|
|
behavior_preset: None,
|
|
};
|
|
let mcp = vec![test_support::tc("mcp.opaque", None)];
|
|
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
|
|
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
|
|
assert!(
|
|
kept_ids.contains(&"baseline.opaque"),
|
|
"baseline kind: None must survive ReadOnly: {kept_ids:?}"
|
|
);
|
|
assert!(
|
|
!kept_ids.contains(&"mcp.opaque"),
|
|
"MCP kind: None MUST be dropped under ReadOnly: {kept_ids:?}"
|
|
);
|
|
assert!(
|
|
kept_ids.contains(&"Kigi:read_file"),
|
|
"baseline Read kind must survive ReadOnly: {kept_ids:?}"
|
|
);
|
|
let _ = factory;
|
|
}
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_mcp_kind_none_kept_under_all() {
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![],
|
|
behavior_preset: None,
|
|
};
|
|
let mcp = vec![test_support::tc("mcp.opaque", None)];
|
|
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::All, "test_session");
|
|
let kept_ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
|
|
assert!(
|
|
kept_ids.contains(&"mcp.opaque"),
|
|
"All mode keeps MCP kind: None: {kept_ids:?}"
|
|
);
|
|
}
|
|
#[tokio::test]
|
|
async fn resolve_session_toolset_mcp_name_override_collision_skipped() {
|
|
let baseline = ToolServerConfig {
|
|
tools: vec![],
|
|
behavior_preset: None,
|
|
};
|
|
let mut mcp_a = test_support::tc("mcp.tool_a", Some(ToolKind::Read));
|
|
mcp_a.name_override = Some("shared_name".into());
|
|
let mut mcp_b = test_support::tc("mcp.tool_b", Some(ToolKind::Read));
|
|
mcp_b.name_override = Some("shared_name".into());
|
|
let mcp = vec![mcp_a, mcp_b];
|
|
let filtered = merge_and_filter(&baseline, &mcp, CapabilityMode::ReadOnly, "test_session");
|
|
let ids: Vec<&str> = filtered.tools.iter().map(|t| t.id.as_str()).collect();
|
|
assert!(ids.contains(&"mcp.tool_a"), "first wins: {ids:?}");
|
|
assert!(
|
|
!ids.contains(&"mcp.tool_b"),
|
|
"duplicate name dropped: {ids:?}"
|
|
);
|
|
}
|
|
#[test]
|
|
fn factory_session_folder_is_tmp_sessions_not_project_cwd() {
|
|
let cwd = PathBuf::from("/workspace");
|
|
let folder = WorkspaceSessionContextFactory::resolve_session_folder("sess-1");
|
|
let expected = PathBuf::from("/tmp/sessions/sess-1");
|
|
assert_eq!(folder, expected);
|
|
assert!(folder.is_dir());
|
|
assert!(!folder.starts_with(&cwd));
|
|
assert_eq!(
|
|
folder.join("terminal").join("call-42.log"),
|
|
PathBuf::from("/tmp/sessions/sess-1/terminal/call-42.log")
|
|
);
|
|
}
|
|
#[test]
|
|
fn factory_session_folder_sanitizes_and_isolates_ids() {
|
|
let sessions = PathBuf::from("/tmp/sessions");
|
|
let hostile = WorkspaceSessionContextFactory::resolve_session_folder("../../etc");
|
|
assert!(hostile.starts_with(&sessions));
|
|
assert_eq!(hostile.parent(), Some(sessions.as_path()));
|
|
assert_ne!(hostile, sessions.join("etc"));
|
|
let a = WorkspaceSessionContextFactory::resolve_session_folder("sess/1");
|
|
let b = WorkspaceSessionContextFactory::resolve_session_folder("sess_1");
|
|
assert_ne!(a, b);
|
|
let empty = WorkspaceSessionContextFactory::resolve_session_folder("");
|
|
assert_eq!(empty.parent(), Some(sessions.as_path()));
|
|
assert!(
|
|
empty
|
|
.file_name()
|
|
.and_then(|n| n.to_str())
|
|
.is_some_and(|s| s.starts_with("anon-"))
|
|
);
|
|
let uuid = "019f3c0a-e2c2-79a3-908d-e8a8f088fe52";
|
|
assert_eq!(
|
|
WorkspaceSessionContextFactory::resolve_session_folder(uuid),
|
|
sessions.join(uuid)
|
|
);
|
|
}
|
|
#[test]
|
|
fn ensure_session_dir_shared_by_state_and_session_folder() {
|
|
let home = tempfile::TempDir::new().unwrap();
|
|
let (under_home, ok) = ensure_session_dir(home.path(), "shared-id");
|
|
assert!(ok.is_ok());
|
|
assert_eq!(under_home, home.path().join("sessions").join("shared-id"));
|
|
assert!(under_home.is_dir());
|
|
let (under_tmp, ok) = ensure_session_dir(std::path::Path::new("/tmp"), "shared-id");
|
|
assert!(ok.is_ok());
|
|
assert_eq!(under_tmp, PathBuf::from("/tmp/sessions/shared-id"));
|
|
assert!(under_tmp.is_dir());
|
|
}
|
|
/// A hostile `session_id` (`../../etc`) is sanitized to a single safe
|
|
/// segment and cannot traverse outside `<home>/sessions/`.
|
|
#[test]
|
|
fn ensure_session_dir_sanitizes_malicious_session_id_no_traversal() {
|
|
let home = tempfile::TempDir::new().unwrap();
|
|
let sessions = home.path().join("sessions");
|
|
let (session_dir, created) = ensure_session_dir(home.path(), "../../etc");
|
|
assert!(created.is_ok());
|
|
assert!(
|
|
session_dir.starts_with(&sessions),
|
|
"session dir escaped sessions/: {}",
|
|
session_dir.display()
|
|
);
|
|
assert_eq!(
|
|
session_dir.parent(),
|
|
Some(sessions.as_path()),
|
|
"session dir must be a direct child of sessions/, got {}",
|
|
session_dir.display()
|
|
);
|
|
let seg = session_dir
|
|
.file_name()
|
|
.and_then(|n| n.to_str())
|
|
.expect("segment is valid utf-8");
|
|
assert!(
|
|
seg.chars()
|
|
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-'),
|
|
"sanitized segment must contain no path separators or dots: {seg:?}"
|
|
);
|
|
assert!(
|
|
session_dir.is_dir(),
|
|
"the confined session dir must be the only thing created"
|
|
);
|
|
}
|
|
/// Sanitization is injective: distinct ids that substitute to the same
|
|
/// base string still map to distinct directories (hash disambiguator),
|
|
/// while already-safe ids map to themselves.
|
|
#[test]
|
|
fn sanitize_session_id_is_injective() {
|
|
assert_eq!(super::sanitize_session_id("sess-1_a"), "sess-1_a");
|
|
assert_ne!(
|
|
super::sanitize_session_id("sess/1"),
|
|
super::sanitize_session_id("sess_1"),
|
|
"substitution collisions must be disambiguated"
|
|
);
|
|
assert_ne!(
|
|
super::sanitize_session_id("sess/1"),
|
|
super::sanitize_session_id("sess.1"),
|
|
);
|
|
assert!(super::sanitize_session_id("").starts_with("anon-"));
|
|
}
|
|
/// A toolset rebuilt for the SAME session rehydrates persisted state from
|
|
/// disk; a DIFFERENT session_id cold-starts with no cross-contamination.
|
|
#[tokio::test]
|
|
async fn tool_state_rehydrates_same_session_and_cold_starts_other() {
|
|
use kigi_tools::types::resources::{State, WebCitationCounter};
|
|
let factory = test_support::TestSessionContextFactory::new();
|
|
let cwd = PathBuf::from("/tmp");
|
|
let (_eff, ts_a, _backend_a) = resolve_session_toolset(
|
|
test_support::baseline_config(),
|
|
CapabilityMode::ReadWrite,
|
|
&[],
|
|
cwd.clone(),
|
|
empty_env(),
|
|
"sess-A",
|
|
&factory,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("build toolset A");
|
|
{
|
|
let mut res = ts_a.resources.lock().await;
|
|
let counter = res.get_or_default::<State<WebCitationCounter>>();
|
|
counter.counter = 123;
|
|
}
|
|
ts_a.save_and_flush_persistence().await;
|
|
drop(ts_a);
|
|
let (_eff, ts_b, _backend_b) = resolve_session_toolset(
|
|
test_support::baseline_config(),
|
|
CapabilityMode::ReadWrite,
|
|
&[],
|
|
cwd.clone(),
|
|
empty_env(),
|
|
"sess-A",
|
|
&factory,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("build toolset B");
|
|
{
|
|
let res = ts_b.resources.lock().await;
|
|
let counter = res
|
|
.get::<State<WebCitationCounter>>()
|
|
.expect("WebCitationCounter must be present after rehydration");
|
|
assert_eq!(
|
|
counter.counter, 123,
|
|
"tool state must survive a rebuild for the same session (rehydration)"
|
|
);
|
|
}
|
|
let (_eff, ts_c, _backend_c) = resolve_session_toolset(
|
|
test_support::baseline_config(),
|
|
CapabilityMode::ReadWrite,
|
|
&[],
|
|
cwd,
|
|
empty_env(),
|
|
"sess-B",
|
|
&factory,
|
|
None,
|
|
None,
|
|
None,
|
|
None,
|
|
)
|
|
.expect("build toolset C");
|
|
{
|
|
let res = ts_c.resources.lock().await;
|
|
let contaminated = res
|
|
.get::<State<WebCitationCounter>>()
|
|
.is_some_and(|c| c.counter == 123);
|
|
assert!(
|
|
!contaminated,
|
|
"a different session_id must cold-start, never inherit sess-A state"
|
|
);
|
|
}
|
|
}
|
|
}
|