diff --git a/hive-runtime/src/acp/mod.rs b/hive-runtime/src/acp/mod.rs index 1f13a7cd..e2b8725f 100644 --- a/hive-runtime/src/acp/mod.rs +++ b/hive-runtime/src/acp/mod.rs @@ -22,6 +22,7 @@ use tokio::sync::futures::Notified; use tokio::sync::{Mutex, Notify}; use tokio::time::Instant; +use self::provider_key::KeyVar; use self::rpc::{Connection, Incoming}; use self::stream::StreamMapper; use crate::spec::AcpCommand; @@ -356,11 +357,14 @@ pub struct AcpRuntime { compact_idle: Duration, choices: Choices, spend: Spend, + resolve_key: provider_key::Resolve, } /// The running agent process and the session loaded into it. struct Live { conn: Connection, + /// The provider key added to the agent's environment at spawn. + key: Option, load_session: bool, /// The recorded session, or a new one whose first prompt is unanswered. loaded: Option, @@ -407,6 +411,7 @@ impl AcpRuntime

{ compact_idle: COMPACT_IDLE, choices: Choices::default(), spend: Spend::default(), + resolve_key: provider_key::resolver(), } } @@ -422,7 +427,7 @@ impl AcpRuntime

{ let (_, servers) = mcp_servers(config)?; let cwd = session_cwd(config); let key = match &self.command.api_key_env { - Some(name) => provider_key::resolve(name).await, + Some(name) => (self.resolve_key)(name.clone()).await, None => None, }; let conn = Connection::spawn( @@ -453,6 +458,7 @@ impl AcpRuntime

{ self.choices.reset(); Ok(Live { conn, + key, load_session: caps["loadSession"] == Value::Bool(true), loaded: None, model: None, @@ -687,9 +693,34 @@ impl AcpRuntime

{ { *guard = None; } + if let Err(Error::Acp(AcpError::EmptyEndTurn | AcpError::IdleTimeout { .. })) = &result + && let Some(live) = guard.as_ref() + && self.key_changed(live).await + { + tracing::warn!( + "the ACP provider key in the store differs from the one the agent was spawned \ + with; respawning the agent" + ); + *guard = None; + } result } + /// Whether the store holds a provider key other than the one `live` was + /// spawned with. No key read is no change, so a store that cannot be read + /// never takes away the key the agent has. + async fn key_changed(&self, live: &Live) -> bool { + let Some(name) = &self.command.api_key_env else { + return false; + }; + let Some(key) = (self.resolve_key)(name.clone()).await else { + return false; + }; + live.key + .as_ref() + .is_none_or(|spawned| spawned.value != key.value) + } + /// Compact the recorded session, if there is one. `checkpoint` runs the /// policy's checkpoint turn before an advertised `compact` command too; /// without one it always runs, since archiving keeps nothing of the @@ -1024,11 +1055,14 @@ mod tests { use std::collections::BTreeMap; use std::path::Path; use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; use hive_claude::{Config, NoopSink, PercentPolicy, Sink}; - use super::{AcpError, AcpRuntime, Choice, ReportedCost, SessionChoices, Spend, TurnKind}; + use super::{ + AcpError, AcpRuntime, Choice, KeyVar, ReportedCost, SessionChoices, Spend, TurnKind, + }; use crate::{AcpCommand, Error, Runtime, TurnStop}; /// An ACP agent in plain `sh`. It answers `initialize`, numbers its @@ -1064,6 +1098,10 @@ mod tests { const AGENT: &str = r#" n=0 prompt= model=m/think effort=low printf 'start\n' >> "$1.methods" +# With TEST_PROVIDER_KEY set: its value at each start, in .keys. +if [ -n "$TEST_PROVIDER_KEY" ]; then + printf '%s\n' "$TEST_PROVIDER_KEY" >> "$1.keys" +fi # With TRAP_INT set: SIGINT creates .int and exits 0. if [ -n "$TRAP_INT" ]; then trap ': > "$1.int"; exit 0' INT @@ -1386,6 +1424,102 @@ done ); } + /// A runtime whose agent gets its provider key in `TEST_PROVIDER_KEY`, + /// resolved to `answers` in turn, the last one for every read after it. + /// Returns how many reads were made. + fn keyed( + dir: &Path, + mode: &str, + answers: &'static [Option<&'static str>], + ) -> (AcpRuntime, Arc) { + let mut runtime = runtime(dir, mode); + runtime.command.api_key_env = Some("TEST_PROVIDER_KEY".to_owned()); + let reads = Arc::new(AtomicUsize::new(0)); + let counted = reads.clone(); + runtime.resolve_key = Arc::new(move |name| { + let read = counted.fetch_add(1, Ordering::SeqCst); + let answer = answers[read.min(answers.len() - 1)]; + Box::pin(async move { + answer.map(|value| KeyVar { + name, + value: value.to_owned(), + }) + }) + }); + (runtime, reads) + } + + /// The provider key each agent start was given. + fn keys(dir: &Path) -> Vec { + std::fs::read_to_string(dir.join("prompts.keys")) + .unwrap() + .lines() + .map(str::to_owned) + .collect() + } + + #[tokio::test] + async fn a_changed_key_after_an_empty_turn_respawns_the_agent_with_it() { + let dir = tempfile::tempdir().unwrap(); + let (runtime, reads) = keyed(dir.path(), "blank", &[Some("old"), Some("new")]); + let config = config(dir.path()); + + let empty = runtime.run(&config, "one", &NoopSink).await; + assert!( + matches!(empty, Err(Error::Acp(AcpError::EmptyEndTurn))), + "{empty:?}" + ); + // Read at spawn and once after the failed turn. + assert_eq!(reads.load(Ordering::SeqCst), 2); + + runtime.run(&config, "two", &NoopSink).await.unwrap(); + assert_eq!(keys(dir.path()), ["old", "new"]); + // The respawned agent carries on with the same session. + assert_eq!(count(dir.path(), "session/new"), 1); + assert_eq!(count(dir.path(), "session/load"), 1); + assert_eq!(recorded(dir.path()), "s1"); + } + + #[tokio::test] + async fn a_changed_key_after_an_idle_timeout_respawns_the_agent_with_it() { + let dir = tempfile::tempdir().unwrap(); + let (runtime, _) = keyed(dir.path(), "silent", &[Some("old"), Some("new")]); + let config = idle_config(dir.path(), 300); + + let stalled = runtime.run(&config, "one", &NoopSink).await; + assert!( + matches!(stalled, Err(Error::Acp(AcpError::IdleTimeout { .. }))), + "{stalled:?}" + ); + runtime.run(&config, "two", &NoopSink).await.unwrap(); + assert_eq!(keys(dir.path()), ["old", "new"]); + } + + #[tokio::test] + async fn an_unchanged_key_after_an_empty_turn_keeps_the_agent() { + let dir = tempfile::tempdir().unwrap(); + let (runtime, reads) = keyed(dir.path(), "blank", &[Some("same")]); + let config = config(dir.path()); + + assert!(runtime.run(&config, "one", &NoopSink).await.is_err()); + runtime.run(&config, "two", &NoopSink).await.unwrap(); + assert_eq!(count(dir.path(), "start"), 1); + assert_eq!(keys(dir.path()), ["same"]); + assert_eq!(reads.load(Ordering::SeqCst), 2); + } + + #[tokio::test] + async fn no_key_read_after_an_empty_turn_keeps_the_agent_and_its_key() { + let dir = tempfile::tempdir().unwrap(); + let (runtime, _) = keyed(dir.path(), "blank", &[Some("old"), None]); + let config = config(dir.path()); + + assert!(runtime.run(&config, "one", &NoopSink).await.is_err()); + runtime.run(&config, "two", &NoopSink).await.unwrap(); + assert_eq!(count(dir.path(), "start"), 1); + assert_eq!(keys(dir.path()), ["old"]); + } + #[tokio::test] async fn an_end_turn_with_output_but_no_usage_succeeds() { let dir = tempfile::tempdir().unwrap(); diff --git a/hive-runtime/src/acp/provider_key.rs b/hive-runtime/src/acp/provider_key.rs index 1b467322..f06417fe 100644 --- a/hive-runtime/src/acp/provider_key.rs +++ b/hive-runtime/src/acp/provider_key.rs @@ -4,8 +4,18 @@ //! //! A value the process already inherited (`backendEnvironmentFile`) wins and //! the store is not asked. Without either, the agent is spawned without it. +//! +//! The key is read when the agent is spawned, and again after a turn fails +//! with [`AcpError::EmptyEndTurn`](super::AcpError::EmptyEndTurn) or +//! [`AcpError::IdleTimeout`](super::AcpError::IdleTimeout), which is how a +//! provider rejecting the key shows. A stored key that differs from the one +//! the agent was spawned with gets the agent respawned with it. Anything else +//! (the same key, an inherited value, no store, a failed or empty read) keeps +//! the running agent. use std::future::Future; +use std::pin::Pin; +use std::sync::Arc; use std::time::Duration; use swarm_secret_client::{ @@ -31,6 +41,15 @@ pub(super) struct KeyVar { pub(super) value: String, } +/// Resolves the variable it is given the name of, as [`resolve`] does. +pub(super) type Resolve = + Arc Pin> + Send>> + Send + Sync>; + +/// A [`Resolve`] calling [`resolve`]. +pub(super) fn resolver() -> Resolve { + Arc::new(|name| Box::pin(async move { resolve(&name).await })) +} + /// Where [`resolve_with`] reads the key from. A trait so the precedence can /// be exercised without a store. trait KeySource { diff --git a/nix/agent-modules/agent-service.nix b/nix/agent-modules/agent-service.nix index 069cb0c9..88c4c81b 100644 --- a/nix/agent-modules/agent-service.nix +++ b/nix/agent-modules/agent-service.nix @@ -206,6 +206,8 @@ in On an ACP agent using the `opencode` preset, a provider key this file leaves unset is read from the swarm secret store at `swarm/agents//acp-provider`, field `api_key`, if one is stored. + After a turn that ends empty or idle, the store is read again, and a + key that changed there gets the agent respawned with it. ''; };