1b4e764e5c
The swarm settings gain `runtime` (`api-key`, the default, or `acp`) and an `acp` provider record: an https base URL, the provider's model, and its prices per million tokens. A PUT now changes only the fields its body names, so the settings form that sends model, effort and limit leaves the runtime alone; `null` clears `effort` or `acp`. The body stays strict, and the stored document keeps #5192's lenient read. The controller sends the runtime in each run request. On `acp` the classifier runs one prompt on a new `opencode acp` session per classification through hive-runtime's AcpRuntime, in a scratch directory, with a config rendered from the settings. opencode runs `read` and other tools without asking, so OPENCODE_PERMISSION {"*":"deny"} takes every tool away from the model, and a deny-all policy refuses any request on top. The title call, which resends the prompt, is turned off. The deadline cancels the turn, and a turn that does not end within 15s of that is dropped, which also covers an agent stuck before its prompt. The answer is the final assistant text, as ACP has no result event. The provider key is read per run from swarm/services/review-classifier/acp-provider, which the bao UI's admins login may create and update, and reaches opencode through AcpRuntime::with_provider_key and the trampoline's allowlist. Each ACP classification's cost is computed from the configured prices and the reported tokens, and added to hyperhive.review_classifier.cost.usage labelled repo, pr, model and currency. Part of #5096
2136 lines
77 KiB
Rust
2136 lines
77 KiB
Rust
//! The ACP backend: one long-lived agent process per runtime, spawned on the
|
|
//! first turn, holding one durable session whose id is kept in a file.
|
|
//!
|
|
//! An ordinary turn that ends with `end_turn` after sending no event and
|
|
//! reporting no usage fails with [`AcpError::EmptyEndTurn`]. An agent that
|
|
//! never reports usage therefore surfaces every such turn as that error.
|
|
|
|
mod provider_key;
|
|
mod rpc;
|
|
mod stream;
|
|
|
|
use std::collections::{BTreeMap, HashMap, VecDeque};
|
|
use std::path::{Path, PathBuf};
|
|
use std::pin::Pin;
|
|
use std::sync::atomic::{AtomicBool, Ordering};
|
|
use std::sync::{Arc, PoisonError};
|
|
use std::time::Duration;
|
|
|
|
use hive_claude::{CompactionPolicy, Config, Progress, Sink};
|
|
use serde_json::{Value, json};
|
|
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;
|
|
use crate::{Error, Result, Runtime};
|
|
|
|
/// ACP protocol version this client speaks.
|
|
const PROTOCOL_VERSION: u64 = 1;
|
|
|
|
/// Stderr lines kept for the error when the agent exits mid-turn.
|
|
const STDERR_TAIL: usize = 20;
|
|
|
|
/// After a turn's reply, how long the agent must stay quiet before the turn
|
|
/// is taken to be over, and the most that wait may take in all.
|
|
const SETTLE_QUIET: Duration = Duration::from_millis(500);
|
|
const SETTLE_MAX: Duration = Duration::from_secs(5);
|
|
|
|
/// How long the agent has to end its turn after `session/cancel` before it is
|
|
/// sent SIGINT.
|
|
const CANCEL_GRACE: Duration = Duration::from_secs(10);
|
|
|
|
/// How long a `compact` command may send nothing before it is cancelled and
|
|
/// the session is replaced instead: room for one summarising model call over a
|
|
/// full window, under a third of the default 600s turn idle window.
|
|
const COMPACT_IDLE: Duration = Duration::from_mins(3);
|
|
|
|
/// The command an agent advertises to compact its session.
|
|
const COMPACT_COMMAND: &str = "compact";
|
|
|
|
/// How long to wait, on a session whose agent has not advertised its
|
|
/// commands yet, for it to do so. An agent may send them only after the
|
|
/// session is attached, or never.
|
|
const COMMANDS_WAIT: Duration = SETTLE_MAX;
|
|
|
|
/// The config option categories [`Config::model`] and [`Config::effort`] set.
|
|
const MODEL_CATEGORY: &str = "model";
|
|
const EFFORT_CATEGORY: &str = "thought_level";
|
|
|
|
/// One `session/request_permission` request, as a [`PermissionPolicy`] sees it.
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct PermissionAsk<'a> {
|
|
/// The tool call's ACP `kind` (`read`, `edit`, `execute`, `fetch`, …;
|
|
/// `other` when the agent gives none).
|
|
pub kind: &'a str,
|
|
/// For a `kind` of `other`: the MCP server, of those handed to the
|
|
/// session, whose tool this is — the tool call's title names it as
|
|
/// `<server>_<tool>`, `<server>__<tool>` or `mcp__<server>__<tool>`.
|
|
/// `None` for every other kind.
|
|
pub mcp_server: Option<&'a str>,
|
|
}
|
|
|
|
/// Decides the agent's `session/request_permission` requests. `true` allows
|
|
/// the call once, `false` rejects it.
|
|
pub type PermissionPolicy = Arc<dyn Fn(&PermissionAsk<'_>) -> bool + Send + Sync>;
|
|
|
|
/// Why an ACP operation failed.
|
|
#[derive(Debug, thiserror::Error)]
|
|
#[non_exhaustive]
|
|
pub enum AcpError {
|
|
#[error("failed to spawn ACP agent `{program}`: {source}")]
|
|
Spawn {
|
|
program: String,
|
|
#[source]
|
|
source: std::io::Error,
|
|
},
|
|
#[error("ACP agent closed its output (exited?)")]
|
|
Closed,
|
|
#[error("ACP agent exited mid-turn; last stderr:\n{stderr_tail}")]
|
|
Exited { stderr_tail: String },
|
|
#[error("ACP agent i/o failed: {0}")]
|
|
Io(#[from] std::io::Error),
|
|
#[error("ACP `{method}` failed: {message} (code {code}) {data}")]
|
|
Rpc {
|
|
method: &'static str,
|
|
code: i64,
|
|
message: String,
|
|
data: String,
|
|
},
|
|
#[error("ACP `{method}` response has no {field}")]
|
|
Malformed {
|
|
method: &'static str,
|
|
field: &'static str,
|
|
},
|
|
#[error(
|
|
"the ACP agent does not accept MCP servers over HTTP (`mcpCapabilities.http` is not true), \
|
|
and the hyperhive tools are only served over HTTP"
|
|
)]
|
|
NoHttpMcp,
|
|
#[error("reading the MCP config {path}: {detail}")]
|
|
McpConfig { path: PathBuf, detail: String },
|
|
#[error(
|
|
"ACP agent sent nothing for {silent_secs}s, so its turn was cancelled; a \
|
|
provider error the agent retries without reporting it, such as HTTP 429, looks like this"
|
|
)]
|
|
IdleTimeout { silent_secs: u64 },
|
|
#[error("agent ended the turn with no output; likely a provider error")]
|
|
EmptyEndTurn,
|
|
#[error(
|
|
"ACP agent sent nothing for {silent_secs}s and did not stop within {grace_secs}s of \
|
|
`session/cancel`, so it was sent SIGINT"
|
|
)]
|
|
IdleKilled { silent_secs: u64, grace_secs: u64 },
|
|
#[error(
|
|
"ACP agent did not stop within {grace_secs}s of `session/cancel`, so it was sent SIGINT"
|
|
)]
|
|
CancelIgnored { grace_secs: u64 },
|
|
}
|
|
|
|
/// Stops the turn an [`AcpRuntime`] has in flight, from outside the call
|
|
/// driving it. Cheap to clone.
|
|
#[derive(Clone)]
|
|
pub struct Canceller(Arc<CancelState>);
|
|
|
|
#[derive(Default)]
|
|
struct CancelState {
|
|
in_turn: AtomicBool,
|
|
requested: Notify,
|
|
}
|
|
|
|
impl Canceller {
|
|
/// Ask the turn in flight to stop. The agent is sent `session/cancel`,
|
|
/// and the turn then ends normally with what it produced so far; an
|
|
/// agent that has not ended it within `CANCEL_GRACE` is sent SIGINT, and
|
|
/// the turn fails with [`AcpError::CancelIgnored`]. Returns whether a turn
|
|
/// was in flight; with none, this does nothing.
|
|
#[must_use]
|
|
pub fn cancel(&self) -> bool {
|
|
let in_turn = self.0.in_turn.load(Ordering::SeqCst);
|
|
if in_turn {
|
|
self.0.requested.notify_waiters();
|
|
}
|
|
in_turn
|
|
}
|
|
}
|
|
|
|
/// A value the session lets the client pick: the one it has, and those it
|
|
/// accepts.
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub struct Choice {
|
|
pub current: String,
|
|
pub values: Vec<String>,
|
|
/// The value the agent last refused to set, until it accepts one, the
|
|
/// value wanted is the one the session holds, or the agent is respawned.
|
|
pub refused: Option<String>,
|
|
}
|
|
|
|
/// The model and effort the loaded session lets the client pick, from its
|
|
/// first `model` and `thought_level` config options. `None` where it offers
|
|
/// none; which effort levels a session offers can change with its model.
|
|
#[derive(Debug, Clone, Default, PartialEq, Eq)]
|
|
pub struct SessionChoices {
|
|
pub model: Option<Choice>,
|
|
pub effort: Option<Choice>,
|
|
}
|
|
|
|
/// Reads the [`SessionChoices`] of an [`AcpRuntime`]'s session from outside
|
|
/// the call driving it. Cheap to clone.
|
|
#[derive(Clone, Default)]
|
|
pub struct Choices(Arc<std::sync::Mutex<Offered>>);
|
|
|
|
#[derive(Default)]
|
|
struct Offered {
|
|
options: Vec<stream::ConfigOption>,
|
|
/// Per category, the value the agent last refused to set.
|
|
refused: HashMap<&'static str, String>,
|
|
}
|
|
|
|
impl Choices {
|
|
/// What the session offers now. Empty until a session is attached, and
|
|
/// after the agent is respawned until one is again.
|
|
#[must_use]
|
|
pub fn get(&self) -> SessionChoices {
|
|
let offered = self.lock();
|
|
let choice = |category| {
|
|
find(&offered.options, category).map(|o| Choice {
|
|
current: o.current.clone(),
|
|
values: o.values.clone(),
|
|
refused: offered.refused.get(category).cloned(),
|
|
})
|
|
};
|
|
SessionChoices {
|
|
model: choice(MODEL_CATEGORY),
|
|
effort: choice(EFFORT_CATEGORY),
|
|
}
|
|
}
|
|
|
|
fn find(&self, category: &str) -> Option<stream::ConfigOption> {
|
|
find(&self.lock().options, category).cloned()
|
|
}
|
|
|
|
fn replace(&self, options: Vec<stream::ConfigOption>) {
|
|
self.lock().options = options;
|
|
}
|
|
|
|
fn reset(&self) {
|
|
*self.lock() = Offered::default();
|
|
}
|
|
|
|
fn refuse(&self, category: &'static str, value: &str) {
|
|
self.lock().refused.insert(category, value.to_owned());
|
|
}
|
|
|
|
fn accept(&self, category: &str) {
|
|
self.lock().refused.remove(category);
|
|
}
|
|
|
|
fn lock(&self) -> std::sync::MutexGuard<'_, Offered> {
|
|
self.0.lock().unwrap_or_else(PoisonError::into_inner)
|
|
}
|
|
}
|
|
|
|
/// An amount of money an agent reported, in an ISO 4217 `currency`.
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct ReportedCost {
|
|
pub amount: f64,
|
|
pub currency: String,
|
|
}
|
|
|
|
/// What the agent has reported spending, shared between the runtime and its
|
|
/// [`Live`] agent so an amount outlives a respawn until it is taken.
|
|
#[derive(Clone, Default)]
|
|
struct Spend(Arc<std::sync::Mutex<Spent>>);
|
|
|
|
#[derive(Default)]
|
|
struct Spent {
|
|
/// The attached session's running total at its last report; `None` until
|
|
/// it reports one.
|
|
total: Option<ReportedCost>,
|
|
/// The attached session is new, so its total started at zero.
|
|
fresh: bool,
|
|
/// Spent since the last [`Spend::take`], per currency.
|
|
untaken: BTreeMap<String, f64>,
|
|
}
|
|
|
|
impl Spend {
|
|
/// Start following a newly attached session. A loaded session's total
|
|
/// before its first report is unknown, so that report only sets the
|
|
/// baseline.
|
|
fn attach(&self, fresh: bool) {
|
|
let mut spent = self.lock();
|
|
spent.total = None;
|
|
spent.fresh = fresh;
|
|
}
|
|
|
|
/// Fold in a running total the attached session reported. A total lower
|
|
/// than the last one means the agent dropped part of the session's
|
|
/// history, so it adds nothing.
|
|
fn observe(&self, total: ReportedCost) {
|
|
let mut spent = self.lock();
|
|
let added = match &spent.total {
|
|
Some(last) if last.currency == total.currency => total.amount - last.amount,
|
|
None if spent.fresh => total.amount,
|
|
_ => 0.0,
|
|
};
|
|
if added > 0.0 {
|
|
*spent.untaken.entry(total.currency.clone()).or_default() += added;
|
|
}
|
|
spent.total = Some(total);
|
|
}
|
|
|
|
fn take(&self) -> Vec<ReportedCost> {
|
|
std::mem::take(&mut self.lock().untaken)
|
|
.into_iter()
|
|
.map(|(currency, amount)| ReportedCost { amount, currency })
|
|
.collect()
|
|
}
|
|
|
|
fn lock(&self) -> std::sync::MutexGuard<'_, Spent> {
|
|
self.0.lock().unwrap_or_else(PoisonError::into_inner)
|
|
}
|
|
}
|
|
|
|
fn find<'a>(
|
|
options: &'a [stream::ConfigOption],
|
|
category: &str,
|
|
) -> Option<&'a stream::ConfigOption> {
|
|
options
|
|
.iter()
|
|
.find(|o| o.category.as_deref() == Some(category))
|
|
}
|
|
|
|
/// Marks a turn in flight for [`Canceller::cancel`] while it lives.
|
|
struct InTurn<'a>(&'a AtomicBool);
|
|
|
|
impl Drop for InTurn<'_> {
|
|
fn drop(&mut self) {
|
|
self.0.store(false, Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
/// Why a turn's prompt is being stopped.
|
|
#[derive(Clone, Copy)]
|
|
enum Stop {
|
|
Operator,
|
|
Idle,
|
|
}
|
|
|
|
/// What a turn is for.
|
|
#[derive(Clone, Copy, PartialEq, Eq)]
|
|
enum TurnKind {
|
|
/// A prompt the runtime was asked to run.
|
|
Ordinary,
|
|
/// A checkpoint or `compact` command turn. A bare `end_turn` is a normal
|
|
/// answer here, so it is not an [`AcpError::EmptyEndTurn`].
|
|
Compaction,
|
|
}
|
|
|
|
/// Turns on an ACP agent's durable session.
|
|
///
|
|
/// From the [`Config`] it reads `cwd` (the agent's working directory),
|
|
/// `mcp_config` (converted to the session's `mcpServers`) and
|
|
/// `system_prompt_file` (prepended to the first prompt of each new session,
|
|
/// since ACP has no system prompt) and `idle_timeout` (with no
|
|
/// `session/update` for that long the turn is cancelled, and fails with
|
|
/// [`AcpError::IdleTimeout`]). Everything else in it is claude's and is
|
|
/// ignored, except `model` and `effort`: before each prompt they are set as
|
|
/// the session's `model` and then `thought_level` config option, when the
|
|
/// session offers that value. [`Runtime::choices`] reads what it offers.
|
|
///
|
|
/// Compaction, proactive once `policy` says so after a turn or on
|
|
/// [`Runtime::compact`], runs the agent's advertised `compact` command as a
|
|
/// prompt on the same session. An agent advertising none, or whose command
|
|
/// fails, gets the policy's checkpoint turn instead (unless it already ran),
|
|
/// and the session is archived, so the next turn starts a new one.
|
|
pub struct AcpRuntime<P: CompactionPolicy> {
|
|
command: AcpCommand,
|
|
session_file: PathBuf,
|
|
permit: PermissionPolicy,
|
|
policy: P,
|
|
live: Mutex<Option<Live>>,
|
|
cancel: Arc<CancelState>,
|
|
cancel_grace: Duration,
|
|
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<KeyVar>,
|
|
load_session: bool,
|
|
/// The recorded session, or a new one whose first prompt is unanswered.
|
|
loaded: Option<String>,
|
|
model: Option<String>,
|
|
/// The commands the agent last advertised for `loaded`; `None` until it
|
|
/// has.
|
|
commands: Option<Vec<String>>,
|
|
/// The config options `loaded` last reported.
|
|
choices: Choices,
|
|
spend: Spend,
|
|
}
|
|
|
|
impl Live {
|
|
fn offer(&mut self, options: Vec<stream::ConfigOption>) {
|
|
if let Some(model) = options
|
|
.iter()
|
|
.find(|o| o.category.as_deref() == Some(MODEL_CATEGORY))
|
|
{
|
|
self.model = Some(model.current.clone());
|
|
}
|
|
self.choices.replace(options);
|
|
}
|
|
}
|
|
|
|
impl<P: CompactionPolicy> AcpRuntime<P> {
|
|
/// A runtime spawning `command`, keeping its session id in
|
|
/// `session_file`, answering its permission requests with `permit`, and
|
|
/// compacting per `policy`.
|
|
#[must_use]
|
|
pub fn new(
|
|
command: AcpCommand,
|
|
session_file: PathBuf,
|
|
permit: PermissionPolicy,
|
|
policy: P,
|
|
) -> Self {
|
|
Self {
|
|
command,
|
|
session_file,
|
|
permit,
|
|
policy,
|
|
live: Mutex::new(None),
|
|
cancel: Arc::default(),
|
|
cancel_grace: CANCEL_GRACE,
|
|
compact_idle: COMPACT_IDLE,
|
|
choices: Choices::default(),
|
|
spend: Spend::default(),
|
|
resolve_key: provider_key::resolver(),
|
|
}
|
|
}
|
|
|
|
/// Spawn the agent with `key` in the variable
|
|
/// [`AcpCommand::api_key_env`] names, in place of the key read from the
|
|
/// agent's own store path. For a caller that is not an agent and has read
|
|
/// its key itself; without `api_key_env` the key is not used.
|
|
#[must_use]
|
|
pub fn with_provider_key(mut self, key: String) -> Self {
|
|
self.resolve_key = Arc::new(move |name| {
|
|
let value = key.clone();
|
|
Box::pin(async move { Some(KeyVar { name, value }) })
|
|
});
|
|
self
|
|
}
|
|
|
|
/// Names a new session until its first prompt has been answered, so a
|
|
/// failed first prompt is retried on the same session.
|
|
fn pending_file(&self) -> PathBuf {
|
|
let mut path = self.session_file.clone().into_os_string();
|
|
path.push(".pending");
|
|
PathBuf::from(path)
|
|
}
|
|
|
|
async fn start(&self, config: &Config) -> Result<Live> {
|
|
let (_, servers) = mcp_servers(config)?;
|
|
let cwd = session_cwd(config);
|
|
let key = match &self.command.api_key_env {
|
|
Some(name) => (self.resolve_key)(name.clone()).await,
|
|
None => None,
|
|
};
|
|
let conn = Connection::spawn(
|
|
&self.command,
|
|
key.as_ref(),
|
|
&cwd,
|
|
self.permit.clone(),
|
|
servers,
|
|
)?;
|
|
let init = conn
|
|
.request(
|
|
"initialize",
|
|
json!({
|
|
"protocolVersion": PROTOCOL_VERSION,
|
|
"clientCapabilities": {
|
|
"fs": { "readTextFile": false, "writeTextFile": false },
|
|
"terminal": false,
|
|
},
|
|
"clientInfo": { "name": "hyperhive", "version": env!("CARGO_PKG_VERSION") },
|
|
}),
|
|
)
|
|
.await?;
|
|
let caps = &init["agentCapabilities"];
|
|
if caps["mcpCapabilities"]["http"] != Value::Bool(true) {
|
|
return Err(AcpError::NoHttpMcp.into());
|
|
}
|
|
tracing::info!(agent = %init["agentInfo"], "ACP agent initialized");
|
|
self.choices.reset();
|
|
Ok(Live {
|
|
conn,
|
|
key,
|
|
load_session: caps["loadSession"] == Value::Bool(true),
|
|
loaded: None,
|
|
model: None,
|
|
commands: None,
|
|
choices: self.choices.clone(),
|
|
spend: self.spend.clone(),
|
|
})
|
|
}
|
|
|
|
/// Attach the durable session: the recorded one (named by the session
|
|
/// file), else a new one whose first prompt was never answered (named by
|
|
/// the pending file), else a new one. Returns its id and whether its first
|
|
/// prompt, the one carrying the system prompt, is still to be answered;
|
|
/// [`Self::turn`] records the session once it has been.
|
|
async fn attach(
|
|
&self,
|
|
live: &mut Live,
|
|
cwd: &Path,
|
|
servers: &[Value],
|
|
) -> Result<(String, bool)> {
|
|
let known = [
|
|
(read_id(&self.session_file), false),
|
|
(read_id(&self.pending_file()), true),
|
|
]
|
|
.into_iter()
|
|
.filter_map(|(id, new)| Some((id?, new)));
|
|
for (id, new) in known {
|
|
if live.loaded.as_deref() == Some(id.as_str()) {
|
|
return Ok((id, new));
|
|
}
|
|
if !live.load_session {
|
|
continue;
|
|
}
|
|
let params = json!({ "sessionId": id, "cwd": cwd, "mcpServers": servers });
|
|
match live.conn.request("session/load", params).await {
|
|
Ok(response) => {
|
|
live.model = stream::session_model(&response);
|
|
live.offer(stream::config_options(&response).unwrap_or_default());
|
|
live.loaded = Some(id.clone());
|
|
live.spend.attach(false);
|
|
live.commands = None;
|
|
discard_stale(live)?;
|
|
return Ok((id, new));
|
|
}
|
|
Err(e @ AcpError::Rpc { .. }) => {
|
|
tracing::warn!(error = %e, session = %id, "ACP session/load failed");
|
|
}
|
|
Err(e) => return Err(e.into()),
|
|
}
|
|
}
|
|
let params = json!({ "cwd": cwd, "mcpServers": servers });
|
|
let response = live.conn.request("session/new", params).await?;
|
|
let id = response["sessionId"]
|
|
.as_str()
|
|
.ok_or(AcpError::Malformed {
|
|
method: "session/new",
|
|
field: "sessionId",
|
|
})?
|
|
.to_owned();
|
|
live.model = stream::session_model(&response);
|
|
live.offer(stream::config_options(&response).unwrap_or_default());
|
|
live.loaded = Some(id.clone());
|
|
live.spend.attach(true);
|
|
live.commands = None;
|
|
write_id(&self.pending_file(), &id)?;
|
|
Ok((id, true))
|
|
}
|
|
|
|
fn record_session(&self, id: &str) -> std::result::Result<(), AcpError> {
|
|
write_id(&self.session_file, id)?;
|
|
remove_if_present(&self.pending_file())
|
|
}
|
|
|
|
async fn turn(
|
|
&self,
|
|
live: &mut Live,
|
|
config: &Config,
|
|
prompt: &str,
|
|
kind: TurnKind,
|
|
sink: &impl Sink,
|
|
mut cancelled: Pin<&mut Notified<'_>>,
|
|
) -> Result<Progress> {
|
|
let cwd = session_cwd(config);
|
|
let (servers, names) = mcp_servers(config)?;
|
|
let (session, created) = self.attach(live, &cwd, &servers).await?;
|
|
choose_model_and_effort(live, &session, config, sink).await?;
|
|
let text = prompt_text(config, prompt, created);
|
|
let params = json!({ "sessionId": session, "prompt": [{ "type": "text", "text": text }] });
|
|
discard_stale(live)?;
|
|
let mut reply = live.conn.send("session/prompt", params).await?;
|
|
let mut mapper = StreamMapper::new(names);
|
|
let mut stderr_tail = VecDeque::new();
|
|
let idle = config.idle_timeout;
|
|
let mut quiet_until = idle.map(|window| Instant::now() + window);
|
|
let mut stop: Option<(Stop, Instant)> = None;
|
|
let response = loop {
|
|
let now = Instant::now();
|
|
tokio::select! {
|
|
// Notifications first: the ones sent before the reply belong
|
|
// to this turn and must reach the sink before it ends.
|
|
biased;
|
|
incoming = live.conn.incoming.recv() => {
|
|
let incoming = incoming.unwrap_or(Incoming::Closed);
|
|
if matches!(incoming, Incoming::Update(_)) {
|
|
quiet_until = idle.map(|window| Instant::now() + window);
|
|
}
|
|
if !deliver(incoming, &mut mapper, live, &mut stderr_tail, sink) {
|
|
let stderr_tail = Vec::from(stderr_tail).join("\n");
|
|
return Err(AcpError::Exited { stderr_tail }.into());
|
|
}
|
|
}
|
|
reply = &mut reply => break reply.unwrap_or(Err(AcpError::Closed))?,
|
|
() = cancelled.as_mut(), if stop.is_none() => {
|
|
tracing::info!("cancelling the ACP turn");
|
|
live.conn.notify("session/cancel", json!({ "sessionId": session })).await?;
|
|
stop = Some((Stop::Operator, Instant::now() + self.cancel_grace));
|
|
}
|
|
() = tokio::time::sleep_until(quiet_until.unwrap_or(now)),
|
|
if stop.is_none() && quiet_until.is_some() =>
|
|
{
|
|
tracing::warn!(?idle, "ACP agent silent; cancelling the turn");
|
|
live.conn.notify("session/cancel", json!({ "sessionId": session })).await?;
|
|
stop = Some((Stop::Idle, Instant::now() + self.cancel_grace));
|
|
}
|
|
() = tokio::time::sleep_until(stop.map_or(now, |(_, deadline)| deadline)),
|
|
if stop.is_some() =>
|
|
{
|
|
let grace_secs = self.cancel_grace.as_secs();
|
|
return Err(match stop {
|
|
Some((Stop::Idle, _)) => AcpError::IdleKilled {
|
|
silent_secs: idle.unwrap_or_default().as_secs(),
|
|
grace_secs,
|
|
},
|
|
_ => AcpError::CancelIgnored { grace_secs },
|
|
}
|
|
.into());
|
|
}
|
|
}
|
|
};
|
|
if created {
|
|
self.record_session(&session)?;
|
|
}
|
|
// Updates may still arrive after the reply: an agent can forward its
|
|
// stream asynchronously. Take them until the agent goes quiet.
|
|
let settled = tokio::time::Instant::now() + SETTLE_MAX;
|
|
while let Ok(Some(incoming)) = tokio::time::timeout_at(
|
|
settled.min(tokio::time::Instant::now() + SETTLE_QUIET),
|
|
live.conn.incoming.recv(),
|
|
)
|
|
.await
|
|
{
|
|
if !deliver(incoming, &mut mapper, live, &mut stderr_tail, sink) {
|
|
break;
|
|
}
|
|
}
|
|
for event in mapper.finish() {
|
|
sink.on_event(&event);
|
|
}
|
|
let stop_reason = match stop {
|
|
Some((Stop::Idle, _)) => {
|
|
return Err(AcpError::IdleTimeout {
|
|
silent_secs: idle.unwrap_or_default().as_secs(),
|
|
}
|
|
.into());
|
|
}
|
|
// Agents differ in the reason they give a cancelled prompt.
|
|
Some((Stop::Operator, _)) => Some("cancelled"),
|
|
None => response["stopReason"].as_str(),
|
|
};
|
|
match stop_reason {
|
|
Some("end_turn") if kind == TurnKind::Ordinary && mapper.sent_nothing(&response) => {
|
|
return Err(AcpError::EmptyEndTurn.into());
|
|
}
|
|
Some("end_turn") => {}
|
|
reason => {
|
|
let reason = reason.unwrap_or("none");
|
|
tracing::warn!(reason, "ACP turn stopped early");
|
|
sink.on_stderr_line(&format!("ACP turn stopped: {reason}"));
|
|
}
|
|
}
|
|
let mut progress = Progress::default();
|
|
progress.created = created;
|
|
progress.telemetry = mapper.telemetry(&response, live.model.as_deref());
|
|
Ok(progress)
|
|
}
|
|
|
|
/// The running agent, spawned if there is none or it has exited.
|
|
async fn running<'g>(
|
|
&self,
|
|
guard: &'g mut Option<Live>,
|
|
config: &Config,
|
|
) -> Result<&'g mut Live> {
|
|
let mut live = guard.take();
|
|
if live.as_mut().is_some_and(|live| live.conn.exited()) {
|
|
tracing::warn!("ACP agent has exited; respawning");
|
|
live = None;
|
|
}
|
|
let live = match live {
|
|
Some(live) => live,
|
|
None => self.start(config).await?,
|
|
};
|
|
Ok(guard.insert(live))
|
|
}
|
|
|
|
/// Run `prompt` as one turn on the durable session.
|
|
async fn prompt(
|
|
&self,
|
|
guard: &mut Option<Live>,
|
|
config: &Config,
|
|
prompt: &str,
|
|
kind: TurnKind,
|
|
sink: &impl Sink,
|
|
) -> Result<Progress> {
|
|
// Registered before the turn is marked in flight, so a cancel is
|
|
// never missed; one that comes before the prompt is sent stops it
|
|
// right after.
|
|
let cancelled = self.cancel.requested.notified();
|
|
tokio::pin!(cancelled);
|
|
cancelled.as_mut().enable();
|
|
self.cancel.in_turn.store(true, Ordering::SeqCst);
|
|
let _in_turn = InTurn(&self.cancel.in_turn);
|
|
let live = self.running(guard, config).await?;
|
|
let result = self.turn(live, config, prompt, kind, sink, cancelled).await;
|
|
if let Err(Error::Acp(
|
|
AcpError::Closed
|
|
| AcpError::Exited { .. }
|
|
| AcpError::Io(_)
|
|
| AcpError::Malformed { .. }
|
|
| AcpError::IdleKilled { .. }
|
|
| AcpError::CancelIgnored { .. },
|
|
)) = &result
|
|
{
|
|
*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
|
|
/// session. A `compact` command that fails, or is silent for
|
|
/// `compact_idle`, is followed by the same archiving, so a compaction
|
|
/// never leaves the session as full as it found it.
|
|
async fn compact_session(
|
|
&self,
|
|
guard: &mut Option<Live>,
|
|
config: &Config,
|
|
sink: &impl Sink,
|
|
checkpoint: bool,
|
|
) -> Result<()> {
|
|
if read_id(&self.session_file).is_none() {
|
|
return Ok(());
|
|
}
|
|
let live = self.running(guard, config).await?;
|
|
let (servers, _) = mcp_servers(config)?;
|
|
let (_, created) = self.attach(live, &session_cwd(config), &servers).await?;
|
|
if created {
|
|
return Ok(());
|
|
}
|
|
let mut checkpointed = false;
|
|
if advertises_compact(live).await? {
|
|
if checkpoint {
|
|
self.checkpoint(guard, config, sink).await;
|
|
checkpointed = true;
|
|
}
|
|
let idle = config
|
|
.idle_timeout
|
|
.map_or(self.compact_idle, |idle| idle.min(self.compact_idle));
|
|
let bounded = Config {
|
|
idle_timeout: Some(idle),
|
|
..config.clone()
|
|
};
|
|
let command = format!("/{COMPACT_COMMAND}");
|
|
let Err(e) = self
|
|
.prompt(guard, &bounded, &command, TurnKind::Compaction, sink)
|
|
.await
|
|
else {
|
|
return Ok(());
|
|
};
|
|
tracing::warn!(error = %e, "ACP compact command failed; starting a new session");
|
|
sink.on_stderr_line(&format!(
|
|
"ACP {command} failed, starting a new session: {e}"
|
|
));
|
|
} else {
|
|
tracing::info!("ACP agent advertises no compact command; starting a new session");
|
|
}
|
|
if !checkpointed {
|
|
self.checkpoint(guard, config, sink).await;
|
|
}
|
|
self.archive()?;
|
|
Ok(())
|
|
}
|
|
|
|
/// The policy's checkpoint turn, if it has one. Best-effort: compaction
|
|
/// goes ahead if it fails.
|
|
async fn checkpoint(&self, guard: &mut Option<Live>, config: &Config, sink: &impl Sink) {
|
|
let Some(prompt) = self.policy.checkpoint_prompt() else {
|
|
return;
|
|
};
|
|
if let Err(e) = self
|
|
.prompt(guard, config, prompt, TurnKind::Compaction, sink)
|
|
.await
|
|
{
|
|
tracing::warn!(error = %e, "ACP checkpoint turn failed");
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<P: CompactionPolicy> Runtime for AcpRuntime<P> {
|
|
async fn run(&self, config: &Config, prompt: &str, sink: &impl Sink) -> Result<Progress> {
|
|
let mut guard = self.live.lock().await;
|
|
let mut progress = self
|
|
.prompt(&mut guard, config, prompt, TurnKind::Ordinary, sink)
|
|
.await?;
|
|
if self.policy.should_compact(progress.telemetry.usage()) {
|
|
match self.compact_session(&mut guard, config, sink, true).await {
|
|
Ok(()) => progress.compacted = true,
|
|
Err(e) => tracing::warn!(error = %e, "ACP compaction failed"),
|
|
}
|
|
}
|
|
Ok(progress)
|
|
}
|
|
|
|
async fn compact(&self, config: &Config, sink: &impl Sink) -> Result<()> {
|
|
let mut guard = self.live.lock().await;
|
|
self.compact_session(&mut guard, config, sink, false).await
|
|
}
|
|
|
|
fn canceller(&self) -> Option<Canceller> {
|
|
Some(Canceller(self.cancel.clone()))
|
|
}
|
|
|
|
fn choices(&self) -> Option<Choices> {
|
|
Some(self.choices.clone())
|
|
}
|
|
|
|
fn take_reported_cost(&self) -> Vec<ReportedCost> {
|
|
self.spend.take()
|
|
}
|
|
|
|
/// Moves the session file aside; the agent keeps its own copy of the
|
|
/// session, so nothing is lost. A new session whose first prompt was never
|
|
/// answered is dropped too.
|
|
fn archive(&self) -> Result<Option<PathBuf>> {
|
|
remove_if_present(&self.pending_file()).map_err(Error::Acp)?;
|
|
if !self.session_file.exists() {
|
|
return Ok(None);
|
|
}
|
|
let stamp = std::time::SystemTime::now()
|
|
.duration_since(std::time::UNIX_EPOCH)
|
|
.map_or(0, |d| d.as_secs());
|
|
let mut archived = self.session_file.clone().into_os_string();
|
|
archived.push(format!(".{stamp}.archived"));
|
|
let archived = PathBuf::from(archived);
|
|
std::fs::rename(&self.session_file, &archived).map_err(AcpError::Io)?;
|
|
Ok(Some(archived))
|
|
}
|
|
}
|
|
|
|
/// Hand one incoming item to the sink, and keep the commands the agent
|
|
/// advertises for the loaded session. Returns `false` once the agent has
|
|
/// closed its output.
|
|
fn deliver(
|
|
incoming: Incoming,
|
|
mapper: &mut StreamMapper,
|
|
agent: &mut Live,
|
|
stderr_tail: &mut VecDeque<String>,
|
|
sink: &impl Sink,
|
|
) -> bool {
|
|
match incoming {
|
|
Incoming::Update(params) => {
|
|
if params["sessionId"].as_str() == agent.loaded.as_deref() {
|
|
keep_advertised(agent, ¶ms["update"]);
|
|
for event in mapper.push(¶ms["update"]) {
|
|
sink.on_event(&event);
|
|
}
|
|
}
|
|
}
|
|
Incoming::Stdout(line) => sink.on_stdout_line(&line),
|
|
Incoming::Stderr(line) => {
|
|
sink.on_stderr_line(&line);
|
|
if stderr_tail.len() == STDERR_TAIL {
|
|
stderr_tail.pop_front();
|
|
}
|
|
stderr_tail.push_back(line);
|
|
}
|
|
Incoming::Closed => return false,
|
|
}
|
|
true
|
|
}
|
|
|
|
/// Keep the commands and config options an `update` for the loaded session
|
|
/// advertises, and the cost it reports.
|
|
fn keep_advertised(agent: &mut Live, update: &Value) {
|
|
if let Some(cost) = stream::reported_cost(update) {
|
|
agent.spend.observe(cost);
|
|
}
|
|
if let Some(advertised) = stream::advertised_commands(update) {
|
|
agent.commands = Some(advertised);
|
|
}
|
|
if let Some(options) = stream::updated_config_options(update) {
|
|
agent.offer(options);
|
|
}
|
|
}
|
|
|
|
/// Drop what the agent sent outside a turn: the history a `session/load`
|
|
/// replays (the caller already has it), and anything that arrived after the
|
|
/// previous turn settled, which must not be shown as part of the next one.
|
|
/// The commands and config options the agent advertises for the loaded
|
|
/// session, and the cost it reports, are kept.
|
|
fn discard_stale(live: &mut Live) -> std::result::Result<(), AcpError> {
|
|
while let Ok(incoming) = live.conn.incoming.try_recv() {
|
|
between_turns(live, incoming)?;
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
fn between_turns(agent: &mut Live, incoming: Incoming) -> std::result::Result<(), AcpError> {
|
|
match incoming {
|
|
Incoming::Update(params) => {
|
|
if params["sessionId"].as_str() == agent.loaded.as_deref() {
|
|
keep_advertised(agent, ¶ms["update"]);
|
|
}
|
|
}
|
|
Incoming::Stdout(line) | Incoming::Stderr(line) => {
|
|
tracing::info!(line = %line, "ACP agent output between turns");
|
|
}
|
|
Incoming::Closed => return Err(AcpError::Closed),
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// Whether the agent advertises a `compact` command for the loaded session,
|
|
/// waiting up to [`COMMANDS_WAIT`] for its commands if it has sent none yet.
|
|
async fn advertises_compact(live: &mut Live) -> std::result::Result<bool, AcpError> {
|
|
discard_stale(live)?;
|
|
let deadline = Instant::now() + COMMANDS_WAIT;
|
|
while live.commands.is_none() {
|
|
match tokio::time::timeout_at(deadline, live.conn.incoming.recv()).await {
|
|
Ok(incoming) => between_turns(live, incoming.unwrap_or(Incoming::Closed))?,
|
|
Err(_) => break,
|
|
}
|
|
}
|
|
Ok(live
|
|
.commands
|
|
.iter()
|
|
.flatten()
|
|
.any(|name| name == COMPACT_COMMAND))
|
|
}
|
|
|
|
/// Set `config`'s model, then its effort, on `session`: the effort levels
|
|
/// on offer depend on the model.
|
|
async fn choose_model_and_effort(
|
|
live: &mut Live,
|
|
session: &str,
|
|
config: &Config,
|
|
sink: &impl Sink,
|
|
) -> std::result::Result<(), AcpError> {
|
|
choose(live, session, MODEL_CATEGORY, config.model.as_deref(), sink).await?;
|
|
choose(
|
|
live,
|
|
session,
|
|
EFFORT_CATEGORY,
|
|
config.effort.as_deref(),
|
|
sink,
|
|
)
|
|
.await
|
|
}
|
|
|
|
/// Set `session`'s config option of `category` to `wanted`, if the session
|
|
/// offers that value and holds another. A refusal is reported to `sink` and
|
|
/// kept as the choice's [`Choice::refused`] until the agent accepts a value
|
|
/// or `wanted` is none or the value the session holds; the turn goes ahead on
|
|
/// the value the session holds.
|
|
async fn choose(
|
|
live: &mut Live,
|
|
session: &str,
|
|
category: &'static str,
|
|
wanted: Option<&str>,
|
|
sink: &impl Sink,
|
|
) -> std::result::Result<(), AcpError> {
|
|
let Some(option) = live.choices.find(category) else {
|
|
return Ok(());
|
|
};
|
|
if wanted.is_none_or(|w| w == option.current) {
|
|
live.choices.accept(category);
|
|
return Ok(());
|
|
}
|
|
let Some(wanted) = wanted.filter(|w| option.values.iter().any(|v| v == w)) else {
|
|
return Ok(());
|
|
};
|
|
let params = json!({ "sessionId": session, "configId": option.id, "value": wanted });
|
|
match live.conn.request("session/set_config_option", params).await {
|
|
Ok(response) => {
|
|
live.choices.accept(category);
|
|
if let Some(options) = stream::config_options(&response) {
|
|
live.offer(options);
|
|
}
|
|
}
|
|
Err(e @ AcpError::Rpc { .. }) => {
|
|
tracing::warn!(error = %e, category, wanted, "ACP session/set_config_option failed");
|
|
sink.on_stderr_line(&format!("ACP agent refused {category} {wanted}: {e}"));
|
|
live.choices.refuse(category, wanted);
|
|
}
|
|
Err(e) => return Err(e),
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// The prompt as sent: the first one of a session carries the system prompt.
|
|
fn prompt_text(config: &Config, prompt: &str, first: bool) -> String {
|
|
match (first, &config.system_prompt_file) {
|
|
(true, Some(path)) => match std::fs::read_to_string(path) {
|
|
Ok(system) => format!("{system}\n\n{prompt}"),
|
|
Err(e) => {
|
|
tracing::warn!(path = %path.display(), error = %e, "system prompt unreadable");
|
|
prompt.to_owned()
|
|
}
|
|
},
|
|
_ => prompt.to_owned(),
|
|
}
|
|
}
|
|
|
|
fn read_id(file: &Path) -> Option<String> {
|
|
std::fs::read_to_string(file)
|
|
.ok()
|
|
.map(|s| s.trim().to_owned())
|
|
.filter(|s| !s.is_empty())
|
|
}
|
|
|
|
fn write_id(file: &Path, id: &str) -> std::result::Result<(), AcpError> {
|
|
if let Some(dir) = file.parent() {
|
|
std::fs::create_dir_all(dir)?;
|
|
}
|
|
std::fs::write(file, id)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn remove_if_present(file: &Path) -> std::result::Result<(), AcpError> {
|
|
match std::fs::remove_file(file) {
|
|
Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e.into()),
|
|
_ => Ok(()),
|
|
}
|
|
}
|
|
|
|
fn session_cwd(config: &Config) -> PathBuf {
|
|
config
|
|
.cwd
|
|
.clone()
|
|
.or_else(|| std::env::current_dir().ok())
|
|
.unwrap_or_else(|| PathBuf::from("."))
|
|
}
|
|
|
|
fn mcp_servers(config: &Config) -> Result<(Vec<Value>, Vec<String>)> {
|
|
let Some(path) = &config.mcp_config else {
|
|
return Ok((Vec::new(), Vec::new()));
|
|
};
|
|
let fail = |detail: String| AcpError::McpConfig {
|
|
path: path.clone(),
|
|
detail,
|
|
};
|
|
let raw = std::fs::read_to_string(path).map_err(|e| fail(e.to_string()))?;
|
|
let parsed: Value = serde_json::from_str(&raw).map_err(|e| fail(e.to_string()))?;
|
|
Ok(stream::mcp_servers(&parsed))
|
|
}
|
|
|
|
#[cfg(test)]
|
|
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, KeyVar, ReportedCost, SessionChoices, Spend, TurnKind,
|
|
};
|
|
use crate::{AcpCommand, Error, Runtime, TurnStop};
|
|
|
|
/// An ACP agent in plain `sh`. It answers `initialize`, numbers its
|
|
/// sessions `s1`, `s2`, …, appends every method it is sent to
|
|
/// `<log>.methods` (and `start` when it starts), and every
|
|
/// `session/prompt` line to `<log>`. It answers every prompt with
|
|
/// `end_turn` and a token `usage`, except the very first one it is sent
|
|
/// across restarts, which it handles per its mode:
|
|
///
|
|
/// - `fail`: an error;
|
|
/// - `silent`: nothing until `session/cancel`, then `end_turn` — the
|
|
/// reply an agent gives a cancelled prompt it was retrying against a
|
|
/// rate-limited (HTTP 429) provider without reporting it;
|
|
/// - `deaf`: nothing, and `session/cancel` is ignored;
|
|
/// - `trickle`: five text chunks 100ms apart, then `end_turn` with no
|
|
/// `usage`;
|
|
/// - `blank`: `end_turn` with no `usage` and nothing before it — how
|
|
/// opencode ends a turn its provider rejected;
|
|
/// - `ok`: `end_turn` and `usage`, like the rest.
|
|
///
|
|
/// With `COMMANDS` set, it advertises that one command on each session it
|
|
/// creates or loads. With `USED` set, it reports that many of 1000 context
|
|
/// tokens used on each prompt to session `s1`. With `STALL_COMPACT` set, it
|
|
/// answers `/compact` only when cancelled, as `silent` does. With `BLANK`
|
|
/// set, it answers a prompt whose text is exactly that as `blank` does.
|
|
/// With `OPTIONS` set, its sessions offer models `m/think` (the default)
|
|
/// and `m/plain`, and effort levels `low` (the default) and `high` on
|
|
/// `m/think` only; it appends each `session/set_config_option` to
|
|
/// `<log>.sets` as `id=value`. With `REFUSE` set, it answers the first
|
|
/// `session/set_config_option` to that value with an error. With `COST`
|
|
/// set, each prompt adds that many dollars to its session's total (from zero,
|
|
/// kept in `<log>.cost.<session>` across restarts) and reports it.
|
|
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 <log>.keys.
|
|
if [ -n "$TEST_PROVIDER_KEY" ]; then
|
|
printf '%s\n' "$TEST_PROVIDER_KEY" >> "$1.keys"
|
|
fi
|
|
# With TRAP_INT set: SIGINT creates <log>.int and exits 0.
|
|
if [ -n "$TRAP_INT" ]; then
|
|
trap ': > "$1.int"; exit 0' INT
|
|
fi
|
|
# With GRANDCHILD set: a background `sleep`, its pid in <log>.grandchild.
|
|
# `&` starts it with SIGINT ignored; `env` puts the default back.
|
|
if [ -n "$GRANDCHILD" ]; then
|
|
sh -c 'echo $$ >"$1.grandchild"; exec env --default-signal=INT sleep 300' sh "$1" >/dev/null 2>&1 &
|
|
fi
|
|
update() {
|
|
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"%s","update":%s}}\n' "$1" "$2"
|
|
}
|
|
advertise() {
|
|
[ -n "$COMMANDS" ] && update "$1" "{\"sessionUpdate\":\"available_commands_update\",\"availableCommands\":[{\"name\":\"$COMMANDS\",\"description\":\"\"}]}"
|
|
}
|
|
ended() {
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn","usage":{"inputTokens":1,"outputTokens":1}}}\n' "$1"
|
|
}
|
|
bare() {
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$1"
|
|
}
|
|
options() {
|
|
[ -n "$OPTIONS" ] || return 0
|
|
o="{\"id\":\"model\",\"name\":\"Model\",\"category\":\"model\",\"type\":\"select\",\"currentValue\":\"$model\",\"options\":[{\"value\":\"m/think\",\"name\":\"Think\"},{\"value\":\"m/plain\",\"name\":\"Plain\"}]}"
|
|
[ "$model" = m/think ] && o="$o,{\"id\":\"effort\",\"name\":\"Effort\",\"category\":\"thought_level\",\"type\":\"select\",\"currentValue\":\"$effort\",\"options\":[{\"value\":\"low\",\"name\":\"Low\"},{\"value\":\"high\",\"name\":\"High\"}]}"
|
|
printf ',"configOptions":[%s]' "$o"
|
|
}
|
|
result() {
|
|
o=$(options)
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{%s}}\n' "$1" "${o#,}"
|
|
}
|
|
while IFS= read -r line; do
|
|
id=${line#*\"id\":}; id=${id%%[,\}]*}
|
|
m=${line#*\"method\":\"}; m=${m%%\"*}
|
|
sid=${line#*\"sessionId\":\"}; sid=${sid%%\"*}
|
|
printf '%s\n' "$m" >> "$1.methods"
|
|
case $m in
|
|
initialize)
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"protocolVersion":1,"agentCapabilities":{"loadSession":true,"mcpCapabilities":{"http":true}}}}\n' "$id" ;;
|
|
session/new)
|
|
n=$((n+1))
|
|
rm -f "$1.cost.s$n"
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"sessionId":"s%s"%s}}\n' "$id" "$n" "$(options)"
|
|
advertise "s$n" ;;
|
|
session/load)
|
|
result "$id"
|
|
advertise "$sid" ;;
|
|
session/set_config_option)
|
|
cid=${line#*\"configId\":\"}; cid=${cid%%\"*}
|
|
val=${line#*\"value\":\"}; val=${val%%\"*}
|
|
printf '%s=%s\n' "$cid" "$val" >> "$1.sets"
|
|
if [ -n "$REFUSE" ] && [ "$val" = "$REFUSE" ] && [ ! -e "$1.refused" ]; then
|
|
: > "$1.refused"
|
|
printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32602,"message":"unknown value"}}\n' "$id"
|
|
continue
|
|
fi
|
|
case $cid in model) model=$val ;; effort) effort=$val ;; esac
|
|
result "$id" ;;
|
|
session/prompt)
|
|
printf '%s\n' "$line" >> "$1"
|
|
[ -n "$USED" ] && [ "$sid" = s1 ] && update "$sid" "{\"sessionUpdate\":\"usage_update\",\"used\":$USED,\"size\":1000}"
|
|
if [ -n "$COST" ]; then
|
|
total=$(( $(cat "$1.cost.$sid" 2>/dev/null || echo 0) + COST ))
|
|
printf '%s' "$total" > "$1.cost.$sid"
|
|
update "$sid" "{\"sessionUpdate\":\"usage_update\",\"cost\":{\"amount\":$total,\"currency\":\"USD\"}}"
|
|
fi
|
|
case $line in
|
|
*'"text":"/compact"'*) [ -n "$STALL_COMPACT" ] && prompt=$id && continue ;;
|
|
esac
|
|
case $line in
|
|
*"\"text\":\"$BLANK\""*) [ -n "$BLANK" ] && bare "$id" && continue ;;
|
|
esac
|
|
if [ -e "$1.once" ]; then
|
|
ended "$id"
|
|
else
|
|
: > "$1.once"
|
|
case $2 in
|
|
fail)
|
|
printf '{"jsonrpc":"2.0","id":%s,"error":{"code":-32603,"message":"provider down"}}\n' "$id" ;;
|
|
ok) ended "$id" ;;
|
|
blank) bare "$id" ;;
|
|
silent) prompt=$id ;;
|
|
trickle)
|
|
for _ in 1 2 3 4 5; do
|
|
sleep 0.1
|
|
printf '{"jsonrpc":"2.0","method":"session/update","params":{"sessionId":"s%s","update":{"sessionUpdate":"agent_message_chunk","content":{"type":"text","text":"x"}}}}\n' "$n"
|
|
done
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$id" ;;
|
|
esac
|
|
fi ;;
|
|
session/cancel)
|
|
if [ -n "$prompt" ]; then
|
|
printf '{"jsonrpc":"2.0","id":%s,"result":{"stopReason":"end_turn"}}\n' "$prompt"
|
|
prompt=
|
|
fi ;;
|
|
esac
|
|
done
|
|
"#;
|
|
|
|
fn runtime(dir: &Path, mode: &str) -> AcpRuntime<PercentPolicy> {
|
|
agent(dir, mode, &[], PercentPolicy::default())
|
|
}
|
|
|
|
fn agent(
|
|
dir: &Path,
|
|
mode: &str,
|
|
env: &[(&str, &str)],
|
|
policy: PercentPolicy,
|
|
) -> AcpRuntime<PercentPolicy> {
|
|
let script = dir.join("agent.sh");
|
|
std::fs::write(&script, AGENT).unwrap();
|
|
let command = AcpCommand {
|
|
command: "/bin/sh".into(),
|
|
args: vec![
|
|
script.display().to_string(),
|
|
dir.join("prompts").display().to_string(),
|
|
mode.to_owned(),
|
|
],
|
|
env: env
|
|
.iter()
|
|
.map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
|
|
.collect::<BTreeMap<_, _>>(),
|
|
api_key_env: None,
|
|
};
|
|
AcpRuntime::new(
|
|
command,
|
|
dir.join("session"),
|
|
Arc::new(|_: &_| false),
|
|
policy,
|
|
)
|
|
}
|
|
|
|
/// Compacts at `percent` of the window, after a `CHECKPOINT` turn.
|
|
fn policy(percent: u8) -> PercentPolicy {
|
|
PercentPolicy {
|
|
percent,
|
|
default_window: None,
|
|
checkpoint_prompt: Some("CHECKPOINT".to_owned()),
|
|
}
|
|
}
|
|
|
|
/// Each prompt the agent was sent: its session and its text.
|
|
fn prompts(dir: &Path) -> Vec<(String, String)> {
|
|
std::fs::read_to_string(dir.join("prompts"))
|
|
.unwrap()
|
|
.lines()
|
|
.map(|line| {
|
|
let request: serde_json::Value = serde_json::from_str(line).unwrap();
|
|
let params = &request["params"];
|
|
let text = params["prompt"][0]["text"].as_str().unwrap();
|
|
(
|
|
params["sessionId"].as_str().unwrap().to_owned(),
|
|
text.to_owned(),
|
|
)
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn config(dir: &Path) -> Config {
|
|
std::fs::write(dir.join("system.md"), "SYSTEM PROMPT").unwrap();
|
|
Config {
|
|
cwd: Some(dir.to_path_buf()),
|
|
system_prompt_file: Some(dir.join("system.md")),
|
|
..Config::default()
|
|
}
|
|
}
|
|
|
|
fn idle_config(dir: &Path, idle_ms: u64) -> Config {
|
|
Config {
|
|
idle_timeout: Some(Duration::from_millis(idle_ms)),
|
|
..config(dir)
|
|
}
|
|
}
|
|
|
|
/// Which of the prompts the agent was sent carried the system prompt.
|
|
fn carried_system_prompt(dir: &Path) -> Vec<bool> {
|
|
std::fs::read_to_string(dir.join("prompts"))
|
|
.unwrap()
|
|
.lines()
|
|
.map(|line| line.contains("SYSTEM PROMPT"))
|
|
.collect()
|
|
}
|
|
|
|
/// How often the agent started, or was sent `method`.
|
|
fn count(dir: &Path, method: &str) -> usize {
|
|
std::fs::read_to_string(dir.join("prompts.methods"))
|
|
.unwrap()
|
|
.lines()
|
|
.filter(|line| *line == method)
|
|
.count()
|
|
}
|
|
|
|
fn recorded(dir: &Path) -> String {
|
|
std::fs::read_to_string(dir.join("session")).unwrap()
|
|
}
|
|
|
|
/// Keeps the stderr lines a turn reports.
|
|
#[derive(Default)]
|
|
struct Stderr(std::sync::Mutex<Vec<String>>);
|
|
|
|
impl Sink for Stderr {
|
|
fn on_stderr_line(&self, line: &str) {
|
|
self.0.lock().unwrap().push(line.to_owned());
|
|
}
|
|
}
|
|
|
|
async fn prompt_sent(dir: &Path) {
|
|
while !dir.join("prompts").exists() {
|
|
tokio::time::sleep(Duration::from_millis(10)).await;
|
|
}
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_session_whose_first_prompt_failed_is_retried_not_resumed() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let (runtime, config) = (runtime(dir.path(), "fail"), config(dir.path()));
|
|
|
|
let failed = runtime.run(&config, "one", &NoopSink).await;
|
|
assert!(
|
|
matches!(failed, Err(Error::Acp(AcpError::Rpc { .. }))),
|
|
"{failed:?}"
|
|
);
|
|
assert!(!dir.path().join("session").exists());
|
|
|
|
let retried = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert!(retried.created);
|
|
let next = runtime.run(&config, "three", &NoopSink).await.unwrap();
|
|
assert!(!next.created);
|
|
|
|
assert_eq!(carried_system_prompt(dir.path()), [true, true, false]);
|
|
assert_eq!(recorded(dir.path()), "s1");
|
|
assert_eq!(count(dir.path(), "session/new"), 1);
|
|
assert!(!dir.path().join("session.pending").exists());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn after_a_restart_a_session_whose_first_prompt_failed_is_retried() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let config = config(dir.path());
|
|
|
|
let before = runtime(dir.path(), "fail");
|
|
assert!(before.run(&config, "one", &NoopSink).await.is_err());
|
|
drop(before);
|
|
|
|
let after = runtime(dir.path(), "fail");
|
|
let retried = after.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert!(retried.created);
|
|
|
|
assert_eq!(carried_system_prompt(dir.path()), [true, true]);
|
|
assert_eq!(count(dir.path(), "start"), 2);
|
|
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_turn_silent_for_the_idle_window_is_cancelled_and_fails() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = runtime(dir.path(), "silent");
|
|
let config = idle_config(dir.path(), 300);
|
|
|
|
let stalled = runtime.run(&config, "one", &NoopSink).await;
|
|
assert!(
|
|
matches!(stalled, Err(Error::Acp(AcpError::IdleTimeout { .. }))),
|
|
"{stalled:?}"
|
|
);
|
|
assert_eq!(count(dir.path(), "session/cancel"), 1);
|
|
|
|
// The agent stopped when asked, so it keeps running and the session
|
|
// it answered on is kept.
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert!(!next.created);
|
|
assert_eq!(count(dir.path(), "start"), 1);
|
|
assert_eq!(recorded(dir.path()), "s1");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_agent_ignoring_the_idle_cancel_is_killed_and_respawned() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut runtime = runtime(dir.path(), "deaf");
|
|
runtime.cancel_grace = Duration::from_millis(200);
|
|
let config = idle_config(dir.path(), 200);
|
|
|
|
let stalled = runtime.run(&config, "one", &NoopSink).await;
|
|
assert!(
|
|
matches!(stalled, Err(Error::Acp(AcpError::IdleKilled { .. }))),
|
|
"{stalled:?}"
|
|
);
|
|
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert!(next.created);
|
|
assert_eq!(count(dir.path(), "start"), 2);
|
|
assert_eq!(count(dir.path(), "session/new"), 1);
|
|
assert_eq!(carried_system_prompt(dir.path()), [true, true]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn updates_keep_the_idle_watchdog_from_firing() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = runtime(dir.path(), "trickle");
|
|
|
|
// Five chunks 100ms apart: longer than the window in all, never
|
|
// silent for as long as it.
|
|
let done = runtime
|
|
.run(&idle_config(dir.path(), 300), "one", &NoopSink)
|
|
.await;
|
|
assert!(done.is_ok(), "{done:?}");
|
|
assert_eq!(count(dir.path(), "session/cancel"), 0);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_end_turn_with_no_output_and_no_usage_fails() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let (runtime, config) = (runtime(dir.path(), "blank"), config(dir.path()));
|
|
|
|
let empty = runtime.run(&config, "one", &NoopSink).await;
|
|
assert!(
|
|
matches!(empty, Err(Error::Acp(AcpError::EmptyEndTurn))),
|
|
"{empty:?}"
|
|
);
|
|
}
|
|
|
|
/// 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<PercentPolicy>, Arc<AtomicUsize>) {
|
|
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<String> {
|
|
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 a_given_provider_key_reaches_the_agent() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut runtime = runtime(dir.path(), "ok").with_provider_key("given".to_owned());
|
|
runtime.command.api_key_env = Some("TEST_PROVIDER_KEY".to_owned());
|
|
|
|
runtime
|
|
.run(&config(dir.path()), "one", &NoopSink)
|
|
.await
|
|
.unwrap();
|
|
assert_eq!(keys(dir.path()), ["given"]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_end_turn_with_output_but_no_usage_succeeds() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let (runtime, config) = (runtime(dir.path(), "trickle"), config(dir.path()));
|
|
|
|
let done = runtime.run(&config, "one", &NoopSink).await;
|
|
assert!(done.is_ok(), "{done:?}");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_cancelled_turn_stops_and_ends_normally() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = runtime(dir.path(), "silent");
|
|
let canceller = runtime.canceller().unwrap();
|
|
assert!(!canceller.cancel(), "no turn is in flight yet");
|
|
|
|
let (config, stderr) = (config(dir.path()), Stderr::default());
|
|
let (done, was_in_flight) = tokio::join!(runtime.run(&config, "one", &stderr), async {
|
|
prompt_sent(dir.path()).await;
|
|
canceller.cancel()
|
|
});
|
|
assert!(was_in_flight);
|
|
assert!(done.is_ok(), "{done:?}");
|
|
assert_eq!(count(dir.path(), "session/cancel"), 1);
|
|
assert!(
|
|
stderr
|
|
.0
|
|
.lock()
|
|
.unwrap()
|
|
.contains(&"ACP turn stopped: cancelled".to_owned()),
|
|
"{stderr:?}",
|
|
stderr = stderr.0.lock().unwrap()
|
|
);
|
|
assert!(!canceller.cancel(), "the turn is over");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_turn_stop_cancels_only_a_turn_in_flight() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = runtime(dir.path(), "silent");
|
|
let stop = TurnStop::from(runtime.canceller().unwrap());
|
|
assert!(!stop.stop(), "no turn is in flight yet");
|
|
|
|
let config = config(dir.path());
|
|
let (done, was_in_flight) = tokio::join!(runtime.run(&config, "one", &NoopSink), async {
|
|
prompt_sent(dir.path()).await;
|
|
stop.stop()
|
|
});
|
|
assert!(was_in_flight);
|
|
assert!(done.is_ok(), "{done:?}");
|
|
assert_eq!(count(dir.path(), "session/cancel"), 1);
|
|
assert!(!stop.stop(), "the turn is over");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_agent_ignoring_a_cancel_is_killed() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut runtime = runtime(dir.path(), "deaf");
|
|
runtime.cancel_grace = Duration::from_millis(200);
|
|
let (canceller, config) = (runtime.canceller().unwrap(), config(dir.path()));
|
|
|
|
let (done, _) = tokio::join!(runtime.run(&config, "one", &NoopSink), async {
|
|
prompt_sent(dir.path()).await;
|
|
canceller.cancel()
|
|
});
|
|
assert!(
|
|
matches!(done, Err(Error::Acp(AcpError::CancelIgnored { .. }))),
|
|
"{done:?}"
|
|
);
|
|
runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(count(dir.path(), "start"), 2);
|
|
}
|
|
|
|
/// Poll `path` until it holds a complete, newline-terminated pid.
|
|
async fn read_pid(path: &Path) -> i32 {
|
|
for _ in 0..250 {
|
|
if let Ok(text) = std::fs::read_to_string(path)
|
|
&& text.ends_with('\n')
|
|
{
|
|
return text.trim().parse().expect("pid file holds a pid");
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(20)).await;
|
|
}
|
|
unreachable!("no pid in {} within 5s", path.display());
|
|
}
|
|
|
|
/// Whether `pid` is still running. A zombie counts as gone: it has
|
|
/// exited, and `kill(pid, 0)` only keeps finding it until it is reaped.
|
|
fn is_running(pid: i32) -> bool {
|
|
// SAFETY: signal 0 delivers nothing; it only checks that `pid` exists.
|
|
if unsafe { libc::kill(pid, 0) } != 0 {
|
|
return false;
|
|
}
|
|
// Field 3 of /proc/<pid>/stat, right after the parenthesised comm, is
|
|
// the state letter. Without /proc the process counts as running.
|
|
!std::fs::read_to_string(format!("/proc/{pid}/stat")).is_ok_and(|stat| {
|
|
stat.rsplit_once(") ")
|
|
.is_some_and(|(_, rest)| rest.starts_with('Z'))
|
|
})
|
|
}
|
|
|
|
/// Require `pid` to be gone within 5s, killing it if it is not.
|
|
async fn assert_gone(pid: i32) {
|
|
let mut alive = is_running(pid);
|
|
for _ in 0..250 {
|
|
if !alive {
|
|
break;
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(20)).await;
|
|
alive = is_running(pid);
|
|
}
|
|
if alive {
|
|
// SAFETY: `kill(2)` on a pid this test started; keeps a failed
|
|
// run from leaving a 300s `sleep` behind.
|
|
unsafe {
|
|
libc::kill(pid, libc::SIGKILL);
|
|
}
|
|
}
|
|
assert!(!alive, "grandchild {pid} outlived the agent");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_agent_ignoring_a_cancel_is_killed_with_the_commands_it_started() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let mut runtime = agent(
|
|
dir.path(),
|
|
"deaf",
|
|
&[("GRANDCHILD", "1")],
|
|
PercentPolicy::default(),
|
|
);
|
|
runtime.cancel_grace = Duration::from_millis(200);
|
|
let (canceller, config) = (runtime.canceller().unwrap(), config(dir.path()));
|
|
|
|
let (done, grandchild) = tokio::join!(runtime.run(&config, "one", &NoopSink), async {
|
|
prompt_sent(dir.path()).await;
|
|
let grandchild = read_pid(&dir.path().join("prompts.grandchild")).await;
|
|
assert!(canceller.cancel());
|
|
grandchild
|
|
});
|
|
assert!(
|
|
matches!(done, Err(Error::Acp(AcpError::CancelIgnored { .. }))),
|
|
"{done:?}"
|
|
);
|
|
assert_gone(grandchild).await;
|
|
}
|
|
|
|
/// A subagent's run drops its runtime when it ends, an interrupted one
|
|
/// included.
|
|
#[tokio::test]
|
|
async fn dropping_the_runtime_kills_the_commands_its_agent_started() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(
|
|
dir.path(),
|
|
"ok",
|
|
&[("GRANDCHILD", "1")],
|
|
PercentPolicy::default(),
|
|
);
|
|
runtime
|
|
.run(&config(dir.path()), "one", &NoopSink)
|
|
.await
|
|
.unwrap();
|
|
let grandchild = read_pid(&dir.path().join("prompts.grandchild")).await;
|
|
assert!(is_running(grandchild));
|
|
|
|
drop(runtime);
|
|
assert_gone(grandchild).await;
|
|
}
|
|
|
|
/// SIGINT rather than SIGKILL: an agent only stops the commands it
|
|
/// started itself if it gets to shut down.
|
|
#[tokio::test]
|
|
async fn dropping_the_runtime_lets_its_agent_shut_down_on_sigint() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(
|
|
dir.path(),
|
|
"ok",
|
|
&[("TRAP_INT", "1")],
|
|
PercentPolicy::default(),
|
|
);
|
|
runtime
|
|
.run(&config(dir.path()), "one", &NoopSink)
|
|
.await
|
|
.unwrap();
|
|
|
|
drop(runtime);
|
|
let marker = dir.path().join("prompts.int");
|
|
for _ in 0..250 {
|
|
if marker.exists() {
|
|
break;
|
|
}
|
|
tokio::time::sleep(Duration::from_millis(20)).await;
|
|
}
|
|
assert!(marker.exists(), "the agent never got SIGINT");
|
|
}
|
|
|
|
fn sent(session: &str, text: &str) -> (String, String) {
|
|
(session.to_owned(), text.to_owned())
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn an_advertised_compact_command_runs_on_the_same_session() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(dir.path(), "ok", &[("COMMANDS", "compact")], policy(0));
|
|
let config = config(dir.path());
|
|
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
runtime.compact(&config, &NoopSink).await.unwrap();
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
|
|
assert!(!next.created);
|
|
assert_eq!(
|
|
prompts(dir.path()),
|
|
[
|
|
sent("s1", "SYSTEM PROMPT\n\none"),
|
|
sent("s1", "/compact"),
|
|
sent("s1", "two"),
|
|
]
|
|
);
|
|
assert_eq!(count(dir.path(), "session/new"), 1);
|
|
assert_eq!(recorded(dir.path()), "s1");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_compact_command_answered_with_a_bare_end_turn_keeps_the_session() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [("COMMANDS", "compact"), ("BLANK", "/compact")];
|
|
let runtime = agent(dir.path(), "ok", &env, policy(0));
|
|
let config = config(dir.path());
|
|
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
runtime.compact(&config, &NoopSink).await.unwrap();
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
|
|
assert!(!next.created);
|
|
assert_eq!(
|
|
prompts(dir.path()),
|
|
[
|
|
sent("s1", "SYSTEM PROMPT\n\none"),
|
|
sent("s1", "/compact"),
|
|
sent("s1", "two"),
|
|
]
|
|
);
|
|
assert_eq!(count(dir.path(), "session/new"), 1);
|
|
assert_eq!(recorded(dir.path()), "s1");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_checkpoint_turn_answered_with_a_bare_end_turn_succeeds() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(dir.path(), "ok", &[("BLANK", "CHECKPOINT")], policy(0));
|
|
let config = config(dir.path());
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
let mut guard = runtime.live.lock().await;
|
|
|
|
let checkpoint = runtime
|
|
.prompt(
|
|
&mut guard,
|
|
&config,
|
|
"CHECKPOINT",
|
|
TurnKind::Compaction,
|
|
&NoopSink,
|
|
)
|
|
.await;
|
|
assert!(checkpoint.is_ok(), "{checkpoint:?}");
|
|
// The same reply to an ordinary turn is the empty-turn error.
|
|
let ordinary = runtime
|
|
.prompt(
|
|
&mut guard,
|
|
&config,
|
|
"CHECKPOINT",
|
|
TurnKind::Ordinary,
|
|
&NoopSink,
|
|
)
|
|
.await;
|
|
assert!(
|
|
matches!(ordinary, Err(Error::Acp(AcpError::EmptyEndTurn))),
|
|
"{ordinary:?}"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn with_no_compact_command_notes_are_written_then_a_new_session_starts() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(dir.path(), "ok", &[("COMMANDS", "init")], policy(0));
|
|
let config = config(dir.path());
|
|
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
runtime.compact(&config, &NoopSink).await.unwrap();
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
|
|
assert!(next.created);
|
|
assert_eq!(
|
|
prompts(dir.path()),
|
|
[
|
|
sent("s1", "SYSTEM PROMPT\n\none"),
|
|
sent("s1", "CHECKPOINT"),
|
|
sent("s2", "SYSTEM PROMPT\n\ntwo"),
|
|
]
|
|
);
|
|
assert_eq!(recorded(dir.path()), "s2");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_turn_past_the_watermark_checkpoints_then_compacts() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [("COMMANDS", "compact"), ("USED", "800")];
|
|
let runtime = agent(dir.path(), "ok", &env, policy(75));
|
|
|
|
let done = runtime
|
|
.run(&config(dir.path()), "one", &NoopSink)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(done.compacted);
|
|
assert_eq!(
|
|
prompts(dir.path()),
|
|
[
|
|
sent("s1", "SYSTEM PROMPT\n\none"),
|
|
sent("s1", "CHECKPOINT"),
|
|
sent("s1", "/compact"),
|
|
]
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_silent_compact_command_is_cut_short_and_a_new_session_started() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [
|
|
("COMMANDS", "compact"),
|
|
("USED", "800"),
|
|
("STALL_COMPACT", "1"),
|
|
];
|
|
let mut runtime = agent(dir.path(), "ok", &env, policy(75));
|
|
runtime.compact_idle = Duration::from_millis(200);
|
|
// The turn's own idle window is far longer than the compact bound.
|
|
let config = idle_config(dir.path(), 5000);
|
|
|
|
let done = tokio::time::timeout(
|
|
Duration::from_secs(3),
|
|
runtime.run(&config, "one", &NoopSink),
|
|
)
|
|
.await
|
|
.expect("the compact bound, not the turn's idle window, ends the /compact")
|
|
.unwrap();
|
|
assert!(done.compacted);
|
|
let next = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
|
|
assert!(next.created && !next.compacted);
|
|
assert_eq!(
|
|
prompts(dir.path()),
|
|
[
|
|
sent("s1", "SYSTEM PROMPT\n\none"),
|
|
sent("s1", "CHECKPOINT"),
|
|
sent("s1", "/compact"),
|
|
sent("s2", "SYSTEM PROMPT\n\ntwo"),
|
|
]
|
|
);
|
|
assert_eq!(recorded(dir.path()), "s2");
|
|
}
|
|
|
|
fn choice(current: &str, values: &[&str]) -> Choice {
|
|
Choice {
|
|
current: current.to_owned(),
|
|
values: values.iter().map(|v| (*v).to_owned()).collect(),
|
|
refused: None,
|
|
}
|
|
}
|
|
|
|
/// Each `session/set_config_option` the agent was sent, as `id=value`.
|
|
fn sets(dir: &Path) -> Vec<String> {
|
|
std::fs::read_to_string(dir.join("prompts.sets"))
|
|
.unwrap_or_default()
|
|
.lines()
|
|
.map(str::to_owned)
|
|
.collect()
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn model_and_effort_are_set_and_effort_goes_with_a_model_without_it() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let runtime = agent(dir.path(), "ok", &[("OPTIONS", "1")], policy(0));
|
|
let choices = runtime.choices().unwrap();
|
|
assert_eq!(choices.get(), SessionChoices::default());
|
|
let mut config = Config {
|
|
model: Some("haiku".into()),
|
|
effort: Some("high".into()),
|
|
..config(dir.path())
|
|
};
|
|
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
// `haiku` is not on offer, so only the effort is set.
|
|
assert_eq!(sets(dir.path()), ["effort=high"]);
|
|
assert_eq!(
|
|
choices.get(),
|
|
SessionChoices {
|
|
model: Some(choice("m/think", &["m/think", "m/plain"])),
|
|
effort: Some(choice("high", &["low", "high"])),
|
|
}
|
|
);
|
|
|
|
config.model = Some("m/plain".into());
|
|
let done = runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(sets(dir.path()), ["effort=high", "model=m/plain"]);
|
|
assert_eq!(
|
|
choices.get(),
|
|
SessionChoices {
|
|
model: Some(choice("m/plain", &["m/think", "m/plain"])),
|
|
effort: None,
|
|
}
|
|
);
|
|
assert_eq!(done.telemetry.model.as_deref(), Some("m/plain"));
|
|
|
|
// Values the session already holds are not set again.
|
|
runtime.run(&config, "three", &NoopSink).await.unwrap();
|
|
assert_eq!(sets(dir.path()).len(), 2);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_refused_value_is_kept_until_the_agent_accepts_one() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [("OPTIONS", "1"), ("REFUSE", "m/plain")];
|
|
let runtime = agent(dir.path(), "ok", &env, policy(0));
|
|
let choices = runtime.choices().unwrap();
|
|
let config = Config {
|
|
model: Some("m/plain".into()),
|
|
..config(dir.path())
|
|
};
|
|
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
assert_eq!(
|
|
choices.get().model,
|
|
Some(Choice {
|
|
refused: Some("m/plain".into()),
|
|
..choice("m/think", &["m/think", "m/plain"])
|
|
})
|
|
);
|
|
|
|
// The next turn asks again, and the agent accepts.
|
|
runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(sets(dir.path()), ["model=m/plain", "model=m/plain"]);
|
|
assert_eq!(
|
|
choices.get().model,
|
|
Some(choice("m/plain", &["m/think", "m/plain"]))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_refusal_is_dropped_when_the_agent_is_respawned() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [("OPTIONS", "1"), ("REFUSE", "m/plain")];
|
|
let runtime = agent(dir.path(), "ok", &env, policy(0));
|
|
let choices = runtime.choices().unwrap();
|
|
let mut config = Config {
|
|
model: Some("m/plain".into()),
|
|
..config(dir.path())
|
|
};
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
assert_eq!(
|
|
choices.get().model.and_then(|c| c.refused).as_deref(),
|
|
Some("m/plain")
|
|
);
|
|
|
|
// `haiku` is not on offer, so the new process is asked to set nothing.
|
|
*runtime.live.lock().await = None;
|
|
config.model = Some("haiku".into());
|
|
runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(count(dir.path(), "start"), 2);
|
|
assert_eq!(sets(dir.path()), ["model=m/plain"]);
|
|
assert_eq!(
|
|
choices.get().model,
|
|
Some(choice("m/think", &["m/think", "m/plain"]))
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn a_refusal_is_dropped_when_the_session_s_own_value_is_wanted() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let env = [("OPTIONS", "1"), ("REFUSE", "m/plain")];
|
|
let runtime = agent(dir.path(), "ok", &env, policy(0));
|
|
let choices = runtime.choices().unwrap();
|
|
let mut config = Config {
|
|
model: Some("m/plain".into()),
|
|
..config(dir.path())
|
|
};
|
|
runtime.run(&config, "one", &NoopSink).await.unwrap();
|
|
assert_eq!(
|
|
choices.get().model.and_then(|c| c.refused).as_deref(),
|
|
Some("m/plain")
|
|
);
|
|
|
|
config.model = Some("m/think".into());
|
|
runtime.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(sets(dir.path()), ["model=m/plain"]);
|
|
assert_eq!(
|
|
choices.get().model,
|
|
Some(choice("m/think", &["m/think", "m/plain"]))
|
|
);
|
|
}
|
|
|
|
fn usd(amount: f64) -> ReportedCost {
|
|
ReportedCost {
|
|
amount,
|
|
currency: "USD".into(),
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn a_new_session_s_first_report_is_all_spent() {
|
|
let spend = Spend::default();
|
|
spend.attach(true);
|
|
spend.observe(usd(0.5));
|
|
spend.observe(usd(1.25));
|
|
assert_eq!(spend.take(), [usd(1.25)]);
|
|
assert_eq!(spend.take(), []);
|
|
spend.observe(usd(2.0));
|
|
assert_eq!(spend.take(), [usd(0.75)]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_loaded_session_s_first_report_only_sets_the_baseline() {
|
|
let spend = Spend::default();
|
|
spend.attach(false);
|
|
spend.observe(usd(10.0));
|
|
assert_eq!(spend.take(), []);
|
|
spend.observe(usd(10.5));
|
|
assert_eq!(spend.take(), [usd(0.5)]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_falling_total_adds_nothing_and_becomes_the_baseline() {
|
|
let spend = Spend::default();
|
|
spend.attach(true);
|
|
spend.observe(usd(3.0));
|
|
spend.observe(usd(1.0));
|
|
spend.observe(usd(1.5));
|
|
assert_eq!(spend.take(), [usd(3.5)]);
|
|
}
|
|
|
|
#[test]
|
|
fn untaken_spend_survives_attaching_another_session() {
|
|
let spend = Spend::default();
|
|
spend.attach(true);
|
|
spend.observe(usd(1.0));
|
|
spend.attach(true);
|
|
spend.observe(usd(0.25));
|
|
assert_eq!(spend.take(), [usd(1.25)]);
|
|
}
|
|
|
|
#[test]
|
|
fn a_changed_currency_only_sets_a_new_baseline() {
|
|
let spend = Spend::default();
|
|
spend.attach(true);
|
|
spend.observe(usd(1.0));
|
|
let eur = |amount| ReportedCost {
|
|
amount,
|
|
currency: "EUR".into(),
|
|
};
|
|
spend.observe(eur(0.75));
|
|
spend.observe(eur(1.25));
|
|
assert_eq!(spend.take(), [eur(0.5), usd(1.0)]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn reported_cost_is_what_the_session_added_since_last_taken() {
|
|
let dir = tempfile::tempdir().unwrap();
|
|
let config = config(dir.path());
|
|
let env = [("COST", "2")];
|
|
|
|
let before = agent(dir.path(), "ok", &env, PercentPolicy::default());
|
|
before.run(&config, "one", &NoopSink).await.unwrap();
|
|
before.run(&config, "two", &NoopSink).await.unwrap();
|
|
assert_eq!(before.take_reported_cost(), [usd(4.0)]);
|
|
assert_eq!(before.take_reported_cost(), []);
|
|
drop(before);
|
|
|
|
let after = agent(dir.path(), "ok", &env, PercentPolicy::default());
|
|
after.run(&config, "three", &NoopSink).await.unwrap();
|
|
assert_eq!(after.take_reported_cost(), []);
|
|
after.run(&config, "four", &NoopSink).await.unwrap();
|
|
assert_eq!(after.take_reported_cost(), [usd(2.0)]);
|
|
|
|
after.archive().unwrap();
|
|
after.run(&config, "five", &NoopSink).await.unwrap();
|
|
assert_eq!(after.take_reported_cost(), [usd(2.0)]);
|
|
assert_eq!(count(dir.path(), "session/load"), 1);
|
|
assert_eq!(count(dir.path(), "session/new"), 2);
|
|
}
|
|
}
|