From 18420266bbc0ae07013b671e1d1e1778d32e0585 Mon Sep 17 00:00:00 2001 From: daniel Date: Sat, 5 Sep 2026 19:57:11 +0100 Subject: [PATCH] refactor: separate test support from production paths --- AGENTS.md | 6 + src/acp_child.rs | 75 +-- src/credentials.rs | 28 +- src/fatal.rs | 30 +- src/plugins.rs | 908 ++++++++++++++------------------ src/protocols/acp.rs | 178 ++----- src/protocols/acp/v2.rs | 852 ++++++++++++++++++++---------- src/provider/chatgpt.rs | 41 +- src/provider/mod.rs | 4 +- src/provider/openai_auth.rs | 333 +++++++----- src/provider/openrouter_auth.rs | 20 +- src/runtime.rs | 108 ++-- src/runtime/tests.rs | 217 ++++++-- src/tools/mcp.rs | 51 +- src/tui/app.rs | 457 ++++------------ src/tui/image.rs | 54 +- src/tui/markdown.rs | 49 +- src/tui/mod.rs | 165 +++++- src/tui/ui.rs | 110 ++-- 19 files changed, 1958 insertions(+), 1728 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 84f8c138..4902fdd0 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -2,4 +2,10 @@ Only when you are using Kit as your agent harness: report issues with Kit's harness at https://github.com/speakeasy-api/kit/issues. Do not report issues with other harnesses there. Do not open an issue on the user's behalf unless the user explicitly requests it; ask the user first when they have not already made that request. Follow [Reporting Kit Issues](docs/user/reporting-kit-issues.md). +## Test boundaries + +Keep test support in test-only modules or files; colocated `#[cfg(test)]` unit tests are fine. Move harmless helper constructors and accessors into child test-support modules. Do not add test-only fields, enum variants, counters, branches, or replacement implementations to production types or executable paths, or disguise instrumentation as telemetry or work statistics. + +Test behavior through real APIs. Use fakes at genuine external or domain boundaries, not spy callbacks, traits, or generics added only to assert implementation details; legitimate dependency injection is fine. Output tests do not prove bounded work: use benchmarks or existing justified iterator boundaries, not flaky wall-clock assertions or hardcoded exact implementation counts. + Do not change release versions in ordinary pull requests. Use a Conventional Commit title for every pull request. Mark a breaking change with `!` after the commit type or scope, or with a `BREAKING CHANGE:` line in the commit body. The release workflow derives the next version from commits since the latest release, advances the version files in a release commit, and applies a minor bump when any commit is breaking or a patch bump otherwise. diff --git a/src/acp_child.rs b/src/acp_child.rs index cd30eeda..91684f62 100644 --- a/src/acp_child.rs +++ b/src/acp_child.rs @@ -648,41 +648,6 @@ impl ChildSession { })? } - #[cfg(test)] - pub(crate) fn closure_probe_for_test() -> (Self, oneshot::Receiver<()>) { - let (tx, mut rx) = mpsc::channel(1); - let (closed_tx, closed_rx) = oneshot::channel(); - tokio::spawn(async move { - while rx.recv().await.is_some() {} - let _ = closed_tx.send(()); - }); - ( - Self { - tx, - session_id: "test".into(), - capabilities: agentkit_acp::AgentCapabilities::default(), - serial: Arc::new(tokio::sync::Mutex::new(())), - closed: watch::channel(false).1, - descendant_parent: None, - }, - closed_rx, - ) - } - - #[cfg(test)] - pub(crate) fn disconnected_for_test() -> Self { - let (tx, rx) = mpsc::channel(1); - drop(rx); - Self { - tx, - session_id: "test".into(), - capabilities: agentkit_acp::AgentCapabilities::default(), - serial: Arc::new(tokio::sync::Mutex::new(())), - closed: watch::channel(false).1, - descendant_parent: None, - } - } - pub async fn fork( &self, model: Option<&str>, @@ -1225,6 +1190,46 @@ fn prompt_outcome( } } +#[cfg(test)] +mod test_support { + use super::*; + + impl ChildSession { + pub(crate) fn closure_probe_for_test() -> (Self, oneshot::Receiver<()>) { + let (tx, mut rx) = mpsc::channel(1); + let (closed_tx, closed_rx) = oneshot::channel(); + tokio::spawn(async move { + while rx.recv().await.is_some() {} + let _ = closed_tx.send(()); + }); + ( + Self { + tx, + session_id: "test".into(), + capabilities: agentkit_acp::AgentCapabilities::default(), + serial: Arc::new(tokio::sync::Mutex::new(())), + closed: watch::channel(false).1, + descendant_parent: None, + }, + closed_rx, + ) + } + + pub(crate) fn disconnected_for_test() -> Self { + let (tx, rx) = mpsc::channel(1); + drop(rx); + Self { + tx, + session_id: "test".into(), + capabilities: agentkit_acp::AgentCapabilities::default(), + serial: Arc::new(tokio::sync::Mutex::new(())), + closed: watch::channel(false).1, + descendant_parent: None, + } + } + } +} + #[cfg(test)] mod tests { use serde_json::json; diff --git a/src/credentials.rs b/src/credentials.rs index 33efe168..a6cdb886 100644 --- a/src/credentials.rs +++ b/src/credentials.rs @@ -119,17 +119,6 @@ impl CredentialStorage { !matches!(self, Self::Memory) } - #[cfg(test)] - pub(crate) fn make_entry_undeletable_for_test(&self, namespace: &str, identity: &str) { - let entry = self.entry(namespace, identity); - entry.save(b"blocked").unwrap(); - let EntryBackend::Filesystem(path) = entry.backend else { - panic!("undeletable credential fixtures require filesystem storage"); - }; - fs::remove_file(&path).unwrap(); - fs::create_dir(&path).unwrap(); - } - pub(crate) async fn lock_refresh(&self) -> Result { let path = match self { Self::Memory => { @@ -584,6 +573,23 @@ fn context(prefix: &str, value: impl std::fmt::Display) -> CredentialStoreError error(format!("{prefix}: {value}")) } +#[cfg(test)] +mod test_support { + use super::*; + + impl CredentialStorage { + pub(crate) fn make_entry_undeletable_for_test(&self, namespace: &str, identity: &str) { + let entry = self.entry(namespace, identity); + entry.save(b"blocked").unwrap(); + let EntryBackend::Filesystem(path) = entry.backend else { + panic!("undeletable credential fixtures require filesystem storage"); + }; + fs::remove_file(&path).unwrap(); + fs::create_dir(&path).unwrap(); + } + } +} + #[cfg(test)] mod tests { use std::{sync::mpsc, time::Duration}; diff --git a/src/fatal.rs b/src/fatal.rs index 996e24db..ff26c3c7 100644 --- a/src/fatal.rs +++ b/src/fatal.rs @@ -432,18 +432,6 @@ fn write_default( ) } -#[cfg(test)] -fn write_in( - base: &Path, - session_id: &str, - surface: Surface, - kind: &str, - code: &str, - message: &str, -) -> Result { - write_in_with_diagnostics(base, session_id, surface, kind, code, message, None) -} - fn write_in_with_diagnostics( base: &Path, session_id: &str, @@ -556,7 +544,10 @@ fn event_order(path: &Path) -> Option<(u64, u32, u64)> { #[cfg(test)] mod tests { - use std::fs; + use std::{ + fs, + path::{Path, PathBuf}, + }; use agentkit_loop::LoopError; use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD}; @@ -566,9 +557,20 @@ mod tests { DIAGNOSTIC_MARKER, FatalRecord, H2Reason, IoClassification, MAX_DIAGNOSTIC_BYTES, MAX_RECORDS_PER_SESSION, ReqwestDiagnostics, Surface, TransportDiagnostics, TransportSource, TransportStage, bounded, classify, event_order, record_loop_error, - render_loop_error, split_diagnostics, write_in, write_in_with_diagnostics, + render_loop_error, split_diagnostics, write_in_with_diagnostics, }; + fn write_in( + base: &Path, + session_id: &str, + surface: Surface, + kind: &str, + code: &str, + message: &str, + ) -> Result { + write_in_with_diagnostics(base, session_id, surface, kind, code, message, None) + } + fn append_diagnostics(message: String, diagnostics: &TransportDiagnostics) -> String { let encoded = serde_json::to_vec(diagnostics).unwrap(); assert!(encoded.len() <= MAX_DIAGNOSTIC_BYTES); diff --git a/src/plugins.rs b/src/plugins.rs index 2c92579c..5e1bd84b 100644 --- a/src/plugins.rs +++ b/src/plugins.rs @@ -110,29 +110,10 @@ struct PluginRuntimeInner { cache_root: PathBuf, skill_cache_root: PathBuf, data_root: PathBuf, - git_mode: GitResolverMode, published: RwLock, generation_barrier: Arc>, } -#[derive(Clone)] -enum GitResolverMode { - Https, - #[cfg(test)] - Local { - repository: PathBuf, - activity: Arc, - }, -} - -#[cfg(test)] -#[derive(Default)] -struct GitActivity { - probes: std::sync::atomic::AtomicUsize, - fetches: std::sync::atomic::AtomicUsize, - archives: std::sync::atomic::AtomicUsize, -} - #[derive(Clone)] struct PublishedPlugins { resolved: Arc, @@ -307,7 +288,6 @@ impl PluginRuntime { cache_root, skill_cache_root, data_root, - git_mode: GitResolverMode::Https, published: RwLock::new(PublishedPlugins { resolved: Arc::new(initial), source_fingerprint: None, @@ -317,17 +297,6 @@ impl PluginRuntime { } } - #[cfg(test)] - fn with_local_git_mode(mut self, repository: &Path, activity: Arc) -> Self { - if let Some(inner) = Arc::get_mut(&mut self.inner) { - inner.git_mode = GitResolverMode::Local { - repository: repository.to_path_buf(), - activity, - }; - } - self - } - pub fn snapshot(&self) -> Arc { match self.inner.published.read() { Ok(published) => published.resolved.clone(), @@ -343,6 +312,14 @@ impl PluginRuntime { } pub(crate) async fn stage(&self) -> Result { + self.stage_with_git_runner(Arc::new(SystemGitRunner::default())) + .await + } + + async fn stage_with_git_runner( + &self, + runner: Arc, + ) -> Result { let contents = match crate::config_files::read_to_string(&self.inner.config_path) { Ok(contents) => contents, Err(error) if error.kind() == io::ErrorKind::NotFound => String::new(), @@ -372,8 +349,8 @@ impl PluginRuntime { &self.inner.cache_root, &self.inner.skill_cache_root, &self.inner.data_root, - self.inner.git_mode.clone(), self.published(), + runner, ) .await .map_err(bounded_diagnostic) @@ -455,8 +432,8 @@ async fn resolve_with_skill_cache( &cache_root, &skill_cache_root, &data_root, - &GitResolverMode::Https, None, + &SystemGitRunner::default(), ) }) .await @@ -487,8 +464,8 @@ async fn stage_with_skill_cache( cache_root: &Path, skill_cache_root: &Path, data_root: &Path, - git_mode: GitResolverMode, published: PublishedPlugins, + runner: Arc, ) -> Result { let configs = configs.clone(); let runtime_root = runtime_root.to_path_buf(); @@ -496,7 +473,7 @@ async fn stage_with_skill_cache( let skill_cache_root = skill_cache_root.to_path_buf(); let data_root = data_root.to_path_buf(); let staged = tokio::task::spawn_blocking(move || -> Result { - let source_plan = source_plan(&configs, &runtime_root, &cache_root, &git_mode)?; + let source_plan = source_plan(&configs, &runtime_root, &cache_root, runner.as_ref())?; if let Some(source_fingerprint) = published .source_fingerprint .filter(|fingerprint| fingerprint.candidate == source_plan.candidate_fingerprint) @@ -512,8 +489,8 @@ async fn stage_with_skill_cache( &cache_root, &skill_cache_root, &data_root, - &git_mode, Some(&source_plan), + runner.as_ref(), )?; source_plan.verify_path_fingerprints(&configs, &runtime_root)?; let source_fingerprint = source_plan.verified_fingerprint(&resolution.git_revisions)?; @@ -555,8 +532,8 @@ fn resolve_blocking( cache_root: &Path, skill_cache_root: &Path, data_root: &Path, - git_mode: &GitResolverMode, source_plan: Option<&SourcePlan>, + runner: &dyn GitRunner, ) -> Result { let mut resolved = ResolvedPlugins::default(); let mut verified_git_revisions = BTreeMap::new(); @@ -577,11 +554,11 @@ fn resolve_blocking( format!("missing staged Git revision for plugin {alias:?}") })?; let verified = - resolve_git_planned(url, subdir.as_deref(), cache_root, git_mode, planned)?; + resolve_git_planned(url, subdir.as_deref(), cache_root, planned, runner)?; verified_git_revisions.insert(alias.clone(), verified.revision); verified.root } - None => resolve_git(url, rev.as_deref(), subdir.as_deref(), cache_root)?, + None => resolve_git(url, rev.as_deref(), subdir.as_deref(), cache_root, runner)?, }, }; let plugin = load_plugin(&root).map_err(|error| { @@ -889,11 +866,6 @@ fn make_tree_writable(root: &Path) { } } -#[cfg(test)] -pub(crate) fn make_tree_writable_for_test(root: &Path) { - make_tree_writable(root); -} - #[cfg(unix)] fn read_only_permissions(metadata: &fs::Metadata) -> fs::Permissions { use std::os::unix::fs::PermissionsExt; @@ -1480,7 +1452,7 @@ fn source_plan( configs: &BTreeMap, runtime_root: &Path, cache_root: &Path, - git_mode: &GitResolverMode, + runner: &dyn GitRunner, ) -> Result { let mut fingerprint = blake3::Hasher::new(); let mut path_sources = BTreeMap::new(); @@ -1544,31 +1516,13 @@ fn source_plan( raw_oid: oid.clone(), commit_oid: oid.clone(), }, - GitRevision::Ref(_) | GitRevision::DefaultBranch => match git_mode { - GitResolverMode::Https => probe_git_revision( - OsStr::new(url.as_str()), - &sha256_text(url.as_str()), - &revision, - cache_root, - GitProtocol::Https, - &SystemGitRunner::default(), - )?, - #[cfg(test)] - GitResolverMode::Local { - repository, - activity, - } => probe_git_revision( - repository.as_os_str(), - &sha256_text(&repository.to_string_lossy()), - &revision, - cache_root, - GitProtocol::Local, - &TrackingGitRunner { - inner: SystemGitRunner::default(), - activity: activity.clone(), - }, - )?, - }, + GitRevision::Ref(_) | GitRevision::DefaultBranch => probe_git_revision( + OsStr::new(url.as_str()), + &sha256_text(url.as_str()), + &revision, + cache_root, + runner, + )?, }; fingerprint.update(&[1]); hash_fingerprint_field(&mut fingerprint, planned.raw_oid.as_bytes()); @@ -1770,13 +1724,6 @@ enum GitRevision { DefaultBranch, } -#[derive(Clone, Copy)] -enum GitProtocol { - Https, - #[cfg(test)] - Local, -} - #[derive(Debug, PartialEq, Eq)] enum GitFailure { Unavailable(io::ErrorKind), @@ -2049,44 +1996,6 @@ impl GitRunner for SystemGitRunner { } } -#[cfg(test)] -struct TrackingGitRunner { - inner: SystemGitRunner, - activity: Arc, -} - -#[cfg(test)] -impl GitRunner for TrackingGitRunner { - fn run(&self, request: GitRunRequest<'_>) -> Result, GitFailure> { - if request - .args - .iter() - .any(|argument| argument == OsStr::new("--exit-code")) - { - self.activity - .probes - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - if request - .args - .iter() - .any(|argument| argument == OsStr::new("fetch")) - { - self.activity - .fetches - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - } - self.inner.run(request) - } - - fn archive(&self, request: GitRunRequest<'_>, destination: &Path) -> Result<(), GitFailure> { - self.activity - .archives - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - self.inner.archive(request, destination) - } -} - fn read_bounded(mut reader: impl Read, limit: u64) -> Result, GitFailure> { let mut bytes = Vec::new(); reader @@ -2282,6 +2191,7 @@ fn resolve_git( rev: Option<&str>, subdir: Option<&str>, cache_root: &Path, + runner: &dyn GitRunner, ) -> Result { let url = validate_git_url(value)?; let revision = rev @@ -2297,8 +2207,7 @@ fn resolve_git( GitSourceRevision::Unplanned(&revision), subdir.as_deref(), cache_root, - GitProtocol::Https, - &SystemGitRunner::default(), + runner, ) .map(|resolved| resolved.root) } @@ -2307,85 +2216,19 @@ fn resolve_git_planned( value: &str, subdir: Option<&str>, cache_root: &Path, - mode: &GitResolverMode, planned: &PlannedGitRevision, + runner: &dyn GitRunner, ) -> Result { let url = validate_git_url(value)?; let subdir = subdir.map(validate_git_subdir).transpose()?; - match mode { - GitResolverMode::Https => resolve_git_source( - OsStr::new(url.as_str()), - &sha256_text(url.as_str()), - GitSourceRevision::Planned(planned), - subdir.as_deref(), - cache_root, - GitProtocol::Https, - &SystemGitRunner::default(), - ), - #[cfg(test)] - GitResolverMode::Local { - repository, - activity, - } => resolve_git_source( - repository.as_os_str(), - &sha256_text(&repository.to_string_lossy()), - GitSourceRevision::Planned(planned), - subdir.as_deref(), - cache_root, - GitProtocol::Local, - &TrackingGitRunner { - inner: SystemGitRunner::default(), - activity: activity.clone(), - }, - ), - } -} - -#[cfg(test)] -fn resolve_git_local( - repository: &Path, - rev: &str, - subdir: Option<&Path>, - cache_root: &Path, - runner: &dyn GitRunner, -) -> Result { - resolve_git_local_revision( - repository, - validate_git_revision(rev)?, - subdir, - cache_root, - runner, - ) -} - -#[cfg(test)] -fn resolve_git_local_revision( - repository: &Path, - revision: GitRevision, - subdir: Option<&Path>, - cache_root: &Path, - runner: &dyn GitRunner, -) -> Result { - let repository = fs::canonicalize(repository) - .map_err(|error| format!("could not resolve test Git repository: {error}"))?; - let subdir = subdir - .map(|path| { - path.to_str() - .ok_or_else(|| "plugin Git subdir must be valid Unicode".to_string()) - .and_then(validate_git_subdir) - }) - .transpose()?; - let source_key = sha256_text(&repository.to_string_lossy()); resolve_git_source( - repository.as_os_str(), - &source_key, - GitSourceRevision::Unplanned(&revision), + OsStr::new(url.as_str()), + &sha256_text(url.as_str()), + GitSourceRevision::Planned(planned), subdir.as_deref(), cache_root, - GitProtocol::Local, runner, ) - .map(|resolved| resolved.root) } fn validate_git_url(value: &str) -> Result { @@ -2488,12 +2331,7 @@ fn sha256_text(value: &str) -> String { .collect() } -fn hardened_git_args( - protocol: GitProtocol, - hooks: &Path, - attributes: &Path, - args: &[&OsStr], -) -> Vec { +fn hardened_git_args(hooks: &Path, attributes: &Path, args: &[&OsStr]) -> Vec { let mut output = vec![ OsString::from("-c"), OsString::from("protocol.allow=never"), @@ -2530,15 +2368,6 @@ fn hardened_git_args( OsString::from("-c"), OsString::from(format!("core.attributesFile={}", attributes.display())), ]; - #[cfg(test)] - if matches!(protocol, GitProtocol::Local) { - output.extend([ - OsString::from("-c"), - OsString::from("protocol.file.allow=always"), - ]); - } - #[cfg(not(test))] - let _ = protocol; output.extend(args.iter().map(|value| (*value).to_os_string())); output } @@ -2568,7 +2397,6 @@ fn hardened_git_config(remote: Option<&str>) -> Vec<(OsString, OsString)> { struct GitCommandContext<'a> { runner: &'a dyn GitRunner, - protocol: GitProtocol, hooks: &'a Path, attributes: &'a Path, remote: Option<&'a str>, @@ -2586,7 +2414,7 @@ impl GitCommandContext<'_> { require_plugin_disk(cwd)?; require_plugin_disk(self.hooks)?; require_plugin_disk(self.attributes)?; - let args = hardened_git_args(self.protocol, self.hooks, self.attributes, args); + let args = hardened_git_args(self.hooks, self.attributes, args); let config = hardened_git_config(self.remote); self.runner .run(GitRunRequest { @@ -2610,7 +2438,7 @@ impl GitCommandContext<'_> { require_plugin_disk(cwd)?; require_plugin_disk(self.hooks)?; require_plugin_disk(self.attributes)?; - let args = hardened_git_args(self.protocol, self.hooks, self.attributes, args); + let args = hardened_git_args(self.hooks, self.attributes, args); let config = hardened_git_config(self.remote); self.runner .archive( @@ -2656,7 +2484,6 @@ fn probe_git_revision( source_key: &str, revision: &GitRevision, cache_root: &Path, - protocol: GitProtocol, runner: &dyn GitRunner, ) -> Result { let source_root = cache_root.join(GIT_CACHE_VERSION).join(source_key); @@ -2676,18 +2503,13 @@ fn probe_git_revision( .map_err(|error| { format!("could not create Git plugin revision probe attributes: {error}") })?; - let remote_config = match protocol { - GitProtocol::Https => Some( - remote - .to_str() - .ok_or("plugin Git URL must be valid Unicode")?, - ), - #[cfg(test)] - GitProtocol::Local => None, - }; + let remote_config = Some( + remote + .to_str() + .ok_or("plugin Git URL must be valid Unicode")?, + ); let git = GitCommandContext { runner, - protocol, hooks: &hooks, attributes: &attributes, remote: remote_config, @@ -2815,7 +2637,6 @@ fn resolve_git_source( source_revision: GitSourceRevision<'_>, subdir: Option<&Path>, cache_root: &Path, - protocol: GitProtocol, runner: &dyn GitRunner, ) -> Result { let (revision, expected_revision) = match source_revision { @@ -2864,15 +2685,11 @@ fn resolve_git_source( } } - let remote_config = match protocol { - GitProtocol::Https => Some( - remote - .to_str() - .ok_or("plugin Git URL must be valid Unicode")?, - ), - #[cfg(test)] - GitProtocol::Local => None, - }; + let remote_config = Some( + remote + .to_str() + .ok_or("plugin Git URL must be valid Unicode")?, + ); let staging_guard = StagingDirectory::create(&source_root, ".staging-", "Git plugin")?; let staging = staging_guard.path(); let hooks = staging.join("hooks"); @@ -2888,7 +2705,6 @@ fn resolve_git_source( .map_err(|error| format!("could not create controlled Git attributes file: {error}"))?; let git = GitCommandContext { runner, - protocol, hooks: &hooks, attributes: &attributes, remote: remote_config, @@ -3737,6 +3553,18 @@ fn select_package_root(extraction: &Path, subdir: Option<&Path>) -> Result, + cache_root: &Path, + runner: &dyn GitRunner, + ) -> Result { + resolve_git_local_revision( + repository, + validate_git_revision(rev)?, + subdir, + cache_root, + runner, + ) + } + + fn resolve_git_local_revision( + repository: &Path, + revision: GitRevision, + subdir: Option<&Path>, + cache_root: &Path, + runner: &dyn GitRunner, + ) -> Result { + let repository = fs::canonicalize(repository) + .map_err(|error| format!("could not resolve test Git repository: {error}"))?; + let subdir = subdir + .map(|path| { + path.to_str() + .ok_or_else(|| "plugin Git subdir must be valid Unicode".to_string()) + .and_then(validate_git_subdir) + }) + .transpose()?; + let source_key = sha256_text(&repository.to_string_lossy()); + resolve_git_source( + repository.as_os_str(), + &source_key, + GitSourceRevision::Unplanned(&revision), + subdir.as_deref(), + cache_root, + &LocalGitRunner { inner: runner }, + ) + .map(|resolved| resolved.root) + } + + // Adapt only the external Git transport for local repository fixtures. The + // resolver and its HTTPS-only command construction remain production code. + struct LocalGitRunner<'a> { + inner: &'a dyn GitRunner, + } + + impl LocalGitRunner<'_> { + fn args(args: &[OsString]) -> Vec { + args.iter() + .map(|arg| { + if arg == "protocol.file.allow=never" { + OsString::from("protocol.file.allow=always") + } else { + arg.clone() + } + }) + .collect() + } + } + + impl GitRunner for LocalGitRunner<'_> { + fn run(&self, request: GitRunRequest<'_>) -> Result, GitFailure> { + let args = Self::args(request.args); + self.inner.run(GitRunRequest { + args: &args, + ..request + }) + } + + fn archive( + &self, + request: GitRunRequest<'_>, + destination: &Path, + ) -> Result<(), GitFailure> { + let args = Self::args(request.args); + self.inner.archive( + GitRunRequest { + args: &args, + ..request + }, + destination, + ) + } + } + + // Route the fixture HTTPS remote to a real local Git subprocess only at + // the transport boundary. URL validation and --get-url still run normally. + struct RepositoryGitRunner { + repository: PathBuf, + } + + impl GitRunner for RepositoryGitRunner { + fn run(&self, request: GitRunRequest<'_>) -> Result, GitFailure> { + let mut args = LocalGitRunner::args(request.args); + if !args.iter().any(|arg| arg == "--get-url") { + for arg in &mut args { + if arg == "https://plugins.example/repository.git" { + *arg = self.repository.as_os_str().to_owned(); + } + } + } + SystemGitRunner::default().run(GitRunRequest { + args: &args, + ..request + }) + } + + fn archive( + &self, + request: GitRunRequest<'_>, + destination: &Path, + ) -> Result<(), GitFailure> { + SystemGitRunner::default().archive(request, destination) + } + } + const MANIFEST: &str = r#"{"$schema":"https://agent-plugins.org/schemas/1.0.0/plugin.schema.json","name":"test-plugin"}"#; fn sha256_hex(bytes: &[u8]) -> String { @@ -3847,20 +3795,6 @@ mod tests { struct RewrittenUrlRunner; - struct RecordingRunner { - inner: SystemGitRunner, - calls: std::sync::Mutex>>, - } - - impl Default for RecordingRunner { - fn default() -> Self { - Self { - inner: SystemGitRunner::default(), - calls: std::sync::Mutex::new(Vec::new()), - } - } - } - impl GitRunner for TimeoutRunner { fn run(&self, _request: GitRunRequest<'_>) -> Result, GitFailure> { Err(GitFailure::Timeout) @@ -3903,22 +3837,6 @@ mod tests { } } - impl GitRunner for RecordingRunner { - fn run(&self, request: GitRunRequest<'_>) -> Result, GitFailure> { - self.calls.lock().unwrap().push(request.args.to_vec()); - self.inner.run(request) - } - - fn archive( - &self, - request: GitRunRequest<'_>, - destination: &Path, - ) -> Result<(), GitFailure> { - self.calls.lock().unwrap().push(request.args.to_vec()); - self.inner.archive(request, destination) - } - } - #[test] fn parses_plugin_sources_with_unknown_fields() { let path: PluginConfig = toml::from_str("source = 'path'\npath = './plugin'").unwrap(); @@ -4056,7 +3974,7 @@ mod tests { &configs, directory.path(), directory.path(), - &GitResolverMode::Https, + &SystemGitRunner::default(), ) .unwrap(); @@ -4089,7 +4007,7 @@ mod tests { &configs, directory.path(), directory.path(), - &GitResolverMode::Https, + &SystemGitRunner::default(), ) .unwrap(); @@ -4127,7 +4045,7 @@ mod tests { &configs, directory.path(), directory.path(), - &GitResolverMode::Https, + &SystemGitRunner::default(), ) .unwrap(); @@ -4346,27 +4264,8 @@ mod tests { assert!(Arc::ptr_eq(&published, &unchanged.resolved)); } - async fn stage_local_git_runtime( - config: PathBuf, - root: &Path, - repository: &Path, - ) -> (PluginRuntime, Arc) { - let activity = Arc::new(GitActivity::default()); - let runtime = PluginRuntime::new( - config, - root.to_path_buf(), - root.join("cache"), - root.join("data"), - ResolvedPlugins::default(), - ) - .with_local_git_mode(repository, activity.clone()); - let staged = runtime.stage().await.unwrap(); - runtime.publish(staged); - (runtime, activity) - } - #[tokio::test] - async fn full_commit_stage_reuses_published_arc_before_git_resolution() { + async fn full_commit_runtime_reuses_cached_and_published_content() { let repository = TestRepository::new(); repository.commit_file("plugin.json", MANIFEST.as_bytes(), "manifest"); let commit = repository.commit_file( @@ -4375,237 +4274,262 @@ mod tests { "skill", ); let directory = tempfile::tempdir().unwrap(); + let cache = directory.path().join("cache"); + let remote = "https://plugins.example/repository.git"; + // Populate the normal cache from a local repository, then use the real + // HTTPS-only runtime without injecting a resolver or subprocess runner. + resolve_git_source( + repository.path().as_os_str(), + &sha256_text(remote), + GitSourceRevision::Unplanned(&GitRevision::Commit(commit.clone())), + None, + &cache, + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, + ) + .unwrap(); let config = directory.path().join("config.toml"); fs::write( &config, - format!( - "[plugins.commit]\nsource = 'git'\nurl = 'https://plugins.example/repository.git'\nrev = '{commit}'\n" - ), + format!("[plugins.commit]\nsource = 'git'\nurl = '{remote}'\nrev = '{commit}'\n"), ) .unwrap(); - let (runtime, activity) = - stage_local_git_runtime(config, directory.path(), repository.path()).await; - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 1 + // No repository remains available to fetch from. + drop(repository); + let runtime = PluginRuntime::new( + config.clone(), + directory.path().to_path_buf(), + cache.clone(), + directory.path().join("data"), + ResolvedPlugins::default(), ); + let staged = runtime.stage().await.unwrap(); + assert_eq!(staged.resolved.skills[0].body, "first"); + runtime.publish(staged); let published = runtime.snapshot(); let unchanged = runtime.stage().await.unwrap(); assert!(Arc::ptr_eq(&published, &unchanged.resolved)); - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 1 + assert_eq!(unchanged.resolved.skills[0].body, "first"); + + let fresh_runtime = PluginRuntime::new( + config, + directory.path().to_path_buf(), + cache, + directory.path().join("data"), + ResolvedPlugins::default(), ); assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 1 + fresh_runtime.stage().await.unwrap().resolved.skills[0].body, + "first" ); } #[tokio::test] - async fn mutable_default_head_stage_reuses_then_refreshes_on_movement() { + async fn mutable_default_head_runtime_reuses_and_invalidates_published_skills() { + assert_mutable_runtime_generations(None).await; + } + + #[tokio::test] + async fn mutable_annotated_tag_runtime_reuses_and_invalidates_published_skills() { + assert_mutable_runtime_generations(Some("refs/tags/stable")).await; + } + + async fn assert_mutable_runtime_generations(rev: Option<&str>) { let repository = TestRepository::new(); repository.commit_file("plugin.json", MANIFEST.as_bytes(), "manifest"); repository.commit_file( "skills/live/SKILL.md", b"---\nname: live\ndescription: Live.\n---\nfirst\n", - "skill", + "first skill", ); repository.git(&["branch", "-M", "main"]); + if rev.is_some() { + repository.git(&["tag", "-a", "stable", "-m", "first tag"]); + } + let runner: Arc = Arc::new(RepositoryGitRunner { + repository: repository.path().to_path_buf(), + }); let directory = tempfile::tempdir().unwrap(); let config = directory.path().join("config.toml"); + let revision = rev + .map(|rev| format!("rev = '{rev}'\n")) + .unwrap_or_default(); fs::write( &config, - "[plugins.head]\nsource = 'git'\nurl = 'https://plugins.example/repository.git'\n", + format!( + "[plugins.live]\nsource = 'git'\nurl = 'https://plugins.example/repository.git'\n{revision}" + ), ) .unwrap(); - let (runtime, activity) = - stage_local_git_runtime(config, directory.path(), repository.path()).await; - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 1 + let runtime = PluginRuntime::new( + config, + directory.path().to_path_buf(), + directory.path().join("cache"), + directory.path().join("data"), + ResolvedPlugins::default(), ); + + let first = runtime.stage_with_git_runner(runner.clone()).await.unwrap(); + assert_eq!(first.resolved.skills[0].body, "first"); + let first_fingerprint = first.source_fingerprint; + runtime.publish(first); let published = runtime.snapshot(); - let unchanged = runtime.stage().await.unwrap(); + let unchanged = runtime.stage_with_git_runner(runner.clone()).await.unwrap(); assert!(Arc::ptr_eq(&published, &unchanged.resolved)); assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 1 + unchanged.source_fingerprint.candidate, + first_fingerprint.candidate ); assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 1 + unchanged.source_fingerprint.resolved, + first_fingerprint.resolved ); + assert_eq!(unchanged.resolved.skills[0].body, "first"); repository.commit_file( "skills/live/SKILL.md", b"---\nname: live\ndescription: Live.\n---\nsecond\n", - "move head", + "second skill", ); - let changed = runtime.stage().await.unwrap(); - assert!(!Arc::ptr_eq(&published, &changed.resolved)); - assert_eq!(changed.resolved.skills[0].body, "second"); - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 2 + if rev.is_some() { + repository.git(&["tag", "-f", "-a", "stable", "-m", "second tag"]); + } + // The config is unchanged: only the remote mutable revision moved. + let changed = runtime.stage_with_git_runner(runner.clone()).await.unwrap(); + assert_ne!( + changed.source_fingerprint.candidate, + first_fingerprint.candidate ); - assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 2 + assert_ne!( + changed.source_fingerprint.resolved, + first_fingerprint.resolved ); + assert!(!Arc::ptr_eq(&published, &changed.resolved)); + assert_eq!(changed.resolved.skills[0].body, "second"); + assert!(Arc::ptr_eq(&published, &runtime.snapshot())); + assert_eq!(published.skills[0].body, "first"); + + let second_fingerprint = changed.source_fingerprint; runtime.publish(changed); let republished = runtime.snapshot(); - let unchanged_again = runtime.stage().await.unwrap(); + assert!(!Arc::ptr_eq(&published, &republished)); + let unchanged_again = runtime.stage_with_git_runner(runner).await.unwrap(); assert!(Arc::ptr_eq(&republished, &unchanged_again.resolved)); assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 2 + unchanged_again.source_fingerprint.candidate, + second_fingerprint.candidate ); assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 2 + unchanged_again.source_fingerprint.resolved, + second_fingerprint.resolved ); + assert_eq!(unchanged_again.resolved.skills[0].body, "second"); } - #[tokio::test] - async fn shared_cache_reuses_probed_mutable_head_then_refreshes_on_movement() { - let repository = TestRepository::new(); - repository.commit_file("plugin.json", MANIFEST.as_bytes(), "manifest"); - repository.commit_file( - "skills/live/SKILL.md", - b"---\nname: live\ndescription: Live.\n---\nfirst\n", - "skill", - ); - repository.git(&["branch", "-M", "main"]); - let directory = tempfile::tempdir().unwrap(); - let config = directory.path().join("config.toml"); - fs::write( - &config, - "[plugins.head]\nsource = 'git'\nurl = 'https://plugins.example/repository.git'\n", + fn resolve_local_plan( + repository: &Path, + plan: &PlannedGitRevision, + cache: &Path, + runner: &dyn GitRunner, + ) -> ResolvedGitSource { + resolve_git_source( + repository.as_os_str(), + &sha256_text(&repository.to_string_lossy()), + GitSourceRevision::Planned(plan), + None, + cache, + &LocalGitRunner { inner: runner }, ) - .unwrap(); - let (_first_runtime, first_activity) = - stage_local_git_runtime(config.clone(), directory.path(), repository.path()).await; - assert_eq!( - first_activity - .probes - .load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - first_activity - .fetches - .load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - first_activity - .archives - .load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - - let (runtime, activity) = - stage_local_git_runtime(config, directory.path(), repository.path()).await; - assert_eq!( - activity.probes.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 0 - ); - assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 0 - ); - assert_eq!(runtime.snapshot().skills[0].body, "first"); - let published = runtime.snapshot(); + .unwrap() + } - repository.commit_file( - "skills/live/SKILL.md", - b"---\nname: live\ndescription: Live.\n---\nsecond\n", - "move head", - ); - let changed = runtime.stage().await.unwrap(); - assert!(!Arc::ptr_eq(&published, &changed.resolved)); - assert_eq!(changed.resolved.skills[0].body, "second"); - assert_eq!( - activity.probes.load(std::sync::atomic::Ordering::Relaxed), - 2 - ); - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); + fn probe_local_revision( + repository: &Path, + revision: &GitRevision, + cache: &Path, + ) -> PlannedGitRevision { + probe_git_revision( + repository.as_os_str(), + &sha256_text(&repository.to_string_lossy()), + revision, + cache, + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, + ) + .unwrap() } - #[tokio::test] - async fn shared_cache_reuses_probed_annotated_tag_without_resolution() { + #[test] + fn planned_default_head_reuses_cache_and_refreshes_content_on_movement() { let repository = TestRepository::new(); repository.commit_file("plugin.json", MANIFEST.as_bytes(), "manifest"); - let commit = repository.commit_file( - "skills/live/SKILL.md", - b"---\nname: live\ndescription: Live.\n---\nfirst\n", - "skill", - ); - repository.git(&["tag", "-a", "stable", "-m", "stable"]); - assert_ne!(repository.git(&["rev-parse", "stable"]), commit); - let directory = tempfile::tempdir().unwrap(); - let config = directory.path().join("config.toml"); - fs::write( - &config, - "[plugins.tag]\nsource = 'git'\nurl = 'https://plugins.example/repository.git'\nrev = 'refs/tags/stable'\n", - ) - .unwrap(); - let (_first_runtime, first_activity) = - stage_local_git_runtime(config.clone(), directory.path(), repository.path()).await; - assert_eq!( - first_activity - .probes - .load(std::sync::atomic::Ordering::Relaxed), - 1 + repository.commit_file("version.txt", b"first", "first"); + repository.git(&["branch", "-M", "main"]); + let cache = tempfile::tempdir().unwrap(); + let plan = + probe_local_revision(repository.path(), &GitRevision::DefaultBranch, cache.path()); + let first = resolve_local_plan( + repository.path(), + &plan, + cache.path(), + &SystemGitRunner::default(), ); - assert_eq!( - first_activity - .fetches - .load(std::sync::atomic::Ordering::Relaxed), - 1 + assert_eq!(fs::read(first.root.join("version.txt")).unwrap(), b"first"); + let unchanged = + probe_local_revision(repository.path(), &GitRevision::DefaultBranch, cache.path()); + // A valid planned cache entry remains usable if the Git process fails. + let reused = + resolve_local_plan(repository.path(), &unchanged, cache.path(), &TimeoutRunner); + assert_eq!(reused.root, first.root); + assert_eq!(fs::read(reused.root.join("version.txt")).unwrap(), b"first"); + + let second_commit = repository.commit_file("version.txt", b"second", "move head"); + let moved = + probe_local_revision(repository.path(), &GitRevision::DefaultBranch, cache.path()); + let second = resolve_local_plan( + repository.path(), + &moved, + cache.path(), + &SystemGitRunner::default(), ); + assert_eq!(second.revision.commit_oid, second_commit); + assert_ne!(second.root, first.root); assert_eq!( - first_activity - .archives - .load(std::sync::atomic::Ordering::Relaxed), - 1 + fs::read(second.root.join("version.txt")).unwrap(), + b"second" ); + assert_eq!(fs::read(first.root.join("version.txt")).unwrap(), b"first"); + let reused = resolve_local_plan(repository.path(), &moved, cache.path(), &TimeoutRunner); + assert_eq!(reused.root, second.root); + } - let (runtime, activity) = - stage_local_git_runtime(config, directory.path(), repository.path()).await; - assert_eq!( - activity.probes.load(std::sync::atomic::Ordering::Relaxed), - 1 - ); - assert_eq!( - activity.fetches.load(std::sync::atomic::Ordering::Relaxed), - 0 + #[test] + fn planned_annotated_tag_reuses_cached_commit_content() { + let repository = TestRepository::new(); + let commit = repository.commit_file("plugin.json", MANIFEST.as_bytes(), "manifest"); + repository.git(&["tag", "-a", "stable", "-m", "stable"]); + let cache = tempfile::tempdir().unwrap(); + let revision = GitRevision::Ref("refs/tags/stable".into()); + let plan = probe_local_revision(repository.path(), &revision, cache.path()); + assert_ne!(plan.raw_oid, commit); + assert_eq!(plan.commit_oid, commit); + let first = resolve_local_plan( + repository.path(), + &plan, + cache.path(), + &SystemGitRunner::default(), ); + let unchanged = probe_local_revision(repository.path(), &revision, cache.path()); + let reused = + resolve_local_plan(repository.path(), &unchanged, cache.path(), &TimeoutRunner); + assert_eq!(reused.root, first.root); + assert_eq!(reused.revision.commit_oid, commit); assert_eq!( - activity.archives.load(std::sync::atomic::Ordering::Relaxed), - 0 + fs::read_to_string(reused.root.join("plugin.json")).unwrap(), + MANIFEST ); - assert_eq!(runtime.snapshot().skills[0].body, "first"); } #[test] @@ -4625,7 +4549,7 @@ mod tests { configs, directory.path(), directory.path(), - &GitResolverMode::Https, + &SystemGitRunner::default(), ) .unwrap() .candidate_fingerprint @@ -4651,8 +4575,9 @@ mod tests { &source_key, &GitRevision::DefaultBranch, cache.path(), - GitProtocol::Local, - &SystemGitRunner::default(), + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, ) .unwrap(); let unchanged = probe_git_revision( @@ -4660,8 +4585,9 @@ mod tests { &source_key, &GitRevision::DefaultBranch, cache.path(), - GitProtocol::Local, - &SystemGitRunner::default(), + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, ) .unwrap(); assert_eq!(first, unchanged); @@ -4671,8 +4597,9 @@ mod tests { &source_key, &GitRevision::DefaultBranch, cache.path(), - GitProtocol::Local, - &SystemGitRunner::default(), + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, ) .unwrap(); assert_ne!(first, changed); @@ -4704,21 +4631,22 @@ mod tests { &source_key, &GitRevision::Ref("refs/tags/stable".into()), cache.path(), - GitProtocol::Local, - &SystemGitRunner::default(), + &LocalGitRunner { + inner: &SystemGitRunner::default(), + }, ) .unwrap(); let second = repository.commit_file("version.txt", b"second", "second"); repository.git(&["tag", "--force", "stable", &second]); - let error = resolve_git_planned( - "https://plugins.example/repository.git", + let error = resolve_git_source( + repository.path().as_os_str(), + &source_key, + GitSourceRevision::Planned(&planned), None, cache.path(), - &GitResolverMode::Local { - repository: repository.path().to_path_buf(), - activity: Arc::new(GitActivity::default()), + &LocalGitRunner { + inner: &SystemGitRunner::default(), }, - &planned, ) .unwrap_err(); assert!(error.contains("moved while")); @@ -4794,12 +4722,7 @@ mod tests { "accepted {invalid}" ); } - let args = hardened_git_args( - GitProtocol::Https, - Path::new("hooks"), - Path::new("attributes"), - &[], - ); + let args = hardened_git_args(Path::new("hooks"), Path::new("attributes"), &[]); let config = hardened_git_config(Some("https://plugins.example/repository=x.git")); assert!(config.contains(&( OsString::from("http.https://plugins.example/repository=x.git.sslVerify"), @@ -4853,12 +4776,7 @@ mod tests { OsStr::new("--get"), OsStr::new(&query), ]; - let args = hardened_git_args( - GitProtocol::Https, - directory.path(), - directory.path(), - &command_args, - ); + let args = hardened_git_args(directory.path(), directory.path(), &command_args); let config = hardened_git_config(Some(remote)); let output = SystemGitRunner { timeout: Duration::from_secs(5), @@ -5186,7 +5104,6 @@ mod tests { GitSourceRevision::Unplanned(&GitRevision::Commit("01".repeat(20))), None, cache.path(), - GitProtocol::Https, &RewrittenUrlRunner, ) .unwrap_err(); @@ -5325,7 +5242,7 @@ mod tests { } #[test] - fn forces_sha1_and_preserves_only_committed_git_attributes() { + fn preserves_only_committed_git_attributes() { let repository = TestRepository::new(); repository.commit_file( ".gitattributes", @@ -5335,49 +5252,11 @@ mod tests { repository.commit_file("plugin.json", MANIFEST.as_bytes(), "plugin"); let commit = repository.commit_file("secret.txt", b"secret", "secret"); let cache = tempfile::tempdir().unwrap(); - let runner = RecordingRunner::default(); + let runner = SystemGitRunner::default(); let root = resolve_git_local(repository.path(), &commit, None, cache.path(), &runner).unwrap(); assert!(root.join(".gitattributes").is_file()); assert!(!root.join("secret.txt").exists()); - - let calls = runner.calls.lock().unwrap(); - assert!( - calls - .iter() - .flatten() - .any(|arg| arg == "--object-format=sha1") - ); - assert_eq!( - calls - .iter() - .filter(|args| args.iter().any(|arg| arg == "--get-url")) - .count(), - 2 - ); - assert!(calls.iter().all(|args| { - args.iter() - .any(|arg| arg.to_string_lossy().starts_with("core.attributesFile=")) - })); - let fetch = calls - .iter() - .find(|args| args.iter().any(|arg| arg == "--no-write-fetch-head")) - .expect("fetch must suppress FETCH_HEAD"); - assert!( - fetch - .iter() - .any(|arg| { arg == &OsString::from(format!("{commit}:{GIT_PRIVATE_FETCH_REF}")) }) - ); - assert!(calls.iter().any(|args| { - args.iter() - .any(|arg| arg == &OsString::from(format!("{GIT_PRIVATE_FETCH_REF}^{{commit}}"))) - })); - assert!( - calls - .iter() - .flatten() - .all(|arg| arg != OsStr::new("FETCH_HEAD^{commit}")) - ); } #[test] @@ -5591,7 +5470,6 @@ mod tests { GitSourceRevision::Unplanned(&GitRevision::Ref("refs/tags/stable".into())), None, cache.path(), - GitProtocol::Https, runner, ) .unwrap_err(); diff --git a/src/protocols/acp.rs b/src/protocols/acp.rs index bf01f4a8..17b4ddb7 100644 --- a/src/protocols/acp.rs +++ b/src/protocols/acp.rs @@ -697,11 +697,6 @@ impl SessionRegistry { .await; } - #[cfg(test)] - async fn reset_authentication(&self) -> bool { - self.reset_authentication_with(async {}).await.0 - } - async fn reset_authentication_with( &self, reset: impl std::future::Future, @@ -710,12 +705,6 @@ impl SessionRegistry { .await } - #[cfg(test)] - async fn shutdown_with_timeout(&self, limit: Duration) { - self.close_sessions_with_timeout(limit, false, || async {}) - .await; - } - async fn close_sessions_with_timeout( &self, limit: Duration, @@ -1970,45 +1959,6 @@ fn record_acp_loop_failure( } } -#[cfg(test)] -#[allow(clippy::too_many_arguments)] -async fn drive_prompt( - session_id: &agentkit_acp::SessionId, - skills: &[agentkit_tool_skills::Skill], - integration: &AcpIntegration, - skill_catalog: &mut skill_catalog::SkillCatalogMonitor, - driver: &mut LoopDriver, - request: PromptRequest, - tasks: &TaskManagerHandle, - background_jobs: &BackgroundJobs, - structured_completion: bool, -) -> Result { - if structured_completion { - let _ = settle_background_jobs(tasks, background_jobs).await?; - } - background_jobs.begin_turn(); - let items = integration.input_port().prompt_to_items(&request)?; - skill_catalog - .submit(skills, items, |items| driver.submit_input(items)) - .map_err(|error| match error { - skill_catalog::SubmitError::Catalog(error) => { - record_acp_runtime_failure(session_id, "skill_catalog", error) - } - skill_catalog::SubmitError::Submit(error) => { - record_acp_loop_failure(session_id, &error) - } - })?; - drive_submitted_prompt( - session_id, - integration, - driver, - tasks, - background_jobs, - structured_completion, - ) - .await -} - #[allow(clippy::too_many_arguments)] async fn drive_runtime_prompt( session_id: &agentkit_acp::SessionId, @@ -2051,26 +2001,6 @@ async fn drive_runtime_prompt( .await } -#[cfg(test)] -async fn drive_submitted_prompt( - session_id: &agentkit_acp::SessionId, - integration: &AcpIntegration, - driver: &mut LoopDriver, - tasks: &TaskManagerHandle, - background_jobs: &BackgroundJobs, - structured_completion: bool, -) -> Result { - drive_until_pause( - session_id, - integration, - driver, - true, - structured_completion.then_some((tasks, background_jobs)), - ) - .await? - .ok_or_else(|| AcpRuntimeError::Loop("prompt ended without a response".into())) -} - async fn drive_unsolicited( session_id: &agentkit_acp::SessionId, integration: &AcpIntegration, @@ -2088,35 +2018,6 @@ async fn drive_unsolicited( .map(|_| ()) } -#[cfg(test)] -async fn drive_autonomous( - session_id: &agentkit_acp::SessionId, - integration: &AcpIntegration, - driver: &mut LoopDriver, -) -> Result<(), AcpRuntimeError> { - let _ = drive_until_pause(session_id, integration, driver, false, None).await?; - Ok(()) -} - -#[cfg(test)] -async fn drive_until_pause( - session_id: &agentkit_acp::SessionId, - integration: &AcpIntegration, - driver: &mut LoopDriver, - answer_prompt: bool, - structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, -) -> Result, AcpRuntimeError> { - let reason = - drive_finalized(session_id, integration, driver, answer_prompt, structured).await?; - if answer_prompt { - Ok(Some(PromptResponse::new( - agentkit_acp::finish_reason_to_stop_reason(&reason)?, - ))) - } else { - Ok(None) - } -} - async fn drive_finalized( session_id: &agentkit_acp::SessionId, integration: &AcpIntegration, @@ -2284,20 +2185,6 @@ pub async fn serve_with_registry( result } -#[cfg(test)] -async fn serve_transport( - runtime: Arc, - transport: impl ConnectTo + 'static, -) -> Result<(), AcpRuntimeError> { - let registry = SessionRegistry::new(); - let result = component(runtime, registry.clone())? - .connect_to(transport) - .await - .map_err(|error| AcpRuntimeError::Sdk(error.to_string())); - registry.shutdown().await; - result -} - pub(crate) fn http_router(runtime: Arc, registry: SessionRegistry) -> axum::Router { agent_client_protocol_http::AcpHttpServer::new(move || { let v1 = component(Arc::clone(&runtime), registry.clone()) @@ -2628,6 +2515,22 @@ async fn drain_client_messages( } } +#[cfg(test)] +mod test_support { + use super::*; + + impl SessionRegistry { + pub(super) async fn reset_authentication(&self) -> bool { + self.reset_authentication_with(async {}).await.0 + } + + pub(super) async fn shutdown_with_timeout(&self, limit: Duration) { + self.close_sessions_with_timeout(limit, false, || async {}) + .await; + } + } +} + #[cfg(test)] pub(super) mod tests { use std::{ @@ -2661,6 +2564,19 @@ pub(super) mod tests { use super::*; use agentkit_acp::StopReason; + async fn serve_transport( + runtime: Arc, + transport: impl ConnectTo + 'static, + ) -> Result<(), AcpRuntimeError> { + let registry = SessionRegistry::new(); + let result = component(runtime, registry.clone())? + .connect_to(transport) + .await + .map_err(|error| AcpRuntimeError::Sdk(error.to_string())); + registry.shutdown().await; + result + } + fn test_activity( id: agentkit_acp::SessionId, tx: mpsc::UnboundedSender, @@ -3988,9 +3904,11 @@ pub(super) mod tests { let tasks = task_manager.handle(); let background_jobs = BackgroundJobs::default(); let mut skill_catalog = skill_catalog::SkillCatalogMonitor::new(&[]).unwrap(); - let response = drive_prompt( + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let response = drive_runtime_prompt( &acp_session_id, - &[], + &runtime, &integration, &mut skill_catalog, &mut driver, @@ -4007,7 +3925,10 @@ pub(super) mod tests { .await .expect("cancellation must be an ACP response, not an RPC error"); - assert_eq!(response.stop_reason, StopReason::Cancelled); + assert_eq!( + agentkit_acp::finish_reason_to_stop_reason(&response).unwrap(), + StopReason::Cancelled + ); drain.abort(); } @@ -4196,9 +4117,11 @@ pub(super) mod tests { agentkit_acp::TextContent::new("start one background call"), )], ); - let prompt = drive_prompt( + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); + let prompt = drive_runtime_prompt( &acp_session_id, - &[], + &runtime, &integration, &mut skill_catalog, &mut driver, @@ -4216,12 +4139,12 @@ pub(super) mod tests { while !entered.load(Ordering::SeqCst) { tokio::task::yield_now().await; } + background_jobs.register_foreground_for_test("background-call"); assert!(detach_compose_call( &tasks, &background_jobs, "background-call", ).await); - background_jobs.register_foreground_for_test("background-call"); while turns.load(Ordering::SeqCst) < 2 { tokio::task::yield_now().await; } @@ -4250,7 +4173,10 @@ pub(super) mod tests { .await .expect("structured prompt did not synthesize the background result") .unwrap(); - assert_eq!(response.stop_reason, StopReason::EndTurn); + assert_eq!( + agentkit_acp::finish_reason_to_stop_reason(&response).unwrap(), + StopReason::EndTurn + ); assert_eq!(turns.load(Ordering::SeqCst), 3); drain.abort(); } @@ -4332,12 +4258,14 @@ pub(super) mod tests { driver .submit_input(vec![Item::text(ItemKind::User, "foreground")]) .unwrap(); - let response = drive_until_pause(&session_id, &integration, &mut driver, true, None) + let response = drive_finalized(&session_id, &integration, &mut driver, true, None) .await - .unwrap() .unwrap(); activity.settle(None, None).unwrap(); - assert_eq!(response.stop_reason, StopReason::EndTurn); + assert_eq!( + agentkit_acp::finish_reason_to_stop_reason(&response).unwrap(), + StopReason::EndTurn + ); assert!(states.try_recv().is_err()); // Multiple logical turns drained within one autonomous interval do not @@ -4345,7 +4273,7 @@ pub(super) mod tests { activity.begin(activity::ExecutionOrigin::Autonomous); for text in ["first continuation", "second continuation"] { driver.submit_input(vec![Item::notification(text)]).unwrap(); - drive_autonomous(&session_id, &integration, &mut driver) + drive_finalized(&session_id, &integration, &mut driver, false, None) .await .unwrap(); } @@ -4483,8 +4411,8 @@ pub(super) mod tests { while !entered.load(Ordering::SeqCst) { tokio::task::yield_now().await; } - assert!(detach_compose_call(&tasks, &background_jobs, "background-call").await); background_jobs.register_foreground_for_test("background-call"); + assert!(detach_compose_call(&tasks, &background_jobs, "background-call").await); assert!(background_jobs.is_detached_for_test("background-call")); timeout(Duration::from_secs(1), reply_rx) .await diff --git a/src/protocols/acp/v2.rs b/src/protocols/acp/v2.rs index 548f8949..7488e195 100644 --- a/src/protocols/acp/v2.rs +++ b/src/protocols/acp/v2.rs @@ -55,18 +55,6 @@ fn available_commands_update(session_id: wire::SessionId) -> wire::UpdateSession ) } -fn complete_new_session( - response: wire::NewSessionResponse, - activation: oneshot::Sender<()>, - respond: impl FnOnce(wire::NewSessionResponse) -> Result<(), E>, - notify: impl FnOnce(wire::UpdateSessionNotification) -> Result<(), E>, -) -> Result<(), E> { - let session_id = response.session_id.clone(); - respond(response)?; - let _ = activation.send(()); - notify(available_commands_update(session_id)) -} - fn sdk_error(error: AcpRuntimeError) -> agent_client_protocol::Error { let detail = error.to_string(); match authentication_method_id(&detail) { @@ -1150,7 +1138,7 @@ async fn session_actor(actor: SessionActor) Some(Command::Prompt(command)) => { let result = prepare_prompt( &session_id, - PromptSkillSource::Runtime(&runtime), + &runtime, &integration, &handle, &mut skill_catalog, @@ -1236,16 +1224,10 @@ async fn session_actor(actor: SessionActor) } } -enum PromptSkillSource<'a> { - #[cfg(test)] - Static(&'a [agentkit_tool_skills::Skill]), - Runtime(&'a Arc), -} - #[allow(clippy::too_many_arguments)] async fn prepare_prompt( session_id: &wire::SessionId, - skill_source: PromptSkillSource<'_>, + runtime: &Arc, integration: &AcpIntegration, handle: &AcpSessionHandle, skill_catalog: &mut skill_catalog::SkillCatalogMonitor, @@ -1269,26 +1251,18 @@ async fn prepare_prompt( let _ = reply.send(Err(error)); return Ok(()); } - let mut current = None; - let skills = match skill_source { - #[cfg(test)] - PromptSkillSource::Static(skills) => skills, - PromptSkillSource::Runtime(runtime) => { - let loaded = match runtime.current_skills().await { - Ok(current) => current, - Err(error) => { - handle.stop_injection_turn(); - let _ = reply.send(Err(AcpRuntimeError::Loop(error))); - return Ok(()); - } - }; - ¤t.insert(loaded).skills + let current = match runtime.current_skills().await { + Ok(current) => current, + Err(error) => { + handle.stop_injection_turn(); + let _ = reply.send(Err(AcpRuntimeError::Loop(error))); + return Ok(()); } }; background_jobs.begin_turn(); let prepared = integration.prompt_to_items(&request).and_then(|items| { skill_catalog - .submit(skills, items, |items| driver.submit_input(items)) + .submit(¤t.skills, items, |items| driver.submit_input(items)) .map_err(|error| match error { skill_catalog::SubmitError::Catalog(error) => { AcpRuntimeError::Loop(format!("skill catalog error: {error}")) @@ -1339,51 +1313,20 @@ async fn prepare_prompt( result } -#[async_trait] -trait TurnControl: Sync { - fn stop_injection_turn(&self); - fn is_cancelled_since(&self, generation: u64) -> bool; - async fn handle_injection_boundary( - &self, - driver: &mut LoopDriver, - terminal: bool, - ) -> Result; -} - -#[async_trait] -impl TurnControl for AcpSessionHandle { - fn stop_injection_turn(&self) { - AcpSessionHandle::stop_injection_turn(self); - } - - fn is_cancelled_since(&self, generation: u64) -> bool { - self.cancellation_handle().is_cancelled_since(generation) - } - - async fn handle_injection_boundary( - &self, - driver: &mut LoopDriver, - terminal: bool, - ) -> Result { - AcpSessionHandle::handle_injection_boundary(self, driver, terminal).await - } -} - -async fn drive_prompt( +async fn drive_prompt( session_id: &wire::SessionId, driver: &mut LoopDriver, - control: &C, + handle: &AcpSessionHandle, cancellation_generation: u64, structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, ) -> Result where S: ModelSession + Send + 'static, - C: TurnControl, { let result = drive_prompt_inner( session_id, driver, - control, + handle, cancellation_generation, structured, ) @@ -1399,29 +1342,34 @@ where result } -async fn drive_prompt_inner( +async fn drive_prompt_inner( session_id: &wire::SessionId, driver: &mut LoopDriver, - control: &C, + handle: &AcpSessionHandle, cancellation_generation: u64, structured: Option<(&TaskManagerHandle, &BackgroundJobs)>, ) -> Result where S: ModelSession + Send + 'static, - C: TurnControl, { loop { let step = match driver.next().await { Ok(step) => step, Err(error) => { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return loop_error_stop_reason(session_id, &error); } }; - if control.is_cancelled_since(cancellation_generation) { + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } match step { @@ -1430,8 +1378,11 @@ where continue; } if result.finish_reason == FinishReason::Error { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return Err(AcpRuntimeError::Loop("model turn failed".into())); @@ -1441,7 +1392,7 @@ where { continue; } - match control.handle_injection_boundary(driver, true).await { + match handle.handle_injection_boundary(driver, true).await { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; } @@ -1452,8 +1403,11 @@ where return Ok(result.finish_reason); } Err(error) => { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return Err(error); @@ -1466,7 +1420,7 @@ where { continue; } - match control.handle_injection_boundary(driver, true).await { + match handle.handle_injection_boundary(driver, true).await { Ok(AcpInjectionBoundary::Delivered | AcpInjectionBoundary::Continue) => { continue; } @@ -1477,8 +1431,11 @@ where return Ok(FinishReason::Completed); } Err(error) => { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return Err(error); @@ -1486,13 +1443,16 @@ where } } LoopStep::Interrupt(LoopInterrupt::AfterToolResult(_)) => { - match control.handle_injection_boundary(driver, false).await { + match handle.handle_injection_boundary(driver, false).await { Ok(AcpInjectionBoundary::Stopped) => { return Ok(FinishReason::Cancelled); } Err(error) => { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return Err(error); @@ -1502,8 +1462,11 @@ where } LoopStep::Interrupt(LoopInterrupt::ApprovalRequest(_)) => { if let Err(error) = driver.cancel_pending_approvals().await { - control.stop_injection_turn(); - if control.is_cancelled_since(cancellation_generation) { + handle.stop_injection_turn(); + if handle + .cancellation_handle() + .is_cancelled_since(cancellation_generation) + { return Ok(FinishReason::Cancelled); } return loop_error_stop_reason(session_id, &error); @@ -1965,12 +1928,12 @@ pub(crate) fn component( let connection = cx.clone(); cx.spawn(async move { match state.new_session(request, connection.clone()).await { - Ok((response, activation)) => complete_new_session( - response, - activation, - |response| responder.respond(response), - |notification| connection.send_notification(notification), - ), + Ok((response, activation)) => { + let session_id = response.session_id.clone(); + responder.respond(response)?; + let _ = activation.send(()); + connection.send_notification(available_commands_update(session_id)) + } Err(error) => responder.respond_with_error(sdk_error(error)), } })?; @@ -2729,7 +2692,7 @@ mod tests { _cancellation: Option, ) -> Result { self.turns.fetch_add(1, Ordering::Relaxed); - if let Some(handle) = &self.interrupt { + if let Some(handle) = self.interrupt.take() { handle.interrupt(); } match self.outcome { @@ -2805,46 +2768,6 @@ mod tests { } } - struct TestTurnControl { - pending_steer: AtomicBool, - boundaries: AtomicU64, - stops: AtomicU64, - } - - impl TestTurnControl { - fn new(pending_steer: bool) -> Self { - Self { - pending_steer: AtomicBool::new(pending_steer), - boundaries: AtomicU64::new(0), - stops: AtomicU64::new(0), - } - } - } - - #[async_trait] - impl TurnControl for TestTurnControl { - fn stop_injection_turn(&self) { - self.stops.fetch_add(1, Ordering::Relaxed); - } - - fn is_cancelled_since(&self, _generation: u64) -> bool { - false - } - - async fn handle_injection_boundary( - &self, - _driver: &mut LoopDriver, - _terminal: bool, - ) -> Result { - self.boundaries.fetch_add(1, Ordering::Relaxed); - if self.pending_steer.swap(false, Ordering::Relaxed) { - Ok(AcpInjectionBoundary::Delivered) - } else { - Ok(AcpInjectionBoundary::Finished) - } - } - } - async fn test_driver( outcome: TestOutcome, session_id: &str, @@ -2923,33 +2846,96 @@ mod tests { )); } - #[test] - fn new_session_response_is_enqueued_before_activation_and_notifications() { - let (activation, activated) = oneshot::channel(); - let activated = std::cell::RefCell::new(activated); - let events = std::cell::RefCell::new(Vec::new()); - let response = wire::NewSessionResponse::new(wire::SessionId::new("session")); + fn send_wire( + channel: &agent_client_protocol::Channel, + method: &str, + id: i64, + params: serde_json::Value, + ) { + channel + .tx + .unbounded_send(agent_client_protocol::TransportFrame::Single( + agent_client_protocol::RawJsonRpcMessage::request(method.into(), params, id.into()) + .unwrap(), + )) + .unwrap(); + } - complete_new_session( - response, - activation, - |_| { - assert!(matches!( - activated.borrow_mut().try_recv(), - Err(oneshot::error::TryRecvError::Empty) - )); - events.borrow_mut().push("response"); - Ok::<(), ()>(()) - }, - |_| { - assert_eq!(activated.borrow_mut().try_recv(), Ok(())); - events.borrow_mut().push("notification"); - Ok::<(), ()>(()) - }, + async fn receive_wire(channel: &mut agent_client_protocol::Channel) -> serde_json::Value { + use futures_util::StreamExt; + let frame = timeout(Duration::from_secs(2), channel.rx.next()) + .await + .expect("ACP frame timed out") + .expect("ACP transport closed"); + let agent_client_protocol::TransportFrame::Single(message) = frame else { + panic!("expected a single ACP message, got {frame:?}"); + }; + serde_json::to_value(message).unwrap() + } + + #[tokio::test] + async fn new_session_response_precedes_notifications_and_activates_actor() { + let root = tempfile::tempdir().unwrap(); + let credentials = crate::credentials::CredentialStorage::Memory; + crate::provider::store_openrouter_test_credentials(&credentials); + let runtime = Runtime::new_with_provider_credentials_and_effort( + root.path(), + "test-model", + crate::ProviderKind::OpenRouter, + credentials, + None, ) .unwrap(); + let (mut client, agent) = agent_client_protocol::Channel::duplex(); + let router = v2_router(runtime, SessionRegistry::new()).unwrap(); + let server = tokio::spawn(async move { router.connect_to(agent).await }); + send_wire( + &client, + "initialize", + 1, + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("ordering-test", "0"), + )) + .unwrap(), + ); + assert_eq!(receive_wire(&mut client).await["id"], 1); + send_wire( + &client, + "session/new", + 2, + serde_json::to_value(wire::NewSessionRequest::new(root.path().to_path_buf())).unwrap(), + ); - assert_eq!(events.into_inner(), ["response", "notification"]); + let response = receive_wire(&mut client).await; + assert_eq!( + response["id"], 2, + "session response must be the first frame: {response}" + ); + let response: wire::NewSessionResponse = + serde_json::from_value(response["result"].clone()).unwrap(); + let notification = receive_wire(&mut client).await; + assert_eq!(notification["method"], "session/update"); + let notification: wire::UpdateSessionNotification = + serde_json::from_value(notification["params"].clone()).unwrap(); + assert_eq!(notification.session_id, response.session_id); + assert!(matches!( + notification.update, + wire::SessionUpdate::AvailableCommandsUpdate(_) + )); + + // Closing requires the actual session actor to process and acknowledge a command. + send_wire( + &client, + "session/close", + 3, + serde_json::to_value(wire::CloseSessionRequest::new(response.session_id)).unwrap(), + ); + let closed = receive_wire(&mut client).await; + assert_eq!(closed["id"], 3); + assert!(closed.get("result").is_some(), "close failed: {closed}"); + server.abort(); + let _ = server.await; } #[test] @@ -2975,6 +2961,8 @@ mod tests { #[tokio::test] async fn foreground_provider_error_after_running_terminalizes_once() { + let root = tempfile::tempdir().unwrap(); + let runtime = Runtime::new(root.path(), "gpt-5.4").unwrap(); let integration = AcpIntegration::default(); let recording = RecordingSink::default(); let sink = ResponseReplacementSink::new(recording.clone()); @@ -3033,7 +3021,7 @@ mod tests { let (result, ()) = tokio::join!( prepare_prompt( &session_id, - PromptSkillSource::Static(&[]), + &runtime, &integration, &handle, &mut skill_catalog, @@ -3163,12 +3151,12 @@ mod tests { while !entered.load(Ordering::SeqCst) { tokio::task::yield_now().await; } + background_jobs.register_foreground_for_test("background-call"); assert!(super::super::detach_compose_call( &tasks, &background_jobs, "background-call", ).await); - background_jobs.register_foreground_for_test("background-call"); while turns.load(Ordering::SeqCst) < 2 { tokio::task::yield_now().await; } @@ -3218,45 +3206,167 @@ mod tests { handle.stop_injection_turn(); } - struct BoundaryCancellationControl { - before_boundary: bool, - boundaries: AtomicU64, + fn bind_test_session( + integration: &AcpIntegration, + session_id: &wire::SessionId, + sink: RecordingSink, + ) -> AcpSessionHandle { + let handle = integration + .bind_session(AcpSessionBinding::new( + session_id.clone(), + SessionId::new(session_id.to_string()), + sink, + )) + .unwrap(); + handle.prepare_injection_turn(); + handle.start_injection_turn(); + handle } - #[async_trait] - impl TurnControl for BoundaryCancellationControl { - fn stop_injection_turn(&self) {} - - fn is_cancelled_since(&self, _generation: u64) -> bool { - self.before_boundary - } + // Hold a real response-committed receipt at the external ACP transport boundary. + // Until activation, the real injection coordinator must wait rather than lose the steer. + async fn staged_injection( + integration: AcpIntegration, + session_id: &wire::SessionId, + ) -> ( + agentkit_acp::v2::AcpInjectAcceptance, + agent_client_protocol::Channel, + tokio::task::JoinHandle>, + ) { + let (mut client, agent) = agent_client_protocol::Channel::duplex(); + let (accepted, mut acceptance) = mpsc::unbounded_channel(); + let server = tokio::spawn(async move { + agent_client_protocol::Agent + .v2() + .on_receive_request( + async move |request: wire::InitializeRequest, responder, _cx| { + responder.respond( + wire::InitializeResponse::new( + request.protocol_version, + wire::Implementation::new("injection-test", "0"), + ) + .capabilities(agentkit_acp::v2::agent_capabilities()), + ) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |request: wire::InjectSessionRequest, responder, cx| { + let integration = integration.clone(); + let accepted = accepted.clone(); + cx.spawn(async move { + let reserved = integration + .reserve_inject_request(request, responder) + .await? + .expect("inject reservation"); + let receipt = reserved.respond_tracked()?.expect("inject acceptance"); + accepted.send(receipt).ok(); + Ok(()) + })?; + Ok(()) + }, + agent_client_protocol::on_receive_request!(), + ) + .connect_to(agent) + .await + }); + send_wire( + &client, + "initialize", + 1, + serde_json::to_value(wire::InitializeRequest::new( + wire::ProtocolVersion::V2, + wire::Implementation::new("injection-test", "0"), + )) + .unwrap(), + ); + assert!(receive_wire(&mut client).await.get("result").is_some()); + send_wire( + &client, + "session/inject", + 2, + serde_json::to_value(wire::InjectSessionRequest::new( + session_id.clone(), + wire::SessionInjectMode::Steer, + vec![wire::ContentBlock::Text(wire::TextContent::new( + "pending steer", + ))], + )) + .unwrap(), + ); + let response = receive_wire(&mut client).await; + assert!( + response.get("result").is_some(), + "injection failed: {response}" + ); + let receipt = timeout(Duration::from_secs(2), acceptance.recv()) + .await + .unwrap() + .unwrap(); + (receipt, client, server) + } - async fn handle_injection_boundary( - &self, - _driver: &mut LoopDriver, - terminal: bool, - ) -> Result { - assert!(!terminal, "cancellation must occur at AfterToolResult"); - self.boundaries.fetch_add(1, Ordering::Relaxed); - Ok(AcpInjectionBoundary::Stopped) - } + #[tokio::test] + async fn cancellation_before_tool_boundary_retires_turn() { + let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("boundary-cancel"); + let handle = bind_test_session(&integration, &session_id, RecordingSink::default()); + let generation = handle.cancellation_handle().generation(); + let (mut driver, _) = test_driver_with_interrupt( + TestOutcome::ToolThenContent, + "boundary-cancel", + Some(handle.clone()), + ) + .await; + driver + .submit_input(vec![Item::text(ItemKind::User, "first")]) + .unwrap(); + assert_eq!( + drive_prompt(&session_id, &mut driver, &handle, generation, None) + .await + .unwrap(), + FinishReason::Cancelled + ); + assert!( + driver + .snapshot() + .transcript + .iter() + .any(|item| item.kind == ItemKind::Tool) + ); + handle.prepare_injection_turn(); + handle.start_injection_turn(); + driver + .submit_input(vec![Item::text(ItemKind::User, "fresh")]) + .unwrap(); + let generation = handle.cancellation_handle().generation(); + assert_eq!( + drive_prompt(&session_id, &mut driver, &handle, generation, None) + .await + .unwrap(), + FinishReason::Completed + ); + assert!( + driver + .snapshot() + .transcript + .iter() + .any(|item| item.kind == ItemKind::Assistant + && item.parts.iter().any( + |part| matches!(part, Part::Text(text) if text.text == "autonomous content") + )) + ); } - async fn assert_boundary_cancellation_retires_turn(before_boundary: bool) { + async fn assert_cancellation_preserves_committed_injection(outcome: TestOutcome) { let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("injection-cancel"); let recording = RecordingSink::default(); + let handle = bind_test_session(&integration, &session_id, recording.clone()); + let (receipt, _client, server) = staged_injection(integration.clone(), &session_id).await; + let message_id = receipt.message_id().clone(); let sink = ResponseReplacementSink::new(recording.clone()); - let session_id = wire::SessionId::new("boundary-cancel"); - let loop_session_id = SessionId::new("boundary-cancel-loop"); let activity = native_activity(session_id.clone(), sink.clone()); - let _handle = integration - .bind_session(AcpSessionBinding::new( - session_id.clone(), - loop_session_id.clone(), - sink.clone(), - )) - .unwrap(); - let turns = Arc::new(AtomicU64::new(0)); let observer = ResponseReplacementObserver::new( integration, sink, @@ -3265,113 +3375,173 @@ mod tests { ); let mut driver = Agent::builder() .model(TestAdapter { - outcome: TestOutcome::ToolThenContent, - turns: turns.clone(), + outcome, + turns: Arc::new(AtomicU64::new(0)), interrupt: None, }) .observer(observer) .build() .unwrap() - .start(SessionConfig::new(loop_session_id).without_cache()) + .start(SessionConfig::new(SessionId::new("injection-cancel")).without_cache()) .await .unwrap(); driver .submit_input(vec![Item::text(ItemKind::User, "first")]) .unwrap(); - let control = BoundaryCancellationControl { - before_boundary, - boundaries: AtomicU64::new(0), - }; - let reason = activity - .execute( + let generation = handle.cancellation_handle().generation(); + { + let turn = activity.execute( ExecutionOrigin::Prompt, - drive_prompt(&session_id, &mut driver, &control, 0, None), + drive_prompt(&session_id, &mut driver, &handle, generation, None), |reason| Some(reason.clone()), - ) - .await - .unwrap(); - assert_eq!(reason, FinishReason::Cancelled); - assert_eq!( - turns.load(Ordering::Relaxed), - 1, - "must not resume cancelled model work" - ); - assert_eq!( - control.boundaries.load(Ordering::Relaxed), - u64::from(!before_boundary) - ); + ); + tokio::pin!(turn); + assert!( + futures_util::poll!(&mut turn).is_pending(), + "unactivated injection must hold the turn open" + ); + handle.interrupt(); + assert_eq!( + timeout(Duration::from_secs(2), &mut turn) + .await + .unwrap() + .unwrap(), + FinishReason::Cancelled + ); + } + let updates = recording.updates.lock().unwrap().clone(); assert!( - driver - .snapshot() - .transcript + !updates .iter() - .any(|item| item.kind == ItemKind::Tool) - ); - assert_running_then_idle( - &recording.updates.lock().unwrap(), - wire::StopReason::Cancelled, + .any(|update| matches!(update.update, wire::SessionUpdate::UserMessage(_))) ); + let states = updates + .into_iter() + .filter(|update| matches!(update.update, wire::SessionUpdate::StateUpdate(_))) + .collect::>(); + assert_running_then_idle(&states, wire::StopReason::Cancelled); recording.updates.lock().unwrap().clear(); - + receipt.activate_after_response().await.unwrap(); + handle.prepare_injection_turn(); + handle.start_injection_turn(); driver .submit_input(vec![Item::text(ItemKind::User, "fresh")]) .unwrap(); - activity - .execute( - ExecutionOrigin::Prompt, - drive_prompt( - &session_id, - &mut driver, - &TestTurnControl::new(false), - 0, - None, - ), - |reason| Some(reason.clone()), - ) - .await + let generation = handle.cancellation_handle().generation(); + assert_eq!( + activity + .execute( + ExecutionOrigin::Prompt, + drive_prompt(&session_id, &mut driver, &handle, generation, None), + |reason| Some(reason.clone()) + ) + .await + .unwrap(), + FinishReason::Completed + ); + let updates = recording.updates.lock().unwrap().clone(); + assert!(updates.iter().any(|update| matches!(&update.update, + wire::SessionUpdate::UserMessage(message) if message.message_id == message_id))); + let states = updates + .into_iter() + .filter(|update| matches!(update.update, wire::SessionUpdate::StateUpdate(_))) + .collect::>(); + assert_running_then_idle(&states, wire::StopReason::EndTurn); + let transcript = driver.snapshot().transcript; + for expected in ["fresh", "pending steer"] { + assert!(transcript.iter().any(|item| { + item.kind == ItemKind::User + && item + .parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text == expected)) + })); + } + server.abort(); + let _ = server.await; + } + + #[tokio::test] + async fn terminal_injection_waits_for_response_activation_and_delivers_input() { + let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("terminal-inject"); + let recording = RecordingSink::default(); + let handle = bind_test_session(&integration, &session_id, recording.clone()); + let (receipt, _client, server) = staged_injection(integration, &session_id).await; + let message_id = receipt.message_id().clone(); + let (mut driver, _) = test_driver(TestOutcome::Content, "terminal-inject").await; + driver + .submit_input(vec![Item::text(ItemKind::User, "first")]) .unwrap(); - assert_eq!(turns.load(Ordering::Relaxed), 2); - assert_running_then_idle( - &recording.updates.lock().unwrap(), - wire::StopReason::EndTurn, + let generation = handle.cancellation_handle().generation(); + { + let turn = drive_prompt(&session_id, &mut driver, &handle, generation, None); + tokio::pin!(turn); + assert!(futures_util::poll!(&mut turn).is_pending()); + receipt.activate_after_response().await.unwrap(); + assert_eq!( + timeout(Duration::from_secs(2), &mut turn) + .await + .unwrap() + .unwrap(), + FinishReason::Completed + ); + } + assert!( + recording + .updates + .lock() + .unwrap() + .iter() + .any(|update| matches!(&update.update, + wire::SessionUpdate::UserMessage(message) if message.message_id == message_id)) ); + assert!(driver.snapshot().transcript.iter().any(|item| { + item.kind == ItemKind::User + && item + .parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text == "pending steer")) + })); + server.abort(); + let _ = server.await; } #[tokio::test] - async fn cancellation_before_tool_boundary_retires_turn() { - assert_boundary_cancellation_retires_turn(true).await; + async fn cancellation_within_tool_boundary_preserves_committed_injection() { + assert_cancellation_preserves_committed_injection(TestOutcome::ToolThenContent).await; } #[tokio::test] - async fn cancellation_within_tool_boundary_retires_turn() { - assert_boundary_cancellation_retires_turn(false).await; + async fn cancellation_at_terminal_injection_boundary_preserves_committed_injection() { + assert_cancellation_preserves_committed_injection(TestOutcome::Content).await; } #[tokio::test] async fn finish_error_stops_before_delivering_pending_steer() { - let (mut driver, turns) = test_driver(TestOutcome::FinishError, "finish-error").await; + let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("finish-error"); + let recording = RecordingSink::default(); + let handle = bind_test_session(&integration, &session_id, recording.clone()); + let (receipt, _client, server) = staged_injection(integration, &session_id).await; + let (mut driver, _) = test_driver(TestOutcome::FinishError, "finish-error").await; driver .submit_input(vec![Item::text(ItemKind::User, "fail")]) .unwrap(); - let control = TestTurnControl::new(true); - - let result = drive_prompt( - &wire::SessionId::new("finish-error"), - &mut driver, - &control, - 0, - None, - ) - .await; - - assert!(matches!( - result, - Err(AcpRuntimeError::Loop(message)) if message == "model turn failed" - )); - assert_eq!(turns.load(Ordering::Relaxed), 1); - assert_eq!(control.boundaries.load(Ordering::Relaxed), 0); - assert!(control.pending_steer.load(Ordering::Relaxed)); - assert_eq!(control.stops.load(Ordering::Relaxed), 1); + let generation = handle.cancellation_handle().generation(); + let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; + assert!( + matches!(result, Err(AcpRuntimeError::Loop(message)) if message == "model turn failed") + ); + assert!(recording.updates.lock().unwrap().is_empty()); + assert!(!driver.snapshot().transcript.iter().any(|item| { + item.parts + .iter() + .any(|part| matches!(part, Part::Text(text) if text.text == "pending steer")) + })); + drop(receipt); + server.abort(); + let _ = server.await; } #[tokio::test] @@ -3403,23 +3573,16 @@ mod tests { #[tokio::test] async fn provider_error_without_cancellation_remains_an_error() { + let integration = AcpIntegration::default(); + let session_id = wire::SessionId::new("provider-error"); + let handle = bind_test_session(&integration, &session_id, RecordingSink::default()); let (mut driver, _) = test_driver(TestOutcome::ProviderError, "provider-error").await; driver .submit_input(vec![Item::text(ItemKind::User, "fail")]) .unwrap(); - let control = TestTurnControl::new(false); - - let result = drive_prompt( - &wire::SessionId::new("provider-error"), - &mut driver, - &control, - 0, - None, - ) - .await; - + let generation = handle.cancellation_handle().generation(); + let result = drive_prompt(&session_id, &mut driver, &handle, generation, None).await; assert!(matches!(result, Err(AcpRuntimeError::Loop(_)))); - assert_eq!(control.stops.load(Ordering::Relaxed), 1); } #[tokio::test] @@ -4253,6 +4416,145 @@ mod tests { )); } + #[test] + fn replay_only_exposes_tagged_developer_compaction_summaries() { + let mut metadata = MetadataMap::new(); + metadata.insert( + crate::compaction::COMPACTION_SUMMARY_METADATA_KEY.into(), + true.into(), + ); + let replay = transcript_replay( + &wire::SessionId::new("saved"), + &[ + Item::text(ItemKind::Developer, "ordinary instruction"), + Item::text(ItemKind::Developer, "summary").with_metadata(metadata), + ], + ); + + let [notification] = replay.as_slice() else { + panic!("only the tagged summary should be replayed"); + }; + let wire::SessionUpdate::AgentMessage(message) = ¬ification.update else { + panic!("compaction summaries replay as agent messages, not notices"); + }; + assert!(matches!( + &message.content, + MaybeUndefined::Value(content) + if matches!(content.as_slice(), [wire::ContentBlock::Text(text)] if text.text == "summary") + )); + } + + #[test] + fn replay_hides_internal_notifications_and_instructions() { + let replay = transcript_replay( + &wire::SessionId::new("saved"), + &[ + Item::text(ItemKind::System, "system instruction"), + Item::text(ItemKind::Context, "internal context"), + Item::text(ItemKind::User, "run the build"), + Item::notification("Background tool call completed: very long raw output"), + Item::text(ItemKind::Assistant, "the build passed"), + ], + ); + + let [user, agent] = replay.as_slice() else { + panic!("only visible conversation messages should be replayed"); + }; + assert!(matches!( + &user.update, + wire::SessionUpdate::UserMessage(message) + if matches!(&message.content, MaybeUndefined::Value(content) + if matches!(content.as_slice(), [wire::ContentBlock::Text(text)] if text.text == "run the build")) + )); + assert!(matches!( + &agent.update, + wire::SessionUpdate::AgentMessage(message) + if matches!(&message.content, MaybeUndefined::Value(content) + if matches!(content.as_slice(), [wire::ContentBlock::Text(text)] if text.text == "the build passed")) + )); + } + + #[test] + fn replay_preserves_media_content_without_inventing_assistant_media() { + let parts = vec![ + Part::text("inspect these"), + Part::media( + Modality::Image, + "image/png", + DataRef::uri("file:///tmp/image.png"), + ), + Part::media( + Modality::Image, + "image/png", + DataRef::uri("https://example.com/result.png"), + ), + Part::media( + Modality::Image, + "image/png", + DataRef::uri("data:image/png;base64,c2VjcmV0"), + ), + ]; + let raw_output = serde_json::to_value(&parts).unwrap(); + let replay = transcript_replay( + &wire::SessionId::new("saved"), + &[ + Item::new(ItemKind::User, parts.clone()), + Item::new(ItemKind::Assistant, parts.clone()), + Item::new( + ItemKind::Tool, + vec![Part::ToolResult(agentkit_core::ToolResultPart::success( + "call-1", + ToolOutput::Parts(parts), + ))], + ), + ], + ); + let [user, agent, tool] = replay.as_slice() else { + panic!("expected user, agent and tool output"); + }; + let wire::SessionUpdate::UserMessage(message) = &user.update else { + panic!("expected user message"); + }; + let MaybeUndefined::Value(content) = &message.content else { + panic!("expected replay content"); + }; + assert!(matches!(content.as_slice(), [ + wire::ContentBlock::Text(text), + wire::ContentBlock::ResourceLink(file), + wire::ContentBlock::ResourceLink(remote), + wire::ContentBlock::Image(image), + ] if text.text == "inspect these" + && file.uri == "file:///tmp/image.png" + && remote.uri == "https://example.com/result.png" + && image.data == "c2VjcmV0" && image.uri.is_none())); + assert!( + matches!(&agent.update, wire::SessionUpdate::AgentMessage(message) + if matches!(&message.content, MaybeUndefined::Value(content) + if matches!(content.as_slice(), [wire::ContentBlock::Text(text)] if text.text == "inspect these"))) + ); + let wire::SessionUpdate::ToolCallUpdate(tool) = &tool.update else { + panic!("expected tool result"); + }; + assert_eq!( + tool.status, + MaybeUndefined::Value(wire::ToolCallStatus::Completed) + ); + assert_eq!(tool.raw_output, MaybeUndefined::Value(raw_output)); + let MaybeUndefined::Value(tool_content) = &tool.content else { + panic!("tool replay must provide visible content, not just raw output"); + }; + let tool_blocks = tool_content + .iter() + .map(|entry| { + let wire::ToolCallContent::Content(content) = entry else { + panic!("expected tool content block"); + }; + content.content.clone() + }) + .collect::>(); + assert_eq!(&tool_blocks, content); + } + #[test] fn replay_preserves_data_url_user_images() { let replay = transcript_replay( diff --git a/src/provider/chatgpt.rs b/src/provider/chatgpt.rs index 99ccbe11..fa693e62 100644 --- a/src/provider/chatgpt.rs +++ b/src/provider/chatgpt.rs @@ -121,8 +121,6 @@ fn subscription_resilience() -> ResilienceConfig { pub struct SubscriptionConfig { pub model: String, pub credential_storage: crate::credentials::CredentialStorage, - #[cfg(test)] - endpoint: Option, } impl SubscriptionConfig { @@ -133,8 +131,6 @@ impl SubscriptionConfig { Ok(Self { model, credential_storage: Default::default(), - #[cfg(test)] - endpoint: None, }) } @@ -145,14 +141,6 @@ impl SubscriptionConfig { self.credential_storage = storage; self } - - fn endpoint(&self) -> &str { - #[cfg(test)] - if let Some(endpoint) = &self.endpoint { - return endpoint; - } - ENDPOINT - } } #[derive(Clone)] @@ -265,7 +253,7 @@ impl ModelAdapter for OpenAiSubscriptionAdapter { }); let mut config = OpenAIResponsesConfig::chatgpt_private(self.config.model.clone(), authentication) - .with_endpoint(self.config.endpoint()) + .with_endpoint(ENDPOINT) .with_originator("kit") .with_user_agent(concat!("kit/", env!("CARGO_PKG_VERSION"))) .with_limits(OpenAIResponsesLimits { @@ -1200,8 +1188,7 @@ mod tests { #[test] fn authentication_attempt_is_bound_and_redacted() { - let record = - auth::TokenRecord::for_test_generation("secret-token", "account-1", "generation-1"); + let record = auth::test_support::token_record("secret-token", "account-1", "generation-1"); let attempt = authentication_attempt(record).unwrap(); assert_eq!( attempt.headers()["ChatGPT-Account-ID"], @@ -1218,20 +1205,20 @@ mod tests { #[test] fn session_binding_rejects_generation_change() { - let expected = auth::TokenRecord::for_test_generation("a", "account", "one") + let expected = auth::test_support::token_record("a", "account", "one") .binding() .unwrap(); assert!( ensure_credential_binding( &expected, - &auth::TokenRecord::for_test_generation("b", "account", "one") + &auth::test_support::token_record("b", "account", "one") ) .is_ok() ); assert!( ensure_credential_binding( &expected, - &auth::TokenRecord::for_test_generation("c", "account", "two") + &auth::test_support::token_record("c", "account", "two") ) .is_err() ); @@ -1240,18 +1227,16 @@ mod tests { #[tokio::test] async fn model_catalog_cache_is_scoped_to_account_and_generation() { let cache = SubscriptionModelCatalogCache::default(); - let first_binding = - auth::TokenRecord::for_test_generation("token", "account-1", "generation-1") - .binding() - .unwrap(); + let first_binding = auth::test_support::token_record("token", "account-1", "generation-1") + .binding() + .unwrap(); let next_generation = - auth::TokenRecord::for_test_generation("token", "account-1", "generation-2") - .binding() - .unwrap(); - let next_account = - auth::TokenRecord::for_test_generation("token", "account-2", "generation-1") + auth::test_support::token_record("token", "account-1", "generation-2") .binding() .unwrap(); + let next_account = auth::test_support::token_record("token", "account-2", "generation-1") + .binding() + .unwrap(); let first = cache .get_or_try_init(&first_binding, || async { @@ -1302,7 +1287,7 @@ mod tests { #[tokio::test] async fn model_catalog_cache_retries_after_failure() { let cache = SubscriptionModelCatalogCache::default(); - let binding = auth::TokenRecord::for_test_generation("token", "account", "generation") + let binding = auth::test_support::token_record("token", "account", "generation") .binding() .unwrap(); diff --git a/src/provider/mod.rs b/src/provider/mod.rs index ee406b2d..a7de6033 100644 --- a/src/provider/mod.rs +++ b/src/provider/mod.rs @@ -150,9 +150,7 @@ pub fn execute_provider_logout( } #[cfg(test)] -pub(crate) fn store_openrouter_test_credentials(storage: &crate::credentials::CredentialStorage) { - openrouter_auth::store_test_credentials(storage); -} +pub(crate) use openrouter_auth::test_support::store_openrouter_test_credentials; #[doc(hidden)] #[derive(Clone, Copy, Debug, Eq, PartialEq)] diff --git a/src/provider/openai_auth.rs b/src/provider/openai_auth.rs index bda6a713..9459c8a9 100644 --- a/src/provider/openai_auth.rs +++ b/src/provider/openai_auth.rs @@ -157,33 +157,6 @@ impl TokenRecord { generation: self.generation.clone(), }) } - - #[cfg(test)] - #[allow(dead_code)] - pub(crate) fn for_test(access_token: &str, account_id: &str) -> Self { - Self { - access_token: access_token.to_owned(), - refresh_token: "test-refresh-token".to_owned(), - id_token: "test-id-token".to_owned(), - expires_at: i64::MAX, - account_id: Some(account_id.to_owned()), - email: None, - plan_type: None, - generation: format!("test-{account_id}"), - } - } - - #[cfg(test)] - #[allow(dead_code)] - pub(crate) fn for_test_generation( - access_token: &str, - account_id: &str, - generation: &str, - ) -> Self { - let mut record = Self::for_test(access_token, account_id); - record.generation = generation.to_owned(); - record - } } #[derive(Clone, Debug, Eq, PartialEq, Zeroize, ZeroizeOnDrop)] @@ -663,15 +636,24 @@ fn refresh_current( rejected_access_token: Option<&str>, token_url: &str, ) -> Result { - refresh_current_at(store, deadline, rejected_access_token, token_url, JWKS_URL) + refresh_with_signing_keys( + store, + deadline, + rejected_access_token, + token_url, + &SigningKeys { + cache: &JWKS_CACHE, + source: &OpenAiSigningKeys, + }, + ) } -fn refresh_current_at( +fn refresh_with_signing_keys( store: &dyn CredentialStore, deadline: Instant, rejected_access_token: Option<&str>, token_url: &str, - jwks_url: &str, + signing_keys: &SigningKeys<'_>, ) -> Result { let current = store.load()?.ok_or_else(|| { AuthError::invalid( @@ -712,7 +694,7 @@ fn refresh_current_at( "token endpoint returned malformed JSON", ) })?; - let next = token_record(response, None, Some(¤t), deadline, jwks_url)?; + let next = token_record(response, None, Some(¤t), deadline, signing_keys)?; store.save(&next)?; Ok(next) } @@ -725,25 +707,28 @@ fn exchange_code( deadline: Instant, token_url: &str, ) -> Result { - exchange_code_at( + exchange_with_signing_keys( code, redirect_uri, verifier, nonce, deadline, token_url, - JWKS_URL, + &SigningKeys { + cache: &JWKS_CACHE, + source: &OpenAiSigningKeys, + }, ) } -fn exchange_code_at( +fn exchange_with_signing_keys( code: &str, redirect_uri: &str, verifier: &str, nonce: &str, deadline: Instant, token_url: &str, - jwks_url: &str, + signing_keys: &SigningKeys<'_>, ) -> Result { if !matches!( redirect_uri, @@ -781,7 +766,7 @@ fn exchange_code_at( "token endpoint returned malformed JSON", ) })?; - token_record(response, Some(nonce), None, deadline, jwks_url) + token_record(response, Some(nonce), None, deadline, signing_keys) } fn token_record( @@ -789,7 +774,7 @@ fn token_record( nonce: Option<&str>, previous: Option<&TokenRecord>, deadline: Instant, - jwks_url: &str, + signing_keys: &SigningKeys<'_>, ) -> Result { if response.access_token.is_empty() || response @@ -807,12 +792,12 @@ fn token_record( "https://api.openai.com/v1", None, deadline, - jwks_url, + signing_keys, )?; let access_account = extract_account_id(&access_claims)?; let (id_token, id_claims) = match response.id_token.take() { Some(id_token) => { - let claims = verify_claims(&id_token, CLIENT_ID, nonce, deadline, jwks_url)?; + let claims = verify_claims(&id_token, CLIENT_ID, nonce, deadline, signing_keys)?; (id_token, Some(claims)) } None if nonce.is_some() => { @@ -911,7 +896,7 @@ fn verify_claims( audience: &str, nonce: Option<&str>, deadline: Instant, - jwks_url: &str, + signing_keys: &SigningKeys<'_>, ) -> Result { if token.len() > MAX_CREDENTIAL_BYTES { return Err(AuthError::invalid("token_invalid", "JWT exceeds 64 KiB")); @@ -928,7 +913,7 @@ fn verify_claims( .kid .filter(|value| !value.is_empty() && value.len() <= 256 && value.is_ascii()) .ok_or_else(|| AuthError::invalid("token_invalid", "JWT omitted a valid key ID"))?; - let key = verification_key(&kid, deadline, jwks_url)?; + let key = verification_key(&kid, deadline, signing_keys)?; let mut validation = Validation::new(Algorithm::RS256); validation.leeway = CLOCK_SKEW_SECONDS as u64; validation.validate_nbf = true; @@ -1019,10 +1004,10 @@ fn valid_email(value: &str) -> bool { fn verification_key( kid: &str, deadline: Instant, - jwks_url: &str, + signing_keys: &SigningKeys<'_>, ) -> Result { for refresh in [false, true] { - let keys = jwks(deadline, jwks_url, refresh)?; + let keys = signing_keys.load(deadline, refresh)?; if let Some(jwk) = keys.find(kid) { if jwk.common.key_algorithm != Some(KeyAlgorithm::RS256) || !matches!(jwk.algorithm, AlgorithmParameters::RSA(_)) @@ -1052,33 +1037,62 @@ fn verification_key( )) } -fn jwks(deadline: Instant, jwks_url: &str, refresh: bool) -> Result { - if !cfg!(test) && jwks_url != JWKS_URL { - return Err(AuthError::invalid( - "jwks_invalid", - "only the pinned OpenAI JWKS endpoint is permitted", - )); - } - if !refresh { - let cache = JWKS_CACHE.lock().unwrap_or_else(|error| error.into_inner()); - if let Some(cached) = cache.get(jwks_url) - && cached.fetched_at.elapsed() < JWKS_TTL - { - return Ok(cached.keys.clone()); +/// Fetches issuer-trusted JWKS documents from an external source. +trait SigningKeySource { + fn fetch(&self, deadline: Instant) -> Result>, AuthError>; +} + +struct OpenAiSigningKeys; + +impl SigningKeySource for OpenAiSigningKeys { + fn fetch(&self, deadline: Instant) -> Result>, AuthError> { + let response = http_client(deadline)? + .get(JWKS_URL) + .send() + .map_err(|_| AuthError::unavailable("jwks_fetch_failed", "JWKS fetch failed"))?; + if !response.status().is_success() { + return Err(AuthError::unavailable( + "jwks_fetch_failed", + "JWKS endpoint rejected the request", + )); } + bounded_body(response) } - let response = http_client(deadline)? - .get(jwks_url) - .send() - .map_err(|_| AuthError::unavailable("jwks_fetch_failed", "JWKS fetch failed"))?; - if !response.status().is_success() { - return Err(AuthError::unavailable( - "jwks_fetch_failed", - "JWKS endpoint rejected the request", - )); +} + +struct SigningKeys<'a> { + cache: &'a Mutex>, + source: &'a dyn SigningKeySource, +} + +impl SigningKeys<'_> { + fn load(&self, deadline: Instant, refresh: bool) -> Result { + if !refresh { + let cache = self.cache.lock().unwrap_or_else(|error| error.into_inner()); + if let Some(cached) = cache.get(JWKS_URL) + && cached.fetched_at.elapsed() < JWKS_TTL + { + return Ok(cached.keys.clone()); + } + } + let body = self.source.fetch(deadline)?; + let keys = parse_jwks(body.as_slice())?; + self.cache + .lock() + .unwrap_or_else(|error| error.into_inner()) + .insert( + JWKS_URL.to_owned(), + CachedJwks { + fetched_at: Instant::now(), + keys: keys.clone(), + }, + ); + Ok(keys) } - let body = bounded_body(response)?; - let keys: JwkSet = serde_json::from_slice(body.as_slice()) +} + +fn parse_jwks(body: &[u8]) -> Result { + let keys: JwkSet = serde_json::from_slice(body) .map_err(|_| AuthError::invalid("jwks_invalid", "JWKS response is malformed"))?; if keys.keys.is_empty() || keys.keys.len() > 64 { return Err(AuthError::invalid( @@ -1097,16 +1111,6 @@ fn jwks(deadline: Instant, jwks_url: &str, refresh: bool) -> Result>, } -#[cfg(test)] -fn process_lock(deadline: Instant) -> Result { - process_lock_scoped(deadline, None) -} - fn process_lock_scoped( deadline: Instant, credential_scope: Option<&std::path::Path>, @@ -1642,9 +1641,36 @@ fn human(stdout: impl Into) -> Output { } } +#[cfg(test)] +pub(crate) mod test_support { + use super::TokenRecord; + + pub(crate) fn token_record( + access_token: &str, + account_id: &str, + generation: &str, + ) -> TokenRecord { + TokenRecord { + access_token: access_token.to_owned(), + refresh_token: "test-refresh-token".to_owned(), + id_token: "test-id-token".to_owned(), + expires_at: i64::MAX, + account_id: Some(account_id.to_owned()), + email: None, + plan_type: None, + generation: generation.to_owned(), + } + } +} + #[cfg(test)] mod tests { use super::*; + + fn process_lock(deadline: Instant) -> Result { + process_lock_scoped(deadline, None) + } + use jsonwebtoken::{EncodingKey, Header, encode, jwk::Jwk}; const TEST_RSA_KEY: &str = "-----BEGIN PRIVATE KEY-----\nMIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC3UQBTeVjOtSY4\nHaZHjSpQPlIIUXSiIq+WRInLoWwYEXmloR41HMmwsCQVV1WFZ7z0wUj14vD3/Bl6\nJG2JTU8ur+RvJojm1gXxg/etp4DG2HVtXong4QE7BKqJufHITMVuEhojkTulHIbW\nXfQQjaxQpGsOIuWcRz3YVB7zpAL7yoeHhvFd7RV+IqG9i4fjN4pzlCTv/TQig+s7\n539MsNx1ZakBfeBhx62JUPhFe6pXdPS2hXVUiTQRPMBm3GimDzyuA3WkVKzPyNMB\n2h+BALRFLslqPaFpul7NIifX36KgUPaimntpvFRxahqyDvJ9ATtq6oMeHaRUMZf5\nkRxjLIXLAgMBAAECggEAAIIV+SVDTMINyrwHo6J4NTlnACTm/jK7FTSNbpC8/E1t\nbpBwGqpAw4pJdKcFqAADSGkSFbRnrJhN+HEKE1uxK3+gp3o43kLw80bFX1Lb4DE7\nahkyp/qXsUfbB9S0dIoEm2srbWElWYN8ZYhkeSNGEKx+q3mx9JPx+kaJa2159flh\nis34maBeEr97gwjAvMjLbdVEpoaEIRC/hmem2ckT5jsDd4HS7RKNXwk/S8O7/PQW\n42xKAvL0APk5J53CDoW4DT78y7t4Rj/dVeRZAhdjDUFP+idZ1r9k6PM8vs5tl1P0\njzcOMzUBFmhnb5MKFvBLc4MKJQYzTT06/qdfAV0M2QKBgQDsoUF+pNQuERUKCI9T\nZey7rFgsbBkK2t0XvgpLwwMbF548HgL+QJhAAONaLe5+2GlZSgb5OgoYYspiuzQT\noz2mqeN2MSMnUtntyUt+Y6IzPEEPg6bVGdoCP3FSvz1L/JDJuJq4cqb3OGPe4yEt\nZDymqUJCDTO52vT0GLdZ6S70twKBgQDGUoYrbButBHX5nwE5XnrjgENGT4RauI76\nQ158MuFmRmpgaWlc37ByVyzMG7x9qxcad4Ry19hsG5KYnL/PNs31a2i/BdfLZyFF\nY0dfNExz6tKf4PWxZhhFhX94f7qseSzXLx8eMQqdds4WQsA13JI9qQJ0pVSNyb/V\nXM/n9XMrjQKBgDsKgSz4M3jLClTWjexhIhAxkE6FKjprIX8rC6abocrAudqGInkN\n5O8TSabWjwtXM/HzZooI0TwEajr4OqYrtNZAzWBQIlVNdtK9xvhiI7Zk8lbMonPJ\nX3vwGHZtAP5Upkuuo+whr0c/6qtSQJTyza9HzCBu6tkUqMm+4QCuDelBAoGAax8S\nF4w6WrcJHj7Tg3BUAmQ6clTrEbGUkPsoov88nmi0dsUZQzAT93681La6lkp+nS4n\nXXzXCnXONh6cwElC8CgHGP8H83cOEpOwbm0qSoZxJCh3rU2PGKYmFyku5JBDNyvd\nrAojSLBuWrnNZopwd1u91tGinT93HcEXD5yVi9UCgYBq1sjl5jlliyHzPWMeV3dn\nkJWDLMpCwrpmQzrhkA02PaZO1BB7QgZeIKTYkzECHT44wHflalVOEEsVZpEn2Ivd\nJz6j2JwX7Ke23MA0MDaV6+7syAwPKx3+pOGwdun2uZNgvS74IWeBEfdMhGrGncX0\nQegKxe+skNhLjXJ5SUTdZg==\n-----END PRIVATE KEY-----\n"; @@ -1794,32 +1820,36 @@ mod tests { (format!("http://{address}/oauth/token"), handle) } - fn serve_jwks(bodies: Vec) -> (String, std::thread::JoinHandle) { - let listener = TcpListener::bind(("127.0.0.1", 0)).unwrap(); - let address = listener.local_addr().unwrap(); - let count = bodies.len(); - let handle = std::thread::spawn(move || { - for body in bodies { - let (mut stream, _) = listener.accept().unwrap(); - stream - .set_read_timeout(Some(Duration::from_secs(2))) - .unwrap(); - let mut request = Vec::new(); - let mut chunk = [0_u8; 1024]; - while !request.windows(4).any(|part| part == b"\r\n\r\n") { - let read = stream.read(&mut chunk).unwrap(); - request.extend_from_slice(&chunk[..read]); - } - write!( - stream, - "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", - body.len(), - ) - .unwrap(); + struct SigningKeyFixture { + bodies: Mutex>, + cache: Mutex>, + } + + impl SigningKeyFixture { + fn new(bodies: Vec) -> Self { + Self { + bodies: Mutex::new(bodies.into()), + cache: Mutex::new(HashMap::new()), } - count - }); - (format!("http://{address}/.well-known/jwks.json"), handle) + } + + fn signing_keys(&self) -> SigningKeys<'_> { + SigningKeys { + cache: &self.cache, + source: self, + } + } + } + + impl SigningKeySource for SigningKeyFixture { + fn fetch(&self, _deadline: Instant) -> Result>, AuthError> { + self.bodies + .lock() + .unwrap() + .pop_front() + .map(|body| Zeroizing::new(body.into_bytes())) + .ok_or_else(|| AuthError::unavailable("jwks_fetch_failed", "fixture is offline")) + } } #[test] @@ -1907,7 +1937,7 @@ mod tests { let directory = tempfile::tempdir().unwrap(); let storage = CredentialStorage::Filesystem(directory.path().to_path_buf()); let store = BackendCredentialStore::new(&storage); - let record = TokenRecord::for_test("access", "account"); + let record = test_support::token_record("access", "account", "test-account"); store.save(&record).unwrap(); let loaded = store.load().unwrap().unwrap(); @@ -1941,23 +1971,22 @@ mod tests { }) .to_string(); let (token_url, server) = serve_once(body); - let (jwks_url, jwks_server) = serve_once(test_jwks(kid)); - let jwks_url = jwks_url.replace("/oauth/token", "/.well-known/jwks.json"); + let signing_keys = std::sync::Arc::new(SigningKeyFixture::new(vec![test_jwks(kid)])); let mut workers = Vec::new(); for _ in 0..8 { let store = std::sync::Arc::clone(&store); let token_url = token_url.clone(); - let jwks_url = jwks_url.clone(); + let signing_keys = std::sync::Arc::clone(&signing_keys); workers.push(std::thread::spawn(move || { let _guard = REFRESH_LOCK .lock() .unwrap_or_else(|error| error.into_inner()); - refresh_current_at( + refresh_with_signing_keys( store.as_ref(), Instant::now() + Duration::from_secs(5), Some("rejected-access-token"), &token_url, - &jwks_url, + &signing_keys.signing_keys(), ) .unwrap() .access_token() @@ -1974,7 +2003,6 @@ mod tests { assert_eq!(refreshed.account_id.as_deref(), Some("account-one")); assert_eq!(refreshed.id_token, old_id_token); let request = server.join().unwrap(); - jwks_server.join().unwrap(); assert!(request.starts_with("POST /oauth/token HTTP/1.1\r\n")); assert!(request.contains("content-type: application/x-www-form-urlencoded")); assert!(request.contains("grant_type=refresh_token")); @@ -2003,19 +2031,18 @@ mod tests { }) .to_string(); let (token_url, token_server) = serve_once(body); - let (jwks_url, jwks_server) = serve_once(test_jwks(kid)); - let record = exchange_code_at( + let signing_keys = SigningKeyFixture::new(vec![test_jwks(kid)]); + let record = exchange_with_signing_keys( "authorization-code", "http://localhost:1455/auth/callback", "pkce-verifier", "expected-nonce", Instant::now() + Duration::from_secs(5), &token_url, - &jwks_url.replace("/oauth/token", "/.well-known/jwks.json"), + &signing_keys.signing_keys(), ) .unwrap(); let request = token_server.join().unwrap(); - jwks_server.join().unwrap(); assert!(request.contains("content-type: application/x-www-form-urlencoded")); assert!(request.contains("grant_type=authorization_code")); assert!(request.contains("code=authorization-code")); @@ -2038,22 +2065,37 @@ mod tests { #[test] fn signed_jwt_rejects_bad_signatures_and_refreshes_unknown_kid_once() { - JWKS_CACHE.lock().unwrap().clear(); let token = jwt("https://api.openai.com/v1", None, "rotated-key", 3600, None); - let (jwks_url, server) = serve_jwks(vec![test_jwks("old-key"), test_jwks("rotated-key")]); + let signing_keys = + SigningKeyFixture::new(vec![test_jwks("old-key"), test_jwks("rotated-key")]); let claims = verify_claims( &token, "https://api.openai.com/v1", None, Instant::now() + Duration::from_secs(5), - &jwks_url, + &signing_keys.signing_keys(), ) .unwrap(); assert_eq!( extract_account_id(&claims).unwrap().as_deref(), Some("account-one") ); - assert_eq!(server.join().unwrap(), 2); + // The refresh must replace, not merge, the cached set. With the source + // now offline, another successful verification must use the published keys. + let cached = signing_keys + .signing_keys() + .load(Instant::now() + Duration::from_secs(1), false) + .unwrap(); + assert!(cached.find("rotated-key").is_some()); + assert!(cached.find("old-key").is_none()); + verify_claims( + &token, + "https://api.openai.com/v1", + None, + Instant::now() + Duration::from_secs(1), + &signing_keys.signing_keys(), + ) + .unwrap(); let mut bad = token.into_bytes(); let signature = bad.iter().rposition(|byte| *byte == b'.').unwrap() + 1; @@ -2064,7 +2106,7 @@ mod tests { "https://api.openai.com/v1", None, Instant::now() + Duration::from_secs(1), - &jwks_url, + &signing_keys.signing_keys(), ) .err() .unwrap(); @@ -2091,7 +2133,7 @@ mod tests { "https://api.openai.com/v1", None, Instant::now() + Duration::from_secs(1), - &jwks_url, + &signing_keys.signing_keys(), ) .is_err() ); @@ -2106,12 +2148,48 @@ mod tests { "https://api.openai.com/v1", None, Instant::now() + Duration::from_secs(1), - &jwks_url, + &signing_keys.signing_keys(), ) .is_err() ); } + #[test] + fn unknown_signing_key_stops_after_one_refresh() { + let keys = SigningKeyFixture::new(vec![test_jwks("old-key"), test_jwks("other-key")]); + let error = verification_key( + "missing-key", + Instant::now() + Duration::from_secs(1), + &keys.signing_keys(), + ) + .err() + .unwrap(); + assert_eq!(error.code, "token_invalid"); + let cached = keys + .signing_keys() + .load(Instant::now() + Duration::from_secs(1), false) + .unwrap(); + assert!(cached.find("other-key").is_some()); + assert!(cached.find("old-key").is_none()); + } + + #[test] + fn jwks_parser_rejects_malformed_empty_oversized_and_duplicate_keys() { + let key: Value = serde_json::from_str(&test_jwks("valid-key")).unwrap(); + for body in [ + "not json".to_owned(), + json!({"keys": []}).to_string(), + json!({"keys": vec![key["keys"][0].clone(); 65]}).to_string(), + json!({"keys": [key["keys"][0].clone(), key["keys"][0].clone()]}).to_string(), + test_jwks(""), + ] { + assert_eq!( + parse_jwks(body.as_bytes()).unwrap_err().code, + "jwks_invalid" + ); + } + } + #[test] fn jwt_not_before_uses_the_bounded_clock_skew() { for (not_before_in, accepted) in [(30, true), (120, false)] { @@ -2123,16 +2201,15 @@ mod tests { 3600, Some(not_before_in), ); - let (jwks_url, server) = serve_once(test_jwks(&kid)); + let signing_keys = SigningKeyFixture::new(vec![test_jwks(&kid)]); let result = verify_claims( &token, "https://api.openai.com/v1", None, Instant::now() + Duration::from_secs(5), - &jwks_url.replace("/oauth/token", "/.well-known/jwks.json"), + &signing_keys.signing_keys(), ); assert_eq!(result.is_ok(), accepted); - server.join().unwrap(); } } diff --git a/src/provider/openrouter_auth.rs b/src/provider/openrouter_auth.rs index bd6e2ef9..c1b33fa1 100644 --- a/src/provider/openrouter_auth.rs +++ b/src/provider/openrouter_auth.rs @@ -483,14 +483,18 @@ fn random_urlsafe() -> Result { } #[cfg(test)] -pub(crate) fn store_test_credentials(storage: &CredentialStorage) { - save( - storage, - &Credentials { - api_key: "test-openrouter-key".into(), - }, - ) - .unwrap(); +pub(crate) mod test_support { + use super::*; + + pub(crate) fn store_openrouter_test_credentials(storage: &CredentialStorage) { + save( + storage, + &Credentials { + api_key: "test-openrouter-key".into(), + }, + ) + .unwrap(); + } } #[cfg(test)] diff --git a/src/runtime.rs b/src/runtime.rs index 5e9ea4c3..6827e59c 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -44,6 +44,58 @@ use crate::{ }, }; +#[cfg(test)] +mod test_support { + use super::*; + + impl Runtime { + pub(crate) fn set_ambient_openrouter_api_key_for_test(&mut self, present: bool) { + self.ambient_openrouter_api_key = present; + } + } + + impl BackgroundJobs { + /// Seeds a foreground job for lifecycle tests, without applying registration policy. + /// Registration-policy tests must invoke `BackgroundableCompose` instead. + pub(crate) fn register_foreground_for_test(&self, call_id: &str) { + let mut jobs = self.state.lock().expect("background jobs poisoned"); + jobs.running.insert( + agentkit_core::ToolCallId::new(call_id), + BackgroundJob { + controller: CancellationController::new(), + foreground_cancellation: None, + cancellation_relay: None, + detached: false, + manual_detach: false, + terminal_published: false, + }, + ); + self.changed(&mut jobs); + } + + pub(crate) fn finish_for_test(&self, call_id: &str) { + self.finish(&agentkit_core::ToolCallId::new(call_id)); + } + + pub(crate) fn is_cancelled_for_test(&self, call_id: &str) -> bool { + let call_id = agentkit_core::ToolCallId::new(call_id); + self.state.lock().is_ok_and(|jobs| { + jobs.running + .get(&call_id) + .is_some_and(|job| job.controller.handle().is_cancelled_since(0)) + }) + } + + pub(crate) fn is_detached_for_test(&self, call_id: &str) -> bool { + let call_id = agentkit_core::ToolCallId::new(call_id); + self.state.lock().is_ok_and(|jobs| { + jobs.running.get(&call_id).is_some_and(|job| job.detached) + || jobs.pending_detaches.contains(&call_id) + }) + } + } +} + #[cfg(test)] mod tests; @@ -552,11 +604,6 @@ impl Runtime { &self.root } - #[cfg(test)] - pub(crate) fn set_ambient_openrouter_api_key_for_test(&mut self, present: bool) { - self.ambient_openrouter_api_key = present; - } - pub(crate) fn supports_terminal_authentication(&self, provider: ProviderKind) -> bool { self.credential_storage.is_persistent() && (provider != ProviderKind::OpenRouter @@ -1913,57 +1960,6 @@ impl BackgroundJobs { self.changed(&mut jobs); } } - - #[cfg(test)] - pub(crate) fn register_foreground_for_test(&self, call_id: &str) { - if let Ok(mut jobs) = self.state.lock() { - let call_id = agentkit_core::ToolCallId::new(call_id); - let manual_detach = jobs.pending_detaches.remove(&call_id); - let controller = CancellationController::new(); - if jobs.cancel_all { - controller.interrupt(); - } - jobs.running.insert( - call_id, - BackgroundJob { - controller, - foreground_cancellation: None, - cancellation_relay: None, - detached: manual_detach, - manual_detach, - terminal_published: false, - }, - ); - if manual_detach { - jobs.background_started = jobs.background_started.wrapping_add(1); - } - self.changed(&mut jobs); - } - } - - #[cfg(test)] - pub(crate) fn finish_for_test(&self, call_id: &str) { - self.finish(&agentkit_core::ToolCallId::new(call_id)); - } - - #[cfg(test)] - pub(crate) fn is_cancelled_for_test(&self, call_id: &str) -> bool { - let call_id = agentkit_core::ToolCallId::new(call_id); - self.state.lock().is_ok_and(|jobs| { - jobs.running - .get(&call_id) - .is_some_and(|job| job.controller.handle().is_cancelled_since(0)) - }) - } - - #[cfg(test)] - pub(crate) fn is_detached_for_test(&self, call_id: &str) -> bool { - let call_id = agentkit_core::ToolCallId::new(call_id); - self.state.lock().is_ok_and(|jobs| { - jobs.running.get(&call_id).is_some_and(|job| job.detached) - || jobs.pending_detaches.contains(&call_id) - }) - } } struct BackgroundJobGuard { diff --git a/src/runtime/tests.rs b/src/runtime/tests.rs index 7d80cb57..730b6790 100644 --- a/src/runtime/tests.rs +++ b/src/runtime/tests.rs @@ -1539,17 +1539,152 @@ async fn compose_background_sanitization_rejects_invalid_and_strips_before_dispa } } -#[test] -fn pending_detach_is_applied_when_compose_registers() { - let jobs = BackgroundJobs::default(); +// Hold execution at the external script-engine boundary, after real compose registration. +struct DelayedComposeBackend { + entered: tokio::sync::mpsc::UnboundedSender, + release: Arc, +} - assert_eq!( - jobs.detach("pending-call"), - Some(DetachRegistration::Registered) +#[async_trait::async_trait] +impl agentkit_tool_compose::ComposeBackend for DelayedComposeBackend { + fn name(&self) -> &'static str { + "delayed" + } + + fn description(&self, _: Option<&[agentkit_tools_core::ToolSpec]>) -> String { + "Controlled script execution".into() + } + + fn script_description(&self) -> &'static str { + "Script held until released or cancelled" + } + + async fn execute( + &self, + run: agentkit_tool_compose::BackendRun, + ) -> Result { + let cancellation = run.cancellation.expect("compose installs cancellation"); + self.entered.send(cancellation.clone()).unwrap(); + tokio::select! { + _ = cancellation.cancelled() => Err(agentkit_tool_compose::ComposeOutcome::Failed( + agentkit_tools_core::ToolError::Cancelled, + )), + _ = self.release.notified() => Ok(json!(7)), + } + } +} + +fn delayed_compose( + root: &std::path::Path, + jobs: BackgroundJobs, +) -> ( + Arc, + tokio::sync::mpsc::UnboundedReceiver, + Arc, +) { + let (entered, entries) = tokio::sync::mpsc::unbounded_channel(); + let release = Arc::new(tokio::sync::Notify::new()); + let inner = agentkit_tool_compose::ComposeTool::new(Default::default()).with_backend( + DelayedComposeBackend { + entered, + release: release.clone(), + }, ); - jobs.register_foreground_for_test("pending-call"); + ( + Arc::new(BackgroundableCompose::new( + inner, + jobs, + root.to_path_buf(), + agentkit_tools_core::ToolRegistry::new(), + )), + entries, + release, + ) +} - assert!(jobs.is_detached_for_test("pending-call")); +fn invoke_delayed_compose( + compose: Arc, + call_id: &'static str, + background: bool, +) -> tokio::task::JoinHandle { + tokio::spawn(async move { + let session_id = SessionId::new("registration-session"); + let turn_id = TurnId::new("registration-turn"); + let permissions = Arc::new(AllowAllPermissions); + let resources: Arc = Arc::new(()); + let owned = OwnedToolContext { + session_id: session_id.clone(), + turn_id: turn_id.clone(), + metadata: MetadataMap::new(), + permissions: permissions.clone(), + resources: resources.clone(), + cancellation: None, + execution_scope: Some(ToolExecutionScope { + executor: Arc::new(BasicToolExecutor::new(Vec::>::new())), + session_id: session_id.clone(), + turn_id: turn_id.clone(), + permissions, + resources, + cancellation: None, + }), + approved_request: None, + }; + compose + .invoke_outcome( + ToolRequest::new( + ToolCallId::new(call_id), + ToolName::new("compose"), + json!({"script": "return 7", "background": background}), + session_id, + turn_id, + ), + &mut owned.borrowed(), + ) + .await + }) +} + +#[tokio::test] +async fn pending_detach_is_applied_when_compose_registers() { + tokio::time::timeout(Duration::from_secs(2), async { + let root = tempfile::tempdir().unwrap(); + let jobs = BackgroundJobs::default(); + let (compose, mut entries, release) = delayed_compose(root.path(), jobs.clone()); + let initial = jobs.activity(); + let call_id = ToolCallId::new("pending-call"); + assert_eq!( + jobs.detach(&call_id.0), + Some(DetachRegistration::Registered) + ); + + let invocation = invoke_delayed_compose(compose, "pending-call", false); + let cancellation = entries.recv().await.unwrap(); + assert!(!cancellation.is_cancelled()); + // The old accessor also returned true for an unconsumed pending detach. + // Inspect the registered job so this cannot pass without applying policy. + { + let state = jobs.state.lock().unwrap(); + let job = state.running.get(&call_id).expect("compose registered"); + assert!(job.detached); + assert!(job.manual_detach); + assert!(!state.pending_detaches.contains(&call_id)); + } + assert!(jobs.activity().active); + assert_eq!( + jobs.activity().background_started, + initial.background_started + 1 + ); + + release.notify_one(); + assert!(matches!( + invocation.await.unwrap(), + ToolExecutionOutcome::Completed(_) + )); + assert!(!jobs.activity().active); + assert!(jobs.activity().unacknowledged_terminals); + }) + .await + .expect("pending detach compose invocation did not finish"); } #[test] @@ -1597,28 +1732,50 @@ fn background_terminal_publication_is_acknowledged_by_call_id() { assert!(!jobs.activity().unacknowledged_terminals); } -#[test] -fn cancel_all_covers_running_and_late_background_registration() { - let jobs = BackgroundJobs::default(); - let initial = jobs.activity(); - jobs.register_foreground_for_test("running"); - assert!(jobs.activity().active); - - jobs.cancel_all(); - assert!(jobs.is_cancelled_for_test("running")); - jobs.register_foreground_for_test("late"); - assert!(jobs.is_cancelled_for_test("late")); - - jobs.finish_for_test("running"); - jobs.finish_for_test("late"); - let quiescent = jobs.activity(); - assert!(!quiescent.active); - assert!(quiescent.generation > initial.generation); - - jobs.begin_turn(); - jobs.register_foreground_for_test("next-turn"); - assert!(!jobs.is_cancelled_for_test("next-turn")); - jobs.finish_for_test("next-turn"); +#[tokio::test] +async fn cancel_all_covers_running_and_late_background_registration() { + tokio::time::timeout(Duration::from_secs(2), async { + for background in [false, true] { + let root = tempfile::tempdir().unwrap(); + let jobs = BackgroundJobs::default(); + let initial = jobs.activity(); + let (compose, mut entries, release) = delayed_compose(root.path(), jobs.clone()); + let running = invoke_delayed_compose(compose.clone(), "running", background); + let cancellation = entries.recv().await.unwrap(); + assert!(jobs.activity().active); + assert!(!cancellation.is_cancelled()); + + jobs.cancel_all(); + assert!(cancellation.is_cancelled()); + assert!(matches!( + running.await.unwrap(), + ToolExecutionOutcome::Failed(agentkit_tools_core::ToolError::Cancelled) + )); + // A late registration must be cancelled before the backend is entered. + let late = invoke_delayed_compose(compose.clone(), "late", background); + assert!(matches!( + late.await.unwrap(), + ToolExecutionOutcome::Failed(agentkit_tools_core::ToolError::Cancelled) + )); + assert!(entries.try_recv().is_err()); + let quiescent = jobs.activity(); + assert!(!quiescent.active); + assert!(quiescent.generation > initial.generation); + + jobs.begin_turn(); + let next = invoke_delayed_compose(compose, "next-turn", background); + let cancellation = entries.recv().await.unwrap(); + assert!(!cancellation.is_cancelled()); + release.notify_one(); + assert!(matches!( + next.await.unwrap(), + ToolExecutionOutcome::Completed(_) + )); + assert!(!jobs.activity().active); + } + }) + .await + .expect("cancel-all compose invocation did not finish"); } #[tokio::test] diff --git a/src/tools/mcp.rs b/src/tools/mcp.rs index 1d8f5569..de28ac07 100644 --- a/src/tools/mcp.rs +++ b/src/tools/mcp.rs @@ -1266,29 +1266,6 @@ impl McpRuntime { } } - #[cfg(test)] - pub(crate) fn publish(&self, session_id: &str, event: McpEvent) { - self.publish_to(session_id, self.event_generation(session_id), event); - } - - #[cfg(test)] - pub(crate) async fn config_source_states(&self) -> Vec<(PathBuf, bool, bool)> { - self.inner - .reload - .lock() - .await - .sources - .iter() - .map(|state| { - ( - state.source.path.clone(), - state.source.required, - state.raw.is_some(), - ) - }) - .collect() - } - pub(crate) async fn refresh(&self) -> Result<(), String> { self.reload_config().await } @@ -2813,6 +2790,34 @@ fn render_search( } } +#[cfg(test)] +mod test_support { + use super::*; + + impl McpRuntime { + pub(crate) fn publish(&self, session_id: &str, event: McpEvent) { + self.publish_to(session_id, self.event_generation(session_id), event); + } + + pub(crate) async fn config_source_states(&self) -> Vec<(PathBuf, bool, bool)> { + self.inner + .reload + .lock() + .await + .sources + .iter() + .map(|state| { + ( + state.source.path.clone(), + state.source.required, + state.raw.is_some(), + ) + }) + .collect() + } + } +} + #[cfg(test)] mod tests { use std::{ diff --git a/src/tui/app.rs b/src/tui/app.rs index dc4b5f9a..bd79a789 100644 --- a/src/tui/app.rs +++ b/src/tui/app.rs @@ -13,21 +13,15 @@ use std::ffi::OsStr; #[cfg(any(target_os = "macos", target_os = "linux"))] use std::process::{Command, Stdio}; -#[cfg(test)] -use agent_client_protocol::schema::v2::RunningStateUpdate; use agent_client_protocol::schema::v2::{ AuthMethodTerminal, StateUpdate, StopReason, ToolCallStatus, ToolKind, }; -#[cfg(test)] -use agentkit_core::{DataRef, Item, ItemKind, Modality, Part, ToolOutput}; use crossterm::event::{ KeyCode, KeyEvent, KeyEventKind, KeyModifiers, MouseButton, MouseEvent, MouseEventKind, }; use ratatui::{layout::Rect, text::Line}; use unicode_segmentation::UnicodeSegmentation; -#[cfg(test)] -use crate::compaction::is_compaction_summary; use crate::events::{GenerationOutcome, RuntimeEvent, SubagentStatus}; use crate::file_search::FileMatch; @@ -98,15 +92,6 @@ pub enum Update { script: Option, backgrounded: bool, }, - /// A tool call changed status or produced output. - #[cfg(test)] - ToolUpdated { - id: String, - status: Option, - script: Option, - output: Vec, - backgrounded: bool, - }, /// A patchable ACP v2 tool call update or content chunk. ToolPatched { id: String, @@ -138,25 +123,6 @@ pub enum Update { ProcessExited(String), } -#[cfg(test)] -impl Update { - pub(super) fn test_text(text: String) -> Self { - Self::AgentMessage { - id: "test-agent".into(), - text, - append: true, - } - } - - pub(super) fn test_thought(text: String) -> Self { - Self::AgentThought { - id: "test-thought".into(), - text, - append: true, - } - } -} - /// Latest provider-reported occupancy of the main model's context window. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct ContextUsage { @@ -870,64 +836,6 @@ fn model_score(choice: &ModelChoice, query: &str) -> Option { ordered_token_score(&all_tokens, &query_tokens).map(|score| ModelScore { tier: 4, ..score }) } -#[cfg(test)] -fn media_label(media: &agentkit_core::MediaPart, index: usize) -> String { - let kind = match media.modality { - Modality::Image => "Image", - Modality::Audio => "Audio", - Modality::Video => "Video", - Modality::Binary => "Media", - }; - match &media.data { - DataRef::Uri(uri) if safe_media_uri(uri) => format!("[{kind} #{index}]({uri})"), - _ => format!("[{kind} #{index}]"), - } -} - -#[cfg(test)] -fn safe_media_uri(uri: &str) -> bool { - uri.len() <= 2_048 - && url::Url::parse(uri).is_ok_and(|uri| matches!(uri.scheme(), "file" | "http" | "https")) -} - -#[cfg(test)] -fn persisted_output(output: &ToolOutput) -> Vec { - let text = match output { - ToolOutput::Text(text) => text.clone(), - ToolOutput::Structured(value) => { - serde_json::to_string_pretty(value).unwrap_or_else(|_| value.to_string()) - } - ToolOutput::Parts(parts) => { - let mut next_media = 0; - let mut next_file = 0; - parts - .iter() - .filter_map(|part| match part { - Part::Text(text) => Some(text.text.clone()), - Part::Media(media) => { - next_media += 1; - Some(media_label(media, next_media)) - } - Part::File(file) => { - next_file += 1; - Some(match &file.data { - DataRef::Uri(uri) if safe_media_uri(uri) => { - format!("[File #{}]({uri})", next_file) - } - _ => format!("[File #{}]", next_file), - }) - } - Part::Structured(value) => Some(value.value.to_string()), - _ => None, - }) - .collect::>() - .join("\n") - } - ToolOutput::Files(files) => format!("{} files", files.len()), - }; - text.lines().map(str::to_string).collect() -} - fn agent_status_rank(status: SubagentStatus) -> u8 { match status { SubagentStatus::Starting => 0, @@ -1270,117 +1178,6 @@ impl App { } } - /// Rebuilds the visible history from the same Items preloaded into the model. - #[cfg(test)] - pub fn restore_transcript(&mut self, session_id: String, transcript: &[Item]) { - self.session_id = Some(session_id); - for item in transcript { - match item.kind { - ItemKind::Developer if is_compaction_summary(item) => { - self.push_block(Block::Notice("context compacted".into())); - } - ItemKind::System - | ItemKind::Developer - | ItemKind::Context - | ItemKind::Notification => continue, - ItemKind::User => { - let mut next_media = 0; - let text = item - .parts - .iter() - .filter_map(|part| match part { - Part::Text(text) => Some(text.text.clone()), - Part::Media(media) => { - next_media += 1; - Some(media_label(media, next_media)) - } - _ => None, - }) - .collect::>() - .join("\n"); - if !text.is_empty() { - if item.kind == ItemKind::User { - self.latest_agent_source.clear(); - self.push_block(Block::User(UserMessage { - text, - images: Vec::new(), - })); - } else { - self.push_block(Block::Notice(text)); - } - } - } - ItemKind::Assistant => { - let mut next_media = 0; - for part in &item.parts { - match part { - Part::Text(text) if !text.text.is_empty() => { - self.latest_agent_source.push_str(&text.text); - self.push_block(Block::Agent(text.text.clone())) - } - Part::Media(media) => { - next_media += 1; - self.push_block(Block::Agent(media_label(media, next_media))); - } - Part::Reasoning(reasoning) if reasoning.summary.is_some() => self - .push_block(Block::Thought { - text: reasoning.summary.clone().unwrap_or_default(), - started: Instant::now(), - millis: Some(0), - }), - Part::ToolCall(call) => self.push_block(Block::Tool(ToolCall { - id: call.id.to_string(), - title: call.name.clone(), - kind: ToolKind::Other, - status: ToolCallStatus::Completed, - started: Instant::now(), - finished: Some(Instant::now()), - script: call - .input - .get("script") - .and_then(serde_json::Value::as_str) - .unwrap_or_default() - .to_string(), - plan: call - .input - .get("script") - .and_then(serde_json::Value::as_str) - .map(parse_plan) - .unwrap_or_default(), - children: Vec::new(), - output: Vec::new(), - intent: call - .input - .get("intent") - .and_then(serde_json::Value::as_str) - .map(str::to_string), - expanded: false, - compose_view: ComposeView::Output, - expansion_explicit: false, - backgrounded: false, - })), - _ => {} - } - } - } - ItemKind::Tool => { - for part in &item.parts { - if let Part::ToolResult(result) = part - && let Some(call) = self.call_mut(&result.call_id.to_string()) - { - call.output = persisted_output(&result.output); - if result.is_error { - call.status = ToolCallStatus::Failed; - } - } - } - } - } - } - self.follow = true; - self.scroll = usize::MAX; - } - /// The tool call the graph pane should show: the running one, else the /// most recent, so a finished program stays readable. pub fn focus_call(&self) -> Option<&ToolCall> { @@ -1434,14 +1231,6 @@ impl App { self.clamp_agents_scroll(); } - #[cfg(test)] - pub fn agents(&self) -> Vec<&AgentRow> { - self.agent_tree_rows() - .into_iter() - .map(|tree_row| tree_row.row) - .collect() - } - pub fn agent_tree_rows(&self) -> Vec> { fn status_rank(status: SubagentStatus) -> u8 { match status { @@ -1938,43 +1727,6 @@ impl App { backgrounded, })); } - #[cfg(test)] - Update::ToolUpdated { - id, - status, - script, - output, - backgrounded, - } => { - let completed_background = { - let Some(call) = self.call_mut(&id) else { - return; - }; - let was_running = call.running(); - if let Some(script) = script { - call.plan = parse_plan(&script); - call.script = script; - } - if !output.is_empty() { - call.output = output; - } - call.backgrounded |= backgrounded; - if let Some(status) = status { - call.status = status; - if !call.running() { - call.finalize_terminal_state(); - } - } - was_running && !call.running() && call.backgrounded - }; - // Autonomous model output follows the terminal update for a - // detached call without a new ACP prompt/TurnEnded pair. Seal - // the current stream so that output starts a new agent block. - self.agent_stream_sealed |= completed_background; - if let Some(index) = self.call_index(&id) { - self.reclassify_dynamic(index); - } - } Update::ToolPatched { id, title, @@ -1998,6 +1750,7 @@ impl App { let Some(call) = self.call_mut(&id) else { return; }; + let was_running = call.running(); if let Some(title) = title { if title == agentkit_tool_compose::COMPOSE_TOOL_NAME && call.title != agentkit_tool_compose::COMPOSE_TOOL_NAME @@ -2033,6 +1786,9 @@ impl App { call.finalize_terminal_state(); } } + let completed_background = was_running && !call.running() && call.backgrounded; + // Autonomous output after a detached call starts a new agent stream. + self.agent_stream_sealed |= completed_background; if let Some(index) = self.call_index(&id) { self.mark_block_dirty(index); self.reclassify_dynamic(index); @@ -2250,21 +2006,6 @@ impl App { self.press = None; } - #[cfg(test)] - pub fn push_user(&mut self, prompt: String) -> u64 { - let id = format!("test-user-{}", self.blocks.len()); - self.apply(Update::UserMessage { - id, - text: prompt, - images: Vec::new(), - append: false, - }); - self.apply(Update::State(StateUpdate::Running( - RunningStateUpdate::new(), - ))); - self.blocks.len() as u64 - } - /// Folds a tool call's raw output open or shut. Completed compose calls /// cycle through their output and source views before folding closed. pub fn toggle_output(&mut self, id: &str) { @@ -2325,17 +2066,6 @@ impl App { self.clamp_agents_scroll(); } - #[cfg(test)] - fn apply_runtime_at(&mut self, event: RuntimeEvent, now_unix_ms: u64) { - match event { - RuntimeEvent::SubagentStateChanged { .. } - | RuntimeEvent::SubagentDescendantsRemoved { .. } => { - self.apply_agent_runtime_at(event, now_unix_ms); - } - event => self.apply_runtime(event), - } - } - fn apply_agent_runtime_at(&mut self, event: RuntimeEvent, now_unix_ms: u64) { match event { RuntimeEvent::SubagentStateChanged { @@ -3853,6 +3583,65 @@ fn has_graphical_session(display: Option<&OsStr>, wayland_display: Option<&OsStr #[cfg(not(any(target_os = "macos", target_os = "linux")))] fn open_url(_url: &str) {} +/// Fixtures and deterministic-clock adapters; not an alternate live event protocol. +#[cfg(test)] +mod test_support { + use agent_client_protocol::schema::v2::RunningStateUpdate; + + use super::*; + + impl Update { + pub(in crate::tui) fn test_text(text: String) -> Self { + Self::AgentMessage { + id: "test-agent".into(), + text, + append: true, + } + } + + pub(in crate::tui) fn test_thought(text: String) -> Self { + Self::AgentThought { + id: "test-thought".into(), + text, + append: true, + } + } + } + + impl App { + pub(super) fn agents(&self) -> Vec<&AgentRow> { + self.agent_tree_rows() + .into_iter() + .map(|tree_row| tree_row.row) + .collect() + } + + pub(in crate::tui) fn push_user(&mut self, prompt: String) -> u64 { + let id = format!("test-user-{}", self.blocks.len()); + self.apply(Update::UserMessage { + id, + text: prompt, + images: Vec::new(), + append: false, + }); + self.apply(Update::State(StateUpdate::Running( + RunningStateUpdate::new(), + ))); + self.blocks.len() as u64 + } + + pub(super) fn apply_runtime_at(&mut self, event: RuntimeEvent, now_unix_ms: u64) { + match event { + RuntimeEvent::SubagentStateChanged { .. } + | RuntimeEvent::SubagentDescendantsRemoved { .. } => { + self.apply_agent_runtime_at(event, now_unix_ms); + } + event => self.apply_runtime(event), + } + } + } +} + #[cfg(test)] mod tests { use std::{ @@ -3864,7 +3653,6 @@ mod tests { AuthMethodTerminal, IdleStateUpdate, RequiresActionStateUpdate, RunningStateUpdate, StateUpdate, StopReason, ToolCallStatus, ToolKind, }; - use agentkit_core::{DataRef, Item, ItemKind, MediaPart, MetadataMap, Modality, Part}; use crossterm::event::{ KeyCode, KeyEvent, KeyEventKind, KeyModifiers, MouseButton, MouseEvent, MouseEventKind, }; @@ -4314,89 +4102,6 @@ mod tests { assert!(!app.compacting); } - #[test] - fn restores_only_tagged_developer_items_as_compaction_markers() { - let mut metadata = MetadataMap::new(); - metadata.insert( - crate::compaction::COMPACTION_SUMMARY_METADATA_KEY.into(), - true.into(), - ); - let transcript = vec![ - Item::text(ItemKind::Developer, "ordinary instruction"), - Item::text(ItemKind::Developer, "summary").with_metadata(metadata), - ]; - let mut app = app(); - app.restore_transcript("session".into(), &transcript); - assert_eq!(app.blocks.len(), 1); - assert!( - matches!(app.blocks.first(), Some(Block::Notice(text)) if text == "context compacted") - ); - } - - #[test] - fn restored_transcript_hides_internal_notifications() { - let transcript = vec![ - Item::text(ItemKind::User, "run the build"), - Item::notification("Background tool call completed: very long raw output"), - Item::text(ItemKind::Assistant, "the build passed"), - ]; - let mut app = app(); - - app.restore_transcript("session".into(), &transcript); - - assert_eq!(app.blocks.len(), 2); - assert!(matches!(&app.blocks[0], Block::User(message) if message.text == "run the build")); - assert!(matches!(&app.blocks[1], Block::Agent(text) if text == "the build passed")); - } - - #[test] - fn restored_media_uses_safe_links_and_never_exposes_data_urls() { - let transcript = vec![ - Item::new( - ItemKind::User, - vec![ - Part::text("inspect these"), - Part::Media(MediaPart::new( - Modality::Image, - "image/png", - DataRef::Uri("file:///tmp/image.png".into()), - )), - Part::Media(MediaPart::new( - Modality::Image, - "image/png", - DataRef::Uri("data:image/png;base64,c2VjcmV0".into()), - )), - ], - ), - Item::new( - ItemKind::Assistant, - vec![ - Part::text("done"), - Part::Media(MediaPart::new( - Modality::Image, - "image/png", - DataRef::Uri("https://example.com/result.png".into()), - )), - ], - ), - ]; - let mut app = app(); - - app.restore_transcript("session".into(), &transcript); - - assert!(matches!( - &app.blocks[0], - Block::User(message) - if message.text == "inspect these\n[Image #1](file:///tmp/image.png)\n[Image #2]" - && !message.text.contains("data:") - )); - assert!(matches!(&app.blocks[1], Block::Agent(text) if text == "done")); - assert!(matches!( - &app.blocks[2], - Block::Agent(text) if text == "[Image #1](https://example.com/result.png)" - )); - } - #[test] fn attributes_nested_calls_to_the_owning_tool_call() { let mut app = app(); @@ -4416,11 +4121,15 @@ mod tests { compose(&mut app, "a = shell({ command: \"sleep 60\" })\nreturn a"); app.apply(Update::Runtime(child("call-1:compose:shell", "shell"))); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { id: "call-1".into(), + title: None, + kind: None, status: Some(ToolCallStatus::Failed), script: None, - output: Vec::new(), + output: None, + append_output: false, + intent: None, backgrounded: false, }); @@ -4542,11 +4251,15 @@ mod tests { RunningStateUpdate::new(), ))); let started = app.turn_started; - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { id: "call-1".into(), + title: None, + kind: None, status: Some(ToolCallStatus::Completed), script: None, - output: Vec::new(), + output: None, + append_output: false, + intent: None, backgrounded: false, }); assert!(app.working()); @@ -6020,11 +5733,15 @@ mod tests { backgrounded: true, }); app.apply(Update::test_text("first completion".into())); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { id: "background".into(), + title: None, + kind: None, status: Some(ToolCallStatus::Completed), script: None, - output: Vec::new(), + output: None, + append_output: false, + intent: None, backgrounded: false, }); app.apply(Update::AgentMessage { @@ -6047,6 +5764,10 @@ mod tests { }) .collect::>(); assert_eq!(agents, ["first completion", "second completion continued"]); + assert_eq!( + app.latest_agent_text().as_deref(), + Some("second completion continued") + ); } fn model_choice(provider: &str, model: &str) -> super::ModelChoice { @@ -6405,11 +6126,15 @@ mod tests { app.phase = Phase::Idle; assert!(app.needs_redraw_tick()); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { id: "background".into(), + title: None, + kind: None, status: Some(ToolCallStatus::Completed), script: None, - output: Vec::new(), + output: None, + append_output: false, + intent: None, backgrounded: true, }); assert!(!app.needs_redraw_tick()); diff --git a/src/tui/image.rs b/src/tui/image.rs index 867e2d15..d927c2bc 100644 --- a/src/tui/image.rs +++ b/src/tui/image.rs @@ -56,26 +56,6 @@ impl ImageRuntime { } } - #[cfg(test)] - pub fn disabled() -> Self { - Self { - picker: None, - cache: HashMap::new(), - decoded_backing_bytes: 0, - clock: 0, - } - } - - #[cfg(test)] - pub fn with_picker(picker: Picker) -> Self { - Self { - picker: Some(picker), - cache: HashMap::new(), - decoded_backing_bytes: 0, - clock: 0, - } - } - pub fn enabled(&self) -> bool { self.picker.is_some() } @@ -85,11 +65,6 @@ impl ImageRuntime { self.decoded_backing_bytes = 0; } - #[cfg(test)] - pub fn cached_entries(&self) -> usize { - self.cache.len() - } - pub fn prepare(&mut self, image: &UserImage, width: u16) -> Option { let picker = self.picker.clone()?; if width == 0 { @@ -195,6 +170,35 @@ fn decode(source: &UserImage) -> Option { reader.decode().ok() } +#[cfg(test)] +mod test_support { + use super::*; + + impl ImageRuntime { + pub fn disabled() -> Self { + Self { + picker: None, + cache: HashMap::new(), + decoded_backing_bytes: 0, + clock: 0, + } + } + + pub fn with_picker(picker: Picker) -> Self { + Self { + picker: Some(picker), + cache: HashMap::new(), + decoded_backing_bytes: 0, + clock: 0, + } + } + + pub fn cached_entries(&self) -> usize { + self.cache.len() + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/tui/markdown.rs b/src/tui/markdown.rs index 6949d667..4722f064 100644 --- a/src/tui/markdown.rs +++ b/src/tui/markdown.rs @@ -17,31 +17,6 @@ use ratatui::{ }; use unicode_width::UnicodeWidthStr; -/// Renders Markdown source into styled transcript lines. -#[cfg(test)] -pub fn render(source: &str) -> Vec> { - render_linked(source) - .into_iter() - .map(|line| line.line()) - .collect() -} - -/// Renders Markdown while attaching each URL directly to the spans it owns. -#[cfg(test)] -pub fn render_linked(source: &str) -> Vec { - render_copyable(source) - .into_iter() - .map(|(line, _)| line) - .collect() -} - -/// Renders Markdown and tags every row of a fenced code block with its exact -/// source content, excluding the fence and language label. -#[cfg(test)] -pub fn render_copyable(source: &str) -> Vec<(LinkedLine, Option>)> { - render_copyable_at_width(source, None) -} - pub(super) fn render_copyable_at_width( source: &str, max_width: Option, @@ -680,7 +655,29 @@ fn inline_with_link_destinations( mod tests { use ratatui::style::Modifier; - use super::{render, render_copyable, render_copyable_at_width, render_linked}; + use super::*; + + /// Renders Markdown source into styled transcript lines. + pub fn render(source: &str) -> Vec> { + render_linked(source) + .into_iter() + .map(|line| line.line()) + .collect() + } + + /// Renders Markdown while attaching each URL directly to the spans it owns. + pub fn render_linked(source: &str) -> Vec { + render_copyable(source) + .into_iter() + .map(|(line, _)| line) + .collect() + } + + /// Renders Markdown and tags every row of a fenced code block with its exact + /// source content, excluding the fence and language label. + pub fn render_copyable(source: &str) -> Vec<(LinkedLine, Option>)> { + render_copyable_at_width(source, None) + } fn line_text(line: &ratatui::text::Line<'_>) -> String { line.spans diff --git a/src/tui/mod.rs b/src/tui/mod.rs index d4c33ab3..66739ec7 100644 --- a/src/tui/mod.rs +++ b/src/tui/mod.rs @@ -374,13 +374,11 @@ enum ConnectedAuthentication { async fn wait_for_connected_authentication( authentication: impl std::future::Future>>, - prepare: impl FnOnce(), app: &mut App, route: &Arc>, updates: &mut mpsc::UnboundedReceiver, exit: &mut oneshot::Receiver>, ) -> ConnectedAuthentication { - prepare(); tokio::pin!(authentication); loop { tokio::select! { @@ -1012,12 +1010,10 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( &method, &mut stop, ); + leave(&mut terminal); + println!("Starting {}…", method.name); let authenticated = wait_for_connected_authentication( authentication, - || { - leave(&mut terminal); - println!("Starting {}…", method.name); - }, &mut app, &transition_session, &mut updates_rx, @@ -1533,12 +1529,10 @@ pub async fn run_with_reasoning_effort_and_openrouter_key( &method, &mut stop, ); + leave(&mut terminal); + println!("Starting {}…", method.name); let authenticated = wait_for_connected_authentication( authentication, - || { - leave(&mut terminal); - println!("Starting {}…", method.name); - }, &mut app, &transition_session, &mut updates_rx, @@ -3000,7 +2994,7 @@ mod tests { } #[tokio::test] - async fn connected_terminal_authentication_prepares_before_observing_agent_exit() { + async fn connected_terminal_authentication_observes_agent_exit() { let root = tempfile::tempdir().unwrap(); let mut app = App::new( root.path().into(), @@ -3019,11 +3013,9 @@ mod tests { .unwrap(); let authentication = std::future::pending::>>(); - let prepared = std::cell::Cell::new(false); let result = wait_for_connected_authentication( authentication, - || prepared.set(true), &mut app, &route, &mut updates_rx, @@ -3031,13 +3023,47 @@ mod tests { ) .await; - assert!(prepared.get()); assert!(matches!( result, ConnectedAuthentication::AgentExited(Ok(Err(_))) )); } + #[cfg(unix)] + #[tokio::test] + async fn connected_terminal_authentication_returns_the_process_exit_status() { + let root = tempfile::tempdir().unwrap(); + let mut app = App::new( + root.path().into(), + "provider".into(), + "model".into(), + "a2a".into(), + ); + let route = Arc::new(Mutex::new(ActiveSessionRoute { + id: "session".into(), + generation: 0, + })); + let (_updates_tx, mut updates_rx) = tokio::sync::mpsc::unbounded_channel(); + let (_exit_tx, mut exit_rx) = tokio::sync::oneshot::channel(); + let mut stop = super::Stop::new().unwrap(); + for code in [0, 7] { + let mut command = tokio::process::Command::new("sh"); + command.args(["-c", &format!("exit {code}")]); + let result = wait_for_connected_authentication( + super::wait_for_terminal_auth(command, &mut stop), + &mut app, + &route, + &mut updates_rx, + &mut exit_rx, + ) + .await; + let ConnectedAuthentication::Completed(Some(Ok(status))) = result else { + panic!("authentication did not return its process status"); + }; + assert_eq!(status.code(), Some(code)); + } + } + #[test] fn translates_available_commands_with_their_session() { let (_, updates) = translate(UpdateSessionNotification::new( @@ -3163,6 +3189,117 @@ mod tests { )); } + #[test] + fn replayed_media_uses_safe_display_content_instead_of_raw_payloads() { + use crate::tui::app::Block; + + let mut app = App::new( + PathBuf::from("/tmp"), + "provider".into(), + "model".into(), + "a2a".into(), + ); + // These are the complete content shapes produced by ACP v2 transcript replay: + // URI media becomes resource links; data URLs become structured image payloads. + let content = vec![ + ContentBlock::Text(TextContent::new("inspect these")), + ContentBlock::ResourceLink(wire::ResourceLink::new( + "file:///tmp/image.png", + "file:///tmp/image.png", + )), + ContentBlock::ResourceLink(wire::ResourceLink::new( + "https://example.com/result.png", + "https://example.com/result.png", + )), + ContentBlock::Image(wire::ImageContent::new("c2VjcmV0", "image/png")), + ]; + let notifications = [ + SessionUpdate::UserMessage(UserMessage::new("user").content(content.clone())), + // Tagged compaction summaries are ordinary agent messages on this path. + SessionUpdate::AgentMessage( + AgentMessage::new("compaction") + .content(vec![ContentBlock::Text(TextContent::new("summary"))]), + ), + SessionUpdate::ToolCallUpdate( + wire::ToolCallUpdate::new("tool-1") + .status(wire::ToolCallStatus::Completed) + .raw_output(Some(json!({"uri": "data:image/png;base64,c2VjcmV0"}))) + .content( + content + .into_iter() + .map(|block| { + wire::ToolCallContent::Content(Box::new(wire::Content::new(block))) + }) + .collect::>(), + ), + ), + ]; + for notification in notifications { + for update in translate_for_session( + UpdateSessionNotification::new("session", notification), + "session", + ) { + app.apply(update); + } + } + + let [Block::User(user), Block::Agent(summary), Block::Tool(tool)] = app.blocks.as_slice() + else { + panic!("expected user, summary and completed tool blocks"); + }; + let links = "inspect these[file:///tmp/image.png](file:///tmp/image.png)[https://example.com/result.png](https://example.com/result.png)"; + assert_eq!(user.text, format!("{links}\n[Image #1]")); + assert_eq!(user.images.len(), 1); + assert_eq!(user.images[0].data, "c2VjcmV0"); + assert_eq!(summary, "summary"); + assert_eq!(tool.status, wire::ToolCallStatus::Completed); + assert_eq!( + tool.output, + [ + "inspect these", + "[file:///tmp/image.png](file:///tmp/image.png)", + "[https://example.com/result.png](https://example.com/result.png)", + "[Image]", + ] + ); + for text in std::iter::once(&user.text).chain(tool.output.iter()) { + assert!(!text.contains("data:")); + assert!(!text.contains("c2VjcmV0")); + } + } + + #[test] + fn replayed_media_translation_rejects_unsafe_and_oversized_links() { + for uri in [ + "data:image/png;base64,c2VjcmV0".to_string(), + "javascript:alert(1)".to_string(), + format!("https://example.com/{}", "x".repeat(2048)), + ] { + let content = vec![ + ContentBlock::ResourceLink(wire::ResourceLink::new(uri.clone(), uri.clone())), + ContentBlock::Image( + wire::ImageContent::new("c2VjcmV0", "image/png").uri(Some(uri)), + ), + ]; + let update = UpdateSessionNotification::new( + "session", + SessionUpdate::UserMessage(UserMessage::new("user").content(content.clone())), + ); + assert!( + matches!(translate_for_session(update, "session").as_slice(), + [Update::UserMessage { text, images, .. }] if text == "[Image #1]" && images.len() == 1) + ); + let update = UpdateSessionNotification::new( + "session", + SessionUpdate::AgentMessage(AgentMessage::new("agent").content(content)), + ); + assert!( + matches!(translate_for_session(update, "session").as_slice(), + [Update::AgentMessage { text, .. }] if text == "[Image]") + ); + } + } + #[test] fn translates_replayed_raw_tool_output() { let update = UpdateSessionNotification::new( diff --git a/src/tui/ui.rs b/src/tui/ui.rs index 700990ac..d1dbc4ff 100644 --- a/src/tui/ui.rs +++ b/src/tui/ui.rs @@ -16,13 +16,6 @@ use ratatui::{ use unicode_segmentation::UnicodeSegmentation; use unicode_width::UnicodeWidthStr; -#[cfg(test)] -thread_local! { - static MATERIALIZED_TRANSCRIPT_ROWS: std::cell::Cell = const { std::cell::Cell::new(0) }; - static REFRESHED_TRANSCRIPT_BLOCKS: std::cell::Cell = const { std::cell::Cell::new(0) }; - static VISITED_TRANSCRIPT_BLOCKS: std::cell::Cell = const { std::cell::Cell::new(0) }; -} - use crate::events::{GenerationOutcome, SubagentStatus}; use super::{ @@ -856,8 +849,6 @@ fn draw_transcript(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRunti ); let mut visible_images: Vec<(usize, usize, i16)> = Vec::new(); let mut materialize = |row: &crate::tui::app::CachedTranscriptRow| { - #[cfg(test)] - MATERIALIZED_TRANSCRIPT_ROWS.with(|count| count.set(count.get() + 1)); visible.push(row.0.clone()); app.row_calls.push(row.1.0.clone()); app.row_code.push(row.1.1.clone()); @@ -870,8 +861,6 @@ fn draw_transcript(frame: &mut Frame<'_>, app: &mut App, images: &mut ImageRunti .saturating_sub(1) .min(app.blocks.len() - 1); while block_index < app.blocks.len() { - #[cfg(test)] - VISITED_TRANSCRIPT_BLOCKS.with(|count| count.set(count.get() + 1)); let span_start = app.transcript_prefixes[block_index]; if span_start >= end { break; @@ -1036,8 +1025,6 @@ fn refresh_transcript_cache_with_images(app: &mut App, images: &mut ImageRuntime let dirty = std::mem::take(&mut app.transcript_dirty); let mut first_changed_count = app.blocks.len(); for block_index in dirty { - #[cfg(test)] - REFRESHED_TRANSCRIPT_BLOCKS.with(|count| count.set(count.get() + 1)); let dynamic = match &app.blocks[block_index] { Block::Thought { millis, .. } => millis.is_none(), Block::Tool(call) => call.running() || call.running_children() > 0, @@ -1091,12 +1078,6 @@ fn refresh_transcript_cache_with_images(app: &mut App, images: &mut ImageRuntime } } -#[cfg(test)] -fn refresh_transcript_cache(app: &mut App, width: usize) { - let mut images = ImageRuntime::disabled(); - refresh_transcript_cache_with_images(app, &mut images, width); -} - fn user_block_rows( message: &UserMessage, width: usize, @@ -2308,8 +2289,8 @@ mod tests { use unicode_width::UnicodeWidthStr; use super::{ - MAX_PROMPT_ROWS, ModelDialogRow, agent_lines, body_layout, draw, draw_agents, - model_dialog_rows, model_dialog_viewport, prompt_lines, refresh_transcript_cache, + ImageRuntime, MAX_PROMPT_ROWS, ModelDialogRow, agent_lines, body_layout, draw, draw_agents, + model_dialog_rows, model_dialog_viewport, prompt_lines, refresh_transcript_cache_with_images, truncate_to_width, user_block_rows, user_line, }; use crate::{ @@ -2322,6 +2303,11 @@ mod tests { }, }; + fn refresh_transcript_cache(app: &mut App, width: usize) { + let mut images = ImageRuntime::disabled(); + refresh_transcript_cache_with_images(app, &mut images, width); + } + fn test_agent( name: &str, status: SubagentStatus, @@ -3334,11 +3320,15 @@ mod tests { { call.attach(format!("extra-{index}"), tool.into(), tool.into()); } - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, - output: vec!["done".into()], + output: Some(vec!["done".into()]), backgrounded: false, }); @@ -3391,11 +3381,15 @@ mod tests { #[test] fn completed_compose_defaults_collapsed_and_cycles_output_script_and_closed() { let mut app = sample(); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, - output: vec!["compose result".into()], + output: Some(vec!["compose result".into()]), backgrounded: false, }); @@ -3429,11 +3423,15 @@ mod tests { let mut app = sample(); app.toggle_last_output(); app.toggle_last_output(); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, - output: vec!["compose result".into()], + output: Some(vec!["compose result".into()]), backgrounded: false, }); @@ -3463,11 +3461,15 @@ mod tests { #[test] fn a_new_tool_collapses_the_previous_compose_output() { let mut app = sample(); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, - output: vec!["compose result".into()], + output: Some(vec!["compose result".into()]), backgrounded: false, }); app.apply(Update::ToolStarted { @@ -3562,11 +3564,15 @@ mod tests { #[test] fn scrolls_back_through_a_long_transcript_with_the_log_pane_open() { let mut app = sample(); - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Failed), script: None, - output: vec!["exit code 1".into()], + output: Some(vec!["exit code 1".into()]), backgrounded: false, }); app.apply(Update::State(StateUpdate::Idle( @@ -3718,13 +3724,19 @@ mod tests { call.title = "shell".into(); call.expanded = false; } - app.apply(Update::ToolUpdated { + app.apply(Update::ToolPatched { + title: None, + kind: None, + append_output: false, + intent: None, id: "call-1".into(), status: Some(agent_client_protocol::schema::v2::ToolCallStatus::Completed), script: None, - output: (0..40) - .map(|index| format!("output line {index}")) - .collect(), + output: Some( + (0..40) + .map(|index| format!("output line {index}")) + .collect(), + ), backgrounded: false, }); let frame = render(&mut app, 100, 24); @@ -4133,7 +4145,7 @@ mod tests { } #[test] - fn materializes_click_metadata_only_for_visible_rows() { + fn click_metadata_matches_visible_rows() { let mut app = App::new( PathBuf::from("/tmp/kit"), "openai-subscription".into(), @@ -4148,15 +4160,8 @@ mod tests { "[visible link](https://example.com/target)".into(), )); - super::MATERIALIZED_TRANSCRIPT_ROWS.with(|count| count.set(0)); - super::VISITED_TRANSCRIPT_BLOCKS.with(|count| count.set(0)); let _ = render(&mut app, 50, 10); - super::MATERIALIZED_TRANSCRIPT_ROWS.with(|count| assert_eq!(count.get(), app.viewport)); - super::VISITED_TRANSCRIPT_BLOCKS.with(|count| { - assert!(count.get() <= app.viewport); - assert!(count.get() < app.blocks.len()); - }); assert_eq!(app.row_links.len(), app.viewport); assert_eq!(app.row_calls.len(), app.viewport); assert_eq!(app.row_code.len(), app.viewport); @@ -4166,10 +4171,17 @@ mod tests { .flatten() .any(|hit| hit.url == "https://example.com/target") ); + + app.scroll_by(-1000); + let frame = render(&mut app, 50, 10); + assert!(frame.contains("old row 0"), "{frame}"); + assert!(!frame.contains("visible link"), "{frame}"); + assert_eq!(app.row_links.len(), app.viewport); + assert!(app.row_links.iter().all(Vec::is_empty)); } #[test] - fn tail_mutation_does_not_inspect_or_rebuild_unchanged_history() { + fn tail_mutation_preserves_cached_history() { let mut app = App::new( PathBuf::from("/tmp/kit"), "openai-subscription".into(), @@ -4185,11 +4197,9 @@ mod tests { let history_revision = app.transcript_cache[0].as_ref().unwrap().revision; let tail_rows = app.transcript_cache[99].as_ref().unwrap().rows.as_ptr(); - super::REFRESHED_TRANSCRIPT_BLOCKS.with(|count| count.set(0)); app.apply(Update::test_text(" changed".into())); refresh_transcript_cache(&mut app, 12); - super::REFRESHED_TRANSCRIPT_BLOCKS.with(|count| assert_eq!(count.get(), 1)); let history = app.transcript_cache[0].as_ref().unwrap(); assert_eq!(history.rows.as_ptr(), history_rows); assert_eq!(history.revision, history_revision); @@ -4197,6 +4207,14 @@ mod tests { app.transcript_cache[99].as_ref().unwrap().rows.as_ptr(), tail_rows ); + let tail = app.transcript_cache[99].as_ref().unwrap(); + let text = tail + .rows + .iter() + .map(|row| line_text(&row.0)) + .collect::(); + assert!(text.contains("history 99"), "{text}"); + assert!(text.contains("changed"), "{text}"); } #[test]