diff --git a/.config/nextest.toml b/.config/nextest.toml index 290d7f9e93..a69c5c6f26 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -4,14 +4,23 @@ failure-output = "immediate-final" # No retries locally by default. retries = 0 -# QMD tests spawn shell scripts that time out under heavy parallelism. -# Run them serially to avoid flaky failures. -[test-groups.qmd-serial] -max-threads = 1 +# QMD tests spawn shell scripts that time out under heavy parallelism. Reserve +# the worker pool because a one-thread test group still competes with other tests. +[[profile.default.overrides]] +filter = "package(moltis-qmd)" +threads-required = "num-test-threads" + +# This test boots the complete headless gateway in a child process. Reserve the +# worker pool so model/provider startup is not starved by the workspace suite. +[[profile.default.overrides]] +filter = "package(moltis) & binary(acp_stdio) & test(without_echo_boots_the_real_backend_and_serves_protocol)" +threads-required = "num-test-threads" +# Synchronous WASM HTTP host tests depend on fixture threads accepting and +# serving loopback requests within millisecond-scale request deadlines. [[profile.default.overrides]] -filter = "package(moltis-qmd)" -test-group = "qmd-serial" +filter = "package(moltis-tools) & test(/^wasm_component::tests::http_host_/)" +threads-required = "num-test-threads" [profile.ci] # In CI, retry flaky tests once before failing the run. @@ -21,3 +30,15 @@ retries = 2 slow-timeout = { period = "60s", terminate-after = 2 } # Produce JUnit XML for GitHub Actions test summary. junit.path = "junit.xml" + +[[profile.ci.overrides]] +filter = "package(moltis) & binary(acp_stdio) & test(without_echo_boots_the_real_backend_and_serves_protocol)" +threads-required = "num-test-threads" + +[[profile.ci.overrides]] +filter = "package(moltis-tools) & test(/^wasm_component::tests::http_host_/)" +threads-required = "num-test-threads" + +[[profile.ci.overrides]] +filter = "package(moltis-qmd)" +threads-required = "num-test-threads" diff --git a/Cargo.lock b/Cargo.lock index 34c2b47361..e55ac6d32d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -127,8 +127,6 @@ dependencies = [ [[package]] name = "agent-client-protocol" version = "0.10.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10eeef5e80864f9c3c148a3f395c3e35a66d37ec7561c7845b2bffae8e841759" dependencies = [ "agent-client-protocol-schema", "anyhow", @@ -235,7 +233,7 @@ version = "1.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" dependencies = [ - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -246,7 +244,7 @@ checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" dependencies = [ "anstyle", "once_cell_polyfill", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -1815,7 +1813,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "117725a109d387c937a1533ce01b450cbde6b88abceea8473c4d7a85853cda3c" dependencies = [ "lazy_static", - "windows-sys 0.48.0", + "windows-sys 0.59.0", ] [[package]] @@ -1824,7 +1822,7 @@ version = "3.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "faf9468729b8cbcea668e36183cb69d317348c2e08e994829fb56ebfdfbaac34" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -2814,7 +2812,7 @@ dependencies = [ "libc", "option-ext", "redox_users 0.5.2", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -3144,7 +3142,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -5367,7 +5365,7 @@ dependencies = [ "js-sys", "log", "wasm-bindgen", - "windows-core", + "windows-core 0.57.0", ] [[package]] @@ -5892,7 +5890,7 @@ dependencies = [ "portable-atomic", "portable-atomic-util", "serde_core", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -7117,11 +7115,15 @@ dependencies = [ name = "moltis" version = "0.1.0" dependencies = [ + "agent-client-protocol", "anyhow", + "async-trait", "clap", "dotenvy", + "moltis-acp", "moltis-agents", "moltis-browser", + "moltis-chat", "moltis-claude-import", "moltis-codex-import", "moltis-common", @@ -7131,6 +7133,8 @@ dependencies = [ "moltis-hermes-import", "moltis-httpd", "moltis-import-core", + "moltis-mcp", + "moltis-mcp-agent-bridge", "moltis-memory", "moltis-node-host", "moltis-oauth", @@ -7139,7 +7143,9 @@ dependencies = [ "moltis-plugins", "moltis-portable", "moltis-projects", + "moltis-protocol", "moltis-providers", + "moltis-service-traits", "moltis-sessions", "moltis-skills", "moltis-tools", @@ -7160,6 +7166,24 @@ dependencies = [ "which 8.0.2", ] +[[package]] +name = "moltis-acp" +version = "0.1.0" +dependencies = [ + "agent-client-protocol", + "anyhow", + "async-trait", + "bytes", + "futures", + "moltis-metrics", + "serde_json", + "tokio", + "tokio-test", + "tokio-util", + "tracing", + "tracing-subscriber", +] + [[package]] name = "moltis-agents" version = "0.1.0" @@ -7413,6 +7437,7 @@ dependencies = [ "futures", "ipnet", "moltis-metrics", + "process-wrap", "reqwest 0.13.2", "secrecy 0.8.0", "serde", @@ -7823,6 +7848,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tokio-test", + "tokio-util", "tracing", "url", "uuid", @@ -8054,6 +8080,7 @@ dependencies = [ "moltis-common", "moltis-config", "moltis-oauth", + "moltis-sessions", "notify-debouncer-full", "secrecy 0.8.0", "serde", @@ -8264,6 +8291,8 @@ name = "moltis-sessions" version = "0.1.0" dependencies = [ "fd-lock", + "hex", + "libc", "moltis-common", "moltis-metrics", "serde", @@ -9055,7 +9084,7 @@ version = "0.50.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" dependencies = [ - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -9191,7 +9220,7 @@ version = "5.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51e219e79014df21a225b1860a479e2dcd7cbd9130f4defd4bd0e191ea31d67d" dependencies = [ - "base64 0.21.7", + "base64 0.22.1", "chrono", "getrandom 0.2.17", "http 1.4.2", @@ -9963,6 +9992,19 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "process-wrap" +version = "9.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e842efad9119158434d193c6682e2ebee4b44d6ad801d7b349623b3f57cdf55" +dependencies = [ + "futures", + "indexmap 2.14.0", + "nix 0.31.2", + "tokio", + "windows 0.62.2", +] + [[package]] name = "prodash" version = "31.0.0" @@ -10000,8 +10042,8 @@ version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ - "heck 0.4.1", - "itertools 0.10.5", + "heck 0.5.0", + "itertools 0.14.0", "log", "multimap", "petgraph", @@ -10018,7 +10060,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -10031,7 +10073,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn 2.0.117", @@ -11024,7 +11066,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys 0.12.1", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -11129,7 +11171,7 @@ dependencies = [ "security-framework 3.7.0", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -12039,7 +12081,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" dependencies = [ "libc", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -12282,7 +12324,7 @@ dependencies = [ "cfg-if", "libc", "psm", - "windows-sys 0.60.2", + "windows-sys 0.61.2", ] [[package]] @@ -12507,7 +12549,7 @@ dependencies = [ "memchr", "ntapi", "objc2-core-foundation", - "windows", + "windows 0.57.0", ] [[package]] @@ -12662,7 +12704,7 @@ dependencies = [ "getrandom 0.4.2", "once_cell", "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -12702,7 +12744,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" dependencies = [ "rustix 1.1.4", - "windows-sys 0.59.0", + "windows-sys 0.61.2", ] [[package]] @@ -15042,7 +15084,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.48.0", + "windows-sys 0.61.2", ] [[package]] @@ -15077,22 +15119,67 @@ version = "0.57.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "12342cb4d8e3b046f3d80effd474a7a02447231330ef77d71daa6fbc40681143" dependencies = [ - "windows-core", + "windows-core 0.57.0", "windows-targets 0.52.6", ] +[[package]] +name = "windows" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "527fadee13e0c05939a6a05d5bd6eec6cd2e3dbd648b9f8e447c6518133d8580" +dependencies = [ + "windows-collections", + "windows-core 0.62.2", + "windows-future", + "windows-numerics", +] + +[[package]] +name = "windows-collections" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b2d95af1a8a14a3c7367e1ed4fc9c20e0a26e79551b1454d72583c97cc6610" +dependencies = [ + "windows-core 0.62.2", +] + [[package]] name = "windows-core" version = "0.57.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d2ed2439a290666cd67ecce2b0ffaad89c2a56b976b736e6ece670297897832d" dependencies = [ - "windows-implement", - "windows-interface", + "windows-implement 0.57.0", + "windows-interface 0.57.0", "windows-result 0.1.2", "windows-targets 0.52.6", ] +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement 0.60.2", + "windows-interface 0.59.3", + "windows-link 0.2.1", + "windows-result 0.4.1", + "windows-strings 0.5.1", +] + +[[package]] +name = "windows-future" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1d6f90251fe18a279739e78025bd6ddc52a7e22f921070ccdc67dde84c605cb" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", + "windows-threading", +] + [[package]] name = "windows-implement" version = "0.57.0" @@ -15104,6 +15191,17 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "windows-interface" version = "0.57.0" @@ -15115,6 +15213,17 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "windows-link" version = "0.1.3" @@ -15127,6 +15236,16 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" +[[package]] +name = "windows-numerics" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e2e40844ac143cdb44aead537bbf727de9b044e107a0f1220392177d15b0f26" +dependencies = [ + "windows-core 0.62.2", + "windows-link 0.2.1", +] + [[package]] name = "windows-registry" version = "0.5.3" @@ -15311,6 +15430,15 @@ dependencies = [ "windows_x86_64_msvc 0.53.1", ] +[[package]] +name = "windows-threading" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3949bd5b99cafdf1c7ca86b43ca564028dfe27d66958f2470940f73d86d75b37" +dependencies = [ + "windows-link 0.2.1", +] + [[package]] name = "windows_aarch64_gnullvm" version = "0.42.2" diff --git a/Cargo.toml b/Cargo.toml index aeb9e23b25..597246f8df 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,7 @@ [workspace] default-members = [ "apps/courier", + "crates/acp", "crates/agents", "crates/auth", "crates/auto-reply", @@ -70,8 +71,10 @@ default-members = [ "crates/webhooks", "crates/whatsapp", ] +exclude = ["vendor/agent-client-protocol"] members = [ "apps/courier", + "crates/acp", "crates/agents", "crates/auth", "crates/auto-reply", @@ -143,6 +146,9 @@ members = [ resolver = "2" [patch.crates-io] +# Keep ACP stdin cancellation and EOF processing independent of blocked stdout, +# while bounding queued wire bytes. Remove once a compatible release includes it. +agent-client-protocol = { path = "vendor/agent-client-protocol" } # Widen libsqlite3-sys pin (0.30.1 → >=0.30.1,<0.36.0) so sqlx-sqlite 0.8.6 # can coexist with matrix-sdk-sqlite's rusqlite 0.37 (libsqlite3-sys 0.35). # Remove once sqlx ≥ 0.9 ships stable. @@ -189,6 +195,7 @@ axum = { features = ["ws"], version = "0.8" } axum-extra = "0.10" # Serialization postcard = { features = ["alloc"], version = "1.1" } +process-wrap = { default-features = false, features = ["job-object", "kill-on-drop", "process-group", "tokio1"], version = "9.1" } schemars = "1.2" serde = { features = ["derive"], version = "1" } serde-big-array = "0.5" @@ -247,6 +254,7 @@ gix = "0.83" hex = "0.4" hmac = "0.12" jsonwebtoken = { features = ["aws_lc_rs"], version = "10" } +libc = "0.2" llama-cpp-2 = "0.1" matrix-sdk = { default-features = false, features = ["e2e-encryption", "markdown", "native-tls", "sqlite"], version = "0.16" } mockito = "1.7" @@ -367,6 +375,7 @@ zeroize = { features = ["derive"], version = "1" } # Workspace crates moltis = { path = "crates/cli" } +moltis-acp = { path = "crates/acp" } moltis-agents = { path = "crates/agents" } moltis-auth = { path = "crates/auth" } moltis-auto-reply = { path = "crates/auto-reply" } diff --git a/README.md b/README.md index 605105c8a7..526f1ee02b 100644 --- a/README.md +++ b/README.md @@ -125,6 +125,7 @@ Verify releases with `gh attestation verify -R moltis-org/moltis` or - **Memory & Recall** — Per-agent memory workspaces, embeddings-powered long-term memory, hybrid vector + full-text search, session persistence with auto-compaction, cross-session recall, Cursor-compatible project context, context-file safety scanning - **Safer Agent Editing** — Automatic checkpoints before built-in skill and memory mutations, restore tooling, session branching - **Extensibility** — MCP servers (stdio + HTTP/SSE), skill system, 15 lifecycle hook events with circuit breaker, destructive command guard +- **Agent Client Protocol** — drives external ACP agents (`codex-acp`, `claude-agent-acp`, Cursor), and serves Moltis to ACP clients over stdio via `moltis acp` - **Security** — Encryption-at-rest vault (XChaCha20-Poly1305 + Argon2id), password + passkey + API key auth, sandbox isolation, SSRF/CSWSH protection - **Operations** — Cron scheduling, OpenTelemetry tracing, Prometheus metrics, cloud deploy (Fly.io, DigitalOcean), Tailscale integration, managed SSH deploy keys, host-pinned remote targets, live tool inventory in Settings, and CLI/web remote-exec doctor flows diff --git a/crates/acp/Cargo.toml b/crates/acp/Cargo.toml new file mode 100644 index 0000000000..a3d33f416c --- /dev/null +++ b/crates/acp/Cargo.toml @@ -0,0 +1,33 @@ +[package] +edition.workspace = true +name = "moltis-acp" +repository.workspace = true +version.workspace = true + +[features] +default = [] +metrics = ["dep:moltis-metrics"] +tracing = ["dep:tracing"] + +[dependencies] +agent-client-protocol = { workspace = true } +anyhow = { workspace = true } +async-trait = { workspace = true } +bytes = { workspace = true } +futures = { workspace = true } +moltis-metrics = { optional = true, workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tokio-util = { features = ["codec", "compat", "io"], workspace = true } +tracing = { optional = true, workspace = true } + +[dev-dependencies] +# `tracing` is here as well as behind the optional feature above: the +# stdout-hygiene test drives a subscriber at TRACE whatever features the crate is +# built with, and the coverage job builds with default features only. +tokio-test = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } + +[lints] +workspace = true diff --git a/crates/acp/src/agent.rs b/crates/acp/src/agent.rs new file mode 100644 index 0000000000..ba03238f6f --- /dev/null +++ b/crates/acp/src/agent.rs @@ -0,0 +1,358 @@ +//! `impl acp::Agent` — the agent side of the Agent Client Protocol. +//! +//! Everything here runs on the [`tokio::task::LocalSet`] that owns the +//! connection, so the types are `Rc`-based and never cross a thread boundary. +//! Real work is delegated to a [`AcpBackend`], which is `Send + Sync` and runs +//! on the ordinary multi-threaded runtime. + +use std::{ + cell::RefCell, + rc::{Rc, Weak}, + sync::Arc, +}; + +use { + agent_client_protocol::{self as acp, Client as _}, + async_trait::async_trait, +}; + +use crate::{ + backend::{AcpBackend, MAX_TURN_UPDATE_BYTES, TurnUpdates, validate_history}, + session::{ACP_SESSION_NAMESPACE, SessionKey, SessionRegistry}, + setup::SessionSetup, +}; + +const MAX_PROMPT_BLOCKS: usize = 256; +const MAX_PROMPT_BYTES: usize = 1024 * 1024; + +/// Protocol handler bridging an ACP client to a Moltis backend. +pub struct MoltisAgent { + backend: Arc, + sessions: SessionRegistry, + /// Weak on purpose: the connection owns the handler, so a strong reference + /// here would form a cycle and leak both for the life of the process. + connection: RefCell>, +} + +impl MoltisAgent { + #[must_use] + pub fn new(backend: Arc) -> Self { + Self { + backend, + sessions: SessionRegistry::new(), + connection: RefCell::new(Weak::new()), + } + } + + /// Attaches the connection used to send `session/update` notifications. + /// + /// `AgentSideConnection::new` consumes the handler, so the handler cannot + /// hold the connection at construction time. Build the agent first, pass a + /// clone into the connection, then call this. The caller must keep the + /// `Rc` alive for as long as the connection is served. + pub fn set_connection(&self, connection: &Rc) { + *self.connection.borrow_mut() = Rc::downgrade(connection); + } + + #[must_use] + pub fn sessions(&self) -> &SessionRegistry { + &self.sessions + } + + /// Clones the connection out of its `RefCell` so no borrow is held across + /// an await point. + fn connection(&self) -> acp::Result> { + self.connection + .borrow() + .upgrade() + .ok_or_else(|| acp::Error::internal_error().data("ACP connection not attached")) + } + + async fn notify( + connection: &Rc, + session_id: &acp::SessionId, + update: acp::SessionUpdate, + ) -> bool { + connection + .session_notification(acp::SessionNotification::new(session_id.clone(), update)) + .await + .is_ok() + } +} + +/// Flattens the supported prompt blocks into the plain text Moltis consumes. +/// Blocks outside the advertised capabilities are rejected instead of being +/// silently replaced or dropped. +pub fn prompt_text(blocks: &[acp::ContentBlock]) -> acp::Result { + if blocks.len() > MAX_PROMPT_BLOCKS { + return Err(acp::Error::invalid_params().data(format!( + "prompt has more than {MAX_PROMPT_BLOCKS} content blocks" + ))); + } + let text = blocks + .iter() + .map(|block| match block { + acp::ContentBlock::Text(text) => Ok(text.text.clone()), + acp::ContentBlock::ResourceLink(link) => Ok(link.uri.to_string()), + _ => Err(acp::Error::invalid_params() + .data("prompt contains content the agent did not advertise")), + }) + .collect::>>()? + .join("\n"); + if text.len() > MAX_PROMPT_BYTES { + return Err( + acp::Error::invalid_params().data(format!("prompt exceeds {MAX_PROMPT_BYTES} bytes")) + ); + } + Ok(text) +} + +#[async_trait(?Send)] +impl acp::Agent for MoltisAgent { + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, args)))] + async fn initialize( + &self, + args: acp::InitializeRequest, + ) -> acp::Result { + let version = if args.protocol_version == acp::ProtocolVersion::V1 { + acp::ProtocolVersion::V1 + } else { + acp::ProtocolVersion::LATEST + }; + let capabilities = self.backend.capabilities(); + Ok(acp::InitializeResponse::new(version) + .agent_capabilities( + acp::AgentCapabilities::new().load_session(capabilities.load_session), + ) + .agent_info( + acp::Implementation::new( + "moltis", + option_env!("MOLTIS_VERSION").unwrap_or(env!("CARGO_PKG_VERSION")), + ) + .title("Moltis"), + )) + } + + async fn authenticate( + &self, + _args: acp::AuthenticateRequest, + ) -> acp::Result { + // The client is the local parent process that spawned us; there is no + // separate identity to establish. + Ok(acp::AuthenticateResponse::new()) + } + + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, args)))] + async fn new_session( + &self, + args: acp::NewSessionRequest, + ) -> acp::Result { + self.sessions.ensure_capacity()?; + let setup_key = SessionKey::namespaced(format!("setup-{}", self.sessions.next_local_id())); + let _setup_guard = self.sessions.begin_setup(&setup_key)?; + let setup = SessionSetup::new(args.cwd, args.mcp_servers).await?; + let key = self.backend.create_session(&setup).await.map_err(|error| { + acp::Error::internal_error().data(format!("failed to create session: {error}")) + })?; + // A backend minting keys outside the namespace would quietly hand the + // client an id that `load_session` must then refuse. That is our bug, + // not the client's, so it is an internal error rather than bad input. + if !key.is_namespaced() { + let _ = self.backend.discard_session(&key).await; + return Err(acp::Error::internal_error().data(format!( + "backend created session {key} outside the `{ACP_SESSION_NAMESPACE}:` namespace" + ))); + } + if let Err(error) = self.sessions.insert(key.clone()) { + let _ = self.backend.discard_session(&key).await; + return Err(error); + } + Ok(acp::NewSessionResponse::new(acp::SessionId::from(key))) + } + + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, args)))] + async fn load_session( + &self, + args: acp::LoadSessionRequest, + ) -> acp::Result { + if !self.backend.capabilities().load_session { + return Err(acp::Error::method_not_found()); + } + // `session_id` is arbitrary client input, and unlike `prompt`/`cancel` + // there is no registry entry to check it against — resuming a session + // this connection never opened is the whole point. The `acp:` namespace + // is therefore the only thing keeping a client from naming a Web UI or + // channel session here and driving it with subsequent prompts. Enforce + // it before the backend sees the key, so no backend has to re-derive + // the invariant to stay isolated. + let key = SessionKey::from(&args.session_id); + if !key.is_namespaced() { + return Err(acp::Error::invalid_params().data(format!( + "session id {key} is outside the `{ACP_SESSION_NAMESPACE}:` namespace" + ))); + } + self.sessions.ensure_capacity_for(&key)?; + let _setup_guard = self.sessions.begin_setup(&key)?; + let setup = SessionSetup::new(args.cwd, args.mcp_servers).await?; + let history = self + .backend + .load_session(&key, &setup) + .await + .map_err(|error| { + let detail = format!("failed to load session {key}: {error}"); + if error.is::() { + acp::Error::invalid_params().data(detail) + } else { + acp::Error::internal_error().data(detail) + } + })?; + + // The spec asks the agent to stream the whole conversation back before + // resolving the request. + let connection = self.connection()?; + if let Err(error) = validate_history(&history) { + let _ = self.backend.discard_session(&key).await; + return Err(acp::Error::internal_error() + .data(format!("cannot replay session history: {error}"))); + } + for update in history { + if !Self::notify(&connection, &args.session_id, update).await { + // The closed connection immediately drives backend shutdown, + // which owns cleanup. Removing the runtime here would race that + // shutdown while its MCP process is still being stopped. + return Err(acp::Error::internal_error() + .data("ACP client disconnected during session replay")); + } + } + if let Err(error) = self.sessions.insert(key.clone()) { + let _ = self.backend.discard_session(&key).await; + return Err(error); + } + Ok(acp::LoadSessionResponse::new()) + } + + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, args)))] + async fn prompt(&self, args: acp::PromptRequest) -> acp::Result { + let key = self.sessions.resolve(&args.session_id)?; + let text = prompt_text(&args.prompt)?; + let _prompt = self.sessions.begin_prompt(&key)?; + let connection = self.connection()?; + self.sessions.clear_cancelled(&key); + + // Memory is bounded by `TurnUpdates`' byte budget. An unbounded item + // queue avoids rejecting a valid burst made up of many tiny deltas. + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + let update_status = updates.clone(); + let turn = self.backend.prompt(&key, text, updates); + let mut turn = std::pin::pin!(turn); + + // Forward updates while the turn runs, on this same task. Doing it here + // rather than in a spawned forwarder means a backend that leaks a + // `TurnUpdates` clone cannot wedge the response: we stop draining the + // moment the turn resolves. + let mut update_bytes = 0usize; + let mut forwarding = true; + let mut update_limit_exceeded = false; + let outcome = loop { + tokio::select! { + biased; + Some(update) = rx.recv() => { + if !forwarding { + continue; + } + update_bytes = update_bytes.saturating_add( + serde_json::to_vec(&update).map_or(MAX_TURN_UPDATE_BYTES + 1, |value| value.len()) + ); + if update_bytes > MAX_TURN_UPDATE_BYTES { + update_limit_exceeded = true; + forwarding = false; + let _ = self.backend.cancel(&key).await; + } else if !Self::notify(&connection, &args.session_id, update).await { + forwarding = false; + let _ = self.backend.cancel(&key).await; + } + }, + result = &mut turn => break result, + } + }; + + // Flush anything already queued so no delta lands after the response. + while forwarding && let Ok(update) = rx.try_recv() { + update_bytes = update_bytes.saturating_add( + serde_json::to_vec(&update).map_or(MAX_TURN_UPDATE_BYTES + 1, |value| value.len()), + ); + if update_bytes > MAX_TURN_UPDATE_BYTES { + update_limit_exceeded = true; + break; + } + if !Self::notify(&connection, &args.session_id, update).await { + break; + } + } + + if update_limit_exceeded || update_status.limit_exceeded() { + return Err(acp::Error::internal_error() + .data(format!("turn updates exceed {MAX_TURN_UPDATE_BYTES} bytes"))); + } + + let cancelled = self.sessions.take_cancelled(&key); + let stop_reason = match outcome { + // A cancel that raced the turn's own completion still has to be + // reported as cancelled: the client is waiting for that signal. + Ok(_) if cancelled => acp::StopReason::Cancelled, + Ok(reason) => reason, + Err(_) if cancelled => acp::StopReason::Cancelled, + Err(error) => { + return Err(acp::Error::internal_error().data(format!("turn failed: {error}"))); + }, + }; + Ok(acp::PromptResponse::new(stop_reason)) + } + + #[cfg_attr(feature = "tracing", tracing::instrument(skip(self, args)))] + async fn cancel(&self, args: acp::CancelNotification) -> acp::Result<()> { + // Arrives out-of-band while `prompt` is still pending. Flag first so the + // pending turn reports `Cancelled` even if the backend finishes first. + let key = self.sessions.resolve(&args.session_id)?; + self.sessions.mark_cancelled(&key); + self.backend.cancel(&key).await.map_err(|error| { + acp::Error::internal_error().data(format!("failed to cancel turn: {error}")) + }) + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use super::*; + + #[test] + fn prompt_text_joins_supported_text_blocks() { + let blocks = vec![ + acp::ContentBlock::from("first"), + acp::ContentBlock::from("second"), + ]; + assert_eq!( + prompt_text(&blocks).expect("supported prompt"), + "first\nsecond" + ); + } + + #[test] + fn prompt_text_rejects_unadvertised_content() { + let image = serde_json::from_value::(serde_json::json!({ + "type": "image", + "data": "AA==", + "mimeType": "image/png" + })) + .expect("valid ACP image block"); + assert!(prompt_text(&[image]).is_err()); + } + + #[test] + fn prompt_text_rejects_oversized_input() { + let oversized = "x".repeat(MAX_PROMPT_BYTES + 1); + assert!(prompt_text(&[acp::ContentBlock::from(oversized)]).is_err()); + } +} diff --git a/crates/acp/src/backend.rs b/crates/acp/src/backend.rs new file mode 100644 index 0000000000..f914fd30c4 --- /dev/null +++ b/crates/acp/src/backend.rs @@ -0,0 +1,214 @@ +//! The seam between the (`!Send`) ACP protocol handler and Moltis's (`Send`) +//! services. +//! +//! `agent_client_protocol` declares its traits with `#[async_trait(?Send)]`, so +//! every future the protocol handler produces is pinned to the thread running +//! the [`tokio::task::LocalSet`]. Moltis's `ChatService` is `Send + Sync` and +//! expects to run on the multi-threaded runtime. +//! +//! [`AcpBackend`] is where the two meet. It is deliberately `Send + Sync`, so +//! implementations live entirely in the threaded world and never learn that a +//! `LocalSet` exists. Streaming flows the other way through [`TurnUpdates`], +//! which is a plain channel sender: the backend pushes updates from whatever +//! task it likes, and the protocol side forwards them as `session/update` +//! notifications from the local thread. + +use {agent_client_protocol as acp, async_trait::async_trait, tokio::sync::mpsc}; + +use crate::{session::SessionKey, setup::SessionSetup}; + +pub const MAX_HISTORY_UPDATES: usize = 10_000; +pub const MAX_HISTORY_BYTES: usize = 8 * 1024 * 1024; +pub const MAX_TURN_UPDATE_BYTES: usize = 8 * 1024 * 1024; + +/// Marker used by backends when a syntactically valid ACP session does not +/// exist. Other load failures are operational and must not be blamed on input. +#[derive(Debug)] +pub struct SessionNotFound; + +impl std::fmt::Display for SessionNotFound { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("session does not exist") + } +} + +impl std::error::Error for SessionNotFound {} + +pub fn validate_history(updates: &[acp::SessionUpdate]) -> anyhow::Result<()> { + if updates.len() > MAX_HISTORY_UPDATES { + anyhow::bail!("session history exceeds {MAX_HISTORY_UPDATES} updates"); + } + let bytes = updates.iter().try_fold(0usize, |total, update| { + Ok::<_, serde_json::Error>(total.saturating_add(serde_json::to_vec(update)?.len())) + })?; + if bytes > MAX_HISTORY_BYTES { + anyhow::bail!("session history exceeds {MAX_HISTORY_BYTES} bytes"); + } + Ok(()) +} + +/// Sink for `session/update` notifications emitted while a turn is running. +/// +/// Cloneable and `Send`, so a backend may hand it to spawned tasks. Sends are +/// non-blocking. A send returns `false` if the bounded buffer is full or the +/// protocol task has stopped listening, allowing the backend to abort rather +/// than growing memory without limit. +#[derive(Clone, Debug)] +pub struct TurnUpdates { + tx: mpsc::UnboundedSender, + bytes_sent: std::sync::Arc, + limit_exceeded: std::sync::Arc, +} + +impl TurnUpdates { + #[must_use] + pub fn new(tx: mpsc::UnboundedSender) -> Self { + Self { + tx, + bytes_sent: std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0)), + limit_exceeded: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)), + } + } + + /// Sends a raw update. Returns `false` when the byte budget is exhausted or + /// the protocol task has stopped listening. + pub fn send(&self, update: acp::SessionUpdate) -> bool { + let Ok(bytes) = serde_json::to_vec(&update).map(|value| value.len()) else { + self.limit_exceeded + .store(true, std::sync::atomic::Ordering::Release); + return false; + }; + let reserved = self.bytes_sent.fetch_update( + std::sync::atomic::Ordering::AcqRel, + std::sync::atomic::Ordering::Acquire, + |used| { + used.checked_add(bytes) + .filter(|total| *total <= MAX_TURN_UPDATE_BYTES) + }, + ); + if reserved.is_err() { + self.limit_exceeded + .store(true, std::sync::atomic::Ordering::Release); + return false; + } + if self.tx.send(update).is_ok() { + return true; + } + self.bytes_sent + .fetch_sub(bytes, std::sync::atomic::Ordering::AcqRel); + false + } + + /// Streams a chunk of the agent's visible reply. + pub fn agent_message(&self, text: impl Into) -> bool { + self.send(acp::SessionUpdate::AgentMessageChunk( + acp::ContentChunk::new(acp::ContentBlock::from(text.into())), + )) + } + + /// Streams a chunk of the agent's reasoning. + pub fn agent_thought(&self, text: impl Into) -> bool { + self.send(acp::SessionUpdate::AgentThoughtChunk( + acp::ContentChunk::new(acp::ContentBlock::from(text.into())), + )) + } + + /// Returns whether the client is still listening. + #[must_use] + pub fn is_open(&self) -> bool { + !self.tx.is_closed() + } + + #[must_use] + pub fn limit_exceeded(&self) -> bool { + self.limit_exceeded + .load(std::sync::atomic::Ordering::Acquire) + } +} + +/// What an [`AcpBackend`] supports, surfaced to the client during `initialize`. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct BackendCapabilities { + /// Whether `session/load` can resume a previously created session. + pub load_session: bool, +} + +/// Moltis-side implementation of a single ACP conversation surface. +/// +/// One process serves one client, matching how ACP harnesses spawn agents, but +/// a backend may hold several sessions at once because a client is free to open +/// more than one. +#[async_trait] +pub trait AcpBackend: Send + Sync + 'static { + /// Creates a new session and returns its Moltis key. + async fn create_session(&self, setup: &SessionSetup) -> anyhow::Result; + + /// Resumes an existing session, returning its history so the protocol layer + /// can replay it to the client as `session/update` notifications. + /// + /// Only called when [`BackendCapabilities::load_session`] is set. + async fn load_session( + &self, + _key: &SessionKey, + _setup: &SessionSetup, + ) -> anyhow::Result> { + Err(anyhow::anyhow!("session/load is not supported")) + } + + /// Rolls back connection-scoped resources for a session setup that the + /// protocol layer could not retain or replay. + async fn discard_session(&self, _key: &SessionKey) -> anyhow::Result<()> { + Ok(()) + } + + /// Runs one turn to completion. + /// + /// Must not return until the turn is over: deltas go out through `updates` + /// while this future is pending, and the returned stop reason is what + /// resolves the client's `session/prompt` call. + async fn prompt( + &self, + key: &SessionKey, + prompt: String, + updates: TurnUpdates, + ) -> anyhow::Result; + + /// Aborts the in-flight turn for `key`, if any. + /// + /// Arrives out-of-band while `prompt` is still pending, so it must not wait + /// on that turn. The pending `prompt` is expected to wind up promptly + /// afterwards. + async fn cancel(&self, key: &SessionKey) -> anyhow::Result<()>; + + /// Releases connection-scoped turns, processes, and registrations. + async fn shutdown(&self) -> anyhow::Result<()> { + Ok(()) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities::default() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn turn_update_budget_is_enforced_before_enqueue() { + let (tx, _rx) = mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + let chunk = "x".repeat(5 * 1024 * 1024); + assert!(updates.agent_message(chunk.clone())); + assert!(!updates.agent_message(chunk)); + assert!(updates.limit_exceeded()); + } + + #[test] + fn many_small_updates_are_limited_by_bytes_not_item_count() { + let (tx, _rx) = mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + assert!((0..2_000).all(|_| updates.agent_message("x"))); + assert!(!updates.limit_exceeded()); + } +} diff --git a/crates/acp/src/echo.rs b/crates/acp/src/echo.rs new file mode 100644 index 0000000000..1da1a2437e --- /dev/null +++ b/crates/acp/src/echo.rs @@ -0,0 +1,180 @@ +//! A backend with no Moltis behind it. +//! +//! Useful for smoke-testing a client's handshake (`moltis acp --echo`) and as +//! the fixture the protocol tests drive, so protocol coverage does not require +//! standing up a gateway. + +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, +}; + +use {agent_client_protocol as acp, async_trait::async_trait}; + +use crate::{ + backend::{AcpBackend, BackendCapabilities, TurnUpdates}, + session::SessionKey, + setup::SessionSetup, +}; + +/// Echoes the prompt back as a streamed agent message. +#[derive(Debug, Default)] +pub struct EchoBackend { + next_id: AtomicU64, + cancels: Mutex>>, +} + +impl EchoBackend { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + fn cancel_flag(&self, key: &SessionKey) -> Arc { + let mut cancels = self + .cancels + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + Arc::clone(cancels.entry(key.clone()).or_default()) + } +} + +#[async_trait] +impl AcpBackend for EchoBackend { + async fn create_session(&self, _setup: &SessionSetup) -> anyhow::Result { + let id = self + .next_id + .fetch_add(1, Ordering::Relaxed) + .saturating_add(1); + Ok(SessionKey::namespaced(format!("echo-{id}"))) + } + + async fn prompt( + &self, + key: &SessionKey, + prompt: String, + updates: TurnUpdates, + ) -> anyhow::Result { + let cancelled = self.cancel_flag(key); + cancelled.store(false, Ordering::SeqCst); + + updates.agent_thought("echoing the prompt back"); + // Chunked so clients exercise their streaming path rather than + // receiving one atomic message. + for word in prompt.split_inclusive(char::is_whitespace) { + if cancelled.load(Ordering::SeqCst) { + return Ok(acp::StopReason::Cancelled); + } + if !updates.agent_message(word.to_string()) { + return Ok(acp::StopReason::Cancelled); + } + } + Ok(acp::StopReason::EndTurn) + } + + async fn cancel(&self, key: &SessionKey) -> anyhow::Result<()> { + self.cancel_flag(key).store(true, Ordering::SeqCst); + Ok(()) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities { + load_session: false, + } + } +} + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use {super::*, tokio::sync::mpsc}; + + async fn setup() -> SessionSetup { + SessionSetup::new(std::env::temp_dir(), Vec::new()) + .await + .expect("setup") + } + + #[tokio::test] + async fn sessions_are_namespaced_and_unique() { + let backend = EchoBackend::new(); + let first = backend + .create_session(&setup().await) + .await + .expect("session"); + let second = backend + .create_session(&setup().await) + .await + .expect("session"); + assert!(first.is_namespaced()); + assert_ne!(first, second); + } + + #[tokio::test] + async fn prompt_streams_chunks_then_ends_turn() { + let backend = EchoBackend::new(); + let key = backend + .create_session(&setup().await) + .await + .expect("session"); + let (tx, mut rx) = mpsc::unbounded_channel(); + let reason = backend + .prompt(&key, "hello there".to_string(), TurnUpdates::new(tx)) + .await + .expect("turn"); + assert_eq!(reason, acp::StopReason::EndTurn); + + let mut text = String::new(); + let mut thoughts = 0; + while let Ok(update) = rx.try_recv() { + match update { + acp::SessionUpdate::AgentMessageChunk(chunk) => { + if let acp::ContentBlock::Text(block) = chunk.content { + text.push_str(&block.text); + } + }, + acp::SessionUpdate::AgentThoughtChunk(_) => thoughts += 1, + _ => {}, + } + } + assert_eq!(text, "hello there"); + assert_eq!(thoughts, 1); + } + + #[tokio::test] + async fn cancel_stops_the_turn() { + let backend = EchoBackend::new(); + let key = backend + .create_session(&setup().await) + .await + .expect("session"); + backend.cancel(&key).await.expect("cancel"); + // The flag is cleared when a turn starts, so a cancel issued before the + // turn must not leak into it. + let (tx, _rx) = mpsc::unbounded_channel(); + let reason = backend + .prompt(&key, "hello".to_string(), TurnUpdates::new(tx)) + .await + .expect("turn"); + assert_eq!(reason, acp::StopReason::EndTurn); + } + + #[tokio::test] + async fn turn_ends_when_the_client_stops_listening() { + let backend = EchoBackend::new(); + let key = backend + .create_session(&setup().await) + .await + .expect("session"); + let (tx, rx) = mpsc::unbounded_channel(); + drop(rx); + let reason = backend + .prompt(&key, "hello".to_string(), TurnUpdates::new(tx)) + .await + .expect("turn"); + assert_eq!(reason, acp::StopReason::Cancelled); + } +} diff --git a/crates/acp/src/lib.rs b/crates/acp/src/lib.rs new file mode 100644 index 0000000000..58698f56bb --- /dev/null +++ b/crates/acp/src/lib.rs @@ -0,0 +1,35 @@ +//! Moltis as an ACP *agent*. +//! +//! Moltis has long been an ACP client (`moltis-external-agents` spawns and +//! drives `codex-acp`, `claude-agent-acp`, and friends). This crate is the +//! inverse: it lets any ACP client — Zed, `buzz-acp`, a bespoke harness — drive +//! Moltis the way it would drive any other agent, with Moltis's sessions, +//! memory, tool policy, and providers staying behind the protocol. +//! +//! ACP is a *surface*, not a channel: there is no correspondent to allowlist +//! and no account to configure, so it sits beside the Web UI and GraphQL rather +//! than alongside Telegram or Nostr. +//! +//! # Layout +//! +//! - [`backend`] — the `Send` seam Moltis implements ([`AcpBackend`]) +//! - [`agent`] — the `!Send` protocol handler ([`MoltisAgent`]) +//! - [`session`] — ACP session ids ↔ namespaced Moltis session keys +//! - [`server`] — transport wiring ([`server::run_stdio`]) +//! - [`echo`] — a backend-free implementation for smoke tests + +pub mod agent; +pub mod backend; +pub mod echo; +pub mod server; +pub mod session; +pub mod setup; + +pub use crate::{ + agent::MoltisAgent, + backend::{AcpBackend, BackendCapabilities, SessionNotFound, TurnUpdates, validate_history}, + echo::EchoBackend, + server::{run_stdio, serve}, + session::{ACP_SESSION_NAMESPACE, SessionKey, SessionRegistry}, + setup::SessionSetup, +}; diff --git a/crates/acp/src/server.rs b/crates/acp/src/server.rs new file mode 100644 index 0000000000..3985ac4441 --- /dev/null +++ b/crates/acp/src/server.rs @@ -0,0 +1,181 @@ +//! Transport wiring: run a [`MoltisAgent`] over a byte stream. +//! +//! # stdout is the wire +//! +//! When served over stdio, stdout carries JSON-RPC framing and nothing else. A +//! stray `println!`, or a `tracing` subscriber whose writer defaults to stdout, +//! corrupts the stream and the client disconnects with a parse error. Callers +//! must point logging at stderr before calling [`run_stdio`] — see +//! `moltis acp` in the CLI, and the `stdout_is_only_protocol_framing` test. + +use std::{io, rc::Rc, sync::Arc}; + +use { + agent_client_protocol as acp, + bytes::Bytes, + futures::TryStreamExt, + tokio::io::{AsyncRead, AsyncWrite}, + tokio_util::{ + codec::{FramedRead, LinesCodec}, + compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}, + io::StreamReader, + }, +}; + +use crate::{agent::MoltisAgent, backend::AcpBackend}; + +/// Maximum size of one newline-delimited JSON-RPC message. +pub const MAX_JSON_RPC_FRAME_BYTES: usize = 4 * 1024 * 1024; + +fn bounded_input(input: R) -> impl AsyncRead + Unpin +where + R: AsyncRead + Unpin, +{ + let frames = FramedRead::new( + input, + LinesCodec::new_with_max_length(MAX_JSON_RPC_FRAME_BYTES), + ) + .map_ok(|line| Bytes::from(format!("{line}\n"))) + .map_err(|error| io::Error::new(io::ErrorKind::InvalidData, error)); + StreamReader::new(frames) +} + +/// Serves one client over the given streams until the connection closes. +/// +/// Must be called from inside a [`tokio::task::LocalSet`]: the protocol handler +/// is `!Send` and its tasks are spawned with `spawn_local`. Use [`run_stdio`] +/// if you do not already have one. +pub async fn serve(backend: Arc, input: R, output: W) -> anyhow::Result<()> +where + R: AsyncRead + Unpin + 'static, + W: AsyncWrite + Unpin + 'static, +{ + let agent = Rc::new(MoltisAgent::new(Arc::clone(&backend))); + let (connection, io_task) = acp::AgentSideConnection::new_with_limits( + Rc::clone(&agent), + output.compat_write(), + bounded_input(input).compat(), + acp::ConnectionLimits::bounded(MAX_JSON_RPC_FRAME_BYTES, MAX_JSON_RPC_FRAME_BYTES, 256), + |future| { + tokio::task::spawn_local(future); + }, + ); + // Held here for the lifetime of the connection; the agent keeps only a weak + // reference so the two do not form a cycle. + let connection = Rc::new(connection); + agent.set_connection(&connection); + + let result = io_task + .await + .map_err(|error| anyhow::anyhow!("ACP connection failed: {error}")); + let shutdown = backend.shutdown().await; + result.and(shutdown) +} + +/// Serves one client over stdio, creating the [`tokio::task::LocalSet`]. +/// +/// Logging must already be pointed away from stdout. +pub async fn run_stdio(backend: Arc) -> anyhow::Result<()> { + let local = tokio::task::LocalSet::new(); + local + .run_until(serve(backend, tokio::io::stdin(), tokio::io::stdout())) + .await +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use std::{ + pin::Pin, + task::{Context, Poll}, + time::Duration, + }; + + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + sync::oneshot, + }; + + use {super::*, crate::echo::EchoBackend}; + + struct BlockingOutput { + started: Option>, + } + + impl AsyncWrite for BlockingOutput { + fn poll_write( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _bytes: &[u8], + ) -> Poll> { + if let Some(started) = self.started.take() { + let _ = started.send(()); + } + Poll::Pending + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + #[tokio::test] + async fn oversized_json_rpc_frames_are_rejected() { + let (mut writer, reader) = tokio::io::duplex(MAX_JSON_RPC_FRAME_BYTES + 2); + let write = tokio::spawn(async move { + writer + .write_all(&vec![b'x'; MAX_JSON_RPC_FRAME_BYTES + 1]) + .await + .expect("write oversized frame"); + writer.write_all(b"\n").await.expect("write newline"); + }); + let mut bounded = bounded_input(reader); + let mut output = Vec::new(); + let error = bounded + .read_to_end(&mut output) + .await + .expect_err("oversized frame must fail"); + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + write.await.expect("writer task"); + } + + #[tokio::test] + async fn blocked_output_does_not_prevent_eof_shutdown() { + let local = tokio::task::LocalSet::new(); + local + .run_until(async { + let (mut input, server_input) = tokio::io::duplex(1024); + let (started_tx, started_rx) = oneshot::channel(); + let server = tokio::task::spawn_local(serve( + Arc::new(EchoBackend::new()), + server_input, + BlockingOutput { + started: Some(started_tx), + }, + )); + + input + .write_all( + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"unknown\",\"params\":{}}\n", + ) + .await + .expect("write request"); + tokio::time::timeout(Duration::from_secs(1), started_rx) + .await + .expect("output write started") + .expect("output signal"); + + input.shutdown().await.expect("close input"); + let result = tokio::time::timeout(Duration::from_secs(1), server) + .await + .expect("server stops after EOF") + .expect("server task"); + assert!(result.is_ok()); + }) + .await; + } +} diff --git a/crates/acp/src/session.rs b/crates/acp/src/session.rs new file mode 100644 index 0000000000..6dec3954d9 --- /dev/null +++ b/crates/acp/src/session.rs @@ -0,0 +1,345 @@ +//! Mapping between ACP session ids and Moltis session keys. +//! +//! The two are the same string. Moltis session keys created for ACP live under +//! a dedicated `acp:` namespace (see [`SessionKey::namespaced`]), which keeps +//! `moltis acp` runs from colliding with Web UI sessions while still letting a +//! client hand a `SessionId` straight back to `session/load`. + +use std::{ + cell::RefCell, + collections::{HashMap, HashSet}, + rc::Rc, +}; + +use agent_client_protocol as acp; + +/// Prefix marking a Moltis session as owned by the ACP surface. +pub const ACP_SESSION_NAMESPACE: &str = "acp"; + +/// Maximum sessions one stdio connection may retain. +pub const MAX_SESSIONS: usize = 64; + +/// Maximum turns or setup operations that may run concurrently per connection. +pub const MAX_CONCURRENT_OPERATIONS: usize = 4; + +const MAX_SESSION_ID_SUFFIX_BYTES: usize = 128; + +/// A Moltis session key. +#[derive(Clone, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct SessionKey(String); + +impl SessionKey { + /// Wraps an existing key verbatim. + #[must_use] + pub fn new(key: impl Into) -> Self { + Self(key.into()) + } + + /// Builds a key inside the ACP namespace. + #[must_use] + pub fn namespaced(id: impl AsRef) -> Self { + Self(format!("{ACP_SESSION_NAMESPACE}:{}", id.as_ref())) + } + + #[must_use] + pub fn as_str(&self) -> &str { + &self.0 + } + + /// Whether this key belongs to the ACP namespace. + #[must_use] + pub fn is_namespaced(&self) -> bool { + self.0.split_once(':').is_some_and(|(prefix, rest)| { + prefix == ACP_SESSION_NAMESPACE + && !rest.is_empty() + && rest.len() <= MAX_SESSION_ID_SUFFIX_BYTES + && rest + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + }) + } +} + +impl std::fmt::Display for SessionKey { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str(&self.0) + } +} + +impl From for acp::SessionId { + fn from(key: SessionKey) -> Self { + Self::from(key.0) + } +} + +impl From<&acp::SessionId> for SessionKey { + fn from(id: &acp::SessionId) -> Self { + Self(id.to_string()) + } +} + +/// Tracks the sessions this connection has opened. +/// +/// Lives on the `LocalSet` thread alongside the protocol handler, so it uses +/// `Rc`/`RefCell` rather than `Arc`/`Mutex` — there is no cross-thread sharing +/// to guard against, and a poisoned-lock path would be dead code. +#[derive(Clone, Debug, Default)] +pub struct SessionRegistry { + inner: Rc>, +} + +#[derive(Debug, Default)] +struct RegistryState { + known: HashSet, + in_flight: HashSet, + /// Sessions whose in-flight turn has been cancelled but not yet observed. + cancelled: HashMap, + next_id: u64, +} + +impl SessionRegistry { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Allocates the next connection-local session identifier. + /// + /// Backends are free to ignore this and mint their own keys; it exists so a + /// backend without its own id source stays collision-free within a run. + pub fn next_local_id(&self) -> u64 { + let mut state = self.inner.borrow_mut(); + state.next_id = state.next_id.saturating_add(1); + state.next_id + } + + pub fn insert(&self, key: SessionKey) -> acp::Result<()> { + self.ensure_capacity_for(&key)?; + let mut state = self.inner.borrow_mut(); + state.known.insert(key); + Ok(()) + } + + pub fn ensure_capacity(&self) -> acp::Result<()> { + if self.inner.borrow().known.len() < MAX_SESSIONS { + return Ok(()); + } + Err(acp::Error::invalid_params().data(format!( + "connection cannot retain more than {MAX_SESSIONS} sessions" + ))) + } + + /// Rejects setup work that would exceed the connection's session budget. + pub fn ensure_capacity_for(&self, key: &SessionKey) -> acp::Result<()> { + let state = self.inner.borrow(); + if state.known.contains(key) || state.known.len() < MAX_SESSIONS { + return Ok(()); + } + Err(acp::Error::invalid_params().data(format!( + "connection cannot retain more than {MAX_SESSIONS} sessions" + ))) + } + + #[must_use] + pub fn contains(&self, key: &SessionKey) -> bool { + self.inner.borrow().known.contains(key) + } + + /// Resolves an incoming `SessionId`, rejecting ids this connection never + /// handed out with `invalid_params` rather than panicking. + pub fn resolve(&self, id: &acp::SessionId) -> acp::Result { + let key = SessionKey::from(id); + if self.contains(&key) { + return Ok(key); + } + Err(acp::Error::invalid_params().data(format!("unknown session id {id}"))) + } + + /// Marks a prompt active, rejecting overlapping turns for one session. + pub fn begin_prompt(&self, key: &SessionKey) -> acp::Result { + let mut state = self.inner.borrow_mut(); + if state.in_flight.len() >= MAX_CONCURRENT_OPERATIONS { + return Err(acp::Error::invalid_params().data(format!( + "connection cannot run more than {MAX_CONCURRENT_OPERATIONS} operations concurrently" + ))); + } + if !state.in_flight.insert(key.clone()) { + return Err(acp::Error::invalid_params() + .data(format!("session {key} already has an active prompt"))); + } + Ok(PromptGuard { + registry: self.clone(), + key: key.clone(), + }) + } + + /// Blocks session setup changes while another setup or prompt is active. + pub fn begin_setup(&self, key: &SessionKey) -> acp::Result { + let mut state = self.inner.borrow_mut(); + if state.in_flight.len() >= MAX_CONCURRENT_OPERATIONS { + return Err(acp::Error::invalid_params().data(format!( + "connection cannot run more than {MAX_CONCURRENT_OPERATIONS} operations concurrently" + ))); + } + if !state.in_flight.insert(key.clone()) { + return Err(acp::Error::invalid_params() + .data(format!("session {key} already has an active operation"))); + } + Ok(PromptGuard { + registry: self.clone(), + key: key.clone(), + }) + } + + /// Marks the session's in-flight turn as cancelled. + pub fn mark_cancelled(&self, key: &SessionKey) { + self.inner.borrow_mut().cancelled.insert(key.clone(), true); + } + + /// Clears any cancellation flag, called when a turn starts. + pub fn clear_cancelled(&self, key: &SessionKey) { + self.inner.borrow_mut().cancelled.remove(key); + } + + /// Takes the cancellation flag, called when a turn finishes. + pub fn take_cancelled(&self, key: &SessionKey) -> bool { + self.inner + .borrow_mut() + .cancelled + .remove(key) + .unwrap_or(false) + } +} + +/// Removes a session's active-prompt marker on every exit path. +pub struct PromptGuard { + registry: SessionRegistry, + key: SessionKey, +} + +impl Drop for PromptGuard { + fn drop(&mut self) { + self.registry.inner.borrow_mut().in_flight.remove(&self.key); + } +} + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn namespaced_keys_are_prefixed() { + let key = SessionKey::namespaced("abc"); + assert_eq!(key.as_str(), "acp:abc"); + assert!(key.is_namespaced()); + } + + #[test] + fn plain_keys_are_not_namespaced() { + assert!(!SessionKey::new("web:abc").is_namespaced()); + assert!(!SessionKey::new("abc").is_namespaced()); + assert!(!SessionKey::new("acp:").is_namespaced()); + assert!(!SessionKey::new("acp:../escape").is_namespaced()); + assert!(!SessionKey::new("acp:line\nbreak").is_namespaced()); + } + + #[test] + fn session_key_round_trips_through_acp_session_id() { + let key = SessionKey::namespaced("round-trip"); + let id = acp::SessionId::from(key.clone()); + assert_eq!(SessionKey::from(&id), key); + } + + #[test] + fn unknown_session_id_is_invalid_params() { + let registry = SessionRegistry::new(); + let error = registry + .resolve(&acp::SessionId::from("acp:nope".to_string())) + .expect_err("unknown session must be rejected"); + assert_eq!(error.code, acp::Error::invalid_params().code); + } + + #[test] + fn known_session_id_resolves() { + let registry = SessionRegistry::new(); + let key = SessionKey::namespaced("known"); + registry.insert(key.clone()).expect("insert known session"); + let id = acp::SessionId::from(key.clone()); + assert_eq!(registry.resolve(&id).expect("known session"), key); + } + + #[test] + fn cancellation_flag_is_take_once() { + let registry = SessionRegistry::new(); + let key = SessionKey::namespaced("cancel"); + registry.insert(key.clone()).expect("insert cancel session"); + assert!(!registry.take_cancelled(&key)); + registry.mark_cancelled(&key); + assert!(registry.take_cancelled(&key)); + assert!(!registry.take_cancelled(&key)); + } + + #[test] + fn local_ids_are_unique() { + let registry = SessionRegistry::new(); + let first = registry.next_local_id(); + let second = registry.next_local_id(); + assert_ne!(first, second); + } + + #[test] + fn setup_and_prompt_operations_are_mutually_exclusive() { + let registry = SessionRegistry::new(); + let key = SessionKey::namespaced("busy"); + registry.insert(key.clone()).expect("insert busy session"); + + let setup = registry.begin_setup(&key).expect("first operation"); + assert!(registry.begin_prompt(&key).is_err()); + drop(setup); + + let prompt = registry.begin_prompt(&key).expect("operation released"); + assert!(registry.begin_setup(&key).is_err()); + drop(prompt); + assert!(registry.begin_setup(&key).is_ok()); + } + + #[test] + fn session_count_is_bounded() { + let registry = SessionRegistry::new(); + for index in 0..MAX_SESSIONS { + registry + .insert(SessionKey::namespaced(format!("session-{index}"))) + .expect("session within limit"); + } + assert!( + registry + .ensure_capacity_for(&SessionKey::namespaced("one-too-many")) + .is_err() + ); + assert!(registry.ensure_capacity().is_err()); + } + + #[test] + fn concurrent_operations_are_bounded() { + let registry = SessionRegistry::new(); + let guards = (0..MAX_CONCURRENT_OPERATIONS) + .map(|index| { + registry + .begin_setup(&SessionKey::namespaced(format!("setup-{index}"))) + .expect("operation within limit") + }) + .collect::>(); + assert!( + registry + .begin_setup(&SessionKey::namespaced("one-too-many")) + .is_err() + ); + drop(guards); + assert!( + registry + .begin_setup(&SessionKey::namespaced("after-release")) + .is_ok() + ); + } +} diff --git a/crates/acp/src/setup.rs b/crates/acp/src/setup.rs new file mode 100644 index 0000000000..fda3dc8637 --- /dev/null +++ b/crates/acp/src/setup.rs @@ -0,0 +1,202 @@ +use std::{collections::HashSet, fmt, path::PathBuf}; + +use agent_client_protocol as acp; + +const MAX_MCP_SERVERS: usize = 16; +const MAX_MCP_ARGS: usize = 256; +const MAX_MCP_ENV_VARS: usize = 256; +const MAX_MCP_SETUP_BYTES: usize = 1024 * 1024; + +/// Validated setup supplied by an ACP client for a new or loaded session. +#[derive(Clone)] +pub struct SessionSetup { + cwd: PathBuf, + mcp_servers: Vec, +} + +impl SessionSetup { + pub async fn new(cwd: PathBuf, mcp_servers: Vec) -> acp::Result { + if !cwd.is_absolute() { + return Err(acp::Error::invalid_params().data("session cwd must be an absolute path")); + } + let cwd = tokio::fs::canonicalize(&cwd).await.map_err(|error| { + acp::Error::invalid_params().data(format!("invalid session cwd: {error}")) + })?; + if !cwd.is_dir() { + return Err(acp::Error::invalid_params().data("session cwd must be a directory")); + } + + if mcp_servers.len() > MAX_MCP_SERVERS { + return Err(acp::Error::invalid_params() + .data(format!("at most {MAX_MCP_SERVERS} MCP servers are allowed"))); + } + + let mut names = HashSet::new(); + let mut setup_bytes = 0usize; + for server in &mcp_servers { + let acp::McpServer::Stdio(server) = server else { + return Err( + acp::Error::invalid_params().data("only stdio MCP servers are supported") + ); + }; + let name = server.name.trim(); + if name.is_empty() { + return Err(acp::Error::invalid_params().data("MCP server name cannot be empty")); + } + if !names.insert(name.to_string()) { + return Err( + acp::Error::invalid_params().data(format!("duplicate MCP server name {name}")) + ); + } + if !server.command.is_absolute() { + return Err(acp::Error::invalid_params().data(format!( + "MCP server {name} command must be an absolute path" + ))); + } + if server.args.len() > MAX_MCP_ARGS { + return Err(acp::Error::invalid_params().data(format!( + "MCP server {name} has more than {MAX_MCP_ARGS} arguments" + ))); + } + if server.env.len() > MAX_MCP_ENV_VARS { + return Err(acp::Error::invalid_params().data(format!( + "MCP server {name} has more than {MAX_MCP_ENV_VARS} environment variables" + ))); + } + setup_bytes = setup_bytes + .saturating_add(name.len()) + .saturating_add(server.command.as_os_str().len()) + .saturating_add(server.args.iter().map(String::len).sum::()); + let mut environment = HashSet::new(); + for variable in &server.env { + if variable.name.is_empty() + || variable.name.contains(['=', '\0']) + || !environment.insert(variable.name.as_str()) + { + return Err(acp::Error::invalid_params().data(format!( + "MCP server {name} has an invalid or duplicate environment variable" + ))); + } + setup_bytes = setup_bytes + .saturating_add(variable.name.len()) + .saturating_add(variable.value.len()); + } + if setup_bytes > MAX_MCP_SETUP_BYTES { + return Err(acp::Error::invalid_params() + .data(format!("MCP setup exceeds {MAX_MCP_SETUP_BYTES} bytes"))); + } + } + + Ok(Self { cwd, mcp_servers }) + } + + #[must_use] + pub fn cwd(&self) -> &std::path::Path { + &self.cwd + } + + #[must_use] + pub fn mcp_servers(&self) -> &[acp::McpServer] { + &self.mcp_servers + } +} + +impl fmt::Debug for SessionSetup { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("SessionSetup") + .field("cwd", &self.cwd) + .field("mcp_server_count", &self.mcp_servers.len()) + .finish() + } +} + +#[cfg(test)] +#[allow(clippy::expect_used)] +mod tests { + use super::*; + + #[tokio::test] + async fn rejects_relative_paths() { + assert!( + SessionSetup::new(PathBuf::from("relative"), Vec::new()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn rejects_missing_paths() { + assert!( + SessionSetup::new(PathBuf::from("/path/that/does/not/exist"), Vec::new()) + .await + .is_err() + ); + } + + #[tokio::test] + async fn rejects_relative_mcp_commands() { + let server = acp::McpServerStdio::new("test", "relative-command"); + assert!( + SessionSetup::new(std::env::temp_dir(), vec![acp::McpServer::Stdio(server)]) + .await + .is_err() + ); + } + + #[tokio::test] + async fn rejects_duplicate_mcp_server_names() { + let first = acp::McpServerStdio::new("test", "/bin/first"); + let second = acp::McpServerStdio::new("test", "/bin/second"); + assert!( + SessionSetup::new(std::env::temp_dir(), vec![ + acp::McpServer::Stdio(first), + acp::McpServer::Stdio(second) + ],) + .await + .is_err() + ); + } + + #[tokio::test] + async fn rejects_duplicate_mcp_environment_names() { + let server = acp::McpServerStdio::new("test", "/bin/test").env(vec![ + acp::EnvVariable::new("TOKEN", "first"), + acp::EnvVariable::new("TOKEN", "second"), + ]); + assert!( + SessionSetup::new(std::env::temp_dir(), vec![acp::McpServer::Stdio(server)]) + .await + .is_err() + ); + } + + #[tokio::test] + async fn debug_does_not_expose_mcp_environment() { + let server = acp::McpServerStdio::new("test", "/bin/test") + .env(vec![acp::EnvVariable::new("TOKEN", "secret")]); + let setup = SessionSetup::new(std::env::temp_dir(), vec![acp::McpServer::Stdio(server)]) + .await + .expect("valid setup"); + let debug = format!("{setup:?}"); + assert!(!debug.contains("secret")); + assert!(!debug.contains("TOKEN")); + } + + #[tokio::test] + async fn rejects_too_many_mcp_servers() { + let servers = (0..=MAX_MCP_SERVERS) + .map(|index| { + acp::McpServer::Stdio(acp::McpServerStdio::new( + format!("server-{index}"), + "/bin/test", + )) + }) + .collect(); + assert!( + SessionSetup::new(std::env::temp_dir(), servers) + .await + .is_err() + ); + } +} diff --git a/crates/acp/tests/protocol.rs b/crates/acp/tests/protocol.rs new file mode 100644 index 0000000000..4d62b490fa --- /dev/null +++ b/crates/acp/tests/protocol.rs @@ -0,0 +1,740 @@ +//! End-to-end protocol coverage over an in-memory loopback. +//! +//! A real `ClientSideConnection` drives a real `AgentSideConnection` over a +//! duplex pipe, so these exercise actual JSON-RPC framing, request/response +//! correlation, and notification ordering — no subprocess, no network. + +#![allow(clippy::unwrap_used, clippy::expect_used)] + +use std::{ + cell::RefCell, + pin::Pin, + rc::Rc, + sync::{ + Arc, Mutex, + atomic::{AtomicBool, Ordering}, + }, + task::{Context, Poll}, + time::Duration, +}; + +use { + agent_client_protocol::{self as acp, Agent as _}, + async_trait::async_trait, + tokio::{ + io::{AsyncRead, AsyncWrite, DuplexStream, ReadBuf}, + sync::Notify, + task::LocalSet, + time::timeout, + }, + tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}, +}; + +use moltis_acp::{ + AcpBackend, BackendCapabilities, EchoBackend, MoltisAgent, SessionKey, SessionSetup, + TurnUpdates, +}; + +/// Every await in these tests should settle in milliseconds; the budget exists +/// to turn a hang into a failure rather than a stuck suite. +const TIMEOUT: Duration = Duration::from_secs(5); + +// ---------------------------------------------------------------- test client + +/// Client that records every `session/update` it receives, in order. +#[derive(Clone, Default)] +struct TestClient { + updates: Rc>>, +} + +impl TestClient { + /// Concatenates the streamed agent message chunks. + fn agent_text(&self) -> String { + self.updates + .borrow() + .iter() + .filter_map(|update| match update { + acp::SessionUpdate::AgentMessageChunk(chunk) => match &chunk.content { + acp::ContentBlock::Text(text) => Some(text.text.clone()), + _ => None, + }, + _ => None, + }) + .collect() + } + + fn kinds(&self) -> Vec<&'static str> { + self.updates + .borrow() + .iter() + .map(|update| match update { + acp::SessionUpdate::UserMessageChunk(_) => "user", + acp::SessionUpdate::AgentMessageChunk(_) => "message", + acp::SessionUpdate::AgentThoughtChunk(_) => "thought", + acp::SessionUpdate::ToolCall(_) => "tool_call", + acp::SessionUpdate::ToolCallUpdate(_) => "tool_call_update", + _ => "other", + }) + .collect() + } +} + +#[async_trait(?Send)] +impl acp::Client for TestClient { + async fn request_permission( + &self, + _args: acp::RequestPermissionRequest, + ) -> acp::Result { + Ok(acp::RequestPermissionResponse::new( + acp::RequestPermissionOutcome::Cancelled, + )) + } + + async fn session_notification(&self, args: acp::SessionNotification) -> acp::Result<()> { + self.updates.borrow_mut().push(args.update); + Ok(()) + } +} + +// -------------------------------------------------------------- test backends + +/// Backend whose turn blocks until released, so a cancel can be delivered while +/// `session/prompt` is genuinely in flight. +#[derive(Default)] +struct BlockingBackend { + started: Arc, + release: Arc, +} + +#[async_trait] +impl AcpBackend for BlockingBackend { + async fn create_session(&self, _setup: &SessionSetup) -> anyhow::Result { + Ok(SessionKey::namespaced("blocking")) + } + + async fn prompt( + &self, + _key: &SessionKey, + _prompt: String, + updates: TurnUpdates, + ) -> anyhow::Result { + updates.agent_message("working"); + self.started.notify_one(); + self.release.notified().await; + // Reports success: the protocol layer is responsible for turning a + // cancelled turn into `StopReason::Cancelled`. + Ok(acp::StopReason::EndTurn) + } + + async fn cancel(&self, _key: &SessionKey) -> anyhow::Result<()> { + self.release.notify_one(); + Ok(()) + } +} + +/// Backend that always fails, to check errors surface as JSON-RPC errors. +struct FailingBackend; + +#[async_trait] +impl AcpBackend for FailingBackend { + async fn create_session(&self, _setup: &SessionSetup) -> anyhow::Result { + Err(anyhow::anyhow!("no sessions today")) + } + + async fn prompt( + &self, + _key: &SessionKey, + _prompt: String, + _updates: TurnUpdates, + ) -> anyhow::Result { + Err(anyhow::anyhow!("turn exploded")) + } + + async fn cancel(&self, _key: &SessionKey) -> anyhow::Result<()> { + Ok(()) + } +} + +/// Backend advertising and serving `session/load`. +struct ResumableBackend; + +#[async_trait] +impl AcpBackend for ResumableBackend { + async fn create_session(&self, _setup: &SessionSetup) -> anyhow::Result { + Ok(SessionKey::namespaced("resumable")) + } + + async fn load_session( + &self, + key: &SessionKey, + _setup: &SessionSetup, + ) -> anyhow::Result> { + if key.as_str() == "acp:missing" { + return Err(moltis_acp::SessionNotFound.into()); + } + if key.as_str() == "acp:broken" { + return Err(anyhow::anyhow!("history store unavailable")); + } + Ok(vec![ + acp::SessionUpdate::UserMessageChunk(acp::ContentChunk::new(acp::ContentBlock::from( + "earlier question".to_string(), + ))), + acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::from( + "earlier answer".to_string(), + ))), + ]) + } + + async fn prompt( + &self, + _key: &SessionKey, + _prompt: String, + _updates: TurnUpdates, + ) -> anyhow::Result { + Ok(acp::StopReason::EndTurn) + } + + async fn cancel(&self, _key: &SessionKey) -> anyhow::Result<()> { + Ok(()) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities { load_session: true } + } +} + +/// Backend that mints keys outside the ACP namespace, and records whether +/// `load_session` was reached. +#[derive(Default)] +struct EscapingBackend { + loaded: Arc, +} + +#[async_trait] +impl AcpBackend for EscapingBackend { + async fn create_session(&self, _setup: &SessionSetup) -> anyhow::Result { + Ok(SessionKey::new("web:escaped")) + } + + async fn load_session( + &self, + _key: &SessionKey, + _setup: &SessionSetup, + ) -> anyhow::Result> { + self.loaded.store(true, Ordering::SeqCst); + Ok(Vec::new()) + } + + async fn prompt( + &self, + _key: &SessionKey, + _prompt: String, + _updates: TurnUpdates, + ) -> anyhow::Result { + Ok(acp::StopReason::EndTurn) + } + + async fn cancel(&self, _key: &SessionKey) -> anyhow::Result<()> { + Ok(()) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities { load_session: true } + } +} + +// ------------------------------------------------------------------- plumbing + +/// Duplex stream that records every byte written through it, so a test can +/// assert on exactly what the agent put on the wire. +struct Tee { + inner: DuplexStream, + seen: Arc>>, +} + +impl AsyncRead for Tee { + fn poll_read( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &mut ReadBuf<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_read(cx, buf) + } +} + +impl AsyncWrite for Tee { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + let written = std::task::ready!(Pin::new(&mut self.inner).poll_write(cx, buf))?; + self.seen + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .extend_from_slice(&buf[..written]); + Poll::Ready(Ok(written)) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +/// A live loopback connection. Holds the agent side alive for the duration of +/// the test: the agent keeps only a weak reference to its own connection. +struct Harness { + client: Rc, + _agent: Rc, + _agent_conn: Rc, + wire: Arc>>, +} + +fn connect(backend: Arc, client: TestClient) -> Harness { + let (agent_io, client_io) = tokio::io::duplex(64 * 1024); + let wire = Arc::new(Mutex::new(Vec::new())); + let (agent_read, agent_write) = tokio::io::split(Tee { + inner: agent_io, + seen: Arc::clone(&wire), + }); + let (client_read, client_write) = tokio::io::split(client_io); + + let agent = Rc::new(MoltisAgent::new(backend)); + let (agent_conn, agent_io_task) = acp::AgentSideConnection::new( + Rc::clone(&agent), + agent_write.compat_write(), + agent_read.compat(), + |future| { + tokio::task::spawn_local(future); + }, + ); + let agent_conn = Rc::new(agent_conn); + agent.set_connection(&agent_conn); + + let (client_conn, client_io_task) = acp::ClientSideConnection::new( + client, + client_write.compat_write(), + client_read.compat(), + |future| { + tokio::task::spawn_local(future); + }, + ); + + tokio::task::spawn_local(async move { + let _ = agent_io_task.await; + }); + tokio::task::spawn_local(async move { + let _ = client_io_task.await; + }); + + Harness { + client: Rc::new(client_conn), + _agent: agent, + _agent_conn: agent_conn, + wire, + } +} + +async fn initialize(harness: &Harness) -> acp::InitializeResponse { + timeout( + TIMEOUT, + harness.client.initialize( + acp::InitializeRequest::new(acp::ProtocolVersion::LATEST) + .client_info(acp::Implementation::new("test-client", "0.0.0").title("Test Client")), + ), + ) + .await + .expect("initialize timed out") + .expect("initialize failed") +} + +async fn new_session(harness: &Harness) -> acp::SessionId { + timeout( + TIMEOUT, + harness + .client + .new_session(acp::NewSessionRequest::new(std::env::temp_dir())), + ) + .await + .expect("session/new timed out") + .expect("session/new failed") + .session_id +} + +/// Runs `body` on a `LocalSet`, since the protocol handler is `!Send`. +fn run_local>(body: F) { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("runtime"); + LocalSet::new().block_on(&runtime, body); +} + +// ---------------------------------------------------------------------- tests + +#[test] +fn initialize_negotiates_version_and_reports_capabilities() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + let response = initialize(&harness).await; + + assert_eq!(response.protocol_version, acp::ProtocolVersion::LATEST); + // The echo backend has nothing to resume, so it must not claim it can. + assert!(!response.agent_capabilities.load_session); + let info = response.agent_info.expect("agent_info advertised"); + assert_eq!(info.name, "moltis"); + }); +} + +#[test] +fn initialize_reports_load_session_when_the_backend_supports_it() { + run_local(async { + let harness = connect(Arc::new(ResumableBackend), TestClient::default()); + let response = initialize(&harness).await; + assert!(response.agent_capabilities.load_session); + }); +} + +#[test] +fn authenticate_succeeds_without_credentials() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + let result = + timeout( + TIMEOUT, + harness.client.authenticate(acp::AuthenticateRequest::new( + acp::AuthMethodId::from("none".to_string()), + )), + ) + .await + .expect("authenticate timed out"); + assert!(result.is_ok(), "local clients need no authentication"); + }); +} + +#[test] +fn prompt_streams_updates_before_returning_a_stop_reason() { + run_local(async { + let client = TestClient::default(); + let harness = connect(Arc::new(EchoBackend::new()), client.clone()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + + let response = timeout( + TIMEOUT, + harness + .client + .prompt(acp::PromptRequest::new(session_id, vec![ + acp::ContentBlock::from("hello world".to_string()), + ])), + ) + .await + .expect("prompt timed out") + .expect("prompt failed"); + + assert_eq!(response.stop_reason, acp::StopReason::EndTurn); + // Every update must already have arrived by the time the call resolves, + // otherwise the agent looks mute and then dumps everything at once. + assert_eq!(client.agent_text(), "hello world"); + let kinds = client.kinds(); + assert_eq!( + kinds.first(), + Some(&"thought"), + "thought precedes the reply: {kinds:?}" + ); + assert!(kinds.iter().filter(|kind| **kind == "message").count() > 1); + }); +} + +#[test] +fn cancel_resolves_the_pending_prompt_as_cancelled() { + run_local(async { + let backend = Arc::new(BlockingBackend::default()); + let started = Arc::clone(&backend.started); + let harness = connect(backend, TestClient::default()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + + let client = Rc::clone(&harness.client); + let prompt_session = session_id.clone(); + let turn = tokio::task::spawn_local(async move { + client + .prompt(acp::PromptRequest::new(prompt_session, vec![ + acp::ContentBlock::from("start working".to_string()), + ])) + .await + }); + + // Cancel only once the turn is genuinely in flight. + timeout(TIMEOUT, started.notified()) + .await + .expect("turn never started"); + timeout( + TIMEOUT, + harness + .client + .cancel(acp::CancelNotification::new(session_id)), + ) + .await + .expect("cancel timed out") + .expect("cancel failed"); + + let response = timeout(TIMEOUT, turn) + .await + .expect("prompt did not resolve promptly after cancel") + .expect("prompt task panicked") + .expect("prompt failed"); + assert_eq!(response.stop_reason, acp::StopReason::Cancelled); + }); +} + +#[test] +fn unknown_session_id_is_rejected_with_invalid_params() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + + let error = timeout( + TIMEOUT, + harness.client.prompt(acp::PromptRequest::new( + acp::SessionId::from("acp:never-created".to_string()), + vec![acp::ContentBlock::from("hello".to_string())], + )), + ) + .await + .expect("prompt timed out") + .expect_err("unknown session must be rejected"); + + assert_eq!(error.code, acp::Error::invalid_params().code); + }); +} + +#[test] +fn cancel_for_unknown_session_does_not_panic() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + + // Notifications carry no reply, so the check here is that the + // connection survives and still serves the next request. + let _ = timeout( + TIMEOUT, + harness + .client + .cancel(acp::CancelNotification::new(acp::SessionId::from( + "acp:never-created".to_string(), + ))), + ) + .await + .expect("cancel timed out"); + + let session_id = new_session(&harness).await; + assert!(session_id.to_string().starts_with("acp:")); + }); +} + +#[test] +fn backend_failures_surface_as_errors_not_panics() { + run_local(async { + let harness = connect(Arc::new(FailingBackend), TestClient::default()); + initialize(&harness).await; + + let error = timeout( + TIMEOUT, + harness + .client + .new_session(acp::NewSessionRequest::new(std::env::temp_dir())), + ) + .await + .expect("session/new timed out") + .expect_err("failing backend must produce an error"); + assert_eq!(error.code, acp::Error::internal_error().code); + }); +} + +#[test] +fn load_session_replays_history_before_responding() { + run_local(async { + let client = TestClient::default(); + let harness = connect(Arc::new(ResumableBackend), client.clone()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + + timeout( + TIMEOUT, + harness.client.load_session(acp::LoadSessionRequest::new( + session_id, + std::env::temp_dir(), + )), + ) + .await + .expect("session/load timed out") + .expect("session/load failed"); + + assert_eq!(client.kinds(), vec!["user", "message"]); + }); +} + +#[test] +fn load_session_is_method_not_found_when_unsupported() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + + let error = timeout( + TIMEOUT, + harness.client.load_session(acp::LoadSessionRequest::new( + session_id, + std::env::temp_dir(), + )), + ) + .await + .expect("session/load timed out") + .expect_err("echo backend cannot resume"); + assert_eq!(error.code, acp::Error::method_not_found().code); + }); +} + +#[test] +fn load_session_classifies_missing_and_operational_failures() { + run_local(async { + let harness = connect(Arc::new(ResumableBackend), TestClient::default()); + initialize(&harness).await; + + let missing = harness + .client + .load_session(acp::LoadSessionRequest::new( + acp::SessionId::from("acp:missing".to_string()), + std::env::temp_dir(), + )) + .await + .expect_err("missing session should fail"); + assert_eq!(missing.code, acp::Error::invalid_params().code); + + let broken = harness + .client + .load_session(acp::LoadSessionRequest::new( + acp::SessionId::from("acp:broken".to_string()), + std::env::temp_dir(), + )) + .await + .expect_err("history read failure should fail"); + assert_eq!(broken.code, acp::Error::internal_error().code); + }); +} + +#[test] +fn load_session_refuses_ids_outside_the_acp_namespace() { + run_local(async { + let backend = Arc::new(EscapingBackend::default()); + let loaded = Arc::clone(&backend.loaded); + let harness = connect(backend, TestClient::default()); + initialize(&harness).await; + + // A Web UI session key, named directly by the client. Nothing in the + // registry rules it out — only the namespace does. + let error = timeout( + TIMEOUT, + harness.client.load_session(acp::LoadSessionRequest::new( + acp::SessionId::from("web:someone-elses-session".to_string()), + std::env::temp_dir(), + )), + ) + .await + .expect("session/load timed out") + .expect_err("foreign session ids must be refused"); + + assert_eq!(error.code, acp::Error::invalid_params().code); + assert!( + !loaded.load(Ordering::SeqCst), + "backend must never be asked to resolve an out-of-namespace key" + ); + }); +} + +#[test] +fn new_session_rejects_a_backend_key_outside_the_acp_namespace() { + run_local(async { + let harness = connect(Arc::new(EscapingBackend::default()), TestClient::default()); + initialize(&harness).await; + + let error = timeout( + TIMEOUT, + harness + .client + .new_session(acp::NewSessionRequest::new(std::env::temp_dir())), + ) + .await + .expect("session/new timed out") + .expect_err("an out-of-namespace backend key must not reach the client"); + assert_eq!(error.code, acp::Error::internal_error().code); + }); +} + +#[test] +fn sessions_are_namespaced_so_they_cannot_collide_with_other_surfaces() { + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + assert!( + session_id.to_string().starts_with("acp:"), + "ACP sessions must live in their own namespace: {session_id}" + ); + }); +} + +/// stdout is the wire: with a subscriber installed and `#[instrument]` spans +/// firing throughout the turn, every byte the agent emits must still be +/// protocol framing. +#[test] +fn wire_carries_only_protocol_framing_while_tracing_is_active() { + let _ = tracing_subscriber::fmt() + .with_writer(std::io::stderr) + .with_max_level(tracing::Level::TRACE) + .try_init(); + + run_local(async { + let harness = connect(Arc::new(EchoBackend::new()), TestClient::default()); + initialize(&harness).await; + let session_id = new_session(&harness).await; + timeout( + TIMEOUT, + harness + .client + .prompt(acp::PromptRequest::new(session_id, vec![ + acp::ContentBlock::from("trace me".to_string()), + ])), + ) + .await + .expect("prompt timed out") + .expect("prompt failed"); + + let bytes = harness + .wire + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let text = String::from_utf8(bytes).expect("wire must be UTF-8"); + assert!(!text.is_empty(), "expected protocol traffic"); + + let mut frames = 0; + for line in text.split('\n').filter(|line| !line.is_empty()) { + let frame: serde_json::Value = serde_json::from_str(line) + .unwrap_or_else(|error| panic!("non-JSON byte on the wire: {error}: {line:?}")); + assert_eq!( + frame.get("jsonrpc").and_then(serde_json::Value::as_str), + Some("2.0"), + "every frame must be JSON-RPC: {line}" + ); + frames += 1; + } + assert!(frames >= 3, "expected initialize + session/new + prompt"); + }); +} diff --git a/crates/agents/src/prompt/formatting.rs b/crates/agents/src/prompt/formatting.rs index bb42308951..1b3669a582 100644 --- a/crates/agents/src/prompt/formatting.rs +++ b/crates/agents/src/prompt/formatting.rs @@ -168,6 +168,7 @@ pub(crate) fn format_host_runtime_line(host: &PromptHostRuntimeContext) -> Optio ("channel_chat_id", host.channel_chat_id.as_deref()), ("channel_chat_type", host.channel_chat_type.as_deref()), ("data_dir", host.data_dir.as_deref()), + ("working_dir", host.working_dir.as_deref()), ("docs_path", host.docs_path.as_deref()), ("config_template", host.config_template_path.as_deref()), ] { diff --git a/crates/agents/src/prompt/tests.rs b/crates/agents/src/prompt/tests.rs index 89d2dd248b..44e346e560 100644 --- a/crates/agents/src/prompt/tests.rs +++ b/crates/agents/src/prompt/tests.rs @@ -360,6 +360,7 @@ fn test_runtime_context_injected_when_provided() { channel_chat_type: None, channel_sender_id: None, data_dir: Some("/home/moltis/.moltis".into()), + working_dir: Some("/workspace/project".into()), docs_path: None, config_template_path: None, sudo_non_interactive: Some(true), @@ -412,6 +413,7 @@ fn test_runtime_context_injected_when_provided() { assert!(prompt.contains("provider=openai")); assert!(prompt.contains("model=gpt-5")); assert!(prompt.contains("data_dir=/home/moltis/.moltis")); + assert!(prompt.contains("working_dir=/workspace/project")); assert!(prompt.contains("sudo_non_interactive=true")); assert!(prompt.contains("sudo_status=passwordless")); assert!(prompt.contains("timezone=Europe/Paris")); diff --git a/crates/agents/src/prompt/types.rs b/crates/agents/src/prompt/types.rs index 04b6539128..f139d20ee8 100644 --- a/crates/agents/src/prompt/types.rs +++ b/crates/agents/src/prompt/types.rs @@ -55,6 +55,7 @@ pub struct PromptHostRuntimeContext { pub channel_chat_type: Option, pub channel_sender_id: Option, pub data_dir: Option, + pub working_dir: Option, pub docs_path: Option, pub config_template_path: Option, pub sudo_non_interactive: Option, diff --git a/crates/agents/src/runner/non_streaming.rs b/crates/agents/src/runner/non_streaming.rs index bbfc375091..d85ba0b437 100644 --- a/crates/agents/src/runner/non_streaming.rs +++ b/crates/agents/src/runner/non_streaming.rs @@ -253,7 +253,11 @@ pub async fn run_agent_loop_with_context_and_limits( messages_count = messages.len(), "calling LLM" ); - trace!(iteration = iterations, messages = ?messages, "LLM request messages"); + trace!( + iteration = iterations, + messages_count = messages.len(), + "LLM request prepared" + ); // Dispatch BeforeLLMCall hook — may block the LLM call. if let Some(ref hooks) = hook_registry { @@ -340,7 +344,11 @@ pub async fn run_agent_loop_with_context_and_limits( "LLM response received" ); if let Some(ref text) = response.text { - trace!(iteration = iterations, text = %text, "LLM response text"); + trace!( + iteration = iterations, + text_bytes = text.len(), + "LLM response text available" + ); } // Fallback: parse tool calls from model text if the provider returned @@ -405,7 +413,10 @@ pub async fn run_agent_loop_with_context_and_limits( && let Some(command) = explicit_shell_command.as_ref() && tools.get("exec").is_some() { - info!(command = %command, "forcing exec tool call from explicit /sh command"); + info!( + command_bytes = command.len(), + "forcing exec tool call from explicit /sh command" + ); // Preserve the model's planning/reasoning text on the assistant // tool-call message. Some providers (e.g. Moonshot thinking mode) // require this history field for follow-up tool turns. @@ -444,7 +455,7 @@ pub async fn run_agent_loop_with_context_and_limits( info!( iteration = iterations, tool_name = %tc.name, - arguments = %tc.arguments, + argument_bytes = serde_json::to_vec(&tc.arguments).map_or(0, |value| value.len()), "LLM requested tool call" ); } @@ -768,7 +779,7 @@ pub async fn run_agent_loop_with_context_and_limits( { if success { info!(tool = %tc.name, id = %tc.id, "tool execution succeeded"); - trace!(tool = %tc.name, result = %result, "tool result"); + trace!(tool = %tc.name, "tool result available"); } else if rejected { warn!( tool = %tc.name, @@ -776,7 +787,7 @@ pub async fn run_agent_loop_with_context_and_limits( "tool call rejected before execution by pre-dispatch validation" ); } else { - warn!(tool = %tc.name, id = %tc.id, error = %error.as_deref().unwrap_or(""), "tool execution failed"); + warn!(tool = %tc.name, id = %tc.id, has_error = error.is_some(), "tool execution failed"); } // Record outcome in the loop detector. Use explicit success/failure @@ -855,7 +866,7 @@ pub async fn run_agent_loop_with_context_and_limits( result_len = tool_result_str.len(), "appending tool result to messages" ); - trace!(tool = %tc.name, content = %tool_result_str, "tool result message content"); + trace!(tool = %tc.name, content_bytes = tool_result_str.len(), "tool result message prepared"); messages.push(ChatMessage::tool(&tc.id, &tool_result_str)); } diff --git a/crates/agents/src/runner/streaming.rs b/crates/agents/src/runner/streaming.rs index 941be50f59..166948e888 100644 --- a/crates/agents/src/runner/streaming.rs +++ b/crates/agents/src/runner/streaming.rs @@ -251,7 +251,11 @@ pub async fn run_agent_loop_streaming_with_limits( messages_count = messages.len(), "calling LLM (streaming)" ); - trace!(iteration = iterations, messages = ?messages, "LLM request messages"); + trace!( + iteration = iterations, + messages_count = messages.len(), + "LLM request prepared" + ); // Dispatch BeforeLLMCall hook — may block the LLM call. if let Some(ref hooks) = hook_registry { @@ -459,12 +463,9 @@ pub async fn run_agent_loop_streaming_with_limits( // Use stream_idx_to_vec_pos to map streaming indices (which may not // start at 0) to the actual position in the tool_calls vec. for (stream_idx, args_str) in &tool_call_args { - // Emit raw accumulated string at debug level so future variants of - // "default to {} because no deltas arrived" can be diagnosed - // without a repro (issue #658). debug!( stream_idx, - args_str = %args_str, + argument_bytes = args_str.len(), "finalizing tool call args" ); if let Some(&vec_pos) = stream_idx_to_vec_pos.get(stream_idx) @@ -537,7 +538,10 @@ pub async fn run_agent_loop_streaming_with_limits( && let Some(command) = explicit_shell_command.as_ref() && tools.get("exec").is_some() { - info!(command = %command, "forcing exec tool call from explicit /sh command"); + info!( + command_bytes = command.len(), + "forcing exec tool call from explicit /sh command" + ); // Preserve streamed reasoning/planning text on the assistant tool // message so providers that validate thinking history accept the // next iteration. @@ -899,7 +903,7 @@ pub async fn run_agent_loop_streaming_with_limits( for (tc, (success, mut result, error, rejected)) in tool_calls.iter().zip(results) { if success { info!(tool = %tc.name, id = %tc.id, "tool execution succeeded"); - trace!(tool = %tc.name, result = %result, "tool result"); + trace!(tool = %tc.name, "tool result available"); } else if rejected { warn!( tool = %tc.name, @@ -907,7 +911,7 @@ pub async fn run_agent_loop_streaming_with_limits( "tool call rejected before execution by pre-dispatch validation" ); } else { - warn!(tool = %tc.name, id = %tc.id, error = %error.as_deref().unwrap_or(""), "tool execution failed"); + warn!(tool = %tc.name, id = %tc.id, has_error = error.is_some(), "tool execution failed"); } // Record outcome in the loop detector (issue #658). @@ -984,7 +988,7 @@ pub async fn run_agent_loop_streaming_with_limits( result_len = tool_result_str.len(), "appending tool result to messages" ); - trace!(tool = %tc.name, content = %tool_result_str, "tool result message content"); + trace!(tool = %tc.name, content_bytes = tool_result_str.len(), "tool result message prepared"); messages.push(ChatMessage::tool(&tc.id, &tool_result_str)); } @@ -1004,7 +1008,7 @@ pub async fn run_agent_loop_streaming_with_limits( let mut guard = inbox.lock().await; if !guard.is_empty() { let combined = guard.drain(..).collect::>().join("\n"); - debug!(steer_text = %combined, "injecting /steer guidance"); + debug!(steer_bytes = combined.len(), "injecting /steer guidance"); messages.push(ChatMessage::system(format!( "[Steering note from the user — adjust your approach accordingly]: {combined}" ))); diff --git a/crates/agents/src/tool_registry.rs b/crates/agents/src/tool_registry.rs index a4464a9cdd..1708b4b6ac 100644 --- a/crates/agents/src/tool_registry.rs +++ b/crates/agents/src/tool_registry.rs @@ -296,6 +296,16 @@ impl ToolRegistry { self.clone_allowed_entries(|name, _| predicate(name)) } + /// Merge another registry while preserving each tool's source metadata. + pub fn extend_from(&mut self, other: &ToolRegistry) { + self.tools.extend(other.tools.iter().map(|(name, entry)| { + (name.clone(), ToolEntry { + tool: Arc::clone(&entry.tool), + source: entry.source.clone(), + }) + })); + } + /// Clone the registry keeping only tools whose name and source metadata match `predicate`. pub fn clone_allowed_entries(&self, mut predicate: F) -> ToolRegistry where diff --git a/crates/benchmarks/benches/boot.rs b/crates/benchmarks/benches/boot.rs index 10c7da0930..38f7e53787 100644 --- a/crates/benchmarks/benches/boot.rs +++ b/crates/benchmarks/benches/boot.rs @@ -87,8 +87,9 @@ const SESSION_KEYS: &[&str] = &[ ]; #[divan::bench(args = SESSION_KEYS)] -fn session_key_to_filename(key: &str) -> String { - divan::black_box(moltis_sessions::store::SessionStore::key_to_filename(key)) +fn session_history_path(bencher: divan::Bencher, key: &str) { + let store = moltis_sessions::store::SessionStore::new(std::path::PathBuf::from("sessions")); + bencher.bench_local(|| divan::black_box(store.history_path_for(divan::black_box(key)))); } fn build_sanitize_input(payload_bytes: usize) -> String { diff --git a/crates/chat/src/agent_loop.rs b/crates/chat/src/agent_loop.rs index 0f62e43a2b..f2ade9f69e 100644 --- a/crates/chat/src/agent_loop.rs +++ b/crates/chat/src/agent_loop.rs @@ -329,7 +329,7 @@ pub(crate) async fn run_explicit_shell_command( state: &Arc, run_id: &str, tool_registry: &Arc>, - session_store: &Arc, + session_store: Option<&Arc>, terminal_runs: &Arc>>, session_key: &str, command: &str, @@ -337,6 +337,7 @@ pub(crate) async fn run_explicit_shell_command( accept_language: Option, conn_id: Option, client_seq: Option, + working_dir: Option, ) -> AssistantTurnOutput { let started = Instant::now(); let tool_call_id = format!("sh_{}", uuid::Uuid::new_v4().simple()); @@ -370,6 +371,9 @@ pub(crate) async fn run_explicit_shell_command( if let Some(cid) = conn_id.as_deref() { exec_params["_conn_id"] = serde_json::json!(cid); } + if let Some(directory) = working_dir { + exec_params["_working_dir"] = serde_json::json!(directory); + } let exec_tool = { let registry = tool_registry.read().await; @@ -407,15 +411,17 @@ pub(crate) async fn run_explicit_shell_command( Some(capped.clone()), None, ); - persist_tool_history_pair( - session_store, - session_key, - assistant_tool_call_msg, - tool_result_msg, - "failed to persist direct /sh assistant tool call", - "failed to persist direct /sh tool result", - ) - .await; + if let Some(session_store) = session_store { + persist_tool_history_pair( + session_store, + session_key, + assistant_tool_call_msg, + tool_result_msg, + "failed to persist direct /sh assistant tool call", + "failed to persist direct /sh tool result", + ) + .await; + } broadcast( state, @@ -460,15 +466,17 @@ pub(crate) async fn run_explicit_shell_command( None, Some(error_text.clone()), ); - persist_tool_history_pair( - session_store, - session_key, - assistant_tool_call_msg, - tool_result_msg, - "failed to persist direct /sh assistant tool call", - "failed to persist direct /sh tool error", - ) - .await; + if let Some(session_store) = session_store { + persist_tool_history_pair( + session_store, + session_key, + assistant_tool_call_msg, + tool_result_msg, + "failed to persist direct /sh assistant tool call", + "failed to persist direct /sh tool error", + ) + .await; + } broadcast( state, diff --git a/crates/chat/src/message.rs b/crates/chat/src/message.rs index d8fd721304..4da1aed91e 100644 --- a/crates/chat/src/message.rs +++ b/crates/chat/src/message.rs @@ -278,8 +278,7 @@ pub(crate) fn user_audio_path_from_params(params: &Value, session_key: &str) -> return None; } - let key = SessionStore::key_to_filename(session_key); - Some(format!("media/{key}/{filename}")) + SessionStore::media_reference(session_key, filename).ok() } pub(crate) fn user_documents_from_params( @@ -288,7 +287,6 @@ pub(crate) fn user_documents_from_params( session_store: &SessionStore, ) -> Option> { let documents = params.get("_document_files")?.as_array()?; - let media_dir_key = SessionStore::key_to_filename(session_key); let mut parsed = Vec::new(); for document in documents { @@ -304,18 +302,19 @@ pub(crate) fn user_documents_from_params( let display_name = sanitize_user_document_display_name(&document.display_name) .unwrap_or_else(|| stored_filename.to_string()); + let Ok(media_ref) = SessionStore::media_reference(session_key, stored_filename) else { + continue; + }; + let Ok(absolute_path) = session_store.media_path_for(session_key, stored_filename) else { + continue; + }; parsed.push(UserDocument { display_name, stored_filename: stored_filename.to_string(), mime_type: mime_type.to_string(), size_bytes: document.size_bytes, - media_ref: format!("media/{media_dir_key}/{stored_filename}"), - absolute_path: Some( - session_store - .media_path_for(session_key, stored_filename) - .to_string_lossy() - .to_string(), - ), + media_ref, + absolute_path: Some(absolute_path.to_string_lossy().to_string()), }); } diff --git a/crates/chat/src/prompt.rs b/crates/chat/src/prompt.rs index 0476fada44..ae9ddb5d33 100644 --- a/crates/chat/src/prompt.rs +++ b/crates/chat/src/prompt.rs @@ -407,6 +407,10 @@ pub(crate) fn build_tool_context( if let Some(cid) = conn_id { tool_context["_conn_id"] = serde_json::json!(cid); } + if let Some(working_dir) = runtime_context.and_then(|context| context.host.working_dir.as_ref()) + { + tool_context["_working_dir"] = serde_json::json!(working_dir); + } tool_context } @@ -754,4 +758,18 @@ mod tests { assert_eq!(filtered.len(), 1); assert_eq!(filtered[0].name, "web-search"); } + + #[test] + fn working_directory_uses_reserved_tool_metadata() { + let mut runtime = PromptRuntimeContext::default(); + runtime.host.working_dir = Some("/workspace/project".into()); + + let context = build_tool_context("acp:test", None, None, Some(&runtime)); + + assert_eq!( + context.get("_working_dir").and_then(Value::as_str), + Some("/workspace/project") + ); + assert!(context.get("working_dir").is_none()); + } } diff --git a/crates/chat/src/run_with_tools.rs b/crates/chat/src/run_with_tools.rs index 22be1d4b31..eb1c51cd6f 100644 --- a/crates/chat/src/run_with_tools.rs +++ b/crates/chat/src/run_with_tools.rs @@ -616,9 +616,12 @@ pub(crate) async fn run_with_tools( warn!("failed to save screenshot media: {e}"); } }); - let sanitized = SessionStore::key_to_filename(&sk_media); - r["screenshot"] = - Value::String(format!("media/{sanitized}/{tool_call_id}.png")); + if let Ok(media_ref) = SessionStore::media_reference( + &sk_media, + &format!("{tool_call_id}.png"), + ) { + r["screenshot"] = Value::String(media_ref); + } } // If screenshot is still a data URI (decode failed), strip it. let strip_screenshot = r @@ -1111,7 +1114,7 @@ pub(crate) async fn run_with_tools( run_id, iterations, tool_calls = tool_calls_made, - response = %display_text, + response_bytes = display_text.len(), silent = is_silent, "agent run complete" ); diff --git a/crates/chat/src/service/chat_impl.rs b/crates/chat/src/service/chat_impl.rs index 72d7c0a5ac..37ad82be4a 100644 --- a/crates/chat/src/service/chat_impl.rs +++ b/crates/chat/src/service/chat_impl.rs @@ -1,23 +1,20 @@ //! `ChatService` trait implementation for `LiveChatService`. +mod queue; mod send; -use std::{ - collections::{HashMap, HashSet}, - path::Path, - sync::Arc, -}; +use std::{path::Path, sync::Arc}; use { async_trait::async_trait, serde_json::Value, - tokio::sync::RwLock, + tokio::sync::oneshot, tracing::{debug, info, warn}, }; use { moltis_agents::{ - ChatMessage, UserContent, + ChatMessage, model::values_to_chat_messages_with_tool_result_limit, prompt::{ build_system_prompt_minimal_runtime_details, @@ -35,19 +32,12 @@ use crate::{ channels::notify_channels_of_compaction, compaction_run, memory_tools::AgentScopedMemoryWriter, - message::{ - infer_reply_medium, user_audio_path_from_params, user_documents_for_persistence, - user_documents_from_params, - }, prompt::{ apply_request_runtime_context, apply_runtime_tool_filters, build_policy_context, build_prompt_runtime_context, clear_prompt_memory_snapshot, discover_skills_if_enabled, filter_skills_for_agent, load_prompt_persona_for_agent, load_prompt_persona_for_session, prompt_build_limits_from_config, resolve_prompt_agent_id, resolve_prompt_mode_context, }, - run_with_tools::run_with_tools, - service::build_persisted_assistant_message, - streaming::run_streaming, types::*, }; @@ -56,293 +46,33 @@ use super::*; #[async_trait] impl ChatService for LiveChatService { async fn send(&self, params: Value) -> ServiceResult { - self.send_impl(params).await + self.send_impl(params, None, true).await } async fn send_sync(&self, params: Value) -> ServiceResult { - let text = params - .get("text") - .and_then(|v| v.as_str()) - .ok_or_else(|| "missing 'text' parameter".to_string())? - .to_string(); - let desired_reply_medium = infer_reply_medium(¶ms, &text); - let requested_agent_id = params + if let Some(agent_id) = params .get("agent_id") - .and_then(|v| v.as_str()) + .and_then(Value::as_str) .map(str::trim) .filter(|value| !value.is_empty()) - .map(str::to_string); - let request_tool_policy = params - .get("_tool_policy") - .cloned() - .map(serde_json::from_value::) - .transpose() - .map_err(|e| format!("invalid '_tool_policy' parameter: {e}"))?; - let ephemeral = params - .get("_ephemeral") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - - let explicit_model = params.get("model").and_then(|v| v.as_str()); - let tool_controls = - moltis_config::schema::AgentToolControls::from_tool_context(Some(¶ms)); - let stream_only = !self.has_tools_sync(); - - // Resolve session key from explicit override. - let session_key = match params.get("_session_key").and_then(|v| v.as_str()) { - Some(sk) => sk.to_string(), - None => "main".to_string(), - }; - - // Resolve provider. - let provider: Arc = { - let reg = self.providers.read().await; - if let Some(id) = explicit_model { - reg.get(id) - .ok_or_else(|| format!("model '{id}' not found"))? - } else if !stream_only { - reg.first_with_tools() - .ok_or_else(|| "no LLM providers configured".to_string())? - } else { - reg.first() - .ok_or_else(|| "no LLM providers configured".to_string())? - } - }; - - let user_audio = user_audio_path_from_params(¶ms, &session_key); - let user_documents = - user_documents_from_params(¶ms, &session_key, self.session_store.as_ref()); - // Persist the user message. - let user_msg = PersistedMessage::User { - content: MessageContent::Text(text.clone()), - created_at: Some(now_ms()), - audio: user_audio, - documents: user_documents - .as_deref() - .and_then(user_documents_for_persistence), - channel: None, - seq: None, - run_id: None, - }; - if !ephemeral { - if let Err(e) = self - .session_store - .append(&session_key, &user_msg.to_value()) - .await - { - warn!("send_sync: failed to persist user message: {e}"); - } - - // Ensure this session appears in the sessions list. - let _ = self.session_metadata.upsert(&session_key, None).await; - } - if let Some(agent_id) = requested_agent_id.as_deref() - && let Err(error) = self - .session_metadata + { + let session_key = self.resolve_session_key_from_params(¶ms).await; + self.session_metadata .set_agent_id(&session_key, Some(agent_id)) .await - { - warn!( - session = %session_key, - agent_id, - error = %error, - "send_sync: failed to assign requested agent to session" - ); + .map_err(ServiceError::message)?; } - if !ephemeral { - self.session_metadata.touch(&session_key, 1).await; + let (completion_tx, completion_rx) = oneshot::channel(); + let started = self.send_impl(params, Some(completion_tx), false).await?; + if started.get("queued").and_then(Value::as_bool) == Some(true) { + return Err("session already has an active turn".into()); } - - let session_entry = self.session_metadata.get(&session_key).await; - let session_agent_id = resolve_prompt_agent_id(session_entry.as_ref()); - let persona = load_prompt_persona_for_session( - &session_key, - session_entry.as_ref(), - self.session_state_store.as_deref(), - ) - .await; - let mut runtime_context = build_prompt_runtime_context( - &self.state, - &persona.config, - &provider, - &session_key, - session_entry.as_ref(), - ) - .await; - runtime_context.mode = resolve_prompt_mode_context(&persona.config, session_entry.as_ref()); - apply_request_runtime_context(&mut runtime_context.host, ¶ms); - - // Load conversation history (excluding the message we just appended). - let mut history = self - .session_store - .read(&session_key) - .await - .unwrap_or_default(); - if !ephemeral && !history.is_empty() { - history.pop(); - } - - let run_id = uuid::Uuid::new_v4().to_string(); - let state = Arc::clone(&self.state); - let tool_registry = if let Some(policy) = request_tool_policy.as_ref() { - let registry_guard = self.tool_registry.read().await; - Arc::new(RwLock::new( - registry_guard.clone_allowed_by(|name| policy.is_allowed(name)), - )) - } else { - Arc::clone(&self.tool_registry) - }; - let hook_registry = self.hook_registry.clone(); - let provider_name = provider.name().to_string(); - let model_id = provider.id().to_string(); - let model_store = Arc::clone(&self.model_store); - let user_message_index = history.len(); - - info!( - run_id = %run_id, - user_message = %text, - model = %model_id, - stream_only, - session = %session_key, - reply_medium = ?desired_reply_medium, - "chat.send_sync" - ); - - if desired_reply_medium == ReplyMedium::Voice { - broadcast( - &state, - "chat", - serde_json::json!({ - "runId": run_id, - "sessionKey": session_key, - "state": "voice_pending", - }), - BroadcastOpts::default(), - ) - .await; + if started.get("rejected").and_then(Value::as_bool) == Some(true) { + return Ok(started); } - - // send_sync is text-only (used by API calls and channels). - let user_content = UserContent::text(&text); - let active_event_forwarders = Arc::new(RwLock::new(HashMap::new())); - let terminal_runs = Arc::new(RwLock::new(HashSet::new())); - let result = if stream_only { - run_streaming( - persona, - &state, - &model_store, - &run_id, - provider, - &model_id, - &user_content, - &provider_name, - &history, - &session_key, - &session_agent_id, - desired_reply_medium, - None, - user_message_index, - &[], - Some(&runtime_context), - None, // send_sync: no sender name - Some(&self.session_store), - None, // send_sync: no client seq - None, // send_sync: no partial assistant tracking - &terminal_runs, - ) - .await - } else { - run_with_tools( - persona, - &state, - &model_store, - &run_id, - provider, - &model_id, - &tool_registry, - &user_content, - &provider_name, - &history, - &session_key, - &session_agent_id, - desired_reply_medium, - None, - Some(&runtime_context), - user_message_index, - &[], - hook_registry, - None, - None, // send_sync: no conn_id - Some(&self.session_store), - false, // send_sync: MCP tools always enabled for API calls - None, // send_sync: no client seq - None, // send_sync: no thinking text tracking - None, // send_sync: no tool call tracking - None, // send_sync: no partial assistant tracking - &active_event_forwarders, - &terminal_runs, - None, // send_sync: no sender name - Some(tool_controls), - ) + completion_rx .await - }; - - // Persist assistant response (even empty ones — needed for LLM history coherence). - if !ephemeral && let Some(ref assistant_output) = result { - let assistant_msg = build_persisted_assistant_message( - assistant_output.clone(), - Some(model_id.clone()), - Some(provider_name.clone()), - None, - Some(run_id.clone()), - ); - if let Err(e) = self - .session_store - .append(&session_key, &assistant_msg.to_value()) - .await - { - warn!("send_sync: failed to persist assistant message: {e}"); - } - // Update metadata message count. - if let Ok(count) = self.session_store.count(&session_key).await { - self.session_metadata.touch(&session_key, count).await; - } - } - - match result { - Some(assistant_output) => Ok(serde_json::json!({ - "text": assistant_output.text, - "inputTokens": assistant_output.input_tokens, - "outputTokens": assistant_output.output_tokens, - "cacheReadTokens": assistant_output.cache_read_tokens, - "cacheWriteTokens": assistant_output.cache_write_tokens, - "durationMs": assistant_output.duration_ms, - "requestInputTokens": assistant_output.request_input_tokens, - "requestOutputTokens": assistant_output.request_output_tokens, - "requestCacheReadTokens": assistant_output.request_cache_read_tokens, - "requestCacheWriteTokens": assistant_output.request_cache_write_tokens, - })), - None => { - // Check the last broadcast for this run to get the actual error message. - let error_msg = state - .last_run_error(&run_id) - .await - .unwrap_or_else(|| "agent run failed (check server logs)".to_string()); - - // Persist the error in the session so it's visible in session history. - let error_entry = PersistedMessage::system(format!("[error] {error_msg}")); - let _ = self - .session_store - .append(&session_key, &error_entry.to_value()) - .await; - // Update metadata so the session shows in the UI. - if let Ok(count) = self.session_store.count(&session_key).await { - self.session_metadata.touch(&session_key, count).await; - } - - Err(error_msg.into()) - }, - } + .map_err(|_| ServiceError::from("chat run was cancelled"))? } async fn abort(&self, params: Value) -> ServiceResult { @@ -1012,9 +742,12 @@ impl ChatService for LiveChatService { apply_request_runtime_context(&mut runtime_context.host, ¶ms); // Resolve project context plus optional command-generated context. - let project_context = self + let (project_context, working_dir) = self .resolve_turn_context(&session_key, conn_id.as_deref()) .await; + runtime_context.host.working_dir = working_dir + .as_ref() + .map(|directory| directory.display().to_string()); // Discover skills (gated on `[skills] enabled` — see #655). let discovered_skills = discover_skills_if_enabled(&persona.config).await; @@ -1152,9 +885,12 @@ impl ChatService for LiveChatService { apply_request_runtime_context(&mut runtime_context.host, ¶ms); // Resolve project context plus optional command-generated context. - let project_context = self + let (project_context, working_dir) = self .resolve_turn_context(&session_key, conn_id.as_deref()) .await; + runtime_context.host.working_dir = working_dir + .as_ref() + .map(|directory| directory.display().to_string()); // Discover skills (gated on `[skills] enabled` — see #655). let discovered_skills = discover_skills_if_enabled(&persona.config).await; diff --git a/crates/chat/src/service/chat_impl/queue.rs b/crates/chat/src/service/chat_impl/queue.rs new file mode 100644 index 0000000000..d812f60349 --- /dev/null +++ b/crates/chat/src/service/chat_impl/queue.rs @@ -0,0 +1,69 @@ +use std::{collections::HashMap, sync::Arc}; + +use { + moltis_config::MessageQueueMode, + tokio::sync::RwLock, + tracing::{info, warn}, +}; + +use crate::{ChatRuntime, service::QueuedMessage}; + +pub(super) async fn drain_queued_messages( + message_queue: &Arc>>>, + state: &Arc, + session_key: &str, + queue_mode: MessageQueueMode, +) { + let queued = message_queue + .write() + .await + .remove(session_key) + .unwrap_or_default(); + if queued.is_empty() { + return; + } + + let chat = state.chat_service().await; + match queue_mode { + MessageQueueMode::Followup => { + let mut iter = queued.into_iter(); + let Some(first) = iter.next() else { + return; + }; + let rest: Vec = iter.collect(); + if !rest.is_empty() { + message_queue + .write() + .await + .entry(session_key.to_string()) + .or_default() + .extend(rest); + } + info!(session = %session_key, "replaying queued message (followup)"); + let mut replay_params = first.params; + replay_params["_queued_replay"] = serde_json::json!(true); + if let Err(error) = chat.send(replay_params).await { + warn!(session = %session_key, %error, "failed to replay queued message"); + } + }, + MessageQueueMode::Collect => { + let combined = queued + .iter() + .filter_map(|message| message.params.get("text").and_then(|value| value.as_str())) + .collect::>(); + if combined.is_empty() { + return; + } + info!(session = %session_key, count = combined.len(), "replaying collected messages"); + let Some(last) = queued.last() else { + return; + }; + let mut merged = last.params.clone(); + merged["text"] = serde_json::json!(combined.join("\n\n")); + merged["_queued_replay"] = serde_json::json!(true); + if let Err(error) = chat.send(merged).await { + warn!(session = %session_key, %error, "failed to replay collected messages"); + } + }, + } +} diff --git a/crates/chat/src/service/chat_impl/send.rs b/crates/chat/src/service/chat_impl/send.rs index b8d33a0e34..d153169901 100644 --- a/crates/chat/src/service/chat_impl/send.rs +++ b/crates/chat/src/service/chat_impl/send.rs @@ -4,11 +4,11 @@ use std::{sync::Arc, time::Duration}; use { serde_json::Value, - tokio::sync::OwnedSemaphorePermit, + tokio::sync::{OwnedSemaphorePermit, oneshot}, tracing::{debug, info, warn}, }; -use {moltis_config::MessageQueueMode, moltis_service_traits::ServiceResult}; +use moltis_service_traits::ServiceResult; #[cfg(feature = "local-llm")] use moltis_providers::model_id::raw_model_id; @@ -30,7 +30,10 @@ use crate::{ types::*, }; -use {super::*, crate::service::build_persisted_assistant_message}; +use { + super::{queue::drain_queued_messages, *}, + crate::service::build_persisted_assistant_message, +}; use { crate::memory_tools::AgentScopedMemoryWriter, @@ -39,7 +42,20 @@ use { impl LiveChatService { #[tracing::instrument(skip(self, params), fields(session_id))] - pub(super) async fn send_impl(&self, mut params: Value) -> ServiceResult { + pub(super) async fn send_impl( + &self, + mut params: Value, + completion: Option>, + queue_if_busy: bool, + ) -> ServiceResult { + let history_limits = params + .get("_history_limits") + .and_then(Value::as_object) + .and_then(|limits| { + let max_messages = usize::try_from(limits.get("max_messages")?.as_u64()?).ok()?; + let max_bytes = usize::try_from(limits.get("max_bytes")?.as_u64()?).ok()?; + Some((max_messages, max_bytes)) + }); // Support both text-only and multimodal content. // - "text": string → plain text message // - "content": array → multimodal content (text + images) @@ -104,6 +120,16 @@ impl LiveChatService { let explicit_model = params.get("model").and_then(|v| v.as_str()); let tool_controls = moltis_config::schema::AgentToolControls::from_tool_context(Some(¶ms)); + let request_tool_policy = params + .get("_tool_policy") + .cloned() + .map(serde_json::from_value::) + .transpose() + .map_err(|error| format!("invalid '_tool_policy' parameter: {error}"))?; + let ephemeral = params + .get("_ephemeral") + .and_then(Value::as_bool) + .unwrap_or(false); // Use streaming-only mode if explicitly requested or if no tools are registered. let explicit_stream_only = params .get("stream_only") @@ -199,6 +225,9 @@ impl LiveChatService { p }, Err(_) => { + if !queue_if_busy { + return Err("session already has an active turn".into()); + } let queue_mode = message_queue_mode; let position = { let mut q = self.message_queue.write().await; @@ -242,6 +271,12 @@ impl LiveChatService { }; if let Some(shell_command) = explicit_shell_command { + if request_tool_policy + .as_ref() + .is_some_and(|policy| !policy.is_allowed("exec")) + { + return Err("exec tool is denied by the request tool policy".into()); + } // Generate run_id early so we can link the user message to this run. let run_id = uuid::Uuid::new_v4().to_string(); let run_id_clone = run_id.clone(); @@ -261,18 +296,16 @@ impl LiveChatService { run_id: Some(run_id.clone()), }; - let history = self - .session_store - .read(&session_key) - .await - .unwrap_or_default(); + let history = self.load_turn_history(&session_key, history_limits).await?; let user_message_index = history.len(); // Ensure the session exists in metadata and counts are up to date. - let _ = self.session_metadata.upsert(&session_key, None).await; - self.session_metadata - .touch(&session_key, history.len() as u32) - .await; + if !ephemeral { + let _ = self.session_metadata.upsert(&session_key, None).await; + self.session_metadata + .touch(&session_key, history.len() as u32) + .await; + } // If this is a web UI message on a channel-bound session, attach the // channel reply target so /sh output can be delivered back to the channel. @@ -334,25 +367,27 @@ impl LiveChatService { info!( run_id = %run_id, - user_message = %text, + user_message_bytes = text.len(), session = %session_key, - command = %shell_command, + command_bytes = shell_command.len(), client_seq = ?client_seq, mode = "explicit_shell", "chat.send" ); // Persist user message now that it will execute immediately. - if let Err(e) = self - .session_store - .append(&session_key, &user_msg.to_value()) - .await + if !ephemeral + && let Err(e) = self + .session_store + .append(&session_key, &user_msg.to_value()) + .await { warn!("failed to persist /sh user message: {e}"); } // Set preview from first user message if not already set. - if let Some(entry) = self.session_metadata.get(&session_key).await + if !ephemeral + && let Some(entry) = self.session_metadata.get(&session_key).await && entry.preview.is_none() { let preview_text = extract_preview_from_value(&user_msg.to_value()); @@ -382,8 +417,16 @@ impl LiveChatService { .and_then(|v| v.as_str()) .map(String::from); let conn_id_for_tool = conn_id.clone(); + let (_, working_dir) = self + .resolve_turn_context(&session_key, conn_id.as_deref()) + .await; + let working_dir = working_dir.map(|directory| directory.display().to_string()); + let (start_tx, start_rx) = oneshot::channel(); let handle = tokio::spawn(async move { + if start_rx.await.is_err() { + return; + } let permit = permit; // hold permit until command run completes if let Some(target) = deferred_channel_target { state.push_channel_reply(&session_key_clone, target).await; @@ -397,7 +440,7 @@ impl LiveChatService { &state, &run_id_clone, &tool_registry, - &session_store, + (!ephemeral).then_some(&session_store), &terminal_runs, &session_key_clone, &shell_command, @@ -405,9 +448,11 @@ impl LiveChatService { accept_language, conn_id_for_tool, client_seq, + working_dir, ) .await; + let completion_result = Ok(turn_result(&assistant_output)); let assistant_msg = build_persisted_assistant_message( assistant_output, None, @@ -415,14 +460,16 @@ impl LiveChatService { client_seq, Some(run_id_clone.clone()), ); - if let Err(e) = session_store - .append(&session_key_clone, &assistant_msg.to_value()) - .await - { - warn!("failed to persist /sh assistant message: {e}"); - } - if let Ok(count) = session_store.count(&session_key_clone).await { - session_metadata.touch(&session_key_clone, count).await; + if !ephemeral { + if let Err(e) = session_store + .append(&session_key_clone, &assistant_msg.to_value()) + .await + { + warn!("failed to persist /sh assistant message: {e}"); + } + if let Ok(count) = session_store.count(&session_key_clone).await { + session_metadata.touch(&session_key_clone, count).await; + } } active_runs.write().await.remove(&run_id_clone); @@ -444,72 +491,25 @@ impl LiveChatService { active_reply_medium.write().await.remove(&session_key_clone); drop(permit); - - // Drain queued messages for this session. - let queued = message_queue - .write() - .await - .remove(&session_key_clone) - .unwrap_or_default(); - if !queued.is_empty() { - let queue_mode = message_queue_mode; - let chat = state_for_drain.chat_service().await; - match queue_mode { - MessageQueueMode::Followup => { - let mut iter = queued.into_iter(); - let Some(first) = iter.next() else { - return; - }; - let rest: Vec = iter.collect(); - if !rest.is_empty() { - message_queue - .write() - .await - .entry(session_key_clone.clone()) - .or_default() - .extend(rest); - } - info!(session = %session_key_clone, "replaying queued message (followup)"); - let mut replay_params = first.params; - replay_params["_queued_replay"] = serde_json::json!(true); - if let Err(e) = chat.send(replay_params).await { - warn!(session = %session_key_clone, error = %e, "failed to replay queued message"); - } - }, - MessageQueueMode::Collect => { - let combined: Vec<&str> = queued - .iter() - .filter_map(|m| m.params.get("text").and_then(|v| v.as_str())) - .collect(); - if !combined.is_empty() { - info!( - session = %session_key_clone, - count = combined.len(), - "replaying collected messages" - ); - let Some(last) = queued.last() else { - return; - }; - let mut merged = last.params.clone(); - merged["text"] = serde_json::json!(combined.join("\n\n")); - merged["_queued_replay"] = serde_json::json!(true); - if let Err(e) = chat.send(merged).await { - warn!(session = %session_key_clone, error = %e, "failed to replay collected messages"); - } - } - }, - } + if let Some(completion) = completion { + let _ = completion.send(completion_result); } + drain_queued_messages( + &message_queue, + &state_for_drain, + &session_key_clone, + message_queue_mode, + ) + .await; }); - self.active_runs - .write() - .await - .insert(run_id.clone(), handle.abort_handle()); - self.active_runs_by_session - .write() - .await - .insert(session_key.clone(), run_id.clone()); + { + let mut runs_by_session = self.active_runs_by_session.write().await; + let mut active_runs = self.active_runs.write().await; + active_runs.insert(run_id.clone(), handle.abort_handle()); + runs_by_session.insert(session_key.clone(), run_id.clone()); + } + let _ = start_tx.send(()); info!( run_id = %run_id, @@ -606,7 +606,7 @@ impl LiveChatService { } // Resolve project context plus optional command-generated context. - let project_context = self + let (project_context, working_dir) = self .resolve_turn_context(&session_key, conn_id.as_deref()) .await; @@ -615,11 +615,7 @@ impl LiveChatService { // Load conversation history (the current user message is NOT yet // persisted — run_streaming / run_agent_loop add it themselves). - let mut history = self - .session_store - .read(&session_key) - .await - .unwrap_or_default(); + let mut history = self.load_turn_history(&session_key, history_limits).await?; info!( session = %session_key, history_len = history.len(), @@ -628,10 +624,12 @@ impl LiveChatService { ); // Update metadata. - let _ = self.session_metadata.upsert(&session_key, None).await; - self.session_metadata - .touch(&session_key, history.len() as u32) - .await; + if !ephemeral { + let _ = self.session_metadata.upsert(&session_key, None).await; + self.session_metadata + .touch(&session_key, history.len() as u32) + .await; + } // If this is a web UI message on a channel-bound session, attach the // channel reply target so the run-start path can route the final @@ -903,6 +901,9 @@ impl LiveChatService { ) .await; runtime_context.mode = resolve_prompt_mode_context(&persona.config, session_entry.as_ref()); + runtime_context.host.working_dir = working_dir + .as_ref() + .map(|directory| directory.display().to_string()); apply_request_runtime_context(&mut runtime_context.host, ¶ms); info!( session = %session_key, @@ -921,7 +922,28 @@ impl LiveChatService { let active_partial_assistant = Arc::clone(&self.active_partial_assistant); let active_reply_medium = Arc::clone(&self.active_reply_medium); let run_id_clone = run_id.clone(); - let tool_registry = Arc::clone(&self.tool_registry); + let overlay = self + .session_tool_overlays + .read() + .await + .get(&session_key) + .cloned(); + let tool_registry = if let Some(overlay) = overlay { + let mut combined = self.tool_registry.read().await.clone_allowed_by(|_| true); + let overlay = overlay.read().await; + combined.extend_from(&overlay); + Arc::new(tokio::sync::RwLock::new(combined)) + } else { + Arc::clone(&self.tool_registry) + }; + let tool_registry = if let Some(policy) = request_tool_policy.as_ref() { + let registry = tool_registry.read().await; + Arc::new(tokio::sync::RwLock::new( + registry.clone_allowed_by(|name| policy.is_allowed(name)), + )) + } else { + tool_registry + }; let hook_registry = self.hook_registry.clone(); // Log if tool mode is active but the provider doesn't support tools. @@ -937,7 +959,7 @@ impl LiveChatService { info!( run_id = %run_id, - user_message = %text, + user_message_bytes = text.len(), model = provider.id(), stream_only, session = %session_key, @@ -952,13 +974,14 @@ impl LiveChatService { let provider_name = provider.name().to_string(); let model_id = provider.id().to_string(); - if self - .session_metadata - .get(&session_key) - .await - .and_then(|entry| entry.model) - .as_deref() - != Some(model_id.as_str()) + if !ephemeral + && self + .session_metadata + .get(&session_key) + .await + .and_then(|entry| entry.model) + .as_deref() + != Some(model_id.as_str()) { self.session_metadata .set_model(&session_key, Some(model_id.clone())) @@ -987,7 +1010,7 @@ impl LiveChatService { let compact_threshold = compute_auto_compact_threshold(context_window, compaction_cfg.threshold_percent); - if estimated_next_input >= compact_threshold { + if !ephemeral && estimated_next_input >= compact_threshold { let pre_compact_msg_count = history.len(); let pre_compact_total = token_usage .current_request_input_tokens @@ -1025,11 +1048,7 @@ impl LiveChatService { match self.compact(compact_params).await { Ok(_) => { // Reload history after compaction. - history = self - .session_store - .read(&session_key) - .await - .unwrap_or_default(); + history = self.load_turn_history(&session_key, history_limits).await?; // This `auto_compact done` event is a lifecycle // signal for subscribers that pre-emptive // auto-compact finished. The mode/token metadata @@ -1078,10 +1097,11 @@ impl LiveChatService { // Persist the user message now that we know it won't be queued. // (Queued messages skip this; they are persisted when replayed.) - if let Err(e) = self - .session_store - .append(&session_key, &user_msg.to_value()) - .await + if !ephemeral + && let Err(e) = self + .session_store + .append(&session_key, &user_msg.to_value()) + .await { warn!("failed to persist user message: {e}"); } @@ -1107,7 +1127,8 @@ impl LiveChatService { .await; // Set preview from the first user message if not already set. - if let Some(entry) = self.session_metadata.get(&session_key).await + if !ephemeral + && let Some(entry) = self.session_metadata.get(&session_key).await && entry.preview.is_none() { let preview_text = extract_preview_from_value(&user_msg.to_value()); @@ -1130,7 +1151,11 @@ impl LiveChatService { let terminal_runs = Arc::clone(&self.terminal_runs); let deferred_channel_target = deferred_channel_target.clone(); + let (start_tx, start_rx) = oneshot::channel(); let handle = tokio::spawn(async move { + if start_rx.await.is_err() { + return; + } let permit = permit; // hold permit until agent run completes let ctx_ref = project_context.as_deref(); if let Some(target) = deferred_channel_target { @@ -1142,10 +1167,12 @@ impl LiveChatService { .write() .await .insert(session_key_clone.clone(), desired_reply_medium); - active_partial_assistant.write().await.insert( - session_key_clone.clone(), - ActiveAssistantDraft::new(&run_id_clone, &model_id, &provider_name, client_seq), - ); + if !ephemeral { + active_partial_assistant.write().await.insert( + session_key_clone.clone(), + ActiveAssistantDraft::new(&run_id_clone, &model_id, &provider_name, client_seq), + ); + } if desired_reply_medium == ReplyMedium::Voice { broadcast( &state, @@ -1187,9 +1214,9 @@ impl LiveChatService { &discovered_skills, Some(&runtime_context), sender_name, - Some(&session_store), + (!ephemeral).then_some(&session_store), client_seq, - Some(Arc::clone(&active_partial_assistant)), + (!ephemeral).then(|| Arc::clone(&active_partial_assistant)), &terminal_runs, ) .await @@ -1215,12 +1242,12 @@ impl LiveChatService { hook_registry, accept_language.clone(), conn_id.clone(), - Some(&session_store), + (!ephemeral).then_some(&session_store), mcp_disabled, client_seq, Some(Arc::clone(&active_thinking_text)), Some(Arc::clone(&active_tool_calls)), - Some(Arc::clone(&active_partial_assistant)), + (!ephemeral).then(|| Arc::clone(&active_partial_assistant)), &active_event_forwarders, &terminal_runs, sender_name, @@ -1271,8 +1298,17 @@ impl LiveChatService { agent_fut.await }; + let completion_result = match assistant_text.as_ref() { + Some(output) => Ok(turn_result(output)), + None => Err(state + .last_run_error(&run_id_clone) + .await + .unwrap_or_else(|| "agent run failed (check server logs)".to_string()) + .into()), + }; + // Persist assistant response (even empty ones — needed for LLM history coherence). - if let Some(assistant_output) = assistant_text { + if !ephemeral && let Some(assistant_output) = assistant_text { let assistant_msg = build_persisted_assistant_message( assistant_output, Some(model_id.clone()), @@ -1366,7 +1402,8 @@ impl LiveChatService { // generation. We check >= 2 (not == 2) because agentic turns // with tool calls produce more than 2 stored messages. // `generate_title_if_needed` guards against duplicate titles. - if auto_title_enabled + if !ephemeral + && auto_title_enabled && let Ok(count) = session_store.count(&session_key_clone).await && count >= 2 && !queued_replay @@ -1402,75 +1439,25 @@ impl LiveChatService { // acquire it. Without this, every replayed `chat.send()` would // fail `try_acquire_owned()` and re-queue the message forever. drop(permit); - - // Drain queued messages for this session. - let queued = message_queue - .write() - .await - .remove(&session_key_clone) - .unwrap_or_default(); - if !queued.is_empty() { - let queue_mode = message_queue_mode; - let chat = state_for_drain.chat_service().await; - match queue_mode { - MessageQueueMode::Followup => { - let mut iter = queued.into_iter(); - let Some(first) = iter.next() else { - return; - }; - // Put remaining messages back so the replayed run's - // own drain loop picks them up after it completes. - let rest: Vec = iter.collect(); - if !rest.is_empty() { - message_queue - .write() - .await - .entry(session_key_clone.clone()) - .or_default() - .extend(rest); - } - info!(session = %session_key_clone, "replaying queued message (followup)"); - let mut replay_params = first.params; - replay_params["_queued_replay"] = serde_json::json!(true); - if let Err(e) = chat.send(replay_params).await { - warn!(session = %session_key_clone, error = %e, "failed to replay queued message"); - } - }, - MessageQueueMode::Collect => { - let combined: Vec<&str> = queued - .iter() - .filter_map(|m| m.params.get("text").and_then(|v| v.as_str())) - .collect(); - if !combined.is_empty() { - info!( - session = %session_key_clone, - count = combined.len(), - "replaying collected messages" - ); - // Use the last queued message as the base params, override text. - let Some(last) = queued.last() else { - return; - }; - let mut merged = last.params.clone(); - merged["text"] = serde_json::json!(combined.join("\n\n")); - merged["_queued_replay"] = serde_json::json!(true); - if let Err(e) = chat.send(merged).await { - warn!(session = %session_key_clone, error = %e, "failed to replay collected messages"); - } - } - }, - } + if let Some(completion) = completion { + let _ = completion.send(completion_result); } + drain_queued_messages( + &message_queue, + &state_for_drain, + &session_key_clone, + message_queue_mode, + ) + .await; }); - self.active_runs - .write() - .await - .insert(run_id.clone(), handle.abort_handle()); - self.active_runs_by_session - .write() - .await - .insert(session_key.clone(), run_id.clone()); + { + let mut runs_by_session = self.active_runs_by_session.write().await; + let mut active_runs = self.active_runs.write().await; + active_runs.insert(run_id.clone(), handle.abort_handle()); + runs_by_session.insert(session_key.clone(), run_id.clone()); + } + let _ = start_tx.send(()); info!( run_id = %run_id, @@ -1481,3 +1468,18 @@ impl LiveChatService { Ok(serde_json::json!({ "ok": true, "runId": run_id })) } } + +fn turn_result(output: &AssistantTurnOutput) -> Value { + serde_json::json!({ + "text": output.text, + "inputTokens": output.input_tokens, + "outputTokens": output.output_tokens, + "cacheReadTokens": output.cache_read_tokens, + "cacheWriteTokens": output.cache_write_tokens, + "durationMs": output.duration_ms, + "requestInputTokens": output.request_input_tokens, + "requestOutputTokens": output.request_output_tokens, + "requestCacheReadTokens": output.request_cache_read_tokens, + "requestCacheWriteTokens": output.request_cache_write_tokens, + }) +} diff --git a/crates/chat/src/service/types.rs b/crates/chat/src/service/types.rs index c9431ffa17..dc4171aec7 100644 --- a/crates/chat/src/service/types.rs +++ b/crates/chat/src/service/types.rs @@ -19,6 +19,7 @@ use { use { moltis_agents::tool_registry::ToolRegistry, moltis_providers::ProviderRegistry, + moltis_service_traits::ServiceError, moltis_sessions::{ PersistedMessage, message::{PersistedFunction, PersistedToolCall}, @@ -232,6 +233,8 @@ pub struct LiveChatService { Arc>>>, pub(in crate::service) terminal_runs: Arc>>, pub(in crate::service) tool_registry: Arc>, + pub(in crate::service) session_tool_overlays: + Arc>>>>, pub(in crate::service) session_store: Arc, pub(in crate::service) session_metadata: Arc, pub(in crate::service) session_state_store: Option>, @@ -261,6 +264,40 @@ pub struct LiveChatService { } impl LiveChatService { + pub(in crate::service) async fn load_turn_history( + &self, + session_key: &str, + limits: Option<(usize, usize)>, + ) -> Result, ServiceError> { + let Some((max_messages, max_bytes)) = limits else { + return Ok(self + .session_store + .read(session_key) + .await + .unwrap_or_default()); + }; + self.session_store + .read_bounded(session_key, max_messages, max_bytes) + .await + .map_err(|error| { + ServiceError::message(format!("failed to read session history: {error}")) + }) + } + + /// Reads persisted history with strict pre-allocation bounds for non-UI + /// protocol surfaces such as ACP session replay. + pub async fn read_session_history_bounded( + &self, + session_key: &str, + max_messages: usize, + max_bytes: usize, + ) -> anyhow::Result> { + self.session_store + .read_bounded(session_key, max_messages, max_bytes) + .await + .map_err(Into::into) + } + pub fn new( providers: Arc>, model_store: Arc>, @@ -277,6 +314,7 @@ impl LiveChatService { active_event_forwarders: Arc::new(RwLock::new(HashMap::new())), terminal_runs: Arc::new(RwLock::new(HashSet::new())), tool_registry: Arc::new(RwLock::new(ToolRegistry::new())), + session_tool_overlays: Arc::new(RwLock::new(HashMap::new())), session_store, session_metadata, session_state_store: None, @@ -308,6 +346,32 @@ impl LiveChatService { self } + pub async fn set_session_tool_overlay( + &self, + session_key: &str, + registry: Arc>, + ) { + self.session_tool_overlays + .write() + .await + .insert(session_key.to_string(), registry); + } + + pub async fn remove_session_tool_overlay(&self, session_key: &str) { + self.session_tool_overlays.write().await.remove(session_key); + } + + pub async fn bind_session_project(&self, session_key: &str, project_id: &str) { + let _ = self.session_metadata.upsert(session_key, None).await; + self.session_metadata + .set_project_id(session_key, Some(project_id.to_string())) + .await; + } + + pub async fn session_exists(&self, session_key: &str) -> bool { + self.session_metadata.get(session_key).await.is_some() + } + pub fn with_session_state_store(mut self, store: Arc) -> Self { self.session_state_store = Some(store); self @@ -616,7 +680,7 @@ impl LiveChatService { &self, session_key: &str, conn_id: Option<&str>, - ) -> Option { + ) -> (Option, Option) { let (project_context, working_dir) = self.resolve_project_context(session_key, conn_id).await; let command_context = moltis_common::context_command::run_context_command( @@ -624,7 +688,10 @@ impl LiveChatService { working_dir.as_deref(), ) .await; - merge_context_sections(project_context, command_context) + ( + merge_context_sections(project_context, command_context), + working_dir, + ) } } diff --git a/crates/chat/src/streaming.rs b/crates/chat/src/streaming.rs index 631954706e..906f90863a 100644 --- a/crates/chat/src/streaming.rs +++ b/crates/chat/src/streaming.rs @@ -347,7 +347,7 @@ pub(crate) async fn run_streaming( run_id, input_tokens = usage.input_tokens, output_tokens = usage.output_tokens, - response = %accumulated, + response_bytes = accumulated.len(), silent = is_silent, "chat stream done" ); diff --git a/crates/cli/Cargo.toml b/crates/cli/Cargo.toml index 6aa8a59e82..f6d0ba7f3d 100644 --- a/crates/cli/Cargo.toml +++ b/crates/cli/Cargo.toml @@ -66,45 +66,53 @@ name = "moltis" path = "src/main.rs" [dependencies] -anyhow = { workspace = true } -clap = { workspace = true } -dotenvy = { workspace = true } -moltis-agents = { workspace = true } -moltis-browser = { workspace = true } -moltis-claude-import = { optional = true, workspace = true } -moltis-codex-import = { optional = true, workspace = true } -moltis-common = { workspace = true } -moltis-config = { workspace = true } -moltis-cron = { workspace = true } -moltis-gateway = { workspace = true } -moltis-hermes-import = { optional = true, workspace = true } -moltis-httpd = { workspace = true } -moltis-import-core = { workspace = true } -moltis-memory = { workspace = true } -moltis-node-host = { workspace = true } -moltis-oauth = { workspace = true } -moltis-onboarding = { workspace = true } -moltis-openclaw-import = { optional = true, workspace = true } -moltis-plugins = { workspace = true } -moltis-portable = { workspace = true } -moltis-projects = { workspace = true } -moltis-providers = { workspace = true } -moltis-sessions = { workspace = true } -moltis-skills = { workspace = true } -moltis-tools = { workspace = true } -moltis-web = { optional = true, workspace = true } -open = { workspace = true } -rand = { workspace = true } -reqwest = { workspace = true } -secrecy = { workspace = true } -serde_json = { workspace = true } -sqlx = { workspace = true } -time = { workspace = true } -tokio = { workspace = true } -tracing = { workspace = true } -tracing-subscriber = { workspace = true } -uuid = { workspace = true } -which = { workspace = true } +agent-client-protocol = { optional = true, workspace = true } +anyhow = { workspace = true } +async-trait = { workspace = true } +clap = { workspace = true } +dotenvy = { workspace = true } +moltis-acp = { optional = true, workspace = true } +moltis-agents = { workspace = true } +moltis-browser = { workspace = true } +moltis-chat = { optional = true, workspace = true } +moltis-claude-import = { optional = true, workspace = true } +moltis-codex-import = { optional = true, workspace = true } +moltis-common = { workspace = true } +moltis-config = { workspace = true } +moltis-cron = { workspace = true } +moltis-gateway = { workspace = true } +moltis-hermes-import = { optional = true, workspace = true } +moltis-httpd = { workspace = true } +moltis-import-core = { workspace = true } +moltis-mcp = { optional = true, workspace = true } +moltis-mcp-agent-bridge = { optional = true, workspace = true } +moltis-memory = { workspace = true } +moltis-node-host = { workspace = true } +moltis-oauth = { workspace = true } +moltis-onboarding = { workspace = true } +moltis-openclaw-import = { optional = true, workspace = true } +moltis-plugins = { workspace = true } +moltis-portable = { workspace = true } +moltis-projects = { workspace = true } +moltis-protocol = { workspace = true } +moltis-providers = { workspace = true } +moltis-service-traits = { optional = true, workspace = true } +moltis-sessions = { workspace = true } +moltis-skills = { workspace = true } +moltis-tools = { workspace = true } +moltis-web = { optional = true, workspace = true } +open = { workspace = true } +rand = { workspace = true } +reqwest = { workspace = true } +secrecy = { workspace = true } +serde_json = { workspace = true } +sqlx = { workspace = true } +time = { workspace = true } +tokio = { workspace = true } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } +uuid = { workspace = true } +which = { workspace = true } [target.'cfg(all(not(target_os = "windows"), not(all(target_os = "linux", target_arch = "aarch64"))))'.dependencies] tikv-jemallocator = { optional = true, workspace = true } @@ -113,10 +121,18 @@ tikv-jemallocator = { optional = true, workspace = true } tempfile = { workspace = true } [features] -agent = ["moltis-gateway/agent", "moltis-web?/agent"] -caldav = ["moltis-gateway/caldav"] +acp = [ + "dep:agent-client-protocol", + "dep:moltis-acp", + "dep:moltis-chat", + "dep:moltis-mcp", + "dep:moltis-mcp-agent-bridge", + "dep:moltis-service-traits", +] +agent = ["moltis-gateway/agent", "moltis-web?/agent"] +caldav = ["moltis-gateway/caldav"] cloudflare-tunnel = ["moltis-httpd/cloudflare-tunnel", "moltis-web?/cloudflare-tunnel"] -home-assistant = ["moltis-gateway/home-assistant"] +home-assistant = ["moltis-gateway/home-assistant"] # User-facing integrations stay default-on in the CLI crate per workspace policy, # even when the underlying implementation remains feature-gated in lower crates. # @@ -127,6 +143,7 @@ home-assistant = ["moltis-gateway/home-assistant"] bundled-skills = ["moltis-gateway/bundled-skills", "moltis-web?/bundled-skills"] default = ["full"] full = [ + "acp", "agent", "bundled-skills", "caldav", diff --git a/crates/cli/src/acp_backend.rs b/crates/cli/src/acp_backend.rs new file mode 100644 index 0000000000..025ab82704 --- /dev/null +++ b/crates/cli/src/acp_backend.rs @@ -0,0 +1,791 @@ +//! Real Moltis turns behind the ACP backend seam. +//! +//! This is the same path the Web UI takes. A prompt becomes `ChatService::send_sync`, +//! which runs the agent loop — providers, tools, memory, session history — and +//! resolves with the final assistant message. The tokens the Web UI renders as +//! they arrive are broadcast as event frames while that call is pending, so to +//! stream them we register a client on the gateway's broadcast registry and +//! forward what it receives. +//! +//! # Why a registered client rather than a bespoke hook +//! +//! `ConnectedClient` is just a bounded `mpsc::Sender` plus subscription +//! metadata, and `broadcast()` fans frames out to every registered client. Using +//! that seam means ACP sees exactly what the Web UI sees, including frames added +//! later, with no parallel notification path in the chat crate to keep in sync. +//! +//! # No server required +//! +//! `prepare_gateway_core` is the transport-agnostic half of startup and binds no +//! socket, so `moltis acp` boots the stack in-process. It does open the databases +//! under `data_dir()`, so it shares session state with a running gateway rather +//! than talking to it. + +mod updates; + +use std::{ + collections::HashMap, + path::{Path, PathBuf}, + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering}, + }, +}; + +use { + agent_client_protocol as acp, + async_trait::async_trait, + moltis_acp::{ + AcpBackend, BackendCapabilities, SessionKey, SessionNotFound, SessionSetup, TurnUpdates, + backend::{MAX_HISTORY_BYTES, MAX_HISTORY_UPDATES}, + }, + moltis_chat::LiveChatService, + moltis_gateway::state::GatewayState, + moltis_protocol::{ClientInfo, ConnectParams, PROTOCOL_VERSION}, + moltis_service_traits::ChatService, + serde_json::{Value, json}, + tokio::sync::{Notify, RwLock, mpsc}, + tracing::{debug, warn}, +}; + +use { + self::updates::{FrameAction, FrameMapper}, + crate::acp_mcp::SessionMcpRuntime, +}; + +/// Frames buffered for one turn before the gateway starts dropping them. +/// +/// Generous because a fast provider can outrun the forwarder briefly; the +/// gateway drops rather than blocks when a client is slow, so an undersized +/// buffer costs tokens rather than backpressure. +const FRAME_BUFFER: usize = 1024; + +/// Serves ACP prompts by running real Moltis turns. +pub struct MoltisBackend { + core: moltis_gateway::server::PreparedGatewayCore, + state: Arc, + chat: Arc, + sessions: RwLock>, + active_prompts: RwLock>>, + lifecycle: Arc, + /// Distinguishes the synthetic clients this backend registers, so two + /// concurrent turns cannot collide on a connection id. + next_conn: AtomicU64, +} + +struct SessionRuntime { + mcp: Option, +} + +#[derive(Default)] +struct PromptSignal { + cancelled: AtomicBool, + notify: Notify, +} + +#[derive(Default)] +struct BackendLifecycle { + closing: AtomicBool, + active: AtomicUsize, + idle: Notify, +} + +impl BackendLifecycle { + fn begin(self: &Arc) -> anyhow::Result { + if self.closing.load(Ordering::Acquire) { + anyhow::bail!("ACP backend is shutting down"); + } + self.active.fetch_add(1, Ordering::AcqRel); + if self.closing.load(Ordering::Acquire) { + self.finish(); + anyhow::bail!("ACP backend is shutting down"); + } + Ok(OperationGuard(Arc::clone(self))) + } + + fn finish(&self) { + if self.active.fetch_sub(1, Ordering::AcqRel) == 1 { + self.idle.notify_waiters(); + } + } + + async fn wait_idle(&self) { + loop { + let notified = self.idle.notified(); + if self.active.load(Ordering::Acquire) == 0 { + break; + } + notified.await; + } + } +} + +struct OperationGuard(Arc); + +impl Drop for OperationGuard { + fn drop(&mut self) { + self.0.finish(); + } +} + +impl MoltisBackend { + #[must_use] + pub fn new(core: moltis_gateway::server::PreparedGatewayCore) -> Self { + let state = Arc::clone(&core.state); + let chat = Arc::clone(&core.live_chat); + Self { + core, + state, + chat, + sessions: RwLock::new(HashMap::new()), + active_prompts: RwLock::new(HashMap::new()), + lifecycle: Arc::new(BackendLifecycle::default()), + next_conn: AtomicU64::new(0), + } + } + + async fn canonical_cwd(setup: &SessionSetup) -> anyhow::Result { + let cwd = tokio::fs::canonicalize(setup.cwd()) + .await + .map_err(|error| anyhow::anyhow!("invalid session cwd: {error}"))?; + if !cwd.is_dir() { + anyhow::bail!("session cwd is not a directory"); + } + Ok(cwd) + } + + async fn bind_project(&self, key: &SessionKey, cwd: &Path) -> anyhow::Result<()> { + let listed = self.state.services.project.list().await?; + let projects = serde_json::from_value::>(listed)?; + let existing = projects.iter().find(|project| { + std::fs::canonicalize(&project.directory).is_ok_and(|directory| directory == cwd) + }); + let project_id = if let Some(project) = existing { + project.id.clone() + } else { + let mut project = moltis_projects::detect::detect_project(cwd) + .ok_or_else(|| anyhow::anyhow!("failed to detect project for {}", cwd.display()))?; + if projects.iter().any(|existing| existing.id == project.id) { + project.id = format!("{}-{}", project.id, uuid::Uuid::new_v4().simple()); + } + self.state + .services + .project + .upsert(serde_json::to_value(&project)?) + .await?; + project.id + }; + self.chat + .bind_session_project(key.as_str(), &project_id) + .await; + Ok(()) + } + + async fn install_session(&self, key: &SessionKey, setup: &SessionSetup) -> anyhow::Result<()> { + if self.lifecycle.closing.load(Ordering::Acquire) { + anyhow::bail!("ACP backend is shutting down"); + } + let cwd = Self::canonical_cwd(setup).await?; + let mcp = SessionMcpRuntime::start(key, setup).await?; + if self.lifecycle.closing.load(Ordering::Acquire) { + if let Some(runtime) = mcp { + runtime.shutdown().await; + } + anyhow::bail!("ACP backend is shutting down"); + } + if let Err(error) = self.bind_project(key, &cwd).await { + if let Some(runtime) = mcp { + runtime.shutdown().await; + } + return Err(error); + } + if self.lifecycle.closing.load(Ordering::Acquire) { + if let Some(runtime) = mcp { + runtime.shutdown().await; + } + anyhow::bail!("ACP backend is shutting down"); + } + if let Some(runtime) = mcp.as_ref() { + self.chat + .set_session_tool_overlay(key.as_str(), runtime.tools()) + .await; + } else { + self.chat.remove_session_tool_overlay(key.as_str()).await; + } + let previous = self + .sessions + .write() + .await + .insert(key.clone(), SessionRuntime { mcp }); + if let Some(previous) = previous + && let Some(mcp) = previous.mcp + { + mcp.shutdown().await; + } + Ok(()) + } + + /// Registers a broadcast client and returns its id plus the frame stream. + /// + /// The caller must pair this with [`Self::unregister`]; [`TurnClient`] does + /// that on drop so an early return cannot leak a registration. + async fn register(&self, key: &SessionKey) -> (String, mpsc::Receiver, Arc) { + let seq = self.next_conn.fetch_add(1, Ordering::Relaxed); + let conn_id = format!("acp-{seq}"); + let (tx, rx) = mpsc::channel::(FRAME_BUFFER); + let delivery_failures = Arc::new(AtomicU64::new(0)); + let client = moltis_gateway::state::ConnectedClient { + conn_id: conn_id.clone(), + connect_params: acp_connect_params(), + sender: tx, + delivery_failures: Arc::clone(&delivery_failures), + connected_at: std::time::Instant::now(), + last_activity_ms: AtomicU64::new(0), + accept_language: None, + remote_ip: None, + timezone: None, + subscriptions: Some(std::collections::HashSet::from(["chat".to_string()])), + session_filter: Some(key.as_str().to_string()), + payload_state_filter: Some(std::collections::HashSet::from([ + "delta".to_string(), + "error".to_string(), + "iteration".to_string(), + "thinking_text".to_string(), + "tool_call_end".to_string(), + "tool_call_start".to_string(), + ])), + joined_channels: std::collections::HashSet::new(), + negotiated_protocol: PROTOCOL_VERSION, + }; + self.state.register_client(client).await; + (conn_id, rx, delivery_failures) + } +} + +fn send_mapped_update( + updates: &TurnUpdates, + mapper: &mut FrameMapper, + update: acp::SessionUpdate, +) -> bool { + if !updates.send(update.clone()) { + return false; + } + mapper.record_sent(&update); + true +} + +fn reconcile_final_text( + updates: &TurnUpdates, + mapper: &mut FrameMapper, + final_text: &str, +) -> anyhow::Result<()> { + let Some(update) = mapper.finish_text(final_text).map_err(anyhow::Error::msg)? else { + return Ok(()); + }; + if !send_mapped_update(updates, mapper, update) { + anyhow::bail!("ACP client stopped reading final turn output"); + } + Ok(()) +} + +/// Connection metadata for the synthetic client backing one ACP turn. +/// +/// The ACP client is the local parent process, already trusted, so this claims +/// the operator role rather than inventing a narrower one it would then have to +/// widen every time a chat frame gained a scope guard. +fn acp_connect_params() -> ConnectParams { + ConnectParams { + min_protocol: PROTOCOL_VERSION, + max_protocol: PROTOCOL_VERSION, + client: ClientInfo { + id: "moltis-acp".to_string(), + display_name: Some("ACP client".to_string()), + version: moltis_config::VERSION.to_string(), + platform: std::env::consts::OS.to_string(), + device_family: None, + model_identifier: None, + mode: "agent".to_string(), + instance_id: None, + }, + caps: None, + commands: None, + permissions: None, + path_env: None, + role: Some("operator".to_string()), + scopes: None, + device: None, + auth: None, + locale: None, + user_agent: None, + timezone: None, + } +} + +/// Keeps a registered broadcast client alive for the duration of a turn. +/// +/// Unregistering matters: a leaked client keeps a channel in the gateway's +/// registry forever, and `broadcast()` walks every registered client on every +/// frame. +struct TurnClient { + state: Arc, + conn_id: String, +} + +impl Drop for TurnClient { + fn drop(&mut self) { + if self.conn_id.is_empty() { + return; + } + let state = Arc::clone(&self.state); + let conn_id = std::mem::take(&mut self.conn_id); + // Drop runs outside async context, and removal takes a write lock. + tokio::spawn(async move { + state.remove_client(&conn_id).await; + }); + } +} + +impl TurnClient { + async fn close(mut self) { + let conn_id = std::mem::take(&mut self.conn_id); + self.state.remove_client(&conn_id).await; + } +} + +#[async_trait] +impl AcpBackend for MoltisBackend { + async fn create_session(&self, setup: &SessionSetup) -> anyhow::Result { + let _operation = self.lifecycle.begin()?; + // Moltis materializes a session on first write, so there is nothing to + // create here beyond choosing the key. The `acp:` namespace is what + // keeps these from colliding with Web UI and channel sessions, and the + // protocol layer rejects anything outside it. + let key = SessionKey::namespaced(uuid::Uuid::new_v4().to_string()); + self.install_session(&key, setup).await?; + debug!(session = %key, "ACP session created"); + Ok(key) + } + + async fn load_session( + &self, + key: &SessionKey, + setup: &SessionSetup, + ) -> anyhow::Result> { + let _operation = self.lifecycle.begin()?; + if !self.chat.session_exists(key.as_str()).await { + return Err(SessionNotFound.into()); + } + let history = self + .chat + .read_session_history_bounded(key.as_str(), MAX_HISTORY_UPDATES, MAX_HISTORY_BYTES) + .await + .map_err(|error| anyhow::anyhow!("failed to read session history: {error}"))?; + let updates = history_to_updates(&history); + moltis_acp::validate_history(&updates)?; + self.install_session(key, setup).await?; + Ok(updates) + } + + async fn discard_session(&self, key: &SessionKey) -> anyhow::Result<()> { + let runtime = self.sessions.write().await.remove(key); + self.chat.remove_session_tool_overlay(key.as_str()).await; + if let Some(SessionRuntime { mcp: Some(mcp), .. }) = runtime { + mcp.shutdown().await; + } + Ok(()) + } + + async fn prompt( + &self, + key: &SessionKey, + prompt: String, + updates: TurnUpdates, + ) -> anyhow::Result { + let _operation = self.lifecycle.begin()?; + let signal = Arc::new(PromptSignal::default()); + let mut active_prompts = self.active_prompts.write().await; + if active_prompts.contains_key(key) { + anyhow::bail!("session {key} already has an active prompt"); + } + active_prompts.insert(key.clone(), Arc::clone(&signal)); + drop(active_prompts); + if self.lifecycle.closing.load(Ordering::Acquire) { + signal.cancelled.store(true, Ordering::Release); + signal.notify.notify_one(); + } + + let result = async { + if signal.cancelled.load(Ordering::Acquire) { + return Ok(acp::StopReason::Cancelled); + } + let (conn_id, mut frames, delivery_failures) = self.register(key).await; + let client = TurnClient { + state: Arc::clone(&self.state), + conn_id, + }; + + let turn = self.chat.send_sync(json!({ + "text": prompt, + "_session_key": key.as_str(), + "_history_limits": { + "max_messages": MAX_HISTORY_UPDATES, + "max_bytes": MAX_HISTORY_BYTES, + }, + })); + let mut turn = std::pin::pin!(turn); + + let mut mapper = FrameMapper::new(); + let mut reported_error: Option = None; + let mut cancelled = false; + let mut delivery_failed = false; + + // Forward frames while the turn runs. `send_sync` resolving is what ends + // the turn — the broadcast has no terminal frame to wait for, and + // waiting for the channel to close would hang, since the gateway holds + // the sender until the client is unregistered. + let outcome = 'turn: loop { + tokio::select! { + biased; + _ = signal.notify.notified() => { + let _ = self + .chat + .abort(json!({ "sessionKey": key.as_str() })) + .await; + cancelled = true; + break Ok(json!({})); + }, + result = &mut turn => break result, + frame = frames.recv() => match frame { + Some(frame) => match mapper.map(&frame, key.as_str()) { + FrameAction::Emit(batch) => { + for update in batch { + if !send_mapped_update(&updates, &mut mapper, update) { + debug!("ACP client stopped reading updates mid-turn"); + let _ = self + .chat + .abort(json!({ "sessionKey": key.as_str() })) + .await; + delivery_failed = true; + break 'turn Ok(json!({})); + } + } + }, + FrameAction::Failed(message) => reported_error = Some(message), + FrameAction::Ignore => {}, + }, + // The gateway dropped our registration; the turn still owns + // the outcome, so wait for it rather than guessing. + None => break (&mut turn).await, + }, + } + }; + + // Frames already queued when the turn resolved are still this turn's + // output; dropping them would truncate the visible reply. + 'flush: while !cancelled + && !delivery_failed + && let Ok(frame) = frames.try_recv() + { + match mapper.map(&frame, key.as_str()) { + FrameAction::Emit(batch) => { + for update in batch { + if !send_mapped_update(&updates, &mut mapper, update) { + delivery_failed = true; + break 'flush; + } + } + }, + FrameAction::Failed(message) => reported_error = Some(message), + FrameAction::Ignore => {}, + } + } + + client.close().await; + + if delivery_failures.load(Ordering::Relaxed) > 0 { + delivery_failed = true; + } + + if cancelled { + Ok(acp::StopReason::Cancelled) + } else if delivery_failed { + Err(anyhow::anyhow!( + "ACP client could not accept all final turn updates" + )) + } else { + match outcome { + Ok(result) if result.get("rejected").and_then(Value::as_bool) == Some(true) => { + let reason = result + .get("reason") + .and_then(Value::as_str) + .unwrap_or("prompt rejected"); + Err(anyhow::anyhow!(reason.to_string())) + }, + Ok(result) => { + if let Some(final_text) = result.get("text").and_then(Value::as_str) { + reconcile_final_text(&updates, &mut mapper, final_text)?; + } + Ok(acp::StopReason::EndTurn) + }, + Err(error) => { + // Prefer the broadcast's message: `send_sync` reports a generic + // failure while the frame carries the provider's own words. + let detail = reported_error.unwrap_or_else(|| error.to_string()); + warn!(session = %key, "ACP turn failed"); + Err(anyhow::anyhow!(detail)) + }, + } + } + } + .await; + + let mut active_prompts = self.active_prompts.write().await; + if active_prompts + .get(key) + .is_some_and(|active| Arc::ptr_eq(active, &signal)) + { + active_prompts.remove(key); + } + result + } + + async fn cancel(&self, key: &SessionKey) -> anyhow::Result<()> { + let signal = self.active_prompts.read().await.get(key).cloned(); + if let Some(signal) = signal.as_ref() { + signal.cancelled.store(true, Ordering::Release); + signal.notify.notify_one(); + } + let result = self + .chat + .abort(json!({ "sessionKey": key.as_str() })) + .await + .map_err(|error| anyhow::anyhow!("failed to abort turn: {error}"))?; + if result.get("aborted").and_then(Value::as_bool) != Some(true) && signal.is_none() { + anyhow::bail!("no active turn found for session {key}"); + } + Ok(()) + } + + async fn shutdown(&self) -> anyhow::Result<()> { + debug!("shutting down ACP backend"); + self.lifecycle.closing.store(true, Ordering::Release); + for signal in self.active_prompts.read().await.values() { + signal.cancelled.store(true, Ordering::Release); + signal.notify.notify_one(); + } + if tokio::time::timeout( + std::time::Duration::from_secs(10), + self.lifecycle.wait_idle(), + ) + .await + .is_err() + { + warn!("timed out waiting for ACP prompts to stop during shutdown"); + } + let sessions = { + let mut sessions = self.sessions.write().await; + std::mem::take(&mut *sessions) + }; + for (key, runtime) in sessions { + let _ = self.chat.abort(json!({ "sessionKey": key.as_str() })).await; + self.chat.remove_session_tool_overlay(key.as_str()).await; + if let Some(mcp) = runtime.mcp { + mcp.shutdown().await; + } + debug!(session = %key, "ACP session shut down"); + } + self.core.mcp_manager.shutdown_all().await; + debug!("ACP backend shut down"); + Ok(()) + } + + fn capabilities(&self) -> BackendCapabilities { + BackendCapabilities { load_session: true } + } +} + +/// Converts persisted session history into replayable ACP updates. +/// +/// `chat.history` returns the same message shapes the Web UI renders. Only user +/// and assistant text carries over: system entries and tool bookkeeping have no +/// ACP representation, and replaying them as messages would put words in the +/// agent's mouth. +fn history_to_updates(messages: &[Value]) -> Vec { + messages + .iter() + .filter_map(|message| { + let text = message_text(message)?; + if text.trim().is_empty() { + return None; + } + match message.get("role").and_then(Value::as_str)? { + "user" => Some(acp::SessionUpdate::UserMessageChunk( + acp::ContentChunk::new(acp::ContentBlock::from(text)), + )), + "assistant" => Some(acp::SessionUpdate::AgentMessageChunk( + acp::ContentChunk::new(acp::ContentBlock::from(text)), + )), + _ => None, + } + }) + .collect() +} + +/// Pulls displayable text out of a persisted message. +/// +/// Content is either a bare string or an array of typed blocks, depending on +/// whether the message carried attachments. +fn message_text(message: &Value) -> Option { + let content = message.get("content")?; + if let Some(text) = content.as_str() { + return Some(text.to_string()); + } + let blocks = content.as_array()?; + let text = blocks + .iter() + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect::>() + .join(""); + (!text.is_empty()).then_some(text) +} + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use {super::*, serde_json::json}; + + #[test] + fn history_replays_only_user_and_assistant_text() { + let history = vec![ + json!({ "role": "user", "content": "hello" }), + json!({ "role": "system", "content": "[error] boom" }), + json!({ "role": "assistant", "content": "hi there" }), + ]; + let updates = history_to_updates(&history); + assert_eq!(updates.len(), 2, "the system entry must not be replayed"); + assert!(matches!( + updates[0], + acp::SessionUpdate::UserMessageChunk(_) + )); + assert!(matches!( + updates[1], + acp::SessionUpdate::AgentMessageChunk(_) + )); + } + + #[test] + fn history_reads_block_structured_content() { + let history = vec![json!({ + "role": "assistant", + "content": [ + { "type": "text", "text": "part one " }, + { "type": "text", "text": "part two" }, + ], + })]; + let updates = history_to_updates(&history); + assert_eq!(updates.len(), 1); + let acp::SessionUpdate::AgentMessageChunk(chunk) = &updates[0] else { + panic!("expected an agent message"); + }; + let acp::ContentBlock::Text(text) = &chunk.content else { + panic!("expected text content"); + }; + assert_eq!(text.text, "part one part two"); + } + + #[test] + fn empty_and_malformed_history_is_not_replayed() { + assert!(history_to_updates(&[]).is_empty()); + // An assistant turn persisted with no text would otherwise replay as an + // empty message. + assert!(history_to_updates(&[json!({ "role": "assistant", "content": " " })]).is_empty()); + assert!(history_to_updates(&[json!({ "role": "assistant" })]).is_empty()); + } + + #[test] + fn created_sessions_are_inside_the_acp_namespace() { + // The protocol layer rejects out-of-namespace keys, so a backend that + // minted one would fail every `session/new`. + let key = SessionKey::namespaced(uuid::Uuid::new_v4().to_string()); + assert!(key.is_namespaced()); + } + + #[test] + fn created_sessions_are_unique() { + let first = SessionKey::namespaced(uuid::Uuid::new_v4().to_string()); + let second = SessionKey::namespaced(uuid::Uuid::new_v4().to_string()); + assert_ne!(first, second); + } + + #[tokio::test] + async fn lifecycle_rejects_new_work_and_waits_for_active_work() { + let lifecycle = Arc::new(BackendLifecycle::default()); + let operation = lifecycle.begin().expect("operation before shutdown"); + lifecycle.closing.store(true, Ordering::Release); + assert!(lifecycle.begin().is_err()); + assert!( + tokio::time::timeout(std::time::Duration::from_millis(10), lifecycle.wait_idle()) + .await + .is_err() + ); + drop(operation); + tokio::time::timeout(std::time::Duration::from_secs(1), lifecycle.wait_idle()) + .await + .expect("active operation should release lifecycle"); + } + + #[test] + fn closed_update_channels_are_not_counted_as_delivered() { + let (tx, rx) = mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + let mut mapper = FrameMapper::new(); + let first = acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new( + acp::ContentBlock::from("first"), + )); + let second = acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new( + acp::ContentBlock::from("second"), + )); + drop(rx); + assert!(!send_mapped_update(&updates, &mut mapper, first)); + assert!(!send_mapped_update(&updates, &mut mapper, second)); + assert!( + mapper + .finish_text("firstsecond") + .is_ok_and(|value| value.is_some()) + ); + } + + #[test] + fn divergent_final_text_fails_the_turn() { + let (tx, _rx) = mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + let mut mapper = FrameMapper::new(); + let streamed = acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new( + acp::ContentBlock::from("draft"), + )); + assert!(send_mapped_update(&updates, &mut mapper, streamed)); + let error = reconcile_final_text(&updates, &mut mapper, "corrected").unwrap_err(); + assert!(error.to_string().contains("diverged")); + } + + #[test] + fn missing_final_suffix_is_delivered() { + let (tx, mut rx) = mpsc::unbounded_channel(); + let updates = TurnUpdates::new(tx); + let mut mapper = FrameMapper::new(); + let streamed = acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new( + acp::ContentBlock::from("partial"), + )); + assert!(send_mapped_update(&updates, &mut mapper, streamed)); + assert!(reconcile_final_text(&updates, &mut mapper, "partial response").is_ok()); + assert!(rx.try_recv().is_ok()); + let update = rx.try_recv().expect("missing final suffix update"); + let acp::SessionUpdate::AgentMessageChunk(chunk) = update else { + panic!("expected agent message"); + }; + let acp::ContentBlock::Text(text) = chunk.content else { + panic!("expected text content"); + }; + assert_eq!(text.text, " response"); + } +} diff --git a/crates/cli/src/acp_backend/updates.rs b/crates/cli/src/acp_backend/updates.rs new file mode 100644 index 0000000000..dcb07931d7 --- /dev/null +++ b/crates/cli/src/acp_backend/updates.rs @@ -0,0 +1,610 @@ +//! Translating Moltis chat broadcasts into ACP `session/update` notifications. +//! +//! The gateway broadcasts a turn's progress as JSON event frames — the same +//! ones the Web UI consumes over its WebSocket. This module is the pure mapping +//! from those frames to ACP updates, kept free of gateway and protocol wiring so +//! it can be tested against literal frames. +//! +//! Two details of the wire format drive the shape here: +//! +//! - `thinking_text` carries the *accumulated* reasoning so far, not a delta, +//! because the Web UI replaces its thinking pane wholesale on each one. ACP +//! thought chunks are incremental, so [`FrameMapper`] keeps the previous text +//! and emits only what was appended. +//! - Frames for every session on the box arrive on the same broadcast, so each +//! one must be matched against the session this turn belongs to before it is +//! forwarded. Skipping that check would leak one client's tokens into another +//! client's turn. + +use {agent_client_protocol as acp, serde_json::Value}; + +const MAX_TOOL_VALUE_BYTES: usize = 64 * 1024; + +/// What a single broadcast frame means for the turn in progress. +/// +/// Not `Eq`: `SessionUpdate` carries content blocks that are only `PartialEq`. +#[derive(Debug, PartialEq)] +pub enum FrameAction { + /// Forward these updates to the client. + Emit(Vec), + /// The run reported an error; the turn should end with this message. + Failed(String), + /// Not for this turn, or carries nothing a client can render. + Ignore, +} + +/// Stateful mapper for one turn's frames. +/// +/// Holds the reasoning text seen so far so accumulated `thinking_text` frames +/// can be converted into incremental thought chunks. +#[derive(Debug, Default)] +pub struct FrameMapper { + reasoning_seen: String, + message_seen: String, +} + +impl FrameMapper { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Maps one raw frame, as serialized onto the broadcast, for `session_key`. + /// + /// Unparseable frames and frames for other sessions are ignored rather than + /// erroring: the broadcast is shared infrastructure and a malformed or + /// unrelated frame is not this turn's problem. + pub fn map(&mut self, raw: &str, session_key: &str) -> FrameAction { + let Ok(frame) = serde_json::from_str::(raw) else { + return FrameAction::Ignore; + }; + if frame.get("event").and_then(Value::as_str) != Some("chat") { + return FrameAction::Ignore; + } + let Some(payload) = frame.get("payload") else { + return FrameAction::Ignore; + }; + if payload.get("sessionKey").and_then(Value::as_str) != Some(session_key) { + return FrameAction::Ignore; + } + self.map_payload(payload) + } + + fn map_payload(&mut self, payload: &Value) -> FrameAction { + let text = || payload.get("text").and_then(Value::as_str).unwrap_or(""); + match payload.get("state").and_then(Value::as_str) { + Some("delta") => { + let delta = text(); + if delta.is_empty() { + return FrameAction::Ignore; + } + FrameAction::Emit(vec![agent_message(delta)]) + }, + Some("iteration") => { + self.message_seen.clear(); + FrameAction::Ignore + }, + Some("thinking_text") => self.reasoning_delta(text()), + Some("tool_call_start") => FrameAction::Emit(vec![tool_call_start(payload)]), + Some("tool_call_end") => FrameAction::Emit(vec![tool_call_end(payload)]), + Some("error") => FrameAction::Failed(error_message(payload)), + // `user_message` echoes what the client just sent us, `queued`, + // `iteration`, `thinking`, and the rest are Web-UI affordances with + // no ACP equivalent. Dropping them keeps the client's transcript to + // what the agent actually said. + _ => FrameAction::Ignore, + } + } + + /// Converts an accumulated reasoning snapshot into the newly added suffix. + fn reasoning_delta(&mut self, accumulated: &str) -> FrameAction { + // A shorter or diverging snapshot means the provider restarted its + // reasoning rather than extended it. Emitting the whole thing again is + // better than emitting a nonsense suffix computed from a stale prefix. + let addition = match accumulated.strip_prefix(self.reasoning_seen.as_str()) { + Some(suffix) => suffix.to_string(), + None => accumulated.to_string(), + }; + self.reasoning_seen = accumulated.to_string(); + if addition.is_empty() { + return FrameAction::Ignore; + } + FrameAction::Emit(vec![agent_thought(addition)]) + } + + /// Reconciles broadcast deltas with the completed turn's authoritative text. + pub fn finish_text(&self, final_text: &str) -> Result, String> { + if final_text.trim() == self.message_seen.trim() { + return Ok(None); + } + let Some(missing) = final_text.strip_prefix(&self.message_seen) else { + return Err("streamed reply diverged from the completed turn".to_string()); + }; + Ok((!missing.is_empty()).then(|| agent_message(missing))) + } + + /// Records only chunks accepted by the bounded protocol update channel. + pub fn record_sent(&mut self, update: &acp::SessionUpdate) { + if let acp::SessionUpdate::AgentMessageChunk(chunk) = update + && let acp::ContentBlock::Text(text) = &chunk.content + { + self.message_seen.push_str(&text.text); + } + } +} + +fn agent_message(text: impl Into) -> acp::SessionUpdate { + acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(acp::ContentBlock::from( + text.into(), + ))) +} + +fn agent_thought(text: impl Into) -> acp::SessionUpdate { + acp::SessionUpdate::AgentThoughtChunk(acp::ContentChunk::new(acp::ContentBlock::from( + text.into(), + ))) +} + +/// Identifier the client uses to correlate a tool call's start with its end. +/// +/// Moltis labels tool calls by name within a run, so the name is what the two +/// frames have in common. +fn tool_call_id(payload: &Value) -> String { + first_non_empty( + &[ + payload.get("toolCallId").and_then(Value::as_str), + payload.get("id").and_then(Value::as_str), + payload.get("toolName").and_then(Value::as_str), + payload.get("tool").and_then(Value::as_str), + payload.get("name").and_then(Value::as_str), + ], + "tool", + ) +} + +fn tool_name(payload: &Value) -> String { + first_non_empty( + &[ + payload.get("toolName").and_then(Value::as_str), + payload.get("tool").and_then(Value::as_str), + payload.get("name").and_then(Value::as_str), + ], + "tool", + ) +} + +fn error_message(payload: &Value) -> String { + let error = payload.get("error"); + first_non_empty( + &[ + error.and_then(Value::as_str), + error + .and_then(|value| value.get("detail")) + .and_then(Value::as_str), + error + .and_then(|value| value.get("message")) + .and_then(Value::as_str), + error + .and_then(|value| value.get("title")) + .and_then(Value::as_str), + payload.get("text").and_then(Value::as_str), + ], + "agent run failed", + ) +} + +fn tool_call_start(payload: &Value) -> acp::SessionUpdate { + let mut call = acp::ToolCall::new( + acp::ToolCallId::from(tool_call_id(payload)), + tool_name(payload), + ) + .status(acp::ToolCallStatus::InProgress); + if let Some(arguments) = payload.get("arguments") { + call = call.raw_input(capped_json(arguments)); + } + acp::SessionUpdate::ToolCall(call) +} + +fn tool_call_end(payload: &Value) -> acp::SessionUpdate { + let failed = payload.get("success").and_then(Value::as_bool) == Some(false) + || payload.get("error").is_some_and(|error| !error.is_null()); + let status = if failed { + acp::ToolCallStatus::Failed + } else { + acp::ToolCallStatus::Completed + }; + let mut fields = acp::ToolCallUpdateFields::default(); + fields.status = Some(status); + let output = if failed { + payload.get("error") + } else { + payload.get("result") + }; + if let Some(output) = output { + let output = capped_json(output); + fields.content = Some(vec![display_json(&output).into()]); + fields.raw_output = Some(output); + } + acp::SessionUpdate::ToolCallUpdate(acp::ToolCallUpdate::new( + acp::ToolCallId::from(tool_call_id(payload)), + fields, + )) +} + +fn capped_json(value: &Value) -> Value { + let encoded_bytes = serde_json::to_vec(value).map_or(MAX_TOOL_VALUE_BYTES + 1, |v| v.len()); + if encoded_bytes <= MAX_TOOL_VALUE_BYTES { + return value.clone(); + } + Value::String(format!( + "[tool value omitted: {encoded_bytes} bytes exceeds {MAX_TOOL_VALUE_BYTES}-byte limit]" + )) +} + +fn display_json(value: &Value) -> String { + value + .as_str() + .map(str::to_owned) + .unwrap_or_else(|| value.to_string()) +} + +/// First non-empty candidate, or `fallback` when every candidate is absent. +fn first_non_empty(candidates: &[Option<&str>], fallback: &str) -> String { + candidates + .iter() + .flatten() + .map(|value| value.trim()) + .find(|value| !value.is_empty()) + .unwrap_or(fallback) + .to_string() +} + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use {super::*, serde_json::json}; + + const SESSION: &str = "acp:s1"; + + fn frame(payload: Value) -> String { + json!({ "type": "event", "event": "chat", "payload": payload }).to_string() + } + + fn chat_frame(state: &str, text: &str) -> String { + frame(json!({ "sessionKey": SESSION, "state": state, "text": text })) + } + + /// Extracts the text of a single emitted content chunk. + fn emitted_text(action: &FrameAction) -> String { + let FrameAction::Emit(updates) = action else { + panic!("expected an emission, got {action:?}"); + }; + assert_eq!(updates.len(), 1, "expected exactly one update"); + match &updates[0] { + acp::SessionUpdate::AgentMessageChunk(chunk) + | acp::SessionUpdate::AgentThoughtChunk(chunk) => match &chunk.content { + acp::ContentBlock::Text(text) => text.text.clone(), + other => panic!("expected text content, got {other:?}"), + }, + other => panic!("expected a content chunk, got {other:?}"), + } + } + + #[test] + fn deltas_become_agent_message_chunks() { + let mut mapper = FrameMapper::new(); + let action = mapper.map(&chat_frame("delta", "hello "), SESSION); + assert_eq!(emitted_text(&action), "hello "); + } + + #[test] + fn frames_for_other_sessions_are_ignored() { + let mut mapper = FrameMapper::new(); + let other = frame(json!({ + "sessionKey": "acp:someone-else", + "state": "delta", + "text": "not yours", + })); + assert_eq!(mapper.map(&other, SESSION), FrameAction::Ignore); + } + + #[test] + fn non_chat_events_are_ignored() { + let mut mapper = FrameMapper::new(); + let tick = json!({ "type": "event", "event": "tick", "payload": {} }).to_string(); + assert_eq!(mapper.map(&tick, SESSION), FrameAction::Ignore); + } + + #[test] + fn malformed_frames_are_ignored_rather_than_erroring() { + let mut mapper = FrameMapper::new(); + assert_eq!(mapper.map("{not json", SESSION), FrameAction::Ignore); + assert_eq!(mapper.map("null", SESSION), FrameAction::Ignore); + } + + #[test] + fn accumulated_reasoning_is_emitted_as_increments() { + let mut mapper = FrameMapper::new(); + // The Web UI receives the whole reasoning each time; ACP wants only the + // part that is new. + let first = mapper.map(&chat_frame("thinking_text", "Let me"), SESSION); + assert_eq!(emitted_text(&first), "Let me"); + let second = mapper.map(&chat_frame("thinking_text", "Let me think"), SESSION); + assert_eq!(emitted_text(&second), " think"); + } + + #[test] + fn unchanged_reasoning_emits_nothing() { + let mut mapper = FrameMapper::new(); + mapper.map(&chat_frame("thinking_text", "same"), SESSION); + assert_eq!( + mapper.map(&chat_frame("thinking_text", "same"), SESSION), + FrameAction::Ignore + ); + } + + #[test] + fn restarted_reasoning_is_emitted_whole() { + let mut mapper = FrameMapper::new(); + mapper.map(&chat_frame("thinking_text", "first attempt"), SESSION); + // Diverges from what we saw: not an extension, so the suffix arithmetic + // would produce garbage. + let action = mapper.map(&chat_frame("thinking_text", "different"), SESSION); + assert_eq!(emitted_text(&action), "different"); + } + + #[test] + fn empty_deltas_are_ignored() { + let mut mapper = FrameMapper::new(); + assert_eq!( + mapper.map(&chat_frame("delta", ""), SESSION), + FrameAction::Ignore + ); + } + + #[test] + fn errors_surface_the_reported_message() { + let mut mapper = FrameMapper::new(); + let raw = frame(json!({ + "sessionKey": SESSION, + "state": "error", + "error": { "detail": "provider exploded" }, + })); + assert_eq!( + mapper.map(&raw, SESSION), + FrameAction::Failed("provider exploded".to_string()) + ); + } + + #[test] + fn errors_without_a_message_still_fail_the_turn() { + let mut mapper = FrameMapper::new(); + let raw = frame(json!({ "sessionKey": SESSION, "state": "error" })); + assert_eq!( + mapper.map(&raw, SESSION), + FrameAction::Failed("agent run failed".to_string()) + ); + } + + #[test] + fn tool_calls_start_in_progress_and_complete() { + let mut mapper = FrameMapper::new(); + let start = frame(json!({ + "sessionKey": SESSION, + "state": "tool_call_start", + "toolName": "read_file", + "arguments": { "path": "README.md" }, + })); + let FrameAction::Emit(updates) = mapper.map(&start, SESSION) else { + panic!("tool call start must be forwarded"); + }; + match &updates[0] { + acp::SessionUpdate::ToolCall(call) => { + assert_eq!(call.title, "read_file"); + assert_eq!(call.status, acp::ToolCallStatus::InProgress); + assert_eq!(call.raw_input, Some(json!({ "path": "README.md" }))); + }, + other => panic!("expected a tool call, got {other:?}"), + } + + let end = frame(json!({ + "sessionKey": SESSION, + "state": "tool_call_end", + "toolName": "read_file", + "result": { "text": "contents" }, + })); + let FrameAction::Emit(updates) = mapper.map(&end, SESSION) else { + panic!("tool call end must be forwarded"); + }; + match &updates[0] { + acp::SessionUpdate::ToolCallUpdate(update) => { + assert_eq!(update.fields.status, Some(acp::ToolCallStatus::Completed)); + assert_eq!( + update.fields.raw_output, + Some(json!({ "text": "contents" })) + ); + assert!( + update + .fields + .content + .as_ref() + .is_some_and(|value| !value.is_empty()) + ); + }, + other => panic!("expected a tool call update, got {other:?}"), + } + } + + #[test] + fn a_failed_tool_call_is_not_reported_as_completed() { + let mut mapper = FrameMapper::new(); + let end = frame(json!({ + "sessionKey": SESSION, + "state": "tool_call_end", + "toolName": "exec", + "error": "non-zero exit", + })); + let FrameAction::Emit(updates) = mapper.map(&end, SESSION) else { + panic!("tool call end must be forwarded"); + }; + match &updates[0] { + acp::SessionUpdate::ToolCallUpdate(update) => { + assert_eq!(update.fields.status, Some(acp::ToolCallStatus::Failed)); + assert_eq!( + update.fields.raw_output, + Some(Value::String("non-zero exit".to_string())) + ); + assert!( + update + .fields + .content + .as_ref() + .is_some_and(|value| !value.is_empty()) + ); + }, + other => panic!("expected a tool call update, got {other:?}"), + } + } + + #[test] + fn unsuccessful_tool_call_is_failed_without_an_error_field() { + let mut mapper = FrameMapper::new(); + let action = mapper.map( + &serde_json::json!({ + "event": "chat", + "payload": { + "sessionKey": SESSION, + "state": "tool_call_end", + "toolName": "exec", + "success": false + } + }) + .to_string(), + SESSION, + ); + let FrameAction::Emit(updates) = action else { + panic!("expected tool update"); + }; + let acp::SessionUpdate::ToolCallUpdate(update) = &updates[0] else { + panic!("expected tool call update"); + }; + assert_eq!(update.fields.status, Some(acp::ToolCallStatus::Failed)); + } + + #[test] + fn null_tool_error_does_not_mark_a_successful_call_failed() { + let mut mapper = FrameMapper::new(); + let action = mapper.map( + &serde_json::json!({ + "event": "chat", + "payload": { + "sessionKey": SESSION, + "state": "tool_call_end", + "toolName": "exec", + "success": true, + "error": null + } + }) + .to_string(), + SESSION, + ); + let FrameAction::Emit(updates) = action else { + panic!("expected tool update"); + }; + let acp::SessionUpdate::ToolCallUpdate(update) = &updates[0] else { + panic!("expected tool call update"); + }; + assert_eq!(update.fields.status, Some(acp::ToolCallStatus::Completed)); + } + + #[test] + fn tool_call_ids_pair_start_with_end() { + let mut mapper = FrameMapper::new(); + let start = frame(json!({ + "sessionKey": SESSION, "state": "tool_call_start", "toolName": "grep", + })); + let end = frame(json!({ + "sessionKey": SESSION, "state": "tool_call_end", "toolName": "grep", + })); + let (FrameAction::Emit(started), FrameAction::Emit(ended)) = + (mapper.map(&start, SESSION), mapper.map(&end, SESSION)) + else { + panic!("both frames must be forwarded"); + }; + let (acp::SessionUpdate::ToolCall(call), acp::SessionUpdate::ToolCallUpdate(update)) = + (&started[0], &ended[0]) + else { + panic!("unexpected update kinds"); + }; + assert_eq!( + call.tool_call_id, update.tool_call_id, + "a client cannot pair the two without a stable id" + ); + } + + #[test] + fn web_ui_only_states_are_dropped() { + let mut mapper = FrameMapper::new(); + for state in [ + "user_message", + "queued", + "iteration", + "thinking", + "thinking_done", + "voice_pending", + "notice", + ] { + assert_eq!( + mapper.map(&chat_frame(state, "x"), SESSION), + FrameAction::Ignore, + "{state} should not reach an ACP client" + ); + } + } + + #[test] + fn completed_text_supplies_a_missing_stream_suffix() { + let mut mapper = FrameMapper::new(); + let FrameAction::Emit(updates) = mapper.map(&chat_frame("delta", "partial"), SESSION) + else { + panic!("delta must emit"); + }; + mapper.record_sent(&updates[0]); + let update = mapper + .finish_text("partial response") + .expect("matching final text") + .expect("missing suffix"); + assert_eq!(emitted_text(&FrameAction::Emit(vec![update])), " response"); + } + + #[test] + fn completed_text_rejects_divergent_streams() { + let mut mapper = FrameMapper::new(); + let FrameAction::Emit(updates) = mapper.map(&chat_frame("delta", "first"), SESSION) else { + panic!("delta must emit"); + }; + mapper.record_sent(&updates[0]); + assert!(mapper.finish_text("different").is_err()); + } + + #[test] + fn new_iterations_reset_final_text_reconciliation() { + let mut mapper = FrameMapper::new(); + let FrameAction::Emit(first) = mapper.map(&chat_frame("delta", "tool preface"), SESSION) + else { + panic!("delta must emit"); + }; + mapper.record_sent(&first[0]); + mapper.map( + &frame(json!({ "sessionKey": SESSION, "state": "iteration", "iteration": 2 })), + SESSION, + ); + let FrameAction::Emit(final_iteration) = + mapper.map(&chat_frame("delta", "final answer"), SESSION) + else { + panic!("delta must emit"); + }; + mapper.record_sent(&final_iteration[0]); + assert_eq!(mapper.finish_text("final answer"), Ok(None)); + } +} diff --git a/crates/cli/src/acp_command.rs b/crates/cli/src/acp_command.rs new file mode 100644 index 0000000000..7504e81d5a --- /dev/null +++ b/crates/cli/src/acp_command.rs @@ -0,0 +1,92 @@ +//! `moltis acp` — serve Moltis to an ACP client over stdio. +//! +//! Any harness that drives ACP agents (Zed, `buzz-acp`, a bespoke runner) spawns +//! the agent as a subprocess and speaks JSON-RPC over its stdin/stdout. One +//! client per process, matching how every ACP harness works. +//! +//! **stdout is the wire.** Callers must have redirected logging to stderr before +//! reaching this module — see `acp_reserves_stdout` in `main.rs`, which is what +//! flips the tracing writer. + +use std::{path::PathBuf, sync::Arc}; + +use {clap::Args, moltis_acp::AcpBackend}; + +use crate::acp_backend::MoltisBackend; + +#[derive(Args, Debug)] +pub struct AcpArgs { + /// Serve a built-in echo agent instead of a real Moltis session. + /// + /// Useful for checking a client's handshake end to end without involving + /// providers, sessions, or tools. + #[arg(long)] + pub echo: bool, + + /// Override the config directory. + #[arg(long)] + pub config_dir: Option, + + /// Override the data directory. + #[arg(long)] + pub data_dir: Option, +} + +/// Resolves the backend to serve and runs the protocol loop until the client +/// disconnects. +pub async fn handle_acp(args: AcpArgs) -> anyhow::Result<()> { + let backend = resolve_backend(&args).await?; + moltis_acp::run_stdio(backend).await +} + +async fn resolve_backend(args: &AcpArgs) -> anyhow::Result> { + if args.echo { + return Ok(Arc::new(moltis_acp::EchoBackend::new())); + } + Ok(Arc::new(MoltisBackend::new(boot_core(args).await?))) +} + +/// Boots the Moltis stack in-process. +/// +/// `prepare_gateway_core_with_profile` is the transport-agnostic half of startup: it wires +/// providers, sessions, memory and tools but binds no socket, so serving ACP +/// does not stand up an HTTP listener or collide with a running gateway on a +/// port. The bind address it takes is only recorded for OAuth callbacks. +async fn boot_core(args: &AcpArgs) -> anyhow::Result { + let core = moltis_gateway::server::prepare_gateway_core_with_profile( + "127.0.0.1", + 0, + true, + None, + args.config_dir.clone(), + args.data_dir.clone(), + None, + None, + None, + moltis_gateway::server::CoreStartupProfile::Headless, + ) + .await?; + Ok(core) +} + +#[allow(clippy::unwrap_used, clippy::expect_used)] +#[cfg(test)] +mod tests { + use super::*; + + fn args(echo: bool) -> AcpArgs { + AcpArgs { + echo, + config_dir: None, + data_dir: None, + } + } + + #[tokio::test] + async fn echo_flag_selects_the_echo_backend() { + // Resolving the echo backend must not boot the Moltis stack: the flag + // exists to check a client's handshake without providers or databases. + let backend = resolve_backend(&args(true)).await.expect("echo backend"); + assert!(!backend.capabilities().load_session); + } +} diff --git a/crates/cli/src/acp_mcp.rs b/crates/cli/src/acp_mcp.rs new file mode 100644 index 0000000000..1c232776e4 --- /dev/null +++ b/crates/cli/src/acp_mcp.rs @@ -0,0 +1,128 @@ +use std::{collections::HashMap, sync::Arc}; + +use { + agent_client_protocol as acp, + moltis_agents::tool_registry::ToolRegistry, + moltis_mcp::{McpManager, McpRegistry, McpServerConfig, StdioLaunchOptions, TransportType}, + secrecy::Secret, + tokio::sync::RwLock, +}; + +use moltis_acp::{SessionKey, SessionSetup}; + +const MAX_PROVIDER_TOOL_NAME_BYTES: usize = 64; + +pub struct SessionMcpRuntime { + manager: Arc, + tools: Arc>, +} + +fn runtime_server_name(namespace: &str, index: usize) -> String { + format!("acp_{namespace}_{index}") +} + +fn runtime_namespace(key: &SessionKey) -> String { + key.as_str() + .strip_prefix("acp:") + .unwrap_or_else(|| key.as_str()) + .bytes() + .filter(u8::is_ascii_alphanumeric) + .take(12) + .map(char::from) + .collect() +} + +impl SessionMcpRuntime { + pub async fn start(key: &SessionKey, setup: &SessionSetup) -> anyhow::Result> { + if setup.mcp_servers().is_empty() { + return Ok(None); + } + + let manager = Arc::new(McpManager::new(McpRegistry::new())); + let options = StdioLaunchOptions { + current_dir: Some(setup.cwd().to_path_buf()), + inherit_parent_env: false, + }; + let namespace = runtime_namespace(key); + for (index, server) in setup.mcp_servers().iter().enumerate() { + let acp::McpServer::Stdio(server) = server else { + manager.shutdown_all().await; + anyhow::bail!("only stdio MCP servers are supported"); + }; + let command = server.command.to_str().map(str::to_owned); + let Some(command) = command else { + manager.shutdown_all().await; + anyhow::bail!("MCP command path is not valid UTF-8"); + }; + let config = McpServerConfig { + command, + args: server.args.clone(), + env: server + .env + .iter() + .map(|variable| (variable.name.clone(), Secret::new(variable.value.clone()))) + .collect::>(), + enabled: true, + transport: TransportType::Stdio, + ..McpServerConfig::default() + }; + // Client MCP identities must not collide with configured servers or + // satisfy an agent preset's allowlist by borrowing a trusted name. + let runtime_name = runtime_server_name(&namespace, index); + if let Err(error) = manager + .start_server_with_options(&runtime_name, &config, &options) + .await + { + manager.shutdown_all().await; + return Err(error.into()); + } + } + + let tools = Arc::new(RwLock::new(ToolRegistry::new())); + moltis_mcp_agent_bridge::sync_mcp_tools(&manager, &tools).await; + let oversized_tool = { + let registry = tools.read().await; + registry.list_schemas().into_iter().find_map(|schema| { + schema + .get("name") + .and_then(serde_json::Value::as_str) + .filter(|name| name.len() > MAX_PROVIDER_TOOL_NAME_BYTES) + .map(str::to_owned) + }) + }; + if let Some(name) = oversized_tool { + manager.shutdown_all().await; + anyhow::bail!( + "client MCP tool name exceeds {MAX_PROVIDER_TOOL_NAME_BYTES} bytes: {name}" + ); + } + Ok(Some(Self { manager, tools })) + } + + pub fn tools(&self) -> Arc> { + Arc::clone(&self.tools) + } + + pub async fn shutdown(self) { + self.manager.shutdown_all().await; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn client_mcp_servers_receive_session_local_names() { + let first_key = SessionKey::new("acp:01234567-89ab-cdef-0123-456789abcdef"); + let second_key = SessionKey::new("acp:abcdef01-2345-6789-abcd-ef0123456789"); + let first_namespace = runtime_namespace(&first_key); + assert_eq!(first_namespace, runtime_namespace(&first_key)); + let first = runtime_server_name(&first_namespace, 0); + let second = runtime_server_name(&runtime_namespace(&second_key), 0); + assert_eq!(first, "acp_0123456789ab_0"); + assert_ne!(first, second); + assert_ne!(first, "github"); + assert!(first.len() <= 18); + } +} diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs index 1b9118db33..e608d46a16 100644 --- a/crates/cli/src/main.rs +++ b/crates/cli/src/main.rs @@ -24,6 +24,12 @@ static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc; #[unsafe(export_name = "malloc_conf")] static malloc_conf: &[u8] = b"dirty_decay_ms:1000,muzzy_decay_ms:1000,background_thread:true\0"; +#[cfg(feature = "acp")] +mod acp_backend; +#[cfg(feature = "acp")] +mod acp_command; +#[cfg(feature = "acp")] +mod acp_mcp; mod auth_commands; mod browser_commands; mod channel_commands; @@ -52,7 +58,13 @@ use { clap::{Parser, Subcommand}, moltis_gateway::logs::{EnabledLogLevels, LogBroadcastLayer, LogBuffer}, tracing::info, - tracing_subscriber::{EnvFilter, fmt, layer::SubscriberExt, util::SubscriberInitExt}, + tracing_subscriber::{ + EnvFilter, Layer as _, + filter::filter_fn, + fmt::{self, writer::BoxMakeWriter}, + layer::SubscriberExt, + util::SubscriberInitExt, + }, }; #[derive(Parser)] @@ -107,6 +119,9 @@ struct Cli { enum Commands { /// Start the gateway server (default when no subcommand is provided). Gateway, + /// Serve Moltis to an ACP client over stdio. + #[cfg(feature = "acp")] + Acp(acp_command::AcpArgs), /// Invoke an agent directly. Agent { #[arg(short, long)] @@ -287,11 +302,48 @@ fn default_telemetry_filter(log_level: &str) -> EnvFilter { filter } +/// Whether the chosen command uses stdout as a machine-readable wire, in which +/// case logs must go to stderr instead. +/// +/// `moltis acp` frames JSON-RPC on stdout: a single log line written there +/// corrupts the stream and the client disconnects with a parse error. +fn command_reserves_stdout(command: Option<&Commands>) -> bool { + #[cfg(feature = "acp")] + if matches!(command, Some(Commands::Acp(_))) { + return true; + } + let _ = command; + false +} + +fn telemetry_target_allowed(target: &str, protocol_payloads_possible: bool) -> bool { + if target.starts_with("agent_client_protocol") { + return false; + } + if !protocol_payloads_possible { + return true; + } + ![ + "moltis_agents", + "moltis_chat", + "moltis_mcp", + "moltis_providers", + "moltis_tools", + ] + .iter() + .any(|prefix| target.starts_with(prefix)) +} + /// Initialise tracing and optionally attach a [`LogBroadcastLayer`] that /// captures events into an in-memory ring buffer for the web UI. fn init_telemetry(cli: &Cli, log_buffer: Option) { - let filter = EnvFilter::try_from_default_env() + let mut filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| default_telemetry_filter(&cli.log_level)); + // The ACP SDK logs complete JSON-RPC payloads, including prompts and MCP + // environment values. Never enable that target, even under RUST_LOG=trace. + if let Ok(directive) = "agent_client_protocol=off".parse() { + filter = filter.add_directive(directive); + } if let Some(ref buffer) = log_buffer { let levels = EnabledLogLevels::from_max_level_hint(filter.max_level_hint()); @@ -300,13 +352,33 @@ fn init_telemetry(cli: &Cli, log_buffer: Option) { let registry = tracing_subscriber::registry().with(filter); - // Optionally attach the in-memory capture layer. - let log_layer = log_buffer.map(LogBroadcastLayer::new); + // Boxed so both destinations share one layer type. + let stdout_reserved = command_reserves_stdout(cli.command.as_ref()); + // Apply the payload filter independently to the in-memory layer. EnvFilter + // specificity must never be able to re-enable protocol payload targets. + let writer = if stdout_reserved { + BoxMakeWriter::new(std::io::stderr) + } else { + BoxMakeWriter::new(std::io::stdout) + }; if cli.json_logs { registry - .with(fmt::layer().json().with_target(true).with_thread_ids(false)) - .with(log_layer) + .with( + fmt::layer() + .json() + .with_target(true) + .with_thread_ids(false) + .with_writer(writer) + .with_filter(filter_fn(move |metadata| { + telemetry_target_allowed(metadata.target(), stdout_reserved) + })), + ) + .with(log_buffer.map(|buffer| { + LogBroadcastLayer::new(buffer).with_filter(filter_fn(move |metadata| { + telemetry_target_allowed(metadata.target(), stdout_reserved) + })) + })) .init(); } else { registry @@ -314,9 +386,18 @@ fn init_telemetry(cli: &Cli, log_buffer: Option) { fmt::layer() .with_target(true) .with_thread_ids(false) - .with_ansi(true), + // ANSI escapes are terminal decoration, never wire content. + .with_ansi(!stdout_reserved) + .with_writer(writer) + .with_filter(filter_fn(move |metadata| { + telemetry_target_allowed(metadata.target(), stdout_reserved) + })), ) - .with(log_layer) + .with(log_buffer.map(|buffer| { + LogBroadcastLayer::new(buffer).with_filter(filter_fn(move |metadata| { + telemetry_target_allowed(metadata.target(), stdout_reserved) + })) + })) .init(); } } @@ -490,6 +571,8 @@ async fn main() -> anyhow::Result<()> { .await .map_err(Into::into) }, + #[cfg(feature = "acp")] + Some(Commands::Acp(args)) => acp_command::handle_acp(args).await, Some(Commands::Agent { message, .. }) => { let result = moltis_agents::runner::run_agent("default", "main", &message).await?; println!("{result}"); @@ -623,7 +706,7 @@ async fn handle_skills(action: SkillAction) -> anyhow::Result<()> { #[cfg(test)] mod tests { - use crate::default_telemetry_filter; + use crate::{default_telemetry_filter, telemetry_target_allowed}; fn manifest_includes_feature(manifest: &str, feature: &str, dependency: &str) -> bool { let Some(start) = manifest.find(&format!("{feature} = [")) else { @@ -645,6 +728,19 @@ mod tests { assert!(filter.contains("matrix_sdk_crypto=error")); } + #[test] + fn protocol_telemetry_filter_cannot_capture_payload_targets() { + assert!(!telemetry_target_allowed( + "agent_client_protocol::rpc", + false + )); + assert!(!telemetry_target_allowed("moltis_providers::openai", true)); + assert!(!telemetry_target_allowed("moltis_mcp::transport", true)); + assert!(!telemetry_target_allowed("moltis_tools::exec", true)); + assert!(telemetry_target_allowed("moltis_cli::acp_backend", true)); + assert!(telemetry_target_allowed("moltis_providers::openai", false)); + } + #[test] fn default_builds_include_all_provider_features() { let cli_manifest = diff --git a/crates/cli/tests/acp_stdio.rs b/crates/cli/tests/acp_stdio.rs new file mode 100644 index 0000000000..597ca5a279 --- /dev/null +++ b/crates/cli/tests/acp_stdio.rs @@ -0,0 +1,165 @@ +//! `moltis acp` speaks JSON-RPC on stdout, so nothing else may write there. +//! +//! The loopback tests in `moltis-acp` prove the protocol layer keeps its own +//! stream clean. They cannot catch the failure that actually bites: a `tracing` +//! subscriber, a stray `println!`, or a startup banner in the *binary* landing +//! on the same file descriptor. That needs the real process, so this drives it +//! end to end with logging turned all the way up. + +#![cfg(feature = "acp")] +#![allow(clippy::unwrap_used, clippy::expect_used)] + +use std::{ + io::{BufRead, BufReader, Read, Write}, + process::{Command, Stdio}, + sync::mpsc, + time::Duration, +}; + +/// Reads `child`'s stdout to EOF after feeding it `input`, returning stdout and +/// stderr. Kills the child if it outlives the deadline. +fn run_acp_with_rust_log(args: &[&str], input: &str, rust_log: &str) -> (String, String) { + let temp = tempfile::tempdir().expect("temp dir"); + std::fs::write( + temp.path().join("moltis.toml"), + "[providers]\noffered = [\"test-disabled\"]\n\n[code_index]\nenabled = false\n", + ) + .expect("write isolated config"); + let mut child = Command::new(env!("CARGO_BIN_EXE_moltis")) + .args(["--log-level", "trace"]) + // Keep the test off the developer's real ~/.moltis. + .args(["--config-dir", &temp.path().to_string_lossy()]) + .args(["--data-dir", &temp.path().to_string_lossy()]) + .args(args) + .env("HOME", temp.path()) + .env("XDG_CONFIG_HOME", temp.path()) + .env("RUST_LOG", rust_log) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("spawn moltis acp"); + let child_stdout = child.stdout.take().expect("stdout"); + let (first_frame_tx, first_frame_rx) = mpsc::sync_channel(1); + let stdout_reader = std::thread::spawn(move || { + let mut reader = BufReader::new(child_stdout); + let mut output = String::new(); + reader.read_line(&mut output).expect("read first frame"); + let _ = first_frame_tx.send(output.clone()); + reader.read_to_string(&mut output).expect("read stdout"); + output + }); + let mut child_stderr = child.stderr.take().expect("stderr"); + let stderr_reader = std::thread::spawn(move || { + let mut output = String::new(); + child_stderr + .read_to_string(&mut output) + .expect("read stderr"); + output + }); + + let mut stdin = child.stdin.take().expect("stdin"); + stdin.write_all(input.as_bytes()).expect("write request"); + stdin.flush().expect("flush"); + // A debug binary links the complete gateway and can be CPU-starved when + // nextest runs the full workspace concurrently. This is a startup deadline, + // not a sleep; focused runs normally answer in about ten seconds. + match first_frame_rx.recv_timeout(Duration::from_secs(180)) { + Ok(frame) if !frame.trim().is_empty() => {}, + result => { + let _ = child.kill(); + let _ = child.wait(); + panic!("moltis acp did not return a response frame: {result:?}"); + }, + } + // The client closes only after receiving its response; dropping stdin then + // signals EOF and exercises graceful backend cleanup. + drop(stdin); + + let deadline = std::time::Instant::now() + Duration::from_secs(60); + let status = loop { + match child.try_wait().expect("wait") { + Some(status) => break status, + None if std::time::Instant::now() >= deadline => { + let _ = child.kill(); + panic!("moltis acp did not exit after stdin closed"); + }, + None => std::thread::sleep(Duration::from_millis(50)), + } + }; + assert!(status.success(), "moltis acp exited with {status}"); + + ( + stdout_reader.join().expect("stdout reader"), + stderr_reader.join().expect("stderr reader"), + ) +} + +fn run_acp(args: &[&str], input: &str) -> (String, String) { + run_acp_with_rust_log(args, input, "trace") +} + +const INITIALIZE: &str = concat!( + r#"{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":1}}"#, + "\n", +); + +#[test] +fn stdout_carries_only_protocol_framing_with_logging_at_trace() { + let (stdout, stderr) = run_acp(&["acp", "--echo"], INITIALIZE); + + let mut frames = 0; + for line in stdout.lines().filter(|line| !line.trim().is_empty()) { + let frame: serde_json::Value = serde_json::from_str(line).unwrap_or_else(|error| { + panic!("stdout must be pure JSON-RPC framing, got {line:?}: {error}") + }); + assert_eq!( + frame.get("jsonrpc").and_then(serde_json::Value::as_str), + Some("2.0"), + "unexpected frame on stdout: {line}" + ); + frames += 1; + } + assert_eq!(frames, 1, "expected exactly one initialize response"); + + // The logs still have to go somewhere, or the redirect is untested. + assert!( + stderr.contains("moltis starting"), + "startup logging should be on stderr, got: {stderr}" + ); +} + +#[test] +fn initialize_reports_moltis_as_the_agent() { + let (stdout, _stderr) = run_acp(&["acp", "--echo"], INITIALIZE); + let line = stdout + .lines() + .find(|line| !line.trim().is_empty()) + .expect("a response frame"); + let frame: serde_json::Value = serde_json::from_str(line).expect("valid JSON"); + assert_eq!(frame["result"]["agentInfo"]["name"], "moltis"); + assert_eq!(frame["result"]["protocolVersion"], 1); +} + +#[test] +fn protocol_payloads_are_not_logged_even_at_trace() { + const SECRET: &str = "ACP-LOG-SENTINEL-DO-NOT-PRINT"; + let input = format!("not-json-{SECRET}\n{INITIALIZE}"); + let (_stdout, stderr) = run_acp_with_rust_log( + &["acp", "--echo"], + &input, + "agent_client_protocol::rpc=trace", + ); + assert!( + !stderr.contains(SECRET), + "malformed protocol payload leaked to stderr: {stderr}" + ); +} + +#[test] +fn without_echo_boots_the_real_backend_and_serves_protocol() { + let (stdout, _stderr) = run_acp(&["acp"], INITIALIZE); + let frame: serde_json::Value = serde_json::from_str(stdout.trim()).expect("valid JSON-RPC"); + assert_eq!(frame["result"]["agentInfo"]["name"], "moltis"); + assert_eq!(frame["result"]["agentCapabilities"]["loadSession"], true); +} diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 563753665b..96e942d58f 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -8,6 +8,7 @@ async-trait = { workspace = true } futures = { workspace = true } ipnet = { workspace = true } moltis-metrics = { optional = true, workspace = true } +process-wrap = { optional = true, workspace = true } reqwest = { workspace = true } secrecy = { workspace = true } serde = { workspace = true } @@ -24,6 +25,7 @@ tempfile = { workspace = true } [features] default = [] metrics = ["dep:moltis-metrics"] +process = ["dep:process-wrap"] [lints] workspace = true diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 5279d4d63e..c9639ad2b6 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -4,6 +4,8 @@ pub mod context_command; pub mod error; pub mod hooks; pub mod http_client; +#[cfg(feature = "process")] +pub mod process_tree; pub mod secret_serde; pub mod ssrf; pub mod types; diff --git a/crates/common/src/process_tree.rs b/crates/common/src/process_tree.rs new file mode 100644 index 0000000000..775f4ac6ed --- /dev/null +++ b/crates/common/src/process_tree.rs @@ -0,0 +1,121 @@ +use std::{io, process::ExitStatus}; + +use { + process_wrap::tokio::{ChildWrapper, CommandWrap, KillOnDrop}, + tokio::process::{ChildStderr, ChildStdin, ChildStdout, Command}, +}; + +/// Owns a spawned process and all descendants that remain in its OS process +/// group (Unix) or Job Object (Windows). +#[derive(Debug)] +pub struct OwnedProcessTree { + child: Box, + terminated: bool, +} + +impl OwnedProcessTree { + pub fn spawn(command: Command) -> io::Result { + let mut command = CommandWrap::from(command); + command.wrap(KillOnDrop); + #[cfg(unix)] + command.wrap(process_wrap::tokio::ProcessGroup::leader()); + #[cfg(windows)] + command.wrap(process_wrap::tokio::JobObject); + + Ok(Self { + child: command.spawn()?, + terminated: false, + }) + } + + pub fn take_stdin(&mut self) -> Option { + self.child.stdin().take() + } + + pub fn take_stdout(&mut self) -> Option { + self.child.stdout().take() + } + + pub fn take_stderr(&mut self) -> Option { + self.child.stderr().take() + } + + pub fn try_wait(&mut self) -> io::Result> { + self.child.try_wait() + } + + pub async fn wait(&mut self) -> io::Result { + let status = self.child.wait().await?; + self.terminated = true; + Ok(status) + } + + pub async fn kill(&mut self) -> io::Result<()> { + Box::into_pin(self.child.kill()).await?; + self.terminated = true; + Ok(()) + } +} + +impl Drop for OwnedProcessTree { + fn drop(&mut self) { + if !self.terminated { + let _ = self.child.start_kill(); + } + } +} + +#[cfg(all(test, unix))] +#[allow(clippy::unwrap_used)] +mod tests { + use std::{process::Stdio, time::Duration}; + + use tokio::io::{AsyncBufReadExt, BufReader}; + + use super::*; + + fn descendant_command(marker: &std::path::Path) -> Command { + let mut command = Command::new("sh"); + command + .arg("-c") + .arg(format!( + "(sleep 1; touch '{}') & echo ready; wait", + marker.display() + )) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .stdin(Stdio::null()); + command + } + + async fn wait_until_ready(tree: &mut OwnedProcessTree) { + let stdout = tree.take_stdout().unwrap(); + let mut line = String::new(); + BufReader::new(stdout).read_line(&mut line).await.unwrap(); + assert_eq!(line.trim(), "ready"); + } + + #[tokio::test] + async fn drop_kills_descendants() { + let temp_dir = tempfile::tempdir().unwrap(); + let marker = temp_dir.path().join("drop-marker"); + let mut tree = OwnedProcessTree::spawn(descendant_command(&marker)).unwrap(); + wait_until_ready(&mut tree).await; + + drop(tree); + tokio::time::sleep(Duration::from_millis(1100)).await; + assert!(!marker.exists(), "descendant survived process-tree drop"); + } + + #[tokio::test] + async fn explicit_kill_terminates_descendants() { + let temp_dir = tempfile::tempdir().unwrap(); + let marker = temp_dir.path().join("kill-marker"); + let mut tree = OwnedProcessTree::spawn(descendant_command(&marker)).unwrap(); + wait_until_ready(&mut tree).await; + + tree.kill().await.unwrap(); + tokio::time::sleep(Duration::from_millis(1100)).await; + assert!(!marker.exists(), "descendant survived process-tree kill"); + } +} diff --git a/crates/external-agents/src/runtimes/acp.rs b/crates/external-agents/src/runtimes/acp.rs index 9a24bf99e4..18bd5b0f37 100644 --- a/crates/external-agents/src/runtimes/acp.rs +++ b/crates/external-agents/src/runtimes/acp.rs @@ -30,6 +30,9 @@ use crate::{ }, }; +const MAX_ACP_FRAME_BYTES: usize = 16 * 1024 * 1024; +const MAX_INFLIGHT_ACP_MESSAGES: usize = 256; + /// Transport for ACP (Agent Client Protocol) agents over JSON-RPC stdio. pub struct AcpTransport { kind: AgentTransportKind, @@ -601,10 +604,15 @@ async fn run_acp_controller( tokio::task::spawn_local(forward_stderr(stderr, Arc::clone(&client))); } - let (conn, io_task) = acp::ClientSideConnection::new( + let (conn, io_task) = acp::ClientSideConnection::new_with_limits( Arc::clone(&client), stdin.compat_write(), stdout.compat(), + acp::ConnectionLimits::bounded( + MAX_ACP_FRAME_BYTES, + MAX_ACP_FRAME_BYTES, + MAX_INFLIGHT_ACP_MESSAGES, + ), |future| { tokio::task::spawn_local(future); }, diff --git a/crates/gateway/src/broadcast.rs b/crates/gateway/src/broadcast.rs index 610277cb04..7839d684c6 100644 --- a/crates/gateway/src/broadcast.rs +++ b/crates/gateway/src/broadcast.rs @@ -143,6 +143,28 @@ pub async fn broadcast( continue; } + if let Some(session_key) = client.session_filter.as_deref() + && frame + .payload + .as_ref() + .and_then(|payload| payload.get("sessionKey")) + .and_then(serde_json::Value::as_str) + != Some(session_key) + { + continue; + } + + if let Some(states) = client.payload_state_filter.as_ref() + && !frame + .payload + .as_ref() + .and_then(|payload| payload.get("state")) + .and_then(serde_json::Value::as_str) + .is_some_and(|state| states.contains(state)) + { + continue; + } + // Channel filter (v4): if event is scoped to a channel, skip clients // that haven't joined it. if let Some(ref ch) = opts.channel diff --git a/crates/gateway/src/channel_events/commands/media.rs b/crates/gateway/src/channel_events/commands/media.rs index 93d7a15633..033f4760dc 100644 --- a/crates/gateway/src/channel_events/commands/media.rs +++ b/crates/gateway/src/channel_events/commands/media.rs @@ -83,10 +83,14 @@ pub(in crate::channel_events) async fn save_channel_attachment( let store = state.services.session_store.as_ref()?; match store.save_media(&session_key, filename, file_data).await { Ok(media_ref) => { - let absolute_path = store - .media_path_for(&session_key, filename) - .to_string_lossy() - .to_string(); + let Ok(absolute_path) = store.media_path_for(&session_key, filename) else { + warn!( + session_key, + filename, "invalid saved channel attachment path" + ); + return None; + }; + let absolute_path = absolute_path.to_string_lossy().to_string(); debug!( session_key, filename, media_ref, absolute_path, "saved channel attachment to session media" diff --git a/crates/gateway/src/external_agents.rs b/crates/gateway/src/external_agents.rs index 18915fe732..e3bb828460 100644 --- a/crates/gateway/src/external_agents.rs +++ b/crates/gateway/src/external_agents.rs @@ -537,6 +537,7 @@ impl ExternalAgentChatService { }, }; let mut assistant_text = String::new(); + let mut thinking_text = String::new(); let mut token_usage = None; let mut external_error = None; while let Some(event) = events.next().await { @@ -558,6 +559,7 @@ impl ExternalAgentChatService { .await; }, ExternalAgentEvent::ThinkingDelta(delta) => { + thinking_text.push_str(&delta); crate::broadcast::broadcast( &self.state, "chat", @@ -565,7 +567,7 @@ impl ExternalAgentChatService { "runId": run_id, "sessionKey": session_key, "state": "thinking_text", - "text": delta, + "text": thinking_text, "seq": seq, }), BroadcastOpts::default(), @@ -645,7 +647,14 @@ impl ExternalAgentChatService { BroadcastOpts::default(), ) .await; - Ok(serde_json::json!({ "ok": true, "runId": run_id })) + Ok(serde_json::json!({ + "ok": true, + "runId": run_id, + "text": assistant_text, + "inputTokens": token_usage.as_ref().map(|usage| usage.input_tokens).unwrap_or(0), + "outputTokens": token_usage.as_ref().map(|usage| usage.output_tokens).unwrap_or(0), + "durationMs": duration_ms, + })) } } @@ -659,7 +668,10 @@ impl ChatService for ExternalAgentChatService { } async fn send_sync(&self, params: Value) -> ServiceResult { - self.send(params).await + if let Some(result) = self.maybe_send_external(¶ms).await { + return result; + } + self.inner.send_sync(params).await } async fn abort(&self, params: Value) -> ServiceResult { diff --git a/crates/gateway/src/server/init_channels.rs b/crates/gateway/src/server/init_channels.rs index c35b2e0a53..bf6a65c76a 100644 --- a/crates/gateway/src/server/init_channels.rs +++ b/crates/gateway/src/server/init_channels.rs @@ -32,6 +32,7 @@ pub(crate) async fn init_channels( session_metadata: Arc, deferred_state: Arc>>, data_dir: &std::path::Path, + start_accounts: bool, ) -> ChannelInitResult { #[cfg(not(feature = "whatsapp"))] let _ = data_dir; @@ -192,70 +193,73 @@ pub(crate) async fn init_channels( let mut pending_starts: Vec<(String, String, serde_json::Value)> = Vec::new(); let mut queued: HashSet<(String, String)> = HashSet::new(); - #[cfg(feature = "telephony")] - if let Some((account_id, account_config)) = crate::methods::phone::phone_channel_account(config) - { - let key = ("telephony".to_string(), account_id.clone()); - if registry.get("telephony").is_some() && queued.insert(key) { - pending_starts.push(("telephony".to_string(), account_id, account_config)); + if start_accounts { + #[cfg(feature = "telephony")] + if let Some((account_id, account_config)) = + crate::methods::phone::phone_channel_account(config) + { + let key = ("telephony".to_string(), account_id.clone()); + if registry.get("telephony").is_some() && queued.insert(key) { + pending_starts.push(("telephony".to_string(), account_id, account_config)); + } } - } - for (channel_type, accounts) in config.channels.all_channel_configs() { - if registry.get(channel_type).is_none() { - if !accounts.is_empty() { - tracing::debug!( - channel_type, - "skipping config — no plugin registered for this channel type" - ); + for (channel_type, accounts) in config.channels.all_channel_configs() { + if registry.get(channel_type).is_none() { + if !accounts.is_empty() { + tracing::debug!( + channel_type, + "skipping config — no plugin registered for this channel type" + ); + } + continue; } - continue; - } - for (account_id, account_config) in accounts { - let key = (channel_type.to_string(), account_id.clone()); - if queued.insert(key) { - pending_starts.push(( - channel_type.to_string(), - account_id.clone(), - account_config.clone(), - )); + for (account_id, account_config) in accounts { + let key = (channel_type.to_string(), account_id.clone()); + if queued.insert(key) { + pending_starts.push(( + channel_type.to_string(), + account_id.clone(), + account_config.clone(), + )); + } } } - } - // Load persisted channels that were not queued from config. - match channel_store.list().await { - Ok(stored) => { - info!("{} stored channel(s) found in database", stored.len()); - for ch in stored { - let key = (ch.channel_type.clone(), ch.account_id.clone()); - if queued.contains(&key) { + // Load persisted channels that were not queued from config. + match channel_store.list().await { + Ok(stored) => { + info!("{} stored channel(s) found in database", stored.len()); + for ch in stored { + let key = (ch.channel_type.clone(), ch.account_id.clone()); + if queued.contains(&key) { + info!( + account_id = ch.account_id, + channel_type = ch.channel_type, + "skipping stored channel (already started from config)" + ); + continue; + } + if registry.get(&ch.channel_type).is_none() { + tracing::warn!( + account_id = ch.account_id, + channel_type = ch.channel_type, + "unsupported channel type, skipping stored account" + ); + continue; + } info!( account_id = ch.account_id, channel_type = ch.channel_type, - "skipping stored channel (already started from config)" - ); - continue; - } - if registry.get(&ch.channel_type).is_none() { - tracing::warn!( - account_id = ch.account_id, - channel_type = ch.channel_type, - "unsupported channel type, skipping stored account" + "starting stored channel" ); - continue; - } - info!( - account_id = ch.account_id, - channel_type = ch.channel_type, - "starting stored channel" - ); - if queued.insert(key) { - pending_starts.push((ch.channel_type, ch.account_id, ch.config)); + if queued.insert(key) { + pending_starts.push((ch.channel_type, ch.account_id, ch.config)); + } } - } - }, - Err(e) => tracing::warn!("failed to load stored channels: {e}"), + }, + Err(e) => tracing::warn!("failed to load stored channels: {e}"), + } } let registry = Arc::new(registry); diff --git a/crates/gateway/src/server/init_memory.rs b/crates/gateway/src/server/init_memory.rs index b1b7a14364..48d221fe18 100644 --- a/crates/gateway/src/server/init_memory.rs +++ b/crates/gateway/src/server/init_memory.rs @@ -21,6 +21,7 @@ pub(crate) async fn init_memory_system( effective_providers: &moltis_config::schema::ProvidersConfig, runtime_env_overrides: &HashMap, db_pool_max_connections: u32, + start_background_tasks: bool, ) -> Option { // Build embedding provider(s) for the fallback chain. let mut embedding_providers: Vec<( @@ -224,7 +225,14 @@ pub(crate) async fn init_memory_system( tracing::warn!("memory migration failed: {e}"); None } else { - build_memory_runtime(mem_cfg, data_dir, embedder, memory_pool).await + build_memory_runtime( + mem_cfg, + data_dir, + embedder, + memory_pool, + start_background_tasks, + ) + .await } }, Err(e) => { @@ -240,6 +248,7 @@ async fn build_memory_runtime( data_dir: &FsPath, embedder: Option>, memory_pool: sqlx::SqlitePool, + start_background_tasks: bool, ) -> Option { let data_memory_file = data_dir.join("MEMORY.md"); let data_memory_file_lower = data_dir.join("memory.md"); @@ -341,6 +350,13 @@ async fn build_memory_runtime( }, }; + if !start_background_tasks { + if let Err(error) = manager.sync().await { + tracing::warn!(%error, "memory: initial headless sync failed"); + } + return Some(manager); + } + // Initial sync + periodic re-sync (15min with watcher, 5min without). let sync_manager = Arc::clone(&manager); tokio::spawn(async move { diff --git a/crates/gateway/src/server/mod.rs b/crates/gateway/src/server/mod.rs index ff32319dc4..dde965e0f4 100644 --- a/crates/gateway/src/server/mod.rs +++ b/crates/gateway/src/server/mod.rs @@ -35,8 +35,8 @@ mod tests_legacy; pub(crate) use hooks::discover_and_build_hooks; pub use { helpers::approval_manager_from_config, - prepare_core::prepare_gateway_core, - prepared::PreparedGatewayCore, + prepare_core::{prepare_gateway_core, prepare_gateway_core_with_profile}, + prepared::{CoreStartupProfile, PreparedGatewayCore}, startup::{ claude_detected_for_ui, codex_detected_for_ui, hermes_detected_for_ui, openclaw_detected_for_ui, start_browser_warmup_after_listener, diff --git a/crates/gateway/src/server/prepare_core.rs b/crates/gateway/src/server/prepare_core.rs index 69e965fb84..2d25a953ff 100644 --- a/crates/gateway/src/server/prepare_core.rs +++ b/crates/gateway/src/server/prepare_core.rs @@ -7,7 +7,7 @@ use { validate_proxy_tls_configuration, }, init_channels, init_code_index, init_memory, - prepared::PreparedGatewayCore, + prepared::{CoreStartupProfile, PreparedGatewayCore}, workspace::{ seed_default_workspace_markdown_files, sync_persona_into_preset, warn_on_workspace_prompt_file_truncation, @@ -38,6 +38,27 @@ mod log_persistence; mod post_state; mod sandbox; mod tool_registration; + +#[cfg(feature = "trusted-network")] +fn apply_web_fetch_network_policy( + tool: moltis_tools::web_fetch::WebFetchTool, + profile: CoreStartupProfile, + sandbox: &moltis_tools::sandbox::SandboxConfig, +) -> moltis_tools::web_fetch::WebFetchTool { + if sandbox.network != moltis_network_filter::NetworkPolicy::Trusted { + return tool; + } + let proxy = if profile.is_headless() { + "http://127.0.0.1:0".to_string() + } else { + format!( + "http://127.0.0.1:{}", + moltis_network_filter::DEFAULT_PROXY_PORT + ) + }; + tool.with_proxy(proxy) +} + /// Prepare the core gateway: load config, run migrations, wire services, /// spawn background tasks, and return the core state without any HTTP layer. /// This is the transport-agnostic initialisation. Non-HTTP consumers (TUI, @@ -54,6 +75,34 @@ pub async fn prepare_gateway_core( tailscale_mode_override: Option, tailscale_reset_on_exit_override: Option, session_event_bus: Option, +) -> anyhow::Result { + prepare_gateway_core_with_profile( + bind, + port, + no_tls, + log_buffer, + config_dir, + data_dir, + tailscale_mode_override, + tailscale_reset_on_exit_override, + session_event_bus, + CoreStartupProfile::Server, + ) + .await +} + +#[allow(clippy::expect_used)] +pub async fn prepare_gateway_core_with_profile( + bind: &str, + port: u16, + no_tls: bool, + log_buffer: Option, + config_dir: Option, + data_dir: Option, + tailscale_mode_override: Option, + tailscale_reset_on_exit_override: Option, + session_event_bus: Option, + profile: CoreStartupProfile, ) -> anyhow::Result { let session_event_bus = session_event_bus.unwrap_or_default(); #[cfg(not(feature = "tailscale"))] @@ -204,7 +253,7 @@ pub async fn prepare_gateway_core( // Refresh dynamic provider model discovery daily. const DYNAMIC_PROVIDER_MODEL_REFRESH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(24 * 60 * 60); - { + if !profile.is_headless() { let registry_for_refresh = Arc::clone(®istry); let provider_config_for_refresh = base_provider_config.clone(); tokio::spawn(async move { @@ -518,15 +567,17 @@ pub async fn prepare_gateway_core( live_mcp .set_credential_store(Arc::clone(&credential_store)) .await; - let mgr = Arc::clone(live_mcp.manager()); - let mcp_for_sync = Arc::clone(&live_mcp); - tokio::spawn(async move { - let started = mgr.start_enabled().await; - if !started.is_empty() { - tracing::info!(servers = ?started, "MCP servers started"); - } - mcp_for_sync.sync_tools_if_ready().await; - }); + if !profile.is_headless() { + let mgr = Arc::clone(live_mcp.manager()); + let mcp_for_sync = Arc::clone(&live_mcp); + tokio::spawn(async move { + let started = mgr.start_enabled().await; + if !started.is_empty() { + tracing::info!(servers = ?started, "MCP servers started"); + } + mcp_for_sync.sync_tools_if_ready().await; + }); + } // If MOLTIS_PASSWORD is set and no password in DB yet, migrate it. if let Some(ref pw) = password @@ -944,7 +995,9 @@ pub async fn prepare_gateway_core( "trusted-network: evaluating network policy" ); - if sandbox_config.network == moltis_network_filter::NetworkPolicy::Trusted { + if !profile.is_headless() + && sandbox_config.network == moltis_network_filter::NetworkPolicy::Trusted + { let domain_mgr = Arc::new( moltis_network_filter::domain_approval::DomainApprovalManager::new( &sandbox_config.trusted_domains, @@ -975,11 +1028,21 @@ pub async fn prepare_gateway_core( moltis_tools::init_shared_http_client(Some(&url)); proxy_shutdown_tx = Some(shutdown_tx); } else { - info!( + debug!( network_policy = ?sandbox_config.network, - "trusted-network proxy not started (policy is not Trusted)" + headless = profile.is_headless(), + "trusted-network proxy not started" ); - moltis_tools::init_shared_http_client(upstream_proxy); + if profile.is_headless() + && sandbox_config.network == moltis_network_filter::NetworkPolicy::Trusted + { + // ACP has no domain-approval UI, so trusted-network mode must + // fail closed rather than bypass policy or bind a proxy port. + info!("trusted-network HTTP access disabled in headless mode"); + moltis_tools::init_shared_http_client(Some("http://127.0.0.1:0")); + } else { + moltis_tools::init_shared_http_client(upstream_proxy); + } proxy_shutdown_tx = None; } @@ -996,10 +1059,13 @@ pub async fn prepare_gateway_core( } // Spawn background sandbox tasks (image pre-build, host provisioning, container GC). - sandbox::spawn_sandbox_background_tasks(&sandbox_router, &deferred_state); + if !profile.is_headless() { + sandbox::spawn_sandbox_background_tasks(&sandbox_router, &deferred_state); + } // Periodic cron session retention pruning. - if let Some(retention_days) = config.cron.session_retention_days + if !profile.is_headless() + && let Some(retention_days) = config.cron.session_retention_days && retention_days > 0 { let prune_store = Arc::clone(&cron_store_for_pruning); @@ -1060,7 +1126,8 @@ pub async fn prepare_gateway_core( } // Pre-pull browser container image. - if config.tools.browser.enabled + if !profile.is_headless() + && config.tools.browser.enabled && !matches!( sandbox_router.config().mode, moltis_tools::sandbox::SandboxMode::Off @@ -1145,6 +1212,7 @@ pub async fn prepare_gateway_core( Arc::clone(&session_metadata), Arc::clone(&deferred_state), &data_dir, + !profile.is_headless(), ) .await; services = channel_result.services; @@ -1225,6 +1293,7 @@ pub async fn prepare_gateway_core( &effective_providers, &runtime_env_overrides, config.server.db_pool_max_connections, + !profile.is_headless(), ) .await; startup_mem_probe.checkpoint("memory_manager.initialized"); @@ -1234,6 +1303,7 @@ pub async fn prepare_gateway_core( startup_mem_probe.checkpoint("code_index.initialized"); post_state::complete_startup(post_state::PostStateInputs { + profile, bind: bind.to_string(), port, config, diff --git a/crates/gateway/src/server/prepare_core/post_state.rs b/crates/gateway/src/server/prepare_core/post_state.rs index 28bf6d7fbb..5af6fe23f9 100644 --- a/crates/gateway/src/server/prepare_core/post_state.rs +++ b/crates/gateway/src/server/prepare_core/post_state.rs @@ -10,7 +10,9 @@ use { mod credential_env; -use credential_env::{CredentialEnvVarProvider, ensure_sandbox_api_key}; +use credential_env::{ + CredentialEnvVarProvider, ensure_sandbox_api_key, gateway_credentials_allowed, +}; use { moltis_providers::{PendingDiscoveries, ProviderRegistry}, @@ -52,6 +54,7 @@ use crate::server::helpers::fs_tools_host_warning_message; use crate::server::helpers::start_skill_hot_reload_watcher; pub(super) struct PostStateInputs { + pub profile: crate::server::CoreStartupProfile, pub bind: String, pub port: u16, pub config: moltis_config::MoltisConfig, @@ -242,6 +245,7 @@ pub(super) async fn complete_startup( inputs: PostStateInputs, ) -> anyhow::Result { let PostStateInputs { + profile, bind, port, config, @@ -405,7 +409,7 @@ pub(super) async fn complete_startup( vault.clone(), ); - { + if !profile.is_headless() { let (webhook_tx, webhook_rx) = tokio::sync::mpsc::channel::(256); let webhook_store: Arc = { let inner: Arc = Arc::new( @@ -484,7 +488,9 @@ pub(super) async fn complete_startup( .or_else(|| config.providers.get("local-llm")) .and_then(|e| e.idle_timeout_secs); svc.populate_lifecycle(global_timeout).await; - svc.lifecycle().spawn_idle_checker(); + if !profile.is_headless() { + svc.lifecycle().spawn_idle_checker(); + } } provider_setup_service.set_broadcaster(Arc::new(crate::provider_setup::GatewayBroadcaster { @@ -558,7 +564,7 @@ pub(super) async fn complete_startup( .get("local") .or_else(|| config.providers.get("local-llm")) .and_then(|e| e.idle_timeout_secs); - tokio::spawn(async move { + let discovery = async move { let startup_discovery_started = std::time::Instant::now(); let prefetched = match tokio::task::spawn_blocking(move || { ProviderRegistry::collect_discoveries(startup_discovery_pending) @@ -634,7 +640,12 @@ pub(super) async fn complete_startup( BroadcastOpts::default(), ) .await; - }); + }; + if profile.is_headless() { + discovery.await; + } else { + tokio::spawn(discovery); + } } { @@ -645,17 +656,15 @@ pub(super) async fn complete_startup( #[cfg(feature = "graphql")] state.set_graphql_enabled(config.graphql.enabled); - { + let live_chat = { let broadcaster: Arc = Arc::new(GatewayApprovalBroadcaster::new(Arc::clone(&state))); // Build gateway URL for sandbox-to-gateway communication. // Only inject when the sandbox network policy allows host access // (Trusted or Bypass). With NetworkPolicy::Blocked the container // has --network=none and host.docker.internal won't resolve. - let sandbox_network_allows_host = !matches!( - sandbox_router.config().network, - moltis_tools::sandbox::NetworkPolicy::Blocked - ); + let sandbox_network_allows_host = + gateway_credentials_allowed(profile, &sandbox_router.config().network); let sandbox_gateway_url = if sandbox_network_allows_host { let scheme = if tls_enabled_for_gateway { "https" @@ -907,6 +916,8 @@ pub(super) async fn complete_startup( { #[cfg(feature = "firecrawl")] let t = t.with_firecrawl(&config.tools.web.firecrawl); + #[cfg(feature = "trusted-network")] + let t = super::apply_web_fetch_network_policy(t, profile, sandbox_router.config()); tool_registry.register(Box::new(t)); } #[cfg(feature = "firecrawl")] @@ -1347,8 +1358,9 @@ pub(super) async fn complete_startup( } let live_chat = Arc::new(chat_service); + let inner_chat: Arc = live_chat.clone(); let chat_with_external_agents = Arc::new(ExternalAgentChatService::new( - live_chat, + inner_chat, external_agent_service, Arc::clone(&state), Arc::clone(&session_store), @@ -1359,15 +1371,22 @@ pub(super) async fn complete_startup( live_mcp .set_tool_registry(Arc::clone(&shared_tool_registry)) .await; + if profile.is_headless() { + let started = live_mcp.manager().start_enabled().await; + if !started.is_empty() { + info!(servers = ?started, "headless MCP servers started"); + } + } crate::mcp_service::sync_mcp_tools(live_mcp.manager(), &shared_tool_registry).await; let schemas = shared_tool_registry.read().await.list_schemas(); let tool_names: Vec<&str> = schemas.iter().filter_map(|s| s["name"].as_str()).collect(); info!(tools = ?tool_names, "agent tools registered"); - } + live_chat + }; #[cfg(feature = "file-watcher")] - { + if !profile.is_headless() { let watcher_state = Arc::clone(&state); tokio::spawn(async move { let (mut watcher, mut rx) = match start_skill_hot_reload_watcher().await { @@ -1421,7 +1440,9 @@ pub(super) async fn complete_startup( let methods = Arc::new(MethodRegistry::new()); #[cfg(feature = "push-notifications")] - let push_service: Option> = { + let push_service: Option> = if profile.is_headless() { + None + } else { match crate::push::PushService::new(&data_dir).await { Ok(svc) => { info!("push notification service initialized"); @@ -1439,6 +1460,8 @@ pub(super) async fn complete_startup( Ok(PreparedGatewayCore { state: Arc::clone(&state), + live_chat, + mcp_manager: Arc::clone(live_mcp.manager()), methods: Arc::clone(&methods), webauthn_registry, #[cfg(feature = "msteams")] diff --git a/crates/gateway/src/server/prepare_core/post_state/credential_env.rs b/crates/gateway/src/server/prepare_core/post_state/credential_env.rs index 57f6f74f3e..6966b0e0d8 100644 --- a/crates/gateway/src/server/prepare_core/post_state/credential_env.rs +++ b/crates/gateway/src/server/prepare_core/post_state/credential_env.rs @@ -8,6 +8,15 @@ use { use crate::auth; +use crate::server::CoreStartupProfile; + +pub(super) fn gateway_credentials_allowed( + profile: CoreStartupProfile, + network: &moltis_tools::sandbox::NetworkPolicy, +) -> bool { + !profile.is_headless() && *network != moltis_tools::sandbox::NetworkPolicy::Blocked +} + pub(super) struct CredentialEnvVarProvider { pub(super) store: Arc, pub(super) gateway_url: Option, @@ -68,3 +77,20 @@ pub(super) async fn ensure_sandbox_api_key(store: &auth::CredentialStore) -> Opt }, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn headless_profiles_never_create_gateway_credentials() { + assert!(!gateway_credentials_allowed( + CoreStartupProfile::Headless, + &moltis_tools::sandbox::NetworkPolicy::Trusted, + )); + assert!(!gateway_credentials_allowed( + CoreStartupProfile::Headless, + &moltis_tools::sandbox::NetworkPolicy::Bypass, + )); + } +} diff --git a/crates/gateway/src/server/prepared.rs b/crates/gateway/src/server/prepared.rs index 2b20108afa..6e6206cb2f 100644 --- a/crates/gateway/src/server/prepared.rs +++ b/crates/gateway/src/server/prepared.rs @@ -5,6 +5,20 @@ use crate::{auth_webauthn::SharedWebAuthnRegistry, methods::MethodRegistry, stat #[cfg(feature = "tailscale")] use crate::tailscale::TailscaleMode; +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum CoreStartupProfile { + #[default] + Server, + Headless, +} + +impl CoreStartupProfile { + #[must_use] + pub fn is_headless(self) -> bool { + self == Self::Headless + } +} + /// Core gateway state produced by [`super::prepare_gateway_core`]. /// /// Contains everything needed to build an HTTP server on top of the core, but @@ -13,6 +27,10 @@ use crate::tailscale::TailscaleMode; pub struct PreparedGatewayCore { /// Shared gateway state (sessions, services, config, etc.). pub state: Arc, + /// Concrete chat service used by non-HTTP surfaces with session-local tools. + pub live_chat: Arc, + /// Configured MCP manager retained for readiness and graceful shutdown. + pub mcp_manager: Arc, /// RPC method registry. pub methods: Arc, /// WebAuthn registry for passkey auth. diff --git a/crates/gateway/src/session/maintenance.rs b/crates/gateway/src/session/maintenance.rs index bc18313826..0dde5c3a4c 100644 --- a/crates/gateway/src/session/maintenance.rs +++ b/crates/gateway/src/session/maintenance.rs @@ -247,9 +247,16 @@ impl LiveSessionService { max.saturating_mul(10).min(200) }; + let known_keys = self + .metadata + .list() + .await + .into_iter() + .map(|entry| entry.key) + .collect::>(); let results = self .store - .search(query, search_limit) + .search_known_keys(query, search_limit, &known_keys) .await .map_err(ServiceError::message)?; diff --git a/crates/gateway/src/session/tests.rs b/crates/gateway/src/session/tests.rs index 78f497d691..04a2ecb815 100644 --- a/crates/gateway/src/session/tests.rs +++ b/crates/gateway/src/session/tests.rs @@ -656,7 +656,7 @@ mod tests { &serde_json::json!({ "role": "assistant", "content": "hi there", - "audio": existing_path, + "audio": existing_path.clone(), "tts_provider": "openai", "run_id": "run-abc", }), @@ -677,7 +677,7 @@ mod tests { .expect("voice generate"); assert_eq!(result["reused"], true); - assert_eq!(result["audio"].as_str(), Some("media/main/voice-msg-1.ogg")); + assert_eq!(result["audio"].as_str(), Some(existing_path.as_str())); assert_eq!(result["ttsProvider"].as_str(), Some("openai")); assert_eq!(mock_tts.convert_calls.load(Ordering::SeqCst), 0); } @@ -726,7 +726,10 @@ mod tests { assert_eq!(result["reused"], false); let audio_path = result["audio"].as_str().unwrap_or_default().to_string(); - assert_eq!(audio_path, "media/main/voice-msg-1.mp3"); + assert_eq!( + audio_path, + SessionStore::media_reference("main", "voice-msg-1.mp3").expect("media reference") + ); assert_eq!(result["ttsProvider"].as_str(), Some("openai")); assert_eq!(mock_tts.convert_calls.load(Ordering::SeqCst), 1); let convert_params = mock_tts @@ -794,7 +797,10 @@ mod tests { assert_eq!(result["reused"], false); let audio_path = result["audio"].as_str().unwrap_or_default().to_string(); - assert_eq!(audio_path, "media/main/voice-msg-1.ogg"); + assert_eq!( + audio_path, + SessionStore::media_reference("main", "voice-msg-1.ogg").expect("media reference") + ); assert_eq!(result["ttsProvider"].as_str(), Some("elevenlabs")); let convert_params = mock_tts .last_convert_params @@ -865,7 +871,7 @@ mod tests { &serde_json::json!({ "role": "assistant", "content": "assistant answer", - "audio": existing_path, + "audio": existing_path.clone(), "run_id": "run-target", }), ) @@ -888,7 +894,7 @@ mod tests { assert_eq!(result["reused"], true); assert_eq!(result["messageIndex"], 2); - assert_eq!(result["audio"].as_str(), Some("media/main/voice-msg-2.ogg")); + assert_eq!(result["audio"].as_str(), Some(existing_path.as_str())); assert_eq!(mock_tts.convert_calls.load(Ordering::SeqCst), 0); } diff --git a/crates/gateway/src/state.rs b/crates/gateway/src/state.rs index 857f9f56a3..4a534b2ea9 100644 --- a/crates/gateway/src/state.rs +++ b/crates/gateway/src/state.rs @@ -103,6 +103,8 @@ pub struct ConnectedClient { pub connect_params: ConnectParams, /// Bounded channel for sending serialized frames to this client's write loop. pub sender: mpsc::Sender, + /// Frames rejected because the client's bounded channel was full. + pub delivery_failures: Arc, pub connected_at: Instant, /// Milliseconds since process start. Updated atomically — no write lock needed. pub last_activity_ms: std::sync::atomic::AtomicU64, @@ -119,6 +121,13 @@ pub struct ConnectedClient { /// `None` = wildcard (receive everything, v3 compat). /// `Some(set)` = only events in the set (or `"*"` = wildcard). pub subscriptions: Option>, + /// Optional session key for in-process clients that consume one session's + /// events. Filtering before enqueue keeps unrelated traffic from filling + /// their bounded delivery channel. + pub session_filter: Option, + /// Optional `payload.state` allowlist for in-process consumers. This makes + /// delivery-failure accounting reflect only states the consumer needs. + pub payload_state_filter: Option>, /// Channels this client has joined (v4 multiplexing). pub joined_channels: HashSet, /// Negotiated protocol version for this connection. @@ -165,7 +174,11 @@ impl ConnectedClient { /// Uses `try_send` to avoid blocking; drops the frame if the client's /// outbound buffer is full (slow consumer protection). pub fn send(&self, frame: &str) -> bool { - self.sender.try_send(frame.to_string()).is_ok() + if self.sender.try_send(frame.to_string()).is_ok() { + return true; + } + self.delivery_failures.fetch_add(1, Ordering::Relaxed); + false } /// Touch the activity timestamp (lock-free). @@ -1098,18 +1111,31 @@ mod tests { timezone: None, }, sender: tx, + delivery_failures: Arc::new(std::sync::atomic::AtomicU64::new(0)), connected_at: Instant::now(), last_activity_ms: std::sync::atomic::AtomicU64::new(0), accept_language: None, remote_ip: None, timezone: None, subscriptions: None, + session_filter: None, + payload_state_filter: None, joined_channels: HashSet::new(), negotiated_protocol: moltis_protocol::PROTOCOL_VERSION, }; (client, rx) } + #[test] + fn connected_client_records_dropped_frames() { + let (client, _rx) = mock_client("slow-client"); + for _ in 0..512 { + assert!(client.send("frame")); + } + assert!(!client.send("overflow")); + assert_eq!(client.delivery_failures.load(Ordering::Relaxed), 1); + } + #[tokio::test] async fn default_channels_offered_include_matrix() { let state = test_state(); @@ -1317,6 +1343,43 @@ mod tests { assert_eq!(frame["event"], "presence"); } + #[tokio::test] + async fn broadcast_skips_other_sessions_for_filtered_clients() { + let state = test_state(); + let (mut client, mut rx) = mock_client("conn-session"); + client.subscriptions = Some(["chat".to_string()].into()); + client.session_filter = Some("acp:expected".to_string()); + client.payload_state_filter = Some(["delta".to_string()].into()); + state.register_client(client).await; + + crate::broadcast::broadcast( + &state, + "chat", + serde_json::json!({"sessionKey": "acp:other", "state": "delta", "text": "secret"}), + crate::broadcast::BroadcastOpts::default(), + ) + .await; + assert!(rx.try_recv().is_err()); + + crate::broadcast::broadcast( + &state, + "chat", + serde_json::json!({"sessionKey": "acp:expected", "state": "final", "text": "hello"}), + crate::broadcast::BroadcastOpts::default(), + ) + .await; + assert!(rx.try_recv().is_err()); + + crate::broadcast::broadcast( + &state, + "chat", + serde_json::json!({"sessionKey": "acp:expected", "state": "delta", "text": "hello"}), + crate::broadcast::BroadcastOpts::default(), + ) + .await; + assert!(rx.try_recv().is_ok()); + } + #[tokio::test] async fn broadcast_channel_filter_skips_non_members() { let state = test_state(); diff --git a/crates/httpd/src/ws.rs b/crates/httpd/src/ws.rs index 35ae272ef8..17e8040571 100644 --- a/crates/httpd/src/ws.rs +++ b/crates/httpd/src/ws.rs @@ -910,12 +910,15 @@ pub async fn handle_connection( conn_id: conn_id.clone(), connect_params: resolved_params, sender: client_tx.clone(), + delivery_failures: Arc::new(std::sync::atomic::AtomicU64::new(0)), connected_at: now, last_activity_ms: std::sync::atomic::AtomicU64::new(0), accept_language, remote_ip, timezone: browser_timezone, subscriptions, + session_filter: None, + payload_state_filter: None, joined_channels: std::collections::HashSet::new(), negotiated_protocol: PROTOCOL_VERSION, }; diff --git a/crates/mcp/Cargo.toml b/crates/mcp/Cargo.toml index 96e0f00995..25a3073e9b 100644 --- a/crates/mcp/Cargo.toml +++ b/crates/mcp/Cargo.toml @@ -10,7 +10,7 @@ metrics = ["dep:moltis-metrics"] [dependencies] async-trait = { workspace = true } futures = { workspace = true } -moltis-common = { workspace = true } +moltis-common = { features = ["process"], workspace = true } moltis-config = { workspace = true } moltis-metrics = { optional = true, workspace = true } moltis-oauth = { workspace = true } @@ -20,6 +20,7 @@ serde = { workspace = true } serde_json = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true } +tokio-util = { features = ["codec"], workspace = true } tracing = { workspace = true } url = { workspace = true } uuid = { workspace = true } diff --git a/crates/mcp/src/client.rs b/crates/mcp/src/client.rs index 5f8eb19a3e..99b7724fe9 100644 --- a/crates/mcp/src/client.rs +++ b/crates/mcp/src/client.rs @@ -20,7 +20,7 @@ use crate::{ remote::ResolvedRemoteConfig, sse_transport::SseTransport, traits::{McpClientTrait, McpTransport}, - transport::StdioTransport, + transport::{StdioLaunchOptions, StdioTransport}, types::{ ClientCapabilities, ClientInfo, InitializeParams, InitializeResult, McpToolDef, McpTransportError, PROTOCOL_VERSION, ToolsCallParams, ToolsCallResult, ToolsListResult, @@ -58,9 +58,29 @@ impl McpClient { env: &HashMap>, request_timeout: Duration, ) -> Result { - info!(server = %server_name, command = %command, args = ?args, "connecting to MCP server"); + Self::connect_with_options( + server_name, + command, + args, + env, + request_timeout, + &StdioLaunchOptions::default(), + ) + .await + } + + pub async fn connect_with_options( + server_name: &str, + command: &str, + args: &[String], + env: &HashMap>, + request_timeout: Duration, + options: &StdioLaunchOptions, + ) -> Result { + info!(server = %server_name, command = %command, arg_count = args.len(), "connecting to MCP server"); let transport = - StdioTransport::spawn_with_timeout(command, args, env, request_timeout).await?; + StdioTransport::spawn_with_options(command, args, env, request_timeout, options) + .await?; let mut client = Self { server_name: server_name.into(), @@ -71,7 +91,7 @@ impl McpClient { }; if let Err(e) = client.initialize().await { - warn!(server = %server_name, error = %e, "MCP initialize handshake failed"); + warn!(server = %server_name, "MCP initialize handshake failed"); return Err(e); } @@ -108,7 +128,7 @@ impl McpClient { }; if let Err(e) = client.initialize().await { - warn!(server = %server_name, error = %e, "remote MCP initialize handshake failed"); + warn!(server = %server_name, "remote MCP initialize handshake failed"); return Err(e); } Ok(client) @@ -139,7 +159,7 @@ impl McpClient { }; if let Err(e) = client.initialize().await { - warn!(server = %server_name, error = %e, "legacy SSE MCP initialize handshake failed"); + warn!(server = %server_name, "legacy SSE MCP initialize handshake failed"); return Err(e); } @@ -177,7 +197,7 @@ impl McpClient { }; if let Err(e) = client.initialize().await { - warn!(server = %server_name, error = %e, "legacy SSE (auth) MCP initialize handshake failed"); + warn!(server = %server_name, "legacy SSE (auth) MCP initialize handshake failed"); return Err(e); } @@ -214,7 +234,7 @@ impl McpClient { }; if let Err(e) = client.initialize().await { - warn!(server = %server_name, error = %e, "MCP SSE (auth) initialize handshake failed"); + warn!(server = %server_name, "MCP SSE (auth) initialize handshake failed"); return Err(e); } @@ -259,8 +279,6 @@ impl McpClient { info!( server = %self.server_name, - protocol = %result.protocol_version, - server_name = %result.server_info.name, "MCP server initialized" ); diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index dd67b934b9..e4f629441e 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -33,5 +33,6 @@ pub use { registry::{McpOAuthConfig, McpRegistry, McpServerConfig, TransportType}, tool_bridge::{McpAgentTool, McpToolBridge}, traits::{McpClientTrait, McpTransport}, + transport::StdioLaunchOptions, types::{McpManagerError, McpTransportError}, }; diff --git a/crates/mcp/src/manager.rs b/crates/mcp/src/manager.rs index e4b2efc4d5..d95b308f67 100644 --- a/crates/mcp/src/manager.rs +++ b/crates/mcp/src/manager.rs @@ -23,6 +23,7 @@ use crate::{ remote::{ResolvedRemoteConfig, header_names, sanitize_url_for_display}, tool_bridge::McpToolBridge, traits::McpClientTrait, + transport::StdioLaunchOptions, types::{McpManagerError, McpToolDef, McpTransportError}, }; @@ -184,7 +185,7 @@ impl McpManager { for (name, config) in enabled { match self.start_server(&name, &config).await { Ok(()) => started.push(name), - Err(e) => warn!(server = %name, error = %e, "failed to start MCP server"), + Err(_) => warn!(server = %name, "failed to start MCP server"), } } started @@ -195,6 +196,16 @@ impl McpManager { /// For SSE servers: attempts unauthenticated first. On 401 Unauthorized, /// stores auth context and returns `McpManagerError::OAuthRequired`. pub async fn start_server(&self, name: &str, config: &McpServerConfig) -> Result<()> { + self.start_server_with_options(name, config, &StdioLaunchOptions::default()) + .await + } + + pub async fn start_server_with_options( + &self, + name: &str, + config: &McpServerConfig, + stdio_options: &StdioLaunchOptions, + ) -> Result<()> { // Shut down existing connection if any. self.stop_server(name).await; @@ -459,12 +470,13 @@ impl McpManager { } }, TransportType::Stdio => { - let client = McpClient::connect( + let client = McpClient::connect_with_options( name, &config.command, &config.args, &config.env, self.effective_timeout_for(config), + stdio_options, ) .await?; (client, None) diff --git a/crates/mcp/src/tool_bridge.rs b/crates/mcp/src/tool_bridge.rs index 8b99630c36..fbfd492506 100644 --- a/crates/mcp/src/tool_bridge.rs +++ b/crates/mcp/src/tool_bridge.rs @@ -262,6 +262,7 @@ mod tests { "_session_key": "abc123", "_accept_language": "en", "_conn_id": "conn-42", + "_working_dir": "/workspace/project", "encoding": "utf-8" }); @@ -282,6 +283,7 @@ mod tests { assert!(!map.contains_key("_session_key")); assert!(!map.contains_key("_accept_language")); assert!(!map.contains_key("_conn_id")); + assert!(!map.contains_key("_working_dir")); } #[tokio::test] diff --git a/crates/mcp/src/transport.rs b/crates/mcp/src/transport.rs index 5b8fbd1bb7..f0ddb679cf 100644 --- a/crates/mcp/src/transport.rs +++ b/crates/mcp/src/transport.rs @@ -2,41 +2,128 @@ use std::{ collections::HashMap, + path::PathBuf, process::Stdio, sync::{ Arc, - atomic::{AtomicU64, Ordering}, + atomic::{AtomicBool, AtomicU64, Ordering}, }, time::Duration, }; +/// Process isolation options for a stdio MCP server launch. +#[derive(Clone, Debug)] +pub struct StdioLaunchOptions { + pub current_dir: Option, + pub inherit_parent_env: bool, +} + +impl Default for StdioLaunchOptions { + fn default() -> Self { + Self { + current_dir: None, + inherit_parent_env: true, + } + } +} + use { + futures::StreamExt, secrecy::{ExposeSecret, Secret}, tokio::{ - io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, - process::{Child, Command}, + io::AsyncWriteExt, + process::Command, sync::{Mutex, oneshot}, }, + tokio_util::codec::{FramedRead, LinesCodec}, tracing::{debug, info, trace, warn}, }; -use crate::{ - error::{Context, Error, Result}, - traits::McpTransport, - types::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse}, +use { + crate::{ + error::{Context, Error, Result}, + traits::McpTransport, + types::{JsonRpcNotification, JsonRpcRequest, JsonRpcResponse}, + }, + moltis_common::process_tree::OwnedProcessTree, }; +const MAX_MCP_STDOUT_LINE_BYTES: usize = 4 * 1024 * 1024; +const MAX_MCP_STDERR_LINE_BYTES: usize = 64 * 1024; + /// Stdio-based transport for an MCP server process. pub struct StdioTransport { - child: Mutex, - stdin: Mutex, + child: Mutex, + stdin: Arc>, pending: Arc>>>, next_id: AtomicU64, request_timeout: Duration, + reader_closed: Arc, /// Handle to the reader task so we can abort on drop. reader_handle: Mutex>>, } +struct PendingRequestGuard { + id_key: Option, + request_id: serde_json::Value, + pending: Arc>>>, + stdin: Arc>, +} + +impl PendingRequestGuard { + fn disarm(&mut self) { + self.id_key = None; + } + + async fn cancel(&mut self) { + let Some(id_key) = self.id_key.as_ref() else { + return; + }; + self.pending.lock().await.remove(id_key); + self.id_key = None; + tokio::spawn(send_cancellation_notification( + self.request_id.clone(), + Arc::clone(&self.stdin), + )); + } +} + +impl Drop for PendingRequestGuard { + fn drop(&mut self) { + let Some(id_key) = self.id_key.take() else { + return; + }; + let request_id = self.request_id.clone(); + let pending = Arc::clone(&self.pending); + let stdin = Arc::clone(&self.stdin); + tokio::spawn(async move { + pending.lock().await.remove(&id_key); + send_cancellation_notification(request_id, stdin).await; + }); + } +} + +async fn send_cancellation_notification( + request_id: serde_json::Value, + stdin: Arc>, +) { + let notification = JsonRpcNotification { + jsonrpc: "2.0".into(), + method: "notifications/cancelled".into(), + params: Some(serde_json::json!({ + "requestId": request_id, + "reason": "request cancelled by client", + })), + }; + let Ok(mut payload) = serde_json::to_string(¬ification) else { + return; + }; + payload.push('\n'); + let mut writer = stdin.lock().await; + let _ = writer.write_all(payload.as_bytes()).await; + let _ = writer.flush().await; +} + impl StdioTransport { /// Spawn the server process and start the reader loop. pub async fn spawn( @@ -54,58 +141,80 @@ impl StdioTransport { env: &HashMap>, request_timeout: Duration, ) -> Result> { - info!( - command = %command, - args = ?args, - "spawning MCP server process" - ); + Self::spawn_with_options( + command, + args, + env, + request_timeout, + &StdioLaunchOptions::default(), + ) + .await + } + + /// Spawn with explicit process isolation and working-directory behavior. + pub async fn spawn_with_options( + command: &str, + args: &[String], + env: &HashMap>, + request_timeout: Duration, + options: &StdioLaunchOptions, + ) -> Result> { + info!(command = %command, arg_count = args.len(), "spawning MCP server process"); let mut cmd = Command::new(command); cmd.args(args) .stdin(Stdio::piped()) .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .kill_on_drop(true); + .stderr(Stdio::piped()); + if !options.inherit_parent_env { + cmd.env_clear(); + } + if let Some(current_dir) = options.current_dir.as_ref() { + cmd.current_dir(current_dir); + } for (name, value) in env { cmd.env(name, value.expose_secret()); } - let mut child = cmd - .spawn() + let mut child = OwnedProcessTree::spawn(cmd) .with_context(|| format!("failed to spawn MCP server: {command}"))?; - let stdin = child.stdin.take().context("failed to capture stdin")?; - let stdout = child.stdout.take().context("failed to capture stdout")?; - let stderr = child.stderr.take(); + let stdin = child.take_stdin().context("failed to capture stdin")?; + let stdout = child.take_stdout().context("failed to capture stdout")?; + let stderr = child.take_stderr(); let pending: Arc>>> = Arc::new(Mutex::new(HashMap::new())); + let reader_closed = Arc::new(AtomicBool::new(false)); let transport = Arc::new(Self { child: Mutex::new(child), - stdin: Mutex::new(stdin), + stdin: Arc::new(Mutex::new(stdin)), pending: Arc::clone(&pending), next_id: AtomicU64::new(1), request_timeout, + reader_closed: Arc::clone(&reader_closed), reader_handle: Mutex::new(None), }); // Start stderr reader task (log server errors). if let Some(stderr) = stderr { tokio::spawn(async move { - let mut reader = BufReader::new(stderr); - let mut line = String::new(); - loop { - line.clear(); - match reader.read_line(&mut line).await { - Ok(0) => break, - Ok(_) => { + let mut lines = FramedRead::new( + stderr, + LinesCodec::new_with_max_length(MAX_MCP_STDERR_LINE_BYTES), + ); + while let Some(line) = lines.next().await { + match line { + Ok(line) => { let trimmed = line.trim(); if !trimmed.is_empty() { - warn!(stderr = %trimmed, "MCP server stderr"); + debug!(bytes = trimmed.len(), "MCP server wrote to stderr"); } }, - Err(_) => break, + Err(error) => { + warn!(%error, "MCP server stderr line exceeded transport limit"); + }, } } }); @@ -113,23 +222,19 @@ impl StdioTransport { // Start stdout reader task. let pending_clone = Arc::clone(&pending); + let reader_closed_clone = Arc::clone(&reader_closed); let handle = tokio::spawn(async move { - let mut reader = BufReader::new(stdout); - let mut line = String::new(); - loop { - line.clear(); - match reader.read_line(&mut line).await { - Ok(0) => { - debug!("MCP server stdout closed"); - break; - }, - Ok(_) => { + let mut lines = FramedRead::new( + stdout, + LinesCodec::new_with_max_length(MAX_MCP_STDOUT_LINE_BYTES), + ); + while let Some(line) = lines.next().await { + match line { + Ok(line) => { let trimmed = line.trim(); if trimmed.is_empty() { continue; } - debug!(raw = %trimmed, "MCP server -> client"); - // Try to parse as response (has id field). match serde_json::from_str::(trimmed) { Ok(resp) => { @@ -142,16 +247,19 @@ impl StdioTransport { } }, Err(e) => { - debug!(error = %e, line = %trimmed, "MCP server sent non-response line"); + debug!(error = %e, bytes = trimmed.len(), "MCP server sent non-response line"); }, } }, - Err(e) => { - warn!(error = %e, "error reading from MCP server stdout"); + Err(error) => { + warn!(%error, "MCP server stdout line exceeded transport limit"); break; }, } } + reader_closed_clone.store(true, Ordering::Release); + pending_clone.lock().await.clear(); + debug!("MCP server stdout closed"); }); *transport.reader_handle.lock().await = Some(handle); @@ -166,6 +274,9 @@ impl McpTransport for StdioTransport { method: &str, params: Option, ) -> Result { + if self.reader_closed.load(Ordering::Acquire) { + return Err(Error::message("MCP server stdout is closed")); + } let id = self.next_id.fetch_add(1, Ordering::SeqCst); let req = JsonRpcRequest::new(id, method, params); let id_key = req.id.to_string(); @@ -173,31 +284,48 @@ impl McpTransport for StdioTransport { let (tx, rx) = oneshot::channel(); { let mut map = self.pending.lock().await; + if self.reader_closed.load(Ordering::Acquire) { + return Err(Error::message("MCP server stdout is closed")); + } map.insert(id_key.clone(), tx); } + let mut pending_guard = PendingRequestGuard { + id_key: Some(id_key.clone()), + request_id: req.id.clone(), + pending: Arc::clone(&self.pending), + stdin: Arc::clone(&self.stdin), + }; let mut payload = serde_json::to_string(&req)?; payload.push('\n'); debug!(method = %method, id = %id, "client -> MCP server"); - { + let write_result = async { let mut stdin = self.stdin.lock().await; stdin.write_all(payload.as_bytes()).await?; - stdin.flush().await?; + stdin.flush().await + } + .await; + if let Err(error) = write_result { + pending_guard.cancel().await; + return Err(error.into()); } - let resp = tokio::time::timeout(self.request_timeout, rx) - .await - .with_context(|| { - format!( + let received = match tokio::time::timeout(self.request_timeout, rx).await { + Ok(received) => received, + Err(_) => { + pending_guard.cancel().await; + return Err(Error::message(format!( "MCP request '{method}' timed out after {}s (no response from server)", self.request_timeout.as_secs() - ) - })? - .with_context(|| { - format!("MCP reader task dropped while waiting for '{method}' response") - })?; + ))); + }, + }; + let resp = received.with_context(|| { + format!("MCP reader task dropped while waiting for '{method}' response") + })?; + pending_guard.disarm(); if let Some(ref err) = resp.error { return Err(Error::message(format!( @@ -228,11 +356,16 @@ impl McpTransport for StdioTransport { } async fn is_alive(&self) -> bool { + if self.reader_closed.load(Ordering::Acquire) { + return false; + } let mut child = self.child.lock().await; matches!(child.try_wait(), Ok(None)) } async fn kill(&self) { + self.reader_closed.store(true, Ordering::Release); + self.pending.lock().await.clear(); if let Some(handle) = self.reader_handle.lock().await.take() { handle.abort(); } @@ -258,6 +391,34 @@ mod tests { assert!(!transport.is_alive().await); } + #[cfg(unix)] + #[tokio::test] + async fn test_kill_terminates_descendants() { + let temp_dir = tempfile::tempdir().unwrap(); + let started = temp_dir.path().join("started"); + let marker = temp_dir.path().join("should-not-exist"); + let script = format!( + "touch '{}'; (sleep 1; touch '{}') & cat", + started.display(), + marker.display() + ); + let args = vec!["-c".to_string(), script]; + let transport = StdioTransport::spawn("sh", &args, &HashMap::new()) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + while !started.exists() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + transport.kill().await; + tokio::time::sleep(Duration::from_millis(1100)).await; + assert!(!marker.exists(), "MCP descendant survived transport kill"); + } + #[tokio::test] async fn test_spawn_nonexistent_command() { let result = @@ -279,7 +440,126 @@ mod tests { let err = transport.request("tools/list", None).await.unwrap_err(); assert!(err.to_string().contains("timed out after 1s")); + assert!(transport.pending.lock().await.is_empty()); transport.kill().await; } + + #[tokio::test] + async fn cancelled_request_removes_pending_entry() { + let temp_dir = tempfile::tempdir().unwrap(); + let messages = temp_dir.path().join("messages.jsonl"); + let script = format!( + "while IFS= read -r line; do printf '%s\\n' \"$line\" >> '{}'; done", + messages.display() + ); + let args = vec!["-c".to_string(), script]; + let transport = StdioTransport::spawn_with_timeout( + "sh", + &args, + &HashMap::new(), + Duration::from_secs(30), + ) + .await + .unwrap(); + let request_transport = Arc::clone(&transport); + let request = + tokio::spawn(async move { request_transport.request("tools/call", None).await }); + + tokio::time::timeout(Duration::from_secs(1), async { + while transport.pending.lock().await.is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + request.abort(); + let _ = request.await; + tokio::time::timeout(Duration::from_secs(1), async { + while !transport.pending.lock().await.is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(1), async { + loop { + let content = tokio::fs::read_to_string(&messages) + .await + .unwrap_or_default(); + if content.contains("notifications/cancelled") { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + transport.kill().await; + } + + #[tokio::test] + async fn cancellation_during_timeout_cleanup_keeps_guard_armed() { + let args = vec!["-c".to_string(), "while read line; do :; done".to_string()]; + let transport = StdioTransport::spawn_with_timeout( + "sh", + &args, + &HashMap::new(), + Duration::from_millis(50), + ) + .await + .unwrap(); + let request_transport = Arc::clone(&transport); + let request = + tokio::spawn(async move { request_transport.request("tools/call", None).await }); + + tokio::time::timeout(Duration::from_secs(1), async { + loop { + if !transport.pending.lock().await.is_empty() { + break; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let pending_guard = transport.pending.lock().await; + tokio::time::sleep(Duration::from_millis(100)).await; + request.abort(); + let _ = request.await; + drop(pending_guard); + tokio::time::timeout(Duration::from_secs(1), async { + while !transport.pending.lock().await.is_empty() { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + transport.kill().await; + } + + #[tokio::test] + async fn test_stdout_closure_fails_pending_request_without_timeout() { + let args = vec!["-c".to_string(), "read line; exit 0".to_string()]; + let transport = StdioTransport::spawn_with_timeout( + "sh", + &args, + &HashMap::new(), + Duration::from_secs(30), + ) + .await + .unwrap(); + + let result = tokio::time::timeout( + Duration::from_secs(1), + transport.request("tools/list", None), + ) + .await + .expect("stdout closure should wake the request"); + assert!(result.is_err()); + assert!(transport.pending.lock().await.is_empty()); + assert!(!transport.is_alive().await); + } } diff --git a/crates/openclaw-import/Cargo.toml b/crates/openclaw-import/Cargo.toml index 3bcbc96d98..6bd5366f3f 100644 --- a/crates/openclaw-import/Cargo.toml +++ b/crates/openclaw-import/Cargo.toml @@ -11,6 +11,7 @@ json5 = { workspace = true } moltis-common = { workspace = true } moltis-config = { workspace = true } moltis-oauth = { workspace = true } +moltis-sessions = { workspace = true } notify-debouncer-full = { optional = true, workspace = true } secrecy = { workspace = true } serde = { workspace = true } diff --git a/crates/openclaw-import/src/sessions.rs b/crates/openclaw-import/src/sessions.rs index f0fadfadfc..c362a26af0 100644 --- a/crates/openclaw-import/src/sessions.rs +++ b/crates/openclaw-import/src/sessions.rs @@ -151,7 +151,15 @@ pub fn import_sessions( .unwrap_or("unknown"); let dest_key = format!("oc:{moltis_agent_id}:{stem}"); let dest_file = - dest_sessions_dir.join(format!("{}.jsonl", sanitize_session_key(&dest_key))); + moltis_sessions::store::SessionStore::new(dest_sessions_dir.to_path_buf()) + .history_path_for(&dest_key); + if let Some(parent) = dest_file.parent() + && let Err(error) = std::fs::create_dir_all(parent) + { + warn!(%error, key = %dest_key, "failed to create session storage directory"); + errors.push(format!("{dest_key}: {error}")); + continue; + } let source_lines = count_lines(&path); @@ -491,10 +499,6 @@ fn count_lines(path: &Path) -> u32 { BufReader::new(file).lines().count() as u32 } -fn sanitize_session_key(key: &str) -> String { - key.replace(':', "_") -} - /// Load existing session metadata from disk. fn load_session_metadata(path: &Path) -> HashMap { if !path.is_file() { @@ -669,6 +673,13 @@ mod tests { std::fs::write(dir.join(format!("{key}.jsonl")), content).unwrap(); } + fn imported_history_path(destination: &Path, key: &str) -> std::path::PathBuf { + let path = moltis_sessions::store::SessionStore::new(destination.to_path_buf()) + .history_path_for(key); + std::fs::create_dir_all(path.parent().unwrap()).unwrap(); + path + } + #[test] fn convert_basic_session() { let tmp = tempfile::tempdir().unwrap(); @@ -690,7 +701,7 @@ mod tests { assert_eq!(report.items_imported, 1); // Verify converted JSONL - let converted_path = dest.join("oc_main_test-session.jsonl"); + let converted_path = imported_history_path(&dest, "oc:main:test-session"); assert!(converted_path.is_file()); let content = std::fs::read_to_string(&converted_path).unwrap(); @@ -724,7 +735,7 @@ mod tests { assert_eq!(report.status, ImportStatus::Success); assert_eq!(report.items_imported, 1); - assert!(dest.join("oc_main_legacy-session.jsonl").is_file()); + assert!(imported_history_path(&dest, "oc:main:legacy-session").is_file()); } #[test] @@ -744,7 +755,8 @@ mod tests { assert_eq!(report.items_imported, 1); - let content = std::fs::read_to_string(dest.join("oc_main_tools.jsonl")).unwrap(); + let content = + std::fs::read_to_string(imported_history_path(&dest, "oc:main:tools")).unwrap(); let lines: Vec<&str> = content.lines().collect(); assert_eq!(lines.len(), 2); @@ -824,8 +836,8 @@ mod tests { assert_eq!(report.items_imported, 2); // Verify both sessions were imported with correct keys - assert!(dest.join("oc_main_s1.jsonl").is_file()); - assert!(dest.join("oc_research_s2.jsonl").is_file()); + assert!(imported_history_path(&dest, "oc:main:s1").is_file()); + assert!(imported_history_path(&dest, "oc:research:s2").is_file()); // Verify metadata has correct agent_id let metadata: HashMap = @@ -876,7 +888,8 @@ mod tests { let report = import_sessions(&detection, &dest, &mem, &default_mapping()); assert_eq!(report.items_imported, 1); - let content = std::fs::read_to_string(dest.join("oc_main_messy.jsonl")).unwrap(); + let content = + std::fs::read_to_string(imported_history_path(&dest, "oc:main:messy")).unwrap(); assert_eq!(content.lines().count(), 1); } @@ -950,7 +963,8 @@ mod tests { assert_eq!(report1.items_imported, 1); assert_eq!(report1.items_updated, 0); - let content1 = std::fs::read_to_string(dest.join("oc_main_growing.jsonl")).unwrap(); + let content1 = + std::fs::read_to_string(imported_history_path(&dest, "oc:main:growing")).unwrap(); assert_eq!(content1.lines().count(), 1); // Append a new message to the source @@ -966,7 +980,8 @@ mod tests { assert_eq!(report2.items_skipped, 0); // Destination should now have 2 messages - let content2 = std::fs::read_to_string(dest.join("oc_main_growing.jsonl")).unwrap(); + let content2 = + std::fs::read_to_string(imported_history_path(&dest, "oc:main:growing")).unwrap(); assert_eq!(content2.lines().count(), 2); } @@ -1102,7 +1117,7 @@ mod tests { // Also write a destination JSONL so it looks like a previous import happened std::fs::write( - dest.join("oc_main_legacy.jsonl"), + imported_history_path(&dest, "oc:main:legacy"), r#"{"role":"user","content":"old message"}"#, ) .unwrap(); @@ -1141,7 +1156,8 @@ mod tests { assert_eq!(report.status, ImportStatus::Success); assert_eq!(report.items_imported, 1); - let converted = std::fs::read_to_string(dest.join("oc_main_timed.jsonl")).unwrap(); + let converted = + std::fs::read_to_string(imported_history_path(&dest, "oc:main:timed")).unwrap(); let mut lines = converted.lines(); let first: serde_json::Value = serde_json::from_str(lines.next().unwrap()).unwrap(); let second: serde_json::Value = serde_json::from_str(lines.next().unwrap()).unwrap(); diff --git a/crates/providers/src/anthropic.rs b/crates/providers/src/anthropic.rs index b171fbb37a..2a71967ac2 100644 --- a/crates/providers/src/anthropic.rs +++ b/crates/providers/src/anthropic.rs @@ -141,7 +141,10 @@ impl AnthropicProvider { } debug!(model = %self.model, "anthropic probe request"); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "anthropic probe request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "anthropic probe request body prepared" + ); let http_resp = self .client @@ -157,7 +160,7 @@ impl AnthropicProvider { if !status.is_success() { let retry_after_ms = retry_after_ms_from_headers(http_resp.headers()); let body_text = http_resp.text().await.unwrap_or_default(); - warn!(status = %status, body = %body_text, "anthropic probe API error"); + warn!(status = %status, body_bytes = body_text.len(), "anthropic probe API error"); anyhow::bail!( "{}", with_retry_after_marker( @@ -779,7 +782,10 @@ impl LlmProvider for AnthropicProvider { reasoning_effort = ?self.reasoning_effort, "anthropic complete request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "anthropic request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "anthropic request body prepared" + ); let http_resp = self .client @@ -795,7 +801,7 @@ impl LlmProvider for AnthropicProvider { if !status.is_success() { let retry_after_ms = retry_after_ms_from_headers(http_resp.headers()); let body_text = http_resp.text().await.unwrap_or_default(); - warn!(status = %status, body = %body_text, "anthropic API error"); + warn!(status = %status, body_bytes = body_text.len(), "anthropic API error"); anyhow::bail!( "{}", with_retry_after_marker( @@ -806,7 +812,10 @@ impl LlmProvider for AnthropicProvider { } let resp = http_resp.json::().await?; - trace!(response = %resp, "anthropic raw response"); + trace!( + response_bytes = serde_json::to_vec(&resp).map_or(0, |value| value.len()), + "anthropic response received" + ); let content = resp["content"].as_array().cloned().unwrap_or_default(); @@ -926,7 +935,7 @@ impl LlmProvider for AnthropicProvider { reasoning_effort = ?self.reasoning_effort, "anthropic stream_with_tools request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "anthropic stream request body"); + trace!(body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), "anthropic stream request body prepared"); let resp = match self .client diff --git a/crates/providers/src/github_copilot/diagnostics.rs b/crates/providers/src/github_copilot/diagnostics.rs index 4e46b208f2..7a5f49d72d 100644 --- a/crates/providers/src/github_copilot/diagnostics.rs +++ b/crates/providers/src/github_copilot/diagnostics.rs @@ -166,7 +166,7 @@ pub(super) fn log_copilot_chat_error( enterprise, stream, status, - error_body = %body, + error_body_bytes = body.len(), "github-copilot chat request failed" ); } diff --git a/crates/providers/src/github_copilot/provider.rs b/crates/providers/src/github_copilot/provider.rs index d8385e5104..ac3fd5e486 100644 --- a/crates/providers/src/github_copilot/provider.rs +++ b/crates/providers/src/github_copilot/provider.rs @@ -534,7 +534,7 @@ async fn collect_streamed_completion( if !status.is_success() { let retry_after_ms = super::super::retry_after_ms_from_headers(http_resp.headers()); let body_text = http_resp.text().await.unwrap_or_default(); - warn!(status = %status, body = %body_text, "github-copilot enterprise API error"); + warn!(status = %status, body_bytes = body_text.len(), "github-copilot enterprise API error"); anyhow::bail!( "{}", super::super::with_retry_after_marker( @@ -673,7 +673,7 @@ async fn collect_streamed_responses_completion( if !status.is_success() { let retry_after_ms = super::super::retry_after_ms_from_headers(http_resp.headers()); let body_text = http_resp.text().await.unwrap_or_default(); - warn!(status = %status, body = %body_text, "github-copilot enterprise responses API error"); + warn!(status = %status, body_bytes = body_text.len(), "github-copilot enterprise responses API error"); anyhow::bail!( "{}", super::super::with_retry_after_marker( @@ -913,7 +913,7 @@ impl LlmProvider for GitHubCopilotProvider { return self.complete_responses(messages, tools).await; } - warn!(status = %status, body = %body_text, "github-copilot API error"); + warn!(status = %status, body_bytes = body_text.len(), "github-copilot API error"); anyhow::bail!( "{}", super::super::with_retry_after_marker( @@ -924,7 +924,10 @@ impl LlmProvider for GitHubCopilotProvider { } let resp = http_resp.json::().await?; - trace!(response = %resp, "github-copilot raw response"); + trace!( + response_bytes = serde_json::to_vec(&resp).map_or(0, |value| value.len()), + "github-copilot response received" + ); let message = &resp["choices"][0]["message"]; @@ -1041,7 +1044,7 @@ impl GitHubCopilotProvider { ) .await; } - warn!(status = %status, body = %body_text, "github-copilot responses API error"); + warn!(status = %status, body_bytes = body_text.len(), "github-copilot responses API error"); anyhow::bail!( "{}", super::super::with_retry_after_marker( @@ -1052,7 +1055,10 @@ impl GitHubCopilotProvider { } let resp = http_resp.json::().await?; - trace!(response = %resp, "github-copilot responses raw response"); + trace!( + response_bytes = serde_json::to_vec(&resp).map_or(0, |value| value.len()), + "github-copilot responses response received" + ); Ok(parse_responses_completion(&resp)) } diff --git a/crates/providers/src/kimi_code.rs b/crates/providers/src/kimi_code.rs index ec8a54c06e..b0e3cd7d51 100644 --- a/crates/providers/src/kimi_code.rs +++ b/crates/providers/src/kimi_code.rs @@ -238,7 +238,10 @@ impl LlmProvider for KimiCodeProvider { tools_count = tools.len(), "kimi-code complete request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "kimi-code request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "kimi-code request body prepared" + ); let mut request = self .client @@ -254,7 +257,7 @@ impl LlmProvider for KimiCodeProvider { if !status.is_success() { let retry_after_ms = super::retry_after_ms_from_headers(http_resp.headers()); let body_text = http_resp.text().await.unwrap_or_default(); - warn!(status = %status, body = %body_text, "kimi-code API error"); + warn!(status = %status, body_bytes = body_text.len(), "kimi-code API error"); let hint = build_access_denied_hint(status, &body_text); if let Some(hint) = hint { anyhow::bail!( @@ -275,7 +278,10 @@ impl LlmProvider for KimiCodeProvider { } let resp = http_resp.json::().await?; - trace!(response = %resp, "kimi-code raw response"); + trace!( + response_bytes = serde_json::to_vec(&resp).map_or(0, |value| value.len()), + "kimi-code response received" + ); let message = &resp["choices"][0]["message"]; @@ -334,7 +340,7 @@ impl LlmProvider for KimiCodeProvider { tools_count = tools.len(), "kimi-code stream_with_tools request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "kimi-code stream request body"); + trace!(body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), "kimi-code stream request body prepared"); let mut request = self .client diff --git a/crates/providers/src/openai/provider/completion.rs b/crates/providers/src/openai/provider/completion.rs index c39fc3dcb2..c04800b348 100644 --- a/crates/providers/src/openai/provider/completion.rs +++ b/crates/providers/src/openai/provider/completion.rs @@ -70,7 +70,10 @@ impl OpenAiProvider { let url = self.chat_completions_url(); debug!(model = %self.model, url = %url, "openai probe request"); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai probe request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "openai probe request body prepared" + ); let http_resp = self.send_chat_completions_request(&body).await?; @@ -84,7 +87,7 @@ impl OpenAiProvider { model = %self.model, provider = %self.provider_name, url = %url, - body = %body_text, + body_bytes = body_text.len(), "openai probe API error" ); } else { @@ -172,7 +175,10 @@ impl OpenAiProvider { self.apply_reasoning_effort_responses(&mut body); debug!(model = %self.model, "openai responses probe request"); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai responses probe request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "openai responses probe request body prepared" + ); let url = self.responses_sse_url(); let http_resp = self @@ -192,7 +198,7 @@ impl OpenAiProvider { status = %status, model = %self.model, provider = %self.provider_name, - body = %body_text, + body_bytes = body_text.len(), "openai responses probe API error" ); anyhow::bail!( @@ -334,7 +340,10 @@ impl OpenAiProvider { reasoning_effort = ?self.reasoning_effort, "openai complete request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "openai request body prepared" + ); let http_resp = self.send_chat_completions_request(&body).await?; @@ -347,7 +356,7 @@ impl OpenAiProvider { status = %status, model = %self.model, provider = %self.provider_name, - body = %body_text, + body_bytes = body_text.len(), "openai API error" ); } else { @@ -368,7 +377,10 @@ impl OpenAiProvider { } let resp = http_resp.json::().await?; - trace!(response = %resp, "openai raw response"); + trace!( + response_bytes = serde_json::to_vec(&resp).map_or(0, |value| value.len()), + "openai response received" + ); let message = &resp["choices"][0]["message"]; @@ -422,7 +434,10 @@ impl OpenAiProvider { tools_count = tools.len(), "openai complete_responses request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai responses request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "openai responses request body prepared" + ); let url = self.responses_sse_url(); let http_resp = self diff --git a/crates/providers/src/openai/provider/streaming.rs b/crates/providers/src/openai/provider/streaming.rs index 2e0dda8827..3e5fe6f20b 100644 --- a/crates/providers/src/openai/provider/streaming.rs +++ b/crates/providers/src/openai/provider/streaming.rs @@ -53,7 +53,7 @@ impl OpenAiProvider { tools_count = tools.len(), "openai stream_responses_sse request" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai responses stream request body"); + trace!(body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), "openai responses stream request body prepared"); let url = self.responses_sse_url(); let resp = match self @@ -195,7 +195,7 @@ impl OpenAiProvider { reasoning_effort = ?self.reasoning_effort, "openai stream_with_tools request (sse)" ); - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai stream request body (sse)"); + trace!(body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), "openai stream request body prepared"); let resp = match self.send_chat_completions_request(&body).await { Ok(r) => { diff --git a/crates/providers/src/openai/provider/websocket.rs b/crates/providers/src/openai/provider/websocket.rs index 87d361b99a..5f10cbc81c 100644 --- a/crates/providers/src/openai/provider/websocket.rs +++ b/crates/providers/src/openai/provider/websocket.rs @@ -151,7 +151,7 @@ impl OpenAiProvider { reasoning_effort = ?self.reasoning_effort, "openai stream_with_tools request (websocket)" ); - trace!(event = %create_event, "openai websocket create event"); + trace!(event_bytes = create_event.to_string().len(), "openai websocket create event"); if let Err(err) = ws_stream .send(Message::Text(create_event.to_string().into())) @@ -192,7 +192,11 @@ impl OpenAiProvider { let Ok(evt) = serde_json::from_str::(&text) else { continue; }; - trace!(event = %evt, "openai websocket event"); + trace!( + event_type = evt["type"].as_str().unwrap_or(""), + event_bytes = text.len(), + "openai websocket event" + ); match evt["type"].as_str().unwrap_or("") { "response.output_text.delta" => { diff --git a/crates/providers/src/openai_codex.rs b/crates/providers/src/openai_codex.rs index 3f8ecfcc02..2f35529a57 100644 --- a/crates/providers/src/openai_codex.rs +++ b/crates/providers/src/openai_codex.rs @@ -494,7 +494,10 @@ impl LlmProvider for OpenAiCodexProvider { } crate::openai::provider::core::apply_openai_responses_tool_choice(&mut body, options)?; - trace!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai-codex request body"); + trace!( + body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), + "openai-codex request body prepared" + ); let http_resp = self .post_responses_request_with_fallback(&token, &account_id, body) @@ -703,7 +706,7 @@ impl LlmProvider for OpenAiCodexProvider { tools_count = tools.len(), "openai-codex stream_with_tools request" ); - debug!(body = %serde_json::to_string(&body).unwrap_or_default(), "openai-codex stream request body"); + debug!(body_bytes = serde_json::to_vec(&body).map_or(0, |value| value.len()), "openai-codex stream request body prepared"); let resp = match self .post_responses_request_with_fallback(&token, &account_id, body) @@ -766,7 +769,7 @@ impl LlmProvider for OpenAiCodexProvider { if let Ok(evt) = serde_json::from_str::(data) { let evt_type = evt["type"].as_str().unwrap_or(""); - trace!(evt_type = %evt_type, evt = %evt, "openai-codex stream event"); + trace!(evt_type = %evt_type, event_bytes = data.len(), "openai-codex stream event"); match evt_type { "response.output_text.delta" => { diff --git a/crates/sessions/Cargo.toml b/crates/sessions/Cargo.toml index 26f3dc41f3..3eb3da77b8 100644 --- a/crates/sessions/Cargo.toml +++ b/crates/sessions/Cargo.toml @@ -9,6 +9,8 @@ metrics = ["dep:moltis-metrics"] [dependencies] fd-lock = { workspace = true } +hex = { workspace = true } +libc = { workspace = true } moltis-common = { workspace = true } moltis-metrics = { optional = true, workspace = true } serde = { workspace = true } diff --git a/crates/sessions/src/lib.rs b/crates/sessions/src/lib.rs index 3173d4160f..73f7633b65 100644 --- a/crates/sessions/src/lib.rs +++ b/crates/sessions/src/lib.rs @@ -1,8 +1,8 @@ //! Session storage and management. //! -//! Sessions are stored as JSONL files (one message per line) at -//! `/agents//sessions/.jsonl` -//! with file locking for concurrent access. +//! Sessions are stored as JSONL files (one message per line) under +//! `/sessions/v1/`, using a reversible filesystem-safe encoding of +//! each session key. Legacy files are copied lazily when their key is accessed. pub mod compaction; pub mod error; @@ -11,6 +11,7 @@ pub mod message; pub mod metadata; pub mod session_events; pub mod state_store; +mod storage_layout; pub mod store; pub use { diff --git a/crates/sessions/src/storage_layout.rs b/crates/sessions/src/storage_layout.rs new file mode 100644 index 0000000000..0831f8b69e --- /dev/null +++ b/crates/sessions/src/storage_layout.rs @@ -0,0 +1,682 @@ +use std::{ + fs::{self, File, OpenOptions}, + io::{Read, Seek, SeekFrom, Write}, + path::{Component, Path, PathBuf}, +}; + +use { + crate::{Error, Result}, + fd_lock::RwLock, + uuid::Uuid, +}; + +pub(crate) const STORAGE_VERSION: &str = "v1"; +const MIGRATED_DIR: &str = ".migrated"; +const MIGRATION_LOCK: &str = ".migration.lock"; +const ENCODED_COMPONENT_BYTES: usize = 120; +const LOWER_HEX: &[u8; 16] = b"0123456789abcdef"; + +#[must_use] +pub(crate) fn encode_key(key: &str) -> String { + let mut encoded = String::with_capacity(key.len().saturating_mul(2).saturating_add(1)); + encoded.push('k'); + for byte in key.bytes() { + encoded.push(char::from(LOWER_HEX[usize::from(byte >> 4)])); + encoded.push(char::from(LOWER_HEX[usize::from(byte & 0x0f)])); + } + encoded +} + +pub(crate) fn decode_key_path(path: &Path) -> Option { + let mut components = path.components(); + let first = match components.next()? { + Component::Normal(component) => component.to_str()?.strip_prefix('k')?.to_string(), + _ => return None, + }; + let mut encoded = first; + for component in components { + let Component::Normal(component) = component else { + return None; + }; + encoded.push_str(component.to_str()?); + } + String::from_utf8(hex::decode(encoded).ok()?).ok() +} + +#[must_use] +pub(crate) fn history_path(base_dir: &Path, key: &str) -> PathBuf { + base_dir + .join(STORAGE_VERSION) + .join(encoded_key_path(key)) + .join("history.jsonl") +} + +#[must_use] +pub(crate) fn media_dir(base_dir: &Path, key: &str) -> PathBuf { + base_dir + .join("media") + .join(STORAGE_VERSION) + .join(encoded_key_path(key)) + .join("files") +} + +pub(crate) fn media_reference(key: &str, filename: &str) -> Result { + validate_media_filename(filename)?; + let encoded_path = encoded_key_components(key).join("/"); + Ok(format!( + "media/{STORAGE_VERSION}/{encoded_path}/files/{filename}" + )) +} + +#[must_use] +pub(crate) fn resolved_media_path(base_dir: &Path, key: &str, filename: &str) -> PathBuf { + let current = media_dir(base_dir, key).join(filename); + if current.exists() || migration_marker(base_dir, key).exists() { + return current; + } + + legacy_media_dir(base_dir, key) + .map(|dir| dir.join(filename)) + .filter(|path| is_regular_file(path)) + .unwrap_or(current) +} + +pub(crate) fn validate_media_filename(filename: &str) -> Result<()> { + if filename.is_empty() || filename.contains('\\') { + return Err(Error::message("invalid session media filename")); + } + let mut components = Path::new(filename).components(); + match (components.next(), components.next()) { + (Some(Component::Normal(_)), None) => Ok(()), + _ => Err(Error::message("invalid session media filename")), + } +} + +pub(crate) fn ensure_migrated(base_dir: &Path, key: &str) -> Result<()> { + let Some((legacy_history, legacy_media)) = legacy_paths(base_dir, key) else { + return Ok(()); + }; + if !legacy_history.exists() && !legacy_media.exists() { + return Ok(()); + } + + with_migration_lock(base_dir, || { + let marker = migration_marker(base_dir, key); + let state = read_migration_state(&marker)?; + if state == Some(MigrationState::Suppressed) { + return Ok(()); + } + + let current_history = history_path(base_dir, key); + let legacy_bytes = if is_regular_file(&legacy_history) { + fs::metadata(&legacy_history)?.len() + } else { + 0 + }; + let copied_bytes = match state { + None if legacy_bytes > 0 && !current_history.exists() => { + if let Some(parent) = current_history.parent() { + create_dir_all_synced(parent)?; + } + copy_file_prefix_atomic(&legacy_history, ¤t_history, legacy_bytes)?; + write_migration_state(&marker, MigrationState::Copied(legacy_bytes))?; + legacy_bytes + }, + None if legacy_bytes > 0 => { + if !file_ranges_match(&legacy_history, 0, ¤t_history, 0, legacy_bytes)? { + return Err(Error::message( + "legacy and current session histories conflict", + )); + } + write_migration_state(&marker, MigrationState::Copied(legacy_bytes))?; + legacy_bytes + }, + None => 0, + Some(MigrationState::Copied(previous_bytes)) => { + if legacy_bytes < previous_bytes { + return Err(Error::message( + "legacy session history changed after migration", + )); + } + if legacy_bytes > previous_bytes { + reconcile_file_range( + &marker, + &legacy_history, + ¤t_history, + previous_bytes, + legacy_bytes, + )?; + } + legacy_bytes + }, + Some(MigrationState::Suppressed) => return Ok(()), + }; + + if is_directory(&legacy_media) { + copy_media_dir(&legacy_media, &media_dir(base_dir, key))?; + } + + if state.is_none() && copied_bytes == 0 { + write_migration_state(&marker, MigrationState::Copied(0))?; + } + Ok(()) + }) +} + +pub(crate) fn suppress_legacy_fallback(base_dir: &Path, key: &str) -> Result<()> { + with_migration_lock(base_dir, || { + write_migration_state(&migration_marker(base_dir, key), MigrationState::Suppressed) + }) +} + +pub(crate) fn legacy_history_path(base_dir: &Path, key: &str) -> Option { + legacy_component(key).map(|component| base_dir.join(format!("{component}.jsonl"))) +} + +pub(crate) fn legacy_media_dir(base_dir: &Path, key: &str) -> Option { + legacy_component(key).map(|component| base_dir.join("media").join(component)) +} + +#[must_use] +pub(crate) fn migration_marker(base_dir: &Path, key: &str) -> PathBuf { + base_dir + .join(STORAGE_VERSION) + .join(MIGRATED_DIR) + .join(encoded_key_path(key)) + .join("state") +} + +fn legacy_paths(base_dir: &Path, key: &str) -> Option<(PathBuf, PathBuf)> { + Some(( + legacy_history_path(base_dir, key)?, + legacy_media_dir(base_dir, key)?, + )) +} + +fn legacy_component(key: &str) -> Option { + let component = key.replace(':', "_"); + let mut components = Path::new(&component).components(); + match (components.next(), components.next()) { + (Some(Component::Normal(_)), None) => Some(component), + _ => None, + } +} + +fn encoded_key_path(key: &str) -> PathBuf { + encoded_key_components(key).into_iter().collect() +} + +fn encoded_key_components(key: &str) -> Vec { + let encoded = hex::encode(key.as_bytes()); + if encoded.is_empty() { + return vec!["k".to_string()]; + } + encoded + .as_bytes() + .chunks(ENCODED_COMPONENT_BYTES) + .enumerate() + .map(|(index, chunk)| { + let chunk = String::from_utf8_lossy(chunk); + if index == 0 { + format!("k{chunk}") + } else { + chunk.into_owned() + } + }) + .collect() +} + +fn with_migration_lock(base_dir: &Path, operation: impl FnOnce() -> Result) -> Result { + let version_dir = base_dir.join(STORAGE_VERSION); + create_dir_all_synced(&version_dir)?; + let lock_file = OpenOptions::new() + .create(true) + .read(true) + .truncate(false) + .write(true) + .open(version_dir.join(MIGRATION_LOCK))?; + let mut lock = RwLock::new(lock_file); + let _guard = lock + .write() + .map_err(|error| Error::lock_failed(error.to_string()))?; + operation() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum MigrationState { + Copied(u64), + Suppressed, +} + +fn read_migration_state(marker: &Path) -> Result> { + let entries = match fs::read_dir(marker) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error.into()), + }; + let mut copied = None; + for entry in entries { + let entry = entry?; + let Some(name) = entry.file_name().to_str().map(str::to_string) else { + continue; + }; + if name == "suppressed" { + return Ok(Some(MigrationState::Suppressed)); + } + if let Some(bytes) = name + .strip_prefix("copied-") + .and_then(|value| value.parse::().ok()) + { + copied = Some(copied.map_or(bytes, |current: u64| current.max(bytes))); + } + } + Ok(copied.map(MigrationState::Copied)) +} + +fn write_migration_state(marker: &Path, state: MigrationState) -> Result<()> { + create_dir_all_synced(marker)?; + let filename = match state { + MigrationState::Copied(bytes) => format!("copied-{bytes}"), + MigrationState::Suppressed => "suppressed".to_string(), + }; + let destination = marker.join(filename); + if destination.exists() { + return Ok(()); + } + let temporary = temporary_path(&destination); + let file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + file.sync_all()?; + fs::rename(&temporary, &destination)?; + sync_parent(&destination)?; + Ok(()) +} + +#[derive(Clone, Copy, Debug)] +struct PendingRange { + source_start: u64, + source_end: u64, + destination_start: u64, +} + +fn pending_range(marker: &Path, source_start: u64) -> Result> { + let entries = match fs::read_dir(marker) { + Ok(entries) => entries, + Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None), + Err(error) => return Err(error.into()), + }; + for entry in entries { + let entry = entry?; + let Some(name) = entry.file_name().to_str().map(str::to_string) else { + continue; + }; + let values = name + .strip_prefix("pending-") + .map(|value| value.split('-').collect::>()); + let Some(values) = values.filter(|values| values.len() == 3) else { + continue; + }; + let Some(start) = values[0].parse::().ok() else { + continue; + }; + if start != source_start { + continue; + } + let Some(source_end) = values[1].parse::().ok() else { + continue; + }; + let Some(destination_start) = values[2].parse::().ok() else { + continue; + }; + return Ok(Some(PendingRange { + source_start: start, + source_end, + destination_start, + })); + } + Ok(None) +} + +fn write_pending_range(marker: &Path, pending: PendingRange) -> Result<()> { + create_dir_all_synced(marker)?; + let destination = marker.join(format!( + "pending-{}-{}-{}", + pending.source_start, pending.source_end, pending.destination_start + )); + if destination.exists() { + return Ok(()); + } + let temporary = temporary_path(&destination); + let file = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + file.sync_all()?; + fs::rename(&temporary, &destination)?; + sync_parent(&destination) +} + +fn copy_media_dir(source: &Path, destination: &Path) -> Result<()> { + create_dir_all_synced(destination)?; + for entry in fs::read_dir(source)? { + let entry = entry?; + let file_type = entry.file_type()?; + if !file_type.is_file() { + continue; + } + let target = destination.join(entry.file_name()); + if !target.exists() { + copy_file_atomic(&entry.path(), &target)?; + } + } + Ok(()) +} + +fn is_regular_file(path: &Path) -> bool { + fs::symlink_metadata(path).is_ok_and(|metadata| metadata.file_type().is_file()) +} + +fn is_directory(path: &Path) -> bool { + fs::symlink_metadata(path).is_ok_and(|metadata| metadata.file_type().is_dir()) +} + +fn copy_file_atomic(source: &Path, destination: &Path) -> Result<()> { + if let Some(parent) = destination.parent() { + create_dir_all_synced(parent)?; + } + let temporary = temporary_path(destination); + fs::copy(source, &temporary)?; + File::open(&temporary)?.sync_all()?; + if destination.exists() { + fs::remove_file(&temporary)?; + return Ok(()); + } + fs::rename(&temporary, destination)?; + sync_parent(destination) +} + +fn copy_file_prefix_atomic(source: &Path, destination: &Path, bytes: u64) -> Result<()> { + if let Some(parent) = destination.parent() { + create_dir_all_synced(parent)?; + } + let temporary = temporary_path(destination); + let source = File::open(source)?; + let mut source = source.take(bytes); + let mut target = OpenOptions::new() + .create_new(true) + .write(true) + .open(&temporary)?; + let copied = std::io::copy(&mut source, &mut target)?; + if copied != bytes { + return Err(Error::message( + "legacy session history changed during migration", + )); + } + target.sync_all()?; + if destination.exists() { + fs::remove_file(&temporary)?; + return Ok(()); + } + fs::rename(&temporary, destination)?; + sync_parent(destination) +} + +fn reconcile_file_range( + marker: &Path, + source: &Path, + destination: &Path, + source_start: u64, + source_end: u64, +) -> Result<()> { + if let Some(parent) = destination.parent() { + create_dir_all_synced(parent)?; + } + let file = OpenOptions::new() + .create(true) + .read(true) + .truncate(false) + .write(true) + .open(destination)?; + let mut lock = RwLock::new(file); + let mut target = lock + .write() + .map_err(|error| Error::lock_failed(error.to_string()))?; + let pending = if let Some(pending) = pending_range(marker, source_start)? { + pending + } else { + let pending = PendingRange { + source_start, + source_end, + destination_start: target.metadata()?.len(), + }; + write_pending_range(marker, pending)?; + pending + }; + if pending.source_end > source_end { + return Err(Error::message( + "legacy session history changed during migration", + )); + } + + let bytes = pending.source_end - source_start; + let expected_end = pending.destination_start.saturating_add(bytes); + let destination_bytes = target.metadata()?.len(); + if destination_bytes >= expected_end + && file_ranges_match( + source, + source_start, + destination, + pending.destination_start, + bytes, + )? + { + write_migration_state(marker, MigrationState::Copied(pending.source_end))?; + drop(target); + if pending.source_end < source_end { + return reconcile_file_range( + marker, + source, + destination, + pending.source_end, + source_end, + ); + } + return Ok(()); + } + if destination_bytes > expected_end { + return Err(Error::message( + "current session history changed during migration", + )); + } + target.set_len(pending.destination_start)?; + target.seek(SeekFrom::Start(pending.destination_start))?; + let mut source_file = File::open(source)?; + source_file.seek(SeekFrom::Start(source_start))?; + let mut source_range = source_file.take(bytes); + let copied = std::io::copy(&mut source_range, &mut *target)?; + if copied != bytes { + return Err(Error::message( + "legacy session history changed during migration", + )); + } + target.flush()?; + target.sync_all()?; + write_migration_state(marker, MigrationState::Copied(pending.source_end))?; + drop(target); + if pending.source_end < source_end { + reconcile_file_range(marker, source, destination, pending.source_end, source_end)?; + } + Ok(()) +} + +fn file_ranges_match( + left: &Path, + left_offset: u64, + right: &Path, + right_offset: u64, + bytes: u64, +) -> Result { + let mut left = File::open(left)?; + let mut right = File::open(right)?; + left.seek(SeekFrom::Start(left_offset))?; + right.seek(SeekFrom::Start(right_offset))?; + let mut left = left.take(bytes); + let mut right = right.take(bytes); + let mut left_buffer = [0_u8; 8192]; + let mut right_buffer = [0_u8; 8192]; + let mut compared = 0_u64; + while compared < bytes { + let remaining = usize::try_from((bytes - compared).min(left_buffer.len() as u64)) + .unwrap_or(left_buffer.len()); + let left_read = left.read(&mut left_buffer[..remaining])?; + let right_read = right.read(&mut right_buffer[..remaining])?; + if left_read != right_read + || left_buffer[..left_read] != right_buffer[..right_read] + || left_read == 0 + { + return Ok(false); + } + compared += left_read as u64; + } + Ok(true) +} + +fn create_dir_all_synced(path: &Path) -> Result<()> { + if path.is_dir() { + return Ok(()); + } + if let Some(parent) = path.parent() + && !parent.as_os_str().is_empty() + { + create_dir_all_synced(parent)?; + } + match fs::create_dir(path) { + Ok(()) => sync_parent(path), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists && path.is_dir() => Ok(()), + Err(error) => Err(error.into()), + } +} + +fn temporary_path(destination: &Path) -> PathBuf { + let filename = destination + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("session"); + destination.with_file_name(format!(".{filename}.{}.tmp", Uuid::new_v4())) +} + +#[cfg(unix)] +fn sync_parent(path: &Path) -> Result<()> { + if let Some(parent) = path.parent() { + File::open(parent)?.sync_all()?; + } + Ok(()) +} + +#[cfg(not(unix))] +fn sync_parent(_path: &Path) -> Result<()> { + Ok(()) +} + +#[cfg(test)] +#[allow(clippy::unwrap_used)] +mod tests { + use super::*; + + #[test] + fn codec_is_reversible_and_injective() { + let keys = [ + "", + "main", + "a:b", + "a_b", + "A", + "a", + "../escape", + "slash/key", + "backslash\\key", + "é", + "e\u{301}", + ]; + let encoded = keys.map(encode_key); + + for (key, component) in keys.iter().zip(&encoded) { + assert_eq!(decode_key_path(Path::new(component)).as_deref(), Some(*key)); + assert!( + component + .chars() + .all(|character| character.is_ascii_lowercase() || character.is_ascii_digit()) + ); + } + for (index, component) in encoded.iter().enumerate() { + assert!(!encoded[..index].contains(component)); + } + } + + #[test] + fn long_keys_are_chunked_and_reversible() { + let key = "long-key".repeat(100); + let path = encoded_key_path(&key); + assert!(path.components().count() > 1); + assert!( + path.components() + .all(|component| component.as_os_str().len() <= ENCODED_COMPONENT_BYTES + 1) + ); + assert_eq!(decode_key_path(&path).as_deref(), Some(key.as_str())); + } + + #[test] + fn unsafe_legacy_keys_cannot_escape_the_store() { + let base = Path::new("/sessions"); + assert!(legacy_history_path(base, "../escape").is_none()); + assert!(legacy_history_path(base, "/absolute").is_none()); + assert!(legacy_history_path(base, "nested/key").is_none()); + } + + #[test] + fn pending_legacy_suffix_is_not_duplicated_on_recovery() { + let directory = tempfile::tempdir().unwrap(); + let base = directory.path(); + let key = "legacy:key"; + let legacy = legacy_history_path(base, key).unwrap(); + let first = b"{\"content\":\"first\"}\n"; + let second = b"{\"content\":\"second\"}\n"; + let third = b"{\"content\":\"third\"}\n"; + fs::write(&legacy, first).unwrap(); + ensure_migrated(base, key).unwrap(); + + OpenOptions::new() + .append(true) + .open(&legacy) + .unwrap() + .write_all(second) + .unwrap(); + let current = history_path(base, key); + OpenOptions::new() + .append(true) + .open(¤t) + .unwrap() + .write_all(second) + .unwrap(); + write_pending_range(&migration_marker(base, key), PendingRange { + source_start: first.len() as u64, + source_end: (first.len() + second.len()) as u64, + destination_start: first.len() as u64, + }) + .unwrap(); + OpenOptions::new() + .append(true) + .open(&legacy) + .unwrap() + .write_all(third) + .unwrap(); + + ensure_migrated(base, key).unwrap(); + assert_eq!( + fs::read(current).unwrap(), + [first.as_slice(), second.as_slice(), third.as_slice()].concat() + ); + } +} diff --git a/crates/sessions/src/store.rs b/crates/sessions/src/store.rs index b36ef47f76..35e1696c04 100644 --- a/crates/sessions/src/store.rs +++ b/crates/sessions/src/store.rs @@ -1,11 +1,11 @@ use std::{ fs::{self, File, OpenOptions}, - io::{BufRead, BufReader, Write}, + io::{BufRead, BufReader, Read, Write}, path::PathBuf, }; use { - crate::{Error, Result}, + crate::{Error, Result, storage_layout}, fd_lock::RwLock, serde::{Deserialize, Serialize}, }; @@ -39,56 +39,78 @@ impl SessionStore { Self { base_dir } } - /// Sanitize a session key for use as a filename. + /// Encode a session key as an injective, filesystem-safe component. pub fn key_to_filename(key: &str) -> String { - key.replace(':', "_") + storage_layout::encode_key(key) + } + + /// Path where new history for a session key is stored. + pub fn history_path_for(&self, key: &str) -> PathBuf { + storage_layout::history_path(&self.base_dir, key) } fn path_for(&self, key: &str) -> PathBuf { - self.base_dir - .join(format!("{}.jsonl", Self::key_to_filename(key))) + self.history_path_for(key) } /// Directory for session media files (screenshots, audio, etc.). fn media_dir_for(&self, key: &str) -> PathBuf { - self.base_dir.join("media").join(Self::key_to_filename(key)) + storage_layout::media_dir(&self.base_dir, key) } /// Absolute path for a session media file. - pub fn media_path_for(&self, key: &str, filename: &str) -> PathBuf { - self.media_dir_for(key).join(filename) + pub fn media_path_for(&self, key: &str, filename: &str) -> Result { + storage_layout::validate_media_filename(filename)?; + Ok(storage_layout::resolved_media_path( + &self.base_dir, + key, + filename, + )) + } + + /// Relative reference for a session media file. + pub fn media_reference(key: &str, filename: &str) -> Result { + storage_layout::media_reference(key, filename) + } + + async fn ensure_migrated(&self, key: &str) -> Result<()> { + let base_dir = self.base_dir.clone(); + let key = key.to_string(); + tokio::task::spawn_blocking(move || storage_layout::ensure_migrated(&base_dir, &key)) + .await? } /// Save a media file for a session. Returns the relative path from base_dir. pub async fn save_media(&self, key: &str, filename: &str, data: &[u8]) -> Result { + storage_layout::validate_media_filename(filename)?; + self.ensure_migrated(key).await?; let dir = self.media_dir_for(key); - let file_path = self.media_path_for(key, filename); + let file_path = dir.join(filename); let data = data.to_vec(); tokio::task::spawn_blocking(move || -> Result<()> { fs::create_dir_all(&dir)?; - fs::write(&file_path, &data)?; + write_media_file(&file_path, &data)?; Ok(()) }) .await??; - let sanitized = Self::key_to_filename(key); - Ok(format!("media/{sanitized}/{filename}")) + Self::media_reference(key, filename) } /// Read a media file. Returns raw bytes. pub async fn read_media(&self, key: &str, filename: &str) -> Result> { - let file_path = self.media_path_for(key, filename); + storage_layout::validate_media_filename(filename)?; + self.ensure_migrated(key).await?; + let file_path = self.media_dir_for(key).join(filename); - tokio::task::spawn_blocking(move || -> Result> { - let data = fs::read(&file_path)?; - Ok(data) - }) - .await? + tokio::task::spawn_blocking(move || -> Result> { read_media_file(&file_path) }) + .await? } /// Append a message (JSON value) as a single line to the session file. pub async fn append(&self, key: &str, message: &serde_json::Value) -> Result<()> { + self.ensure_migrated(key).await?; let path = self.path_for(key); let line = serde_json::to_string(message)?; @@ -111,6 +133,7 @@ impl SessionStore { /// Read all messages from a session file. pub async fn read(&self, key: &str) -> Result> { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result> { @@ -138,6 +161,66 @@ impl SessionStore { .await? } + /// Read a session without allocating beyond the caller's input bounds. + /// Limits apply to JSONL file bytes and parsed messages, including entries + /// that a higher-level consumer may later filter out. + pub async fn read_bounded( + &self, + key: &str, + max_messages: usize, + max_bytes: usize, + ) -> Result> { + self.ensure_migrated(key).await?; + let path = self.path_for(key); + + tokio::task::spawn_blocking(move || -> Result> { + if !path.exists() { + return Ok(vec![]); + } + let file = File::open(&path)?; + let mut reader = BufReader::new(file); + let mut messages = Vec::new(); + let mut line = String::new(); + let mut bytes_read = 0usize; + loop { + line.clear(); + let remaining = max_bytes.saturating_sub(bytes_read); + let read = reader + .by_ref() + .take(remaining.saturating_add(1) as u64) + .read_line(&mut line)?; + if read == 0 { + break; + } + bytes_read = bytes_read.saturating_add(read); + if bytes_read > max_bytes { + return Err(Error::message(format!( + "session history exceeds {max_bytes} bytes" + ))); + } + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + match serde_json::from_str(trimmed) { + Ok(value) => { + if messages.len() >= max_messages { + return Err(Error::message(format!( + "session history exceeds {max_messages} messages" + ))); + } + messages.push(value); + }, + Err(error) => { + tracing::warn!(%error, "skipping malformed JSONL line"); + }, + } + } + Ok(messages) + }) + .await? + } + /// Read all messages from a session that match a given `run_id`. pub async fn read_by_run_id(&self, key: &str, run_id: &str) -> Result> { let all = self.read(key).await?; @@ -150,6 +233,7 @@ impl SessionStore { /// Read the last N messages from a session file. pub async fn read_last_n(&self, key: &str, n: usize) -> Result> { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result> { @@ -177,15 +261,19 @@ impl SessionStore { /// Delete the session file and its media directory. pub async fn clear(&self, key: &str) -> Result<()> { + self.ensure_migrated(key).await?; let path = self.path_for(key); let media_dir = self.media_dir_for(key); + let base_dir = self.base_dir.clone(); + let key = key.to_string(); tokio::task::spawn_blocking(move || -> Result<()> { + storage_layout::suppress_legacy_fallback(&base_dir, &key)?; if path.exists() { fs::remove_file(&path)?; } if media_dir.exists() { - let _ = fs::remove_dir_all(&media_dir); + fs::remove_dir_all(&media_dir)?; } Ok(()) }) @@ -196,40 +284,40 @@ impl SessionStore { /// List all session keys by scanning JSONL files in the base directory. pub fn list_keys(&self) -> Vec { - let Ok(entries) = fs::read_dir(&self.base_dir) else { - return vec![]; - }; - entries - .filter_map(|e| e.ok()) - .filter_map(|e| { - let name = e.file_name().to_string_lossy().to_string(); - name.strip_suffix(".jsonl").map(|s| s.replace('_', ":")) - }) + stored_sessions(&self.base_dir) + .into_iter() + .map(|(key, _)| key) .collect() } /// Search all sessions for messages containing `query` (case-insensitive). /// Returns up to `max_results` hits, at most one per session. pub async fn search(&self, query: &str, max_results: usize) -> Result> { + self.search_known_keys(query, max_results, &[]).await + } + + /// Search sessions after migrating the supplied authoritative keys. + pub async fn search_known_keys( + &self, + query: &str, + max_results: usize, + known_keys: &[String], + ) -> Result> { let base = self.base_dir.clone(); let query = query.to_lowercase(); + let known_keys = known_keys.to_vec(); tokio::task::spawn_blocking(move || { + for key in known_keys { + if let Err(error) = storage_layout::ensure_migrated(&base, &key) { + tracing::warn!(session_key = %key, %error, "failed to migrate session for search"); + } + } let mut results = Vec::new(); - let entries = fs::read_dir(&base)?; - - for entry in entries.flatten() { + for (session_key, path) in stored_sessions(&base) { if results.len() >= max_results { break; } - let path = entry.path(); - let Some(name) = path.file_name().and_then(|n| n.to_str()) else { - continue; - }; - let Some(key_raw) = name.strip_suffix(".jsonl") else { - continue; - }; - let session_key = key_raw.replace('_', ":"); let Ok(file) = File::open(&path) else { continue; @@ -280,6 +368,7 @@ impl SessionStore { /// Replace the entire session history with the given messages. pub async fn replace_history(&self, key: &str, messages: Vec) -> Result<()> { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result<()> { @@ -311,6 +400,7 @@ impl SessionStore { /// Lines that fail to deserialize into `PersistedMessage` are skipped /// (with a warning), matching the behavior of [`read`]. pub async fn read_typed(&self, key: &str) -> Result> { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result> { @@ -344,6 +434,7 @@ impl SessionStore { key: &str, n: usize, ) -> Result> { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result> { @@ -375,6 +466,7 @@ impl SessionStore { key: &str, messages: &[crate::message::PersistedMessage], ) -> Result<()> { + self.ensure_migrated(key).await?; let path = self.path_for(key); let values: Vec = messages.iter().map(|m| m.to_value()).collect(); @@ -413,6 +505,7 @@ impl SessionStore { /// Count messages in a session file without parsing them. pub async fn count(&self, key: &str) -> Result { + self.ensure_migrated(key).await?; let path = self.path_for(key); tokio::task::spawn_blocking(move || -> Result { @@ -432,6 +525,77 @@ impl SessionStore { } } +fn stored_sessions(base_dir: &std::path::Path) -> Vec<(String, PathBuf)> { + let mut sessions = Vec::new(); + let version_dir = base_dir.join(storage_layout::STORAGE_VERSION); + collect_stored_sessions(&version_dir, &version_dir, &mut sessions); + + sessions.sort_by(|left, right| left.0.cmp(&right.0)); + sessions +} + +fn collect_stored_sessions( + version_dir: &std::path::Path, + directory: &std::path::Path, + sessions: &mut Vec<(String, PathBuf)>, +) { + let Ok(entries) = fs::read_dir(directory) else { + return; + }; + for entry in entries.flatten() { + let Ok(file_type) = entry.file_type() else { + continue; + }; + let path = entry.path(); + if file_type.is_dir() { + if entry.file_name() != ".migrated" { + collect_stored_sessions(version_dir, &path, sessions); + } + continue; + } + if !file_type.is_file() || entry.file_name() != "history.jsonl" { + continue; + } + let Some(key_dir) = path.parent() else { + continue; + }; + let Ok(encoded_path) = key_dir.strip_prefix(version_dir) else { + continue; + }; + let Some(key) = storage_layout::decode_key_path(encoded_path) else { + continue; + }; + sessions.push((key, path)); + } +} + +fn write_media_file(path: &std::path::Path, data: &[u8]) -> Result<()> { + let mut options = OpenOptions::new(); + options.create(true).write(true).truncate(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NOFOLLOW); + } + let mut file = options.open(path)?; + file.write_all(data)?; + Ok(()) +} + +fn read_media_file(path: &std::path::Path) -> Result> { + let mut options = OpenOptions::new(); + options.read(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + options.custom_flags(libc::O_NOFOLLOW); + } + let mut file = options.open(path)?; + let mut data = Vec::new(); + file.read_to_end(&mut data)?; + Ok(data) +} + #[allow(clippy::unwrap_used, clippy::expect_used)] #[cfg(test)] mod tests { @@ -668,7 +832,7 @@ mod tests { #[tokio::test] async fn test_key_sanitization() { - let (store, _dir) = temp_store(); + let (store, dir) = temp_store(); store .append("session:abc-123", &json!({"role": "user"})) @@ -676,6 +840,13 @@ mod tests { .unwrap(); let msgs = store.read("session:abc-123").await.unwrap(); assert_eq!(msgs.len(), 1); + assert_eq!( + store.history_path_for("session:abc-123"), + dir.path() + .join("v1") + .join(SessionStore::key_to_filename("session:abc-123")) + .join("history.jsonl") + ); } #[tokio::test] @@ -684,7 +855,10 @@ mod tests { let data = b"fake png data"; let path = store.save_media("main", "call_1.png", data).await.unwrap(); - assert_eq!(path, "media/main/call_1.png"); + assert_eq!( + path, + SessionStore::media_reference("main", "call_1.png").unwrap() + ); let read_back = store.read_media("main", "call_1.png").await.unwrap(); assert_eq!(read_back, data); @@ -699,7 +873,10 @@ mod tests { .save_media("session:abc", "shot.png", data) .await .unwrap(); - assert_eq!(path, "media/session_abc/shot.png"); + assert_eq!( + path, + SessionStore::media_reference("session:abc", "shot.png").unwrap() + ); let read_back = store.read_media("session:abc", "shot.png").await.unwrap(); assert_eq!(read_back, data); @@ -708,12 +885,14 @@ mod tests { #[test] fn test_media_path_for_uses_session_media_dir() { let (store, dir) = temp_store(); - let path = store.media_path_for("session:abc", "report.pdf"); + let path = store.media_path_for("session:abc", "report.pdf").unwrap(); assert_eq!( path, dir.path() .join("media") - .join("session_abc") + .join("v1") + .join(SessionStore::key_to_filename("session:abc")) + .join("files") .join("report.pdf") ); } @@ -725,6 +904,39 @@ mod tests { assert!(result.is_err()); } + #[tokio::test] + async fn media_filenames_cannot_escape_the_session_directory() { + let (store, _dir) = temp_store(); + for filename in ["", ".", "..", "../escape", "nested/file", "nested\\file"] { + assert!(store.save_media("main", filename, b"data").await.is_err()); + assert!(store.read_media("main", filename).await.is_err()); + assert!(store.media_path_for("main", filename).is_err()); + assert!(SessionStore::media_reference("main", filename).is_err()); + } + } + + #[cfg(unix)] + #[tokio::test] + async fn media_io_does_not_follow_symlinks() { + use std::os::unix::fs::symlink; + + let (store, dir) = temp_store(); + let outside = dir.path().join("outside.bin"); + fs::write(&outside, b"outside").unwrap(); + let media_path = store.media_path_for("main", "linked.bin").unwrap(); + fs::create_dir_all(media_path.parent().unwrap()).unwrap(); + symlink(&outside, &media_path).unwrap(); + + assert!( + store + .save_media("main", "linked.bin", b"changed") + .await + .is_err() + ); + assert!(store.read_media("main", "linked.bin").await.is_err()); + assert_eq!(fs::read(outside).unwrap(), b"outside"); + } + #[tokio::test] async fn test_clear_removes_media_dir() { let (store, dir) = temp_store(); @@ -739,7 +951,12 @@ mod tests { .await .unwrap(); - let media_dir = dir.path().join("media").join("main"); + let media_dir = dir + .path() + .join("media") + .join("v1") + .join(SessionStore::key_to_filename("main")) + .join("files"); assert!(media_dir.exists()); store.clear("main").await.unwrap(); @@ -748,6 +965,222 @@ mod tests { assert!(store.read("main").await.unwrap().is_empty()); } + #[tokio::test] + async fn formerly_colliding_keys_have_isolated_history_media_and_search() { + let (store, _dir) = temp_store(); + store + .append("a:b", &json!({"role": "user", "content": "colon value"})) + .await + .unwrap(); + store + .append( + "a_b", + &json!({"role": "user", "content": "underscore value"}), + ) + .await + .unwrap(); + store.save_media("a:b", "same.bin", b"colon").await.unwrap(); + store + .save_media("a_b", "same.bin", b"underscore") + .await + .unwrap(); + + assert_eq!( + store.read("a:b").await.unwrap()[0]["content"], + "colon value" + ); + assert_eq!( + store.read("a_b").await.unwrap()[0]["content"], + "underscore value" + ); + assert_eq!(store.read_media("a:b", "same.bin").await.unwrap(), b"colon"); + assert_eq!( + store.read_media("a_b", "same.bin").await.unwrap(), + b"underscore" + ); + + let colon_results = store.search("colon value", 10).await.unwrap(); + let underscore_results = store.search("underscore value", 10).await.unwrap(); + assert_eq!(colon_results[0].session_key, "a:b"); + assert_eq!(underscore_results[0].session_key, "a_b"); + assert_eq!(store.list_keys(), vec!["a:b", "a_b"]); + } + + #[tokio::test] + async fn chunk_prefix_keys_have_isolated_media_and_clear() { + let (store, _dir) = temp_store(); + let shorter = "a".repeat(60); + let longer = "a".repeat(61); + store + .save_media(&shorter, "same.bin", b"short") + .await + .unwrap(); + store + .save_media(&longer, "same.bin", b"long") + .await + .unwrap(); + store + .append(&shorter, &json!({"content": "short"})) + .await + .unwrap(); + store + .append(&longer, &json!({"content": "long"})) + .await + .unwrap(); + + store.clear(&shorter).await.unwrap(); + + assert_eq!(store.read(&longer).await.unwrap()[0]["content"], "long"); + assert_eq!( + store.read_media(&longer, "same.bin").await.unwrap(), + b"long" + ); + assert!(store.read(&shorter).await.unwrap().is_empty()); + } + + #[tokio::test] + async fn legacy_collision_is_copied_per_key_and_clear_does_not_resurrect_it() { + let (store, dir) = temp_store(); + fs::write( + dir.path().join("a_b.jsonl"), + "{\"role\":\"user\",\"content\":\"legacy\"}\n", + ) + .unwrap(); + let legacy_media = dir.path().join("media").join("a_b"); + fs::create_dir_all(&legacy_media).unwrap(); + fs::write(legacy_media.join("old.bin"), b"legacy media").unwrap(); + + assert_eq!(store.read("a:b").await.unwrap()[0]["content"], "legacy"); + assert_eq!(store.read("a_b").await.unwrap()[0]["content"], "legacy"); + assert_eq!( + store.read_media("a:b", "old.bin").await.unwrap(), + b"legacy media" + ); + assert_eq!( + store.read_media("a_b", "old.bin").await.unwrap(), + b"legacy media" + ); + + store + .append("a:b", &json!({"role": "assistant", "content": "colon"})) + .await + .unwrap(); + store + .append( + "a_b", + &json!({"role": "assistant", "content": "underscore"}), + ) + .await + .unwrap(); + assert_eq!(store.read("a:b").await.unwrap()[1]["content"], "colon"); + assert_eq!(store.read("a_b").await.unwrap()[1]["content"], "underscore"); + + store.clear("a:b").await.unwrap(); + assert!(store.read("a:b").await.unwrap().is_empty()); + assert_eq!(store.read("a_b").await.unwrap().len(), 2); + assert_eq!( + store.read_media("a_b", "old.bin").await.unwrap(), + b"legacy media" + ); + assert!(dir.path().join("a_b.jsonl").exists()); + } + + #[tokio::test] + async fn authoritative_search_migrates_legacy_underscore_keys_without_guessing() { + let (store, dir) = temp_store(); + fs::write( + dir.path().join("agent_private.jsonl"), + "{\"role\":\"user\",\"content\":\"private history\"}\n", + ) + .unwrap(); + + assert!( + store + .search("private history", 10) + .await + .unwrap() + .is_empty() + ); + let results = store + .search_known_keys("private history", 10, &["agent_private".to_string()]) + .await + .unwrap(); + + assert_eq!(results.len(), 1); + assert_eq!(results[0].session_key, "agent_private"); + } + + #[tokio::test] + async fn migrated_sessions_reconcile_later_legacy_appends_and_media() { + let (store, dir) = temp_store(); + let legacy_history = dir.path().join("old_key.jsonl"); + fs::write( + &legacy_history, + "{\"role\":\"user\",\"content\":\"first\"}\n", + ) + .unwrap(); + let legacy_media = dir.path().join("media").join("old_key"); + fs::create_dir_all(&legacy_media).unwrap(); + + assert_eq!(store.read("old:key").await.unwrap().len(), 1); + OpenOptions::new() + .append(true) + .open(&legacy_history) + .unwrap() + .write_all(b"{\"role\":\"user\",\"content\":\"second\"}\n") + .unwrap(); + fs::write(legacy_media.join("late.bin"), b"late media").unwrap(); + + let messages = store.read("old:key").await.unwrap(); + assert_eq!(messages.len(), 2); + assert_eq!(messages[1]["content"], "second"); + assert_eq!( + store.read_media("old:key", "late.bin").await.unwrap(), + b"late media" + ); + } + + #[tokio::test] + async fn late_created_conflicting_legacy_history_is_not_marked_as_copied() { + let (store, dir) = temp_store(); + store + .append("late:key", &json!({"content": "current"})) + .await + .unwrap(); + fs::write( + dir.path().join("late_key.jsonl"), + "{\"content\":\"legacy\"}\n", + ) + .unwrap(); + + let error = store.read("late:key").await.unwrap_err(); + assert!(error.to_string().contains("histories conflict")); + assert!( + storage_layout::migration_marker(dir.path(), "late:key") + .read_dir() + .is_err() + ); + } + + #[tokio::test] + async fn traversal_like_key_stays_inside_versioned_storage() { + let root = tempfile::tempdir().unwrap(); + let base = root.path().join("sessions"); + let store = SessionStore::new(base.clone()); + store + .append("../escape", &json!({"role": "user"})) + .await + .unwrap(); + + assert!( + store + .history_path_for("../escape") + .starts_with(base.join("v1")) + ); + assert!(!root.path().join("escape.jsonl").exists()); + assert_eq!(store.read("../escape").await.unwrap().len(), 1); + } + // --- Typed API tests --- #[tokio::test] @@ -785,6 +1218,25 @@ mod tests { } } + #[tokio::test] + async fn read_bounded_rejects_message_and_byte_overflow() { + let (store, _dir) = temp_store(); + store + .append("main", &json!({"content": "one"})) + .await + .unwrap(); + store + .append("main", &json!({"content": "two"})) + .await + .unwrap(); + + let message_error = store.read_bounded("main", 1, 1024).await.unwrap_err(); + assert!(message_error.to_string().contains("exceeds 1 messages")); + + let byte_error = store.read_bounded("main", 10, 8).await.unwrap_err(); + assert!(byte_error.to_string().contains("exceeds 8 bytes")); + } + #[tokio::test] async fn test_read_typed_empty() { let (store, _dir) = temp_store(); diff --git a/crates/tools/Cargo.toml b/crates/tools/Cargo.toml index 53f3e62f7f..71c082248b 100644 --- a/crates/tools/Cargo.toml +++ b/crates/tools/Cargo.toml @@ -41,7 +41,7 @@ image = { workspace = true } ipnet = { workspace = true } moltis-agents = { workspace = true } moltis-browser = { workspace = true } -moltis-common = { workspace = true } +moltis-common = { features = ["process"], workspace = true } moltis-config = { workspace = true } moltis-cron = { workspace = true } moltis-media = { workspace = true } diff --git a/crates/tools/src/exec.rs b/crates/tools/src/exec.rs index 691c9c9d5d..8e9d8b9ee2 100644 --- a/crates/tools/src/exec.rs +++ b/crates/tools/src/exec.rs @@ -6,11 +6,17 @@ use std::time::Instant; use { async_trait::async_trait, serde::{Deserialize, Serialize}, - tokio::process::Command, + tokio::{ + io::{AsyncRead, AsyncReadExt}, + process::Command, + }, tracing::{debug, info, warn}, }; -use crate::{Result, error::Error}; +use { + crate::{Result, error::Error}, + moltis_common::process_tree::OwnedProcessTree, +}; #[cfg(feature = "metrics")] use moltis_metrics::{ @@ -112,19 +118,40 @@ impl Default for ExecOpts { } } -fn truncate_output_for_display(output: &mut String, max_output_bytes: usize) { - if output.len() <= max_output_bytes { - return; +async fn read_output_limited( + mut reader: impl AsyncRead + Unpin, + max_output_bytes: usize, +) -> std::io::Result { + let mut retained = Vec::with_capacity(max_output_bytes.min(64 * 1024)); + let mut buffer = [0_u8; 8 * 1024]; + let mut truncated = false; + + loop { + let read = reader.read(&mut buffer).await?; + if read == 0 { + break; + } + let keep = read.min(max_output_bytes.saturating_sub(retained.len())); + retained.extend_from_slice(&buffer[..keep]); + truncated |= keep < read; + } + + let mut output = String::from_utf8_lossy(&retained).into_owned(); + if output.len() > max_output_bytes { + output.truncate(output.floor_char_boundary(max_output_bytes)); + truncated = true; + } + if truncated { + output.push_str("\n... [output truncated]"); } - output.truncate(output.floor_char_boundary(max_output_bytes)); - output.push_str("\n... [output truncated]"); + Ok(output) } /// Execute a shell command with timeout and output limits. -#[tracing::instrument(skip(opts), fields(timeout_secs = opts.timeout.as_secs()))] +#[tracing::instrument(skip(command, opts), fields(timeout_secs = opts.timeout.as_secs()))] pub async fn exec_command(command: &str, opts: &ExecOpts) -> Result { debug!( - command, + command_bytes = command.len(), timeout_secs = opts.timeout.as_secs(), "exec_command" ); @@ -143,8 +170,7 @@ pub async fn exec_command(command: &str, opts: &ExecOpts) -> Result cmd.stderr(std::process::Stdio::piped()); // Prevent the child from inheriting stdin. cmd.stdin(std::process::Stdio::null()); - - let child = cmd.spawn().map_err(|e| { + let mut child = OwnedProcessTree::spawn(cmd).map_err(|e| { if e.kind() == std::io::ErrorKind::NotFound { if let Some(ref dir) = opts.working_dir { Error::message(format!( @@ -159,18 +185,24 @@ pub async fn exec_command(command: &str, opts: &ExecOpts) -> Result } })?; - let result = tokio::time::timeout(opts.timeout, child.wait_with_output()).await; + let stdout = child + .take_stdout() + .ok_or_else(|| Error::message("failed to capture command stdout"))?; + let stderr = child + .take_stderr() + .ok_or_else(|| Error::message("failed to capture command stderr"))?; + let execution = async { + tokio::try_join!( + child.wait(), + read_output_limited(stdout, opts.max_output_bytes), + read_output_limited(stderr, opts.max_output_bytes), + ) + }; + let result = tokio::time::timeout(opts.timeout, execution).await; match result { - Ok(Ok(output)) => { - let mut stdout = String::from_utf8_lossy(&output.stdout).into_owned(); - let mut stderr = String::from_utf8_lossy(&output.stderr).into_owned(); - - // Truncate if exceeding limit. - truncate_output_for_display(&mut stdout, opts.max_output_bytes); - truncate_output_for_display(&mut stderr, opts.max_output_bytes); - - let exit_code = output.status.code().unwrap_or(-1); + Ok(Ok((status, stdout, stderr))) => { + let exit_code = status.code().unwrap_or(-1); debug!( exit_code, stdout_len = stdout.len(), @@ -186,7 +218,8 @@ pub async fn exec_command(command: &str, opts: &ExecOpts) -> Result }, Ok(Err(e)) => Err(Error::message(format!("failed to run command: {e}"))), Err(_) => { - warn!(command, "exec timeout"); + let _ = child.kill().await; + warn!(command_bytes = command.len(), "exec timeout"); Err(Error::message(format!( "command timed out after {}s", opts.timeout.as_secs() @@ -418,7 +451,7 @@ impl AgentTool for ExecTool { let cwd = params.get("working_dir").and_then(|v| v.as_str()); info!( - command, + command_bytes = command.len(), node_id = %node_id, timeout_secs, "exec forwarding to remote node" @@ -472,6 +505,13 @@ impl AgentTool for ExecTool { .and_then(|v| v.as_str()) .filter(|s| !s.is_empty()) .map(PathBuf::from) + .or_else(|| { + params + .get("_working_dir") + .and_then(|v| v.as_str()) + .filter(|s| !s.is_empty()) + .map(PathBuf::from) + }) .or_else(|| self.working_dir.clone()); let runs_on_host = !(is_sandboxed && has_container_backend); @@ -528,7 +568,7 @@ impl AgentTool for ExecTool { } info!( - command, + command_bytes = command.len(), timeout_secs, ?working_dir, is_sandboxed, @@ -544,7 +584,10 @@ impl AgentTool for ExecTool { if needs_approval && let Some(ref mgr) = self.approval_manager { let action = mgr.check_command(command).await?; if action == ApprovalAction::NeedsApproval { - info!(command, "command needs approval, waiting..."); + info!( + command_bytes = command.len(), + "command needs approval, waiting..." + ); let (req_id, rx) = mgr.create_request(command, session_key).await; // Broadcast to connected clients. @@ -557,7 +600,7 @@ impl AgentTool for ExecTool { let decision = mgr.wait_for_decision(rx).await; match decision { ApprovalDecision::Approved => { - info!(command, "command approved"); + info!(command_bytes = command.len(), "command approved"); }, ApprovalDecision::Denied => { return Err( @@ -712,7 +755,7 @@ impl AgentTool for ExecTool { Error::message(format!("sandbox preparation failed: {error}")).into(), ); } - debug!(session = sk, sandbox_id = %id, command, "sandbox running command"); + debug!(session = sk, sandbox_id = %id, command_bytes = command.len(), "sandbox running command"); let mut sandbox_result = backend.exec(&id, command, &opts).await?; for retry_idx in 1..=MAX_SANDBOX_RECOVERY_RETRIES { if sandbox_result.exit_code == 0 @@ -724,7 +767,7 @@ impl AgentTool for ExecTool { warn!( session = sk, sandbox_id = %id, - command, + command_bytes = command.len(), retry_idx, max_retries = MAX_SANDBOX_RECOVERY_RETRIES, "sandbox exec failed because container is unavailable, reinitializing and retrying" @@ -742,11 +785,15 @@ impl AgentTool for ExecTool { } sandbox_result } else { - debug!(session = sk, command, "running unsandboxed"); + debug!( + session = sk, + command_bytes = command.len(), + "running unsandboxed" + ); exec_command(command, &opts).await? } } else if let Some(ref id) = self.sandbox_id { - debug!(sandbox_id = %id, command, "static sandbox running command"); + debug!(sandbox_id = %id, command_bytes = command.len(), "static sandbox running command"); self.sandbox.ensure_ready(id, None).await?; let mut sandbox_result = self.sandbox.exec(id, command, &opts).await?; for retry_idx in 1..=MAX_SANDBOX_RECOVERY_RETRIES { @@ -758,7 +805,7 @@ impl AgentTool for ExecTool { warn!( sandbox_id = %id, - command, + command_bytes = command.len(), retry_idx, max_retries = MAX_SANDBOX_RECOVERY_RETRIES, "sandbox exec failed because container is unavailable, reinitializing and retrying" @@ -792,7 +839,7 @@ impl AgentTool for ExecTool { } info!( - command, + command_bytes = command.len(), exit_code = result.exit_code, stdout_len = result.stdout.len(), stderr_len = result.stderr.len(), diff --git a/crates/tools/src/exec/tests.rs b/crates/tools/src/exec/tests.rs index b9af2dfd1d..a6db8778d1 100644 --- a/crates/tools/src/exec/tests.rs +++ b/crates/tools/src/exec/tests.rs @@ -19,14 +19,24 @@ impl TestBroadcaster { } } -#[test] -fn truncate_output_for_display_handles_multibyte_boundary() { - let mut output = format!("{}л{}", "a".repeat(1999), "z".repeat(10)); - truncate_output_for_display(&mut output, 2000); +#[tokio::test] +async fn limited_output_handles_multibyte_boundary() { + let input = format!("{}л{}", "a".repeat(1999), "z".repeat(10)); + let output = read_output_limited(input.as_bytes(), 2000).await.unwrap(); assert!(output.contains("[output truncated]")); assert!(!output.contains('л')); } +#[tokio::test] +async fn limited_output_caps_lossy_utf8_expansion() { + const LIMIT: usize = 10; + const MARKER: &str = "\n... [output truncated]"; + let input = [0xff_u8; 100]; + let output = read_output_limited(input.as_slice(), LIMIT).await.unwrap(); + assert!(output.ends_with(MARKER)); + assert!(output.len() <= LIMIT + MARKER.len()); +} + #[async_trait] impl ApprovalBroadcaster for TestBroadcaster { async fn broadcast_request( @@ -74,6 +84,21 @@ async fn test_exec_timeout() { assert!(result.is_err()); } +#[tokio::test] +async fn test_exec_timeout_kills_before_later_side_effect() { + let temp_dir = tempfile::tempdir().unwrap(); + let marker = temp_dir.path().join("should-not-exist"); + let command = format!("(sleep 1; touch '{}') & wait", marker.display()); + let opts = ExecOpts { + timeout: Duration::from_millis(50), + ..Default::default() + }; + + assert!(exec_command(&command, &opts).await.is_err()); + tokio::time::sleep(Duration::from_millis(1100)).await; + assert!(!marker.exists(), "timed-out descendant continued running"); +} + #[tokio::test] async fn test_exec_tool() { let temp_dir = tempfile::tempdir().unwrap(); @@ -104,6 +129,21 @@ async fn test_exec_tool_empty_working_dir() { assert!(!result["stdout"].as_str().unwrap().trim().is_empty()); } +#[tokio::test] +async fn test_exec_tool_uses_internal_working_dir_default() { + let temp_dir = tempfile::tempdir().unwrap(); + let tool = ExecTool::default(); + let result = tool + .execute(serde_json::json!({ + "command": "pwd", + "_working_dir": temp_dir.path(), + })) + .await + .unwrap(); + let reported = std::fs::canonicalize(result["stdout"].as_str().unwrap().trim()).unwrap(); + assert_eq!(reported, temp_dir.path().canonicalize().unwrap()); +} + #[tokio::test] async fn test_exec_tool_safe_command_no_approval_needed() { let mgr = Arc::new(ApprovalManager::default()); diff --git a/crates/tools/src/sessions_communicate.rs b/crates/tools/src/sessions_communicate.rs index 7c56e9069c..2b9d4a71b8 100644 --- a/crates/tools/src/sessions_communicate.rs +++ b/crates/tools/src/sessions_communicate.rs @@ -288,12 +288,6 @@ impl AgentTool for SessionsSearchTool { None }; - let search_limit = limit.saturating_mul(4).max(limit); - let hits = - self.store.search(query, search_limit).await.map_err(|e| { - Error::message(format!("failed to search sessions for '{query}': {e}")) - })?; - let entries: HashMap = self .metadata .list() @@ -301,6 +295,13 @@ impl AgentTool for SessionsSearchTool { .into_iter() .map(|entry| (entry.key.clone(), entry)) .collect(); + let known_keys = entries.keys().cloned().collect::>(); + let search_limit = limit.saturating_mul(4).max(limit); + let hits = self + .store + .search_known_keys(query, search_limit, &known_keys) + .await + .map_err(|e| Error::message(format!("failed to search sessions for '{query}': {e}")))?; let mut results = Vec::with_capacity(limit); for hit in hits { diff --git a/crates/tools/src/wasm_component.rs b/crates/tools/src/wasm_component.rs index b4bd1382cb..5d311c000c 100644 --- a/crates/tools/src/wasm_component.rs +++ b/crates/tools/src/wasm_component.rs @@ -298,10 +298,10 @@ mod tests { super::{HttpError, HttpHostImpl, HttpRequest, PureToolValue, marshal_tool_result}, std::{ collections::HashMap, - io::{Read, Write}, + io::{ErrorKind, Read, Write}, net::TcpListener, thread::{self, JoinHandle}, - time::Duration, + time::{Duration, Instant}, }, }; @@ -338,19 +338,33 @@ mod tests { fn spawn_http_server(body: &'static [u8], delay: Option) -> (String, JoinHandle<()>) { let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let addr = listener.local_addr().unwrap(); + listener.set_nonblocking(true).unwrap(); let handle = thread::spawn(move || { - if let Ok((mut stream, _)) = listener.accept() { - let mut request_buffer = [0_u8; 2048]; - let _ = stream.read(&mut request_buffer); - if let Some(wait) = delay { - thread::sleep(wait); + let deadline = Instant::now() + Duration::from_secs(5); + while Instant::now() < deadline { + match listener.accept() { + Ok((mut stream, _)) => { + stream + .set_read_timeout(Some(Duration::from_secs(5))) + .unwrap(); + let mut request_buffer = [0_u8; 2048]; + let _ = stream.read(&mut request_buffer); + if let Some(wait) = delay { + thread::sleep(wait); + } + let headers = format!( + "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + let _ = stream.write_all(headers.as_bytes()); + let _ = stream.write_all(body); + return; + }, + Err(error) if error.kind() == ErrorKind::WouldBlock => { + thread::sleep(Duration::from_millis(5)); + }, + Err(_) => return, } - let headers = format!( - "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", - body.len() - ); - let _ = stream.write_all(headers.as_bytes()); - let _ = stream.write_all(body); } }); (format!("http://{addr}/"), handle) diff --git a/crates/web/ui/e2e/specs/chat-api-message.spec.js b/crates/web/ui/e2e/specs/chat-api-message.spec.js index b76b27206d..10af9fbb46 100644 --- a/crates/web/ui/e2e/specs/chat-api-message.spec.js +++ b/crates/web/ui/e2e/specs/chat-api-message.spec.js @@ -149,4 +149,35 @@ test.describe("API-sent user messages (GH #729)", () => { await expect(page.locator(".msg.user")).toHaveCount(0, { timeout: 2_000 }); expect(pageErrors).toEqual([]); }); + + test("channel audio cache does not reconstruct storage paths", async ({ page }) => { + await expectRpcOk(page, "system-event", { + event: "chat", + payload: { + sessionKey: "agent:private", + state: "channel_user", + text: "Voice message", + channel: { audio_filename: "voice.ogg" }, + }, + }); + + const cachedHistory = await page.evaluate(async () => { + var appScript = document.querySelector('script[type="module"][src*="js/app.js"]'); + if (!appScript) throw new Error("app module script not found"); + var appUrl = new URL(appScript.src, window.location.origin); + var prefix = appUrl.href.slice(0, appUrl.href.length - "js/app.js".length); + var cache = await import(`${prefix}js/stores/session-history-cache.js`); + return cache.getSessionHistory("agent:private"); + }); + + expect(cachedHistory).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + role: "user", + audio: "voice.ogg", + }), + ]), + ); + expect(pageErrors).toEqual([]); + }); }); diff --git a/crates/web/ui/src/components/SessionHeader.tsx b/crates/web/ui/src/components/SessionHeader.tsx index 95413976a9..e7ea1fa294 100644 --- a/crates/web/ui/src/components/SessionHeader.tsx +++ b/crates/web/ui/src/components/SessionHeader.tsx @@ -511,7 +511,7 @@ export function SessionHeader({ /> ) : ( **Tool calls are reported, not gated.** A client sees tools as they run but +> cannot veto them: `session/request_permission` is not yet wired to Moltis's +> tool policy, so tool approval remains governed by your Moltis configuration, +> not by the ACP client. If you need a client to approve individual tool use, +> that is not available yet. + +## Architecture notes + +The protocol crate declares its traits with `#[async_trait(?Send)]`, so the +handler is pinned to the thread running a `tokio::task::LocalSet`, while +Moltis's services are `Send + Sync`. The `AcpBackend` trait in `crates/acp` is +where the two meet: implementations are `Send + Sync` and never learn a +`LocalSet` exists, and streaming flows back through a plain channel that the +protocol layer drains while the turn runs. diff --git a/docs/src/openclaw-import.md b/docs/src/openclaw-import.md index 7a1cb86b2d..562a3ffcd9 100644 --- a/docs/src/openclaw-import.md +++ b/docs/src/openclaw-import.md @@ -22,7 +22,7 @@ If the directory exists and contains recognizable OpenClaw files (`openclaw.json | **Skills** | `skills/` directories with `SKILL.md` | `~/.moltis/skills/` | Copies entire skill directories; skips duplicates | | **Memory** | `MEMORY.md` and all `memory/*.md` files | `~/.moltis/MEMORY.md` and `~/.moltis/memory/` | Imports daily logs, project notes, and all other markdown memory files. Appends with `` separator for idempotency | | **Channels** | Telegram and Discord bot configuration in `openclaw.json` | `moltis.toml` channels section | Supports both flat and multi-account Telegram configs | -| **Sessions** | JSONL conversation files under `agents/*/sessions/` | `~/.moltis/sessions/` and `~/.moltis/memory/sessions/` | Converts OpenClaw message format to Moltis format; prefixes keys with `oc:`. Also generates markdown transcripts for memory search indexing | +| **Sessions** | JSONL conversation files under `agents/*/sessions/` | `~/.moltis/sessions/v1/` and `~/.moltis/memory/sessions/` | Converts OpenClaw message format to Moltis format; prefixes keys with `oc:`. Also generates markdown transcripts for memory search indexing | | **MCP Servers** | `mcp-servers.json` | `~/.moltis/mcp-servers.json` | Merges with existing servers; skips duplicates by name | | **Workspace Files** | `SOUL.md`, `IDENTITY.md`, `USER.md`, `TOOLS.md`, `AGENTS.md`, `HEARTBEAT.md`, `BOOT.md` | `~/.moltis/` (root) or `~/.moltis/agents//` | Copies raw workspace files; skips if destination already has user content. Replaces auto-seeded defaults | | **Agent Presets** | Non-default agents in `agents.list` | `moltis.toml` `[agents.presets.*]` | Creates `spawn_agent` presets with name, theme, and model. Existing presets are preserved | diff --git a/vendor/agent-client-protocol/.cargo-ok b/vendor/agent-client-protocol/.cargo-ok new file mode 100644 index 0000000000..5f8b795830 --- /dev/null +++ b/vendor/agent-client-protocol/.cargo-ok @@ -0,0 +1 @@ +{"v":1} \ No newline at end of file diff --git a/vendor/agent-client-protocol/.cargo_vcs_info.json b/vendor/agent-client-protocol/.cargo_vcs_info.json new file mode 100644 index 0000000000..e2de8d613b --- /dev/null +++ b/vendor/agent-client-protocol/.cargo_vcs_info.json @@ -0,0 +1,6 @@ +{ + "git": { + "sha1": "23ab9bf555484f262ab4b19fdd812dea1a2fc961" + }, + "path_in_vcs": "src/agent-client-protocol" +} \ No newline at end of file diff --git a/vendor/agent-client-protocol/.gitignore b/vendor/agent-client-protocol/.gitignore new file mode 100644 index 0000000000..e9e21997b1 --- /dev/null +++ b/vendor/agent-client-protocol/.gitignore @@ -0,0 +1,2 @@ +/target/ +/Cargo.lock diff --git a/vendor/agent-client-protocol/CHANGELOG.md b/vendor/agent-client-protocol/CHANGELOG.md new file mode 100644 index 0000000000..9fa515ae4c --- /dev/null +++ b/vendor/agent-client-protocol/CHANGELOG.md @@ -0,0 +1,148 @@ +# Changelog + +## [0.10.4](https://github.com/agentclientprotocol/rust-sdk/compare/v0.10.3...v0.10.4) - 2026-03-31 + +### Added + +- *(schema)* Update schema to 0.11.4 ([#95](https://github.com/agentclientprotocol/rust-sdk/pull/95)) + +### Fixed + +- add warning logs for silent failures in RPC message handling ([#92](https://github.com/agentclientprotocol/rust-sdk/pull/92)) +- Clearer error message when connection is broken before messages are sent ([#89](https://github.com/agentclientprotocol/rust-sdk/pull/89)) + +### Other + +- Fix the rpc_test and example use following the new schema api ([#88](https://github.com/agentclientprotocol/rust-sdk/pull/88)) + +## [0.10.3](https://github.com/agentclientprotocol/rust-sdk/compare/v0.10.2...v0.10.3) - 2026-03-25 + +### Added + +- *(unstable)* Add logout support ([#84](https://github.com/agentclientprotocol/rust-sdk/pull/84)) +- *(schema)* Update schema to 0.11.3 ([#82](https://github.com/agentclientprotocol/rust-sdk/pull/82)) + +## [0.10.2](https://github.com/agentclientprotocol/rust-sdk/compare/v0.10.1...v0.10.2) - 2026-03-11 + +### Added + +- *(unstable)* Add support for session/close methods ([#77](https://github.com/agentclientprotocol/rust-sdk/pull/77)) + +## [0.10.1](https://github.com/agentclientprotocol/rust-sdk/compare/v0.10.0...v0.10.1) - 2026-03-10 + +### Added + +- Stabilize session_list and session_info_update ([#74](https://github.com/agentclientprotocol/rust-sdk/pull/74)) + +### Fixed + +- Make examples compile again ([#76](https://github.com/agentclientprotocol/rust-sdk/pull/76)) + +## [0.10.0](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.5...v0.10.0) - 2026-03-05 + +### Added + +- Add more unstable feature flags from schema ([#71](https://github.com/agentclientprotocol/rust-sdk/pull/71)) +- [**breaking**] Update to schema crate v0.11.0 ([#69](https://github.com/agentclientprotocol/rust-sdk/pull/69)) + +## [0.9.5](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.4...v0.9.5) - 2026-03-03 + +### Fixed + +- handle escaped forward slashes in JSON-RPC method names ([#65](https://github.com/agentclientprotocol/rust-sdk/pull/65)) + +## [0.9.4](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.3...v0.9.4) - 2026-02-04 + +### Added + +- Update to 0.10.8 of the schema ([#51](https://github.com/agentclientprotocol/rust-sdk/pull/51)) + +## [0.9.3](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.2...v0.9.3) - 2026-01-09 + +### Other + +- update Cargo.toml dependencies + +## [0.9.2](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.1...v0.9.2) - 2025-12-17 + +### Added + +- *(unstable)* Add initial support for session config options ([#36](https://github.com/agentclientprotocol/rust-sdk/pull/36)) + +## [0.9.1](https://github.com/agentclientprotocol/rust-sdk/compare/v0.9.0...v0.9.1) - 2025-12-17 + +### Added + +- *(unstable)* Add initial support for resuming sessions ([#34](https://github.com/agentclientprotocol/rust-sdk/pull/34)) +- *(unstable)* Add initial support for forking sessions ([#33](https://github.com/agentclientprotocol/rust-sdk/pull/33)) +- *(unstable)* Add initial support for listing sessions ([#31](https://github.com/agentclientprotocol/rust-sdk/pull/31)) + +### Other + +- Add test for unstable session info feature ([#35](https://github.com/agentclientprotocol/rust-sdk/pull/35)) + +## [0.9.0](https://github.com/agentclientprotocol/rust-sdk/compare/v0.8.0...v0.9.0) - 2025-12-08 + +Update to v0.10.0 of agent-client-protocol-schema + +## 0.8.0 (2025-12-01) + +The types from the Rust crate, `agent-client-protocol-schema` has major breaking changes. All exported type are now marked as `#[non_exhaustive]`. Since the schema itself is JSON, and we can introduce new fields and variants in a non-breaking way, we wanted to allow for the same behavior in the Rust library. + +All enum variants are also tuple variants now, with their own structs. This made it nicer to represent in the JSON Schema, and also made sure we have `_meta` fields on all variants. + +This upgrade will likely come with a lot of compilation errors, but ideally upgrading will be more painless in the future. + +## 0.7.0 (2025-10-24) + +- Add ability for agents and clients to provide information about their implementation +- Fix incorrectly serialized `_meta` field on `SetSessionModeResponse` + +## 0.6.0 (2025-10-23) + +- Provide missing `_meta` fields on certain enum variants. +- More consistent enum usage. Enums are always either newtype or struct variants within a single enum, not mixed. + +## 0.5.0 (2025-10-20) + +- Export necessary RPC types. Fixes an issue where certain fields weren't public enough. +- Make id types easier to create and add `PartialEq` and `Eq` impls for as many types as possible. +- Export `acp::Result` for easier indication of ACP errors. +- Use `acp::Error`/`acp::Result` instead of `anyhow::Error`/`anyhow::Result` for all return types. + +## 0.4.7 (2025-10-13) + +- Depend on `agent-client-protocol-schema` for schema types + +## 0.4.6 (2025-10-10) + +### Rust + +- Fix: support all valid JSON-RPC ids (int, string, null) + +## 0.4.5 (2025-10-02) + +- No changes + +## 0.4.4 (2025-09-30) + +- Provide default trait implementations for optional capability-based `Agent` and `Client` methods. + +## 0.4.3 (2025-09-25) + +- impl `Agent` and `Client` for `Rc` and `Arc` where `T` implements either trait. + +## 0.4.2 (2025-09-22) + +**Unstable** fix missing method for model selection in Rust library. + +## 0.4.1 (2025-09-22) + +**Unstable** initial support for model selection. + +## 0.4.0 (2025-09-17) + +- Make `Agent` and `Client` dyn compatible (you'll need to annotate them with `#[async_trait]`) [#97](https://github.com/agentclientprotocol/agent-client-protocol/pull/97) +- `ext_method` and `ext_notification` methods are now more consistent with the other trait methods [#95](https://github.com/agentclientprotocol/agent-client-protocol/pull/95) + - There are also distinct types for `ExtRequest`, `ExtResponse`, and `ExtNotification` +- Rexport `serde_json::RawValue` for easier use [#95](https://github.com/agentclientprotocol/agent-client-protocol/pull/95) diff --git a/vendor/agent-client-protocol/Cargo.toml b/vendor/agent-client-protocol/Cargo.toml new file mode 100644 index 0000000000..0a88bca9c0 --- /dev/null +++ b/vendor/agent-client-protocol/Cargo.toml @@ -0,0 +1,172 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2024" +name = "agent-client-protocol" +version = "0.10.4" +authors = ["Zed "] +build = false +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "A protocol for standardizing communication between code editors and AI coding agents" +homepage = "https://github.com/agentclientprotocol/rust-sdk" +documentation = "https://docs.rs/agent-client-protocol" +readme = "README.md" +keywords = [ + "agent", + "client", + "protocol", + "ai", + "editor", +] +categories = [ + "development-tools", + "api-bindings", +] +license = "Apache-2.0" +repository = "https://github.com/agentclientprotocol/rust-sdk" + +[features] +unstable = [ + "unstable_auth_methods", + "unstable_cancel_request", + "unstable_elicitation", + "unstable_logout", + "unstable_nes", + "unstable_session_additional_directories", + "unstable_session_fork", + "unstable_session_close", + "unstable_session_model", + "unstable_session_resume", + "unstable_session_usage", + "unstable_message_id", + "unstable_boolean_config", +] +unstable_auth_methods = ["agent-client-protocol-schema/unstable_auth_methods"] +unstable_boolean_config = ["agent-client-protocol-schema/unstable_boolean_config"] +unstable_cancel_request = ["agent-client-protocol-schema/unstable_cancel_request"] +unstable_elicitation = ["agent-client-protocol-schema/unstable_elicitation"] +unstable_logout = ["agent-client-protocol-schema/unstable_logout"] +unstable_message_id = ["agent-client-protocol-schema/unstable_message_id"] +unstable_nes = ["agent-client-protocol-schema/unstable_nes"] +unstable_session_additional_directories = ["agent-client-protocol-schema/unstable_session_additional_directories"] +unstable_session_close = ["agent-client-protocol-schema/unstable_session_close"] +unstable_session_fork = ["agent-client-protocol-schema/unstable_session_fork"] +unstable_session_model = ["agent-client-protocol-schema/unstable_session_model"] +unstable_session_resume = ["agent-client-protocol-schema/unstable_session_resume"] +unstable_session_usage = ["agent-client-protocol-schema/unstable_session_usage"] + +[lib] +name = "agent_client_protocol" +path = "src/lib.rs" + +[[example]] +name = "agent" +path = "examples/agent.rs" + +[[example]] +name = "client" +path = "examples/client.rs" + +[dependencies.agent-client-protocol-schema] +version = "=0.11.4" + +[dependencies.anyhow] +version = "1.0" + +[dependencies.async-broadcast] +version = "0.7" + +[dependencies.async-trait] +version = "0.1" + +[dependencies.derive_more] +version = "2" +features = ["from"] + +[dependencies.futures] +version = "0.3.31" + +[dependencies.log] +version = "0.4" + +[dependencies.serde] +version = "1.0" +features = [ + "derive", + "rc", +] + +[dependencies.serde_json] +version = "1" +features = ["raw_value"] + +[dev-dependencies.env_logger] +version = "0.11" + +[dev-dependencies.futures-util] +version = "0.3" +features = ["io"] + +[dev-dependencies.piper] +version = "0.2" + +[dev-dependencies.pretty_assertions] +version = "1" + +[dev-dependencies.rustyline] +version = "17" + +[dev-dependencies.tokio] +version = "1.48" +features = ["full"] + +[dev-dependencies.tokio-util] +version = "0.7" +features = ["compat"] + +[lints.clippy] +doc_markdown = "allow" +missing_errors_doc = "allow" +missing_panics_doc = "allow" +needless_pass_by_value = "allow" +similar_names = "allow" +struct_field_names = "allow" +too_many_lines = "allow" +wildcard_imports = "allow" + +[lints.clippy.pedantic] +level = "warn" +priority = -1 + +[lints.rust] +let-underscore = "warn" +missing_debug_implementations = "warn" + +[lints.rust.future_incompatible] +level = "warn" +priority = -1 + +[lints.rust.nonstandard_style] +level = "warn" +priority = -1 + +[lints.rust.rust_2018_idioms] +level = "warn" +priority = -1 + +[lints.rust.unused] +level = "warn" +priority = -1 diff --git a/vendor/agent-client-protocol/Cargo.toml.orig b/vendor/agent-client-protocol/Cargo.toml.orig new file mode 100644 index 0000000000..30f19ab07e --- /dev/null +++ b/vendor/agent-client-protocol/Cargo.toml.orig @@ -0,0 +1,67 @@ +[package] +name = "agent-client-protocol" +authors = ["Zed "] +version = "0.10.4" +edition = "2024" +license = "Apache-2.0" +description = "A protocol for standardizing communication between code editors and AI coding agents" +repository = "https://github.com/agentclientprotocol/rust-sdk" +homepage = "https://github.com/agentclientprotocol/rust-sdk" +documentation = "https://docs.rs/agent-client-protocol" +readme = "README.md" +keywords = ["agent", "client", "protocol", "ai", "editor"] +categories = ["development-tools", "api-bindings"] + +[features] +unstable = [ + "unstable_auth_methods", + "unstable_cancel_request", + "unstable_elicitation", + "unstable_logout", + "unstable_nes", + "unstable_session_additional_directories", + "unstable_session_fork", + "unstable_session_close", + "unstable_session_model", + "unstable_session_resume", + "unstable_session_usage", + "unstable_message_id", + "unstable_boolean_config", +] +unstable_auth_methods = ["agent-client-protocol-schema/unstable_auth_methods"] +unstable_cancel_request = ["agent-client-protocol-schema/unstable_cancel_request"] +unstable_elicitation = ["agent-client-protocol-schema/unstable_elicitation"] +unstable_logout = ["agent-client-protocol-schema/unstable_logout"] +unstable_nes = ["agent-client-protocol-schema/unstable_nes"] +unstable_session_additional_directories = ["agent-client-protocol-schema/unstable_session_additional_directories"] +unstable_session_fork = ["agent-client-protocol-schema/unstable_session_fork"] +unstable_session_model = ["agent-client-protocol-schema/unstable_session_model"] +unstable_session_resume = ["agent-client-protocol-schema/unstable_session_resume"] +unstable_session_usage = ["agent-client-protocol-schema/unstable_session_usage"] +unstable_session_close = ["agent-client-protocol-schema/unstable_session_close"] +unstable_message_id = ["agent-client-protocol-schema/unstable_message_id"] +unstable_boolean_config = ["agent-client-protocol-schema/unstable_boolean_config"] + + +[dependencies] +agent-client-protocol-schema.workspace = true +anyhow.workspace = true +async-broadcast.workspace = true +async-trait.workspace = true +derive_more.workspace = true +futures.workspace = true +log.workspace = true +serde.workspace = true +serde_json.workspace = true + +[dev-dependencies] +env_logger.workspace = true +futures-util.workspace = true +piper.workspace = true +pretty_assertions.workspace = true +rustyline.workspace = true +tokio.workspace = true +tokio-util.workspace = true + +[lints] +workspace = true diff --git a/vendor/agent-client-protocol/README.md b/vendor/agent-client-protocol/README.md new file mode 100644 index 0000000000..ef510fdf21 --- /dev/null +++ b/vendor/agent-client-protocol/README.md @@ -0,0 +1,31 @@ + + Agent Client Protocol + + +# Agent Client Protocol + +The Agent Client Protocol (ACP) standardizes communication between _code editors_ (interactive programs for viewing and editing source code) and _coding agents_ (programs that use generative AI to autonomously modify code). + +Learn more at [agentclientprotocol.com](https://agentclientprotocol.com/). + +## Integrations + +- [Schema](./schema/schema.json) +- [Agents](https://agentclientprotocol.com/overview/agents) +- [Clients](https://agentclientprotocol.com/overview/clients) +- Official Libraries + - **Kotlin**: [`acp-kotlin`](https://github.com/agentclientprotocol/kotlin-sdk) - Supports JVM, other targets are in progress, see [samples](https://github.com/agentclientprotocol/kotlin-sdk/tree/master/samples/kotlin-acp-client-sample/src/main/kotlin/com/agentclientprotocol/samples) + - **Python**: [`python-sdk`](https://github.com/agentclientprotocol/python-sdk) - See [examples](https://github.com/agentclientprotocol/python-sdk/tree/main/examples) + - **Rust**: [`agent-client-protocol`](https://crates.io/crates/agent-client-protocol) - See [examples/agent.rs](https://github.com/agentclientprotocol/rust-sdk/blob/main/examples/agent.rs) and [examples/client.rs](https://github.com/agentclientprotocol/rust-sdk/blob/main/examples/client.rs) + - **TypeScript**: [`@agentclientprotocol/sdk`](https://www.npmjs.com/package/@agentclientprotocol/sdk) - See [examples/](https://github.com/agentclientprotocol/typescript-sdk/tree/main/src/examples) +- [Community Libraries](https://agentclientprotocol.com/libraries/community) + +## Contributing + +ACP is a protocol intended for broad adoption across the ecosystem; we follow a structured process to ensure changes are well-considered. Read the [Contributing Guide](./CONTRIBUTING.md) for more information. + +## Contribution Policy + +This project does not require a Contributor License Agreement (CLA). Instead, contributions are accepted under the following terms: + +> By contributing to this project, you agree that your contributions will be licensed under the [Apache License, Version 2.0](https://www.apache.org/licenses/LICENSE-2.0). You affirm that you have the legal right to submit your work, that you are not including code you do not have rights to, and that you understand contributions are made without requiring a Contributor License Agreement (CLA). diff --git a/vendor/agent-client-protocol/examples/agent.rs b/vendor/agent-client-protocol/examples/agent.rs new file mode 100644 index 0000000000..926865417b --- /dev/null +++ b/vendor/agent-client-protocol/examples/agent.rs @@ -0,0 +1,193 @@ +//! A simple ACP agent server for educational purposes. +//! +//! The agent communicates with clients over stdio. To run it with logging: +//! +//! ```bash +//! RUST_LOG=info cargo run --example agent +//! ``` +//! +//! To connect it to the example client from this crate: +//! +//! ```bash +//! cargo build --example agent && cargo run --example client -- target/debug/examples/agent +//! ``` + +use std::cell::Cell; + +use { + agent_client_protocol::{self as acp, Client as _}, + serde_json::json, + tokio::sync::{mpsc, oneshot}, + tokio_util::compat::{TokioAsyncReadCompatExt as _, TokioAsyncWriteCompatExt as _}, +}; + +struct ExampleAgent { + session_update_tx: mpsc::UnboundedSender<(acp::SessionNotification, oneshot::Sender<()>)>, + next_session_id: Cell, +} + +impl ExampleAgent { + fn new( + session_update_tx: mpsc::UnboundedSender<(acp::SessionNotification, oneshot::Sender<()>)>, + ) -> Self { + Self { + session_update_tx, + next_session_id: Cell::new(0), + } + } +} + +#[async_trait::async_trait(?Send)] +impl acp::Agent for ExampleAgent { + async fn initialize( + &self, + arguments: acp::InitializeRequest, + ) -> Result { + log::info!("Received initialize request {arguments:?}"); + Ok(acp::InitializeResponse::new(acp::ProtocolVersion::V1) + .agent_info(acp::Implementation::new("example-agent", "0.1.0").title("Example Agent"))) + } + + async fn authenticate( + &self, + arguments: acp::AuthenticateRequest, + ) -> Result { + log::info!("Received authenticate request {arguments:?}"); + Ok(acp::AuthenticateResponse::default()) + } + + async fn new_session( + &self, + arguments: acp::NewSessionRequest, + ) -> Result { + log::info!("Received new session request {arguments:?}"); + let session_id = self.next_session_id.get(); + self.next_session_id.set(session_id + 1); + Ok(acp::NewSessionResponse::new(session_id.to_string())) + } + + async fn load_session( + &self, + arguments: acp::LoadSessionRequest, + ) -> Result { + log::info!("Received load session request {arguments:?}"); + Ok(acp::LoadSessionResponse::new()) + } + + async fn prompt( + &self, + arguments: acp::PromptRequest, + ) -> Result { + log::info!("Received prompt request {arguments:?}"); + for content in ["Client sent: ".into()].into_iter().chain(arguments.prompt) { + let (tx, rx) = oneshot::channel(); + self.session_update_tx + .send(( + acp::SessionNotification::new( + arguments.session_id.clone(), + acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk::new(content)), + ), + tx, + )) + .map_err(|_| acp::Error::internal_error())?; + rx.await.map_err(|_| acp::Error::internal_error())?; + } + Ok(acp::PromptResponse::new(acp::StopReason::EndTurn)) + } + + async fn cancel(&self, args: acp::CancelNotification) -> Result<(), acp::Error> { + log::info!("Received cancel request {args:?}"); + Ok(()) + } + + async fn set_session_mode( + &self, + args: acp::SetSessionModeRequest, + ) -> Result { + log::info!("Received set session mode request {args:?}"); + Ok(acp::SetSessionModeResponse::default()) + } + + #[cfg(feature = "unstable_session_model")] + async fn set_session_model( + &self, + args: acp::SetSessionModelRequest, + ) -> Result { + log::info!("Received select model request {args:?}"); + Ok(acp::SetSessionModelResponse::default()) + } + + async fn set_session_config_option( + &self, + args: acp::SetSessionConfigOptionRequest, + ) -> Result { + log::info!("Received set session config option request {args:?}"); + let value = args + .value + .as_value_id() + .ok_or(acp::Error::invalid_params())? + .clone(); + let option = + acp::SessionConfigOption::select(args.config_id, "Example Option", value, vec![ + acp::SessionConfigSelectOption::new("option1", "Option 1"), + acp::SessionConfigSelectOption::new("option2", "Option 2"), + ]); + Ok(acp::SetSessionConfigOptionResponse::new(vec![option])) + } + + async fn ext_method(&self, args: acp::ExtRequest) -> Result { + log::info!( + "Received extension method call: method={}, params={:?}", + args.method, + args.params + ); + Ok(acp::ExtResponse::new( + serde_json::value::to_raw_value(&json!({"example": "response"}))?.into(), + )) + } + + async fn ext_notification(&self, args: acp::ExtNotification) -> Result<(), acp::Error> { + log::info!( + "Received extension notification: method={}, params={:?}", + args.method, + args.params + ); + Ok(()) + } +} + +#[tokio::main(flavor = "current_thread")] +async fn main() -> acp::Result<()> { + env_logger::init(); + + let outgoing = tokio::io::stdout().compat_write(); + let incoming = tokio::io::stdin().compat(); + + // The AgentSideConnection will spawn futures onto our Tokio runtime. + // LocalSet and spawn_local are used because the futures from the + // agent-client-protocol crate are not Send. + let local_set = tokio::task::LocalSet::new(); + local_set + .run_until(async move { + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + // Start up the ExampleAgent connected to stdio. + let (conn, handle_io) = + acp::AgentSideConnection::new(ExampleAgent::new(tx), outgoing, incoming, |fut| { + tokio::task::spawn_local(fut); + }); + // Kick off a background task to send the ExampleAgent's session notifications to the client. + tokio::task::spawn_local(async move { + while let Some((session_notification, tx)) = rx.recv().await { + let result = conn.session_notification(session_notification).await; + if let Err(e) = result { + log::error!("{e}"); + break; + } + tx.send(()).ok(); + } + }); + // Run until stdin/stdout are closed. + handle_io.await + }) + .await +} diff --git a/vendor/agent-client-protocol/examples/client.rs b/vendor/agent-client-protocol/examples/client.rs new file mode 100644 index 0000000000..9df4b1ff7f --- /dev/null +++ b/vendor/agent-client-protocol/examples/client.rs @@ -0,0 +1,174 @@ +//! A simple ACP client for educational purposes. +//! +//! The client starts an agent as a subprocess and communicates with it over stdio. Run the client like this: +//! +//! ```bash +//! cargo run --example client -- path/to/agent --agent-arg +//! ``` +//! +//! To connect it to the example agent from this crate: +//! +//! ```bash +//! cargo build --example agent && cargo run --example client -- target/debug/examples/agent +//! ``` + +use { + agent_client_protocol::{self as acp, Agent as _}, + tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}, +}; + +struct ExampleClient {} + +#[async_trait::async_trait(?Send)] +impl acp::Client for ExampleClient { + async fn request_permission( + &self, + _args: acp::RequestPermissionRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn write_text_file( + &self, + _args: acp::WriteTextFileRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn read_text_file( + &self, + _args: acp::ReadTextFileRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn create_terminal( + &self, + _args: acp::CreateTerminalRequest, + ) -> Result { + Err(acp::Error::method_not_found()) + } + + async fn terminal_output( + &self, + _args: acp::TerminalOutputRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn release_terminal( + &self, + _args: acp::ReleaseTerminalRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn wait_for_terminal_exit( + &self, + _args: acp::WaitForTerminalExitRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn kill_terminal( + &self, + _args: acp::KillTerminalRequest, + ) -> acp::Result { + Err(acp::Error::method_not_found()) + } + + async fn session_notification( + &self, + args: acp::SessionNotification, + ) -> acp::Result<(), acp::Error> { + if let acp::SessionUpdate::AgentMessageChunk(acp::ContentChunk { content, .. }) = + args.update + { + let text = match content { + acp::ContentBlock::Text(text_content) => text_content.text, + acp::ContentBlock::Image(_) => "".into(), + acp::ContentBlock::Audio(_) => "