Author SHA1 Message Date
Hare 62ada5eaa4 fix: bound safe-boundary pause escalation 2026-08-27 23:57:20 +09:00
Hare f5ff0b7c13 fix: confirm bash cancellation cleanup 2026-08-27 23:01:06 +09:00
Hare 3337cafcdf feat: own cancellable tool execution lifecycle 2026-08-27 23:01:00 +09:00
Hare e87784118b fix: preserve resumable paused tool calls 2026-08-27 21:48:59 +09:00
Hare 40fada28ea fix: preserve confirmed output on interceptor abort 2026-08-27 21:30:04 +09:00
Hare 58cc94d4b7 fix: terminalize interrupted tool executions 2026-08-27 20:54:13 +09:00
Hare ccabea59c9 fix: terminalize parallel tool outputs on completion 2026-08-27 19:16:49 +09:00
Hare 183c37446e fix: fence running snapshots on input commit 2026-08-27 15:42:20 +09:00
Hare 7aa06afc45 chore: preserve reviewed source lineage
# Conflicts:
#	crates/agen/README.md
#	crates/agen/examples/engine_cancel_demo.rs
#	crates/agen/examples/engine_cli.rs
#	crates/agen/src/engine.rs
#	crates/agen/tests/annotated_history_test.rs
#	crates/agen/tests/callback_test.rs
#	crates/agen/tests/engine_fixtures.rs
#	crates/agen/tests/engine_state_test.rs
#	crates/agen/tests/parallel_execution_test.rs
#	crates/agen/tests/reasoning_round_trip_test.rs
#	crates/session-store/tests/session_test.rs
#	crates/worker/src/worker.rs
2026-08-27 15:13:27 +09:00
Hare 1515a2fb86 fix: reconcile typed history with run exits 2026-08-27 15:12:05 +09:00
Hare ec798c58d7 fix: preserve annotated history through rewind 2026-08-27 14:54:24 +09:00
Hare e365189276 feat: add provenance-aware worker history 2026-08-27 14:54:24 +09:00
Hare 116d610ad0 fix: project Ticket mutation outputs to human keys 2026-08-27 14:18:48 +09:00
Hare 75c570962d Merge commit '7edc588202dfbfd4c834f677f510ddda7f3d6451' into work/00001M10HW6BV-model-facing-resource-projection 2026-08-27 13:46:19 +09:00
Hare cae8ac1799 fix: allow missing Objective query snippets 2026-08-27 13:46:10 +09:00
Hare 917cc222a3 fix: resolve relation summaries to Ticket keys 2026-08-27 13:22:08 +09:00
Hare 7edc588202 chore: refresh T-528 after T-541 2026-08-27 13:09:44 +09:00
Hare c83461508b Merge commit '21b3dd1da1b1bbf18799a0623bf67dbe6266067c' into work/00001M10HW6BV-model-facing-resource-projection 2026-08-27 13:09:15 +09:00
Hare 374449e663 feat: add provenance-aware worker history 2026-08-27 12:51:01 +09:00
Hare 4c876a201b fix: validate projected resource keys canonically 2026-08-27 12:49:59 +09:00
Hare 21b3dd1da1 Merge commit '0496cd907bc7bb96e9aa1c6d385bedb616bf3233' into work/00001M10FJVA2-orchestrator-queue-notice 2026-08-27 12:46:27 +09:00
Hare 5ca0ea9228 fix: sanitize orchestrator queue attention 2026-08-27 12:46:10 +09:00
Hare d5c3a68a37 fix: use Ticket resource keys in handoffs 2026-08-27 12:45:08 +09:00
Hare 2b33b9158d chore: refresh T-528 against develop 2026-08-27 12:45:01 +09:00
Hare 3a7a3307ef fix: map internal worker terminal lifecycles 2026-08-27 12:27:55 +09:00
Hare 0496cd907b fix: use Ticket resource keys in handoffs 2026-08-27 12:24:19 +09:00
Hare 975b4fa700 feat: add typed engine run exits 2026-08-27 11:42:49 +09:00
79 changed files with 6643 additions and 1256 deletions
+6 -5
View File
@@ -21,20 +21,21 @@ agen = { version = "0.2.1", features = ["codex"] }
## Quick start ## Quick start
Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns. Supply an implementation of [`LlmClient`](https://docs.rs/agen/latest/agen/llm_client/trait.LlmClient.html), keep conversation history in your application, then run a turn. The first call consumes the mutable engine and returns a cache-locked engine for later turns.
```no_run ```no_run
use agen::{Engine, EngineError}; use agen::{Engine, EngineError, History};
use agen::llm_client::LlmClient; use agen::llm_client::LlmClient;
async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> { async fn conversation<C: LlmClient>(client: C) -> Result<(), EngineError> {
let mut history = History::new();
let output = Engine::new(client) let output = Engine::new(client)
.system_prompt("You are a concise assistant.") .system_prompt("You are a concise assistant.")
.run("Explain typed state in one sentence.") .run(&mut history, "Explain typed state in one sentence.")
.await?; .await;
let mut engine = output.engine; let mut engine = output.engine;
let _result = engine.run("Give a Rust example.").await?; let _result = engine.run(&mut history, "Give a Rust example.").await;
Ok(()) Ok(())
} }
``` ```
+10 -10
View File
@@ -4,7 +4,7 @@
use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme}; use agen::llm_client::scheme::{Scheme, anthropic::AnthropicScheme};
use agen::llm_client::transport::{HttpTransport, ResolvedAuth}; use agen::llm_client::transport::{HttpTransport, ResolvedAuth};
use agen::{Engine, EngineResult}; use agen::{Engine, EngineRunExit, StopReason};
use std::time::Duration; use std::time::Duration;
#[tokio::main] #[tokio::main]
@@ -29,6 +29,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let base_url = scheme.default_base_url().to_string(); let base_url = scheme.default_base_url().to_string();
let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap); let client = HttpTransport::new(scheme, model, base_url, ResolvedAuth::ApiKey(api_key), cap);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history = agen::History::new();
println!("🚀 Starting Engine..."); println!("🚀 Starting Engine...");
println!("💡 Will cancel after 2 seconds\n"); println!("💡 Will cancel after 2 seconds\n");
@@ -45,16 +46,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
println!("📡 Sending request to LLM..."); println!("📡 Sending request to LLM...");
match engine.run("Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await { let output = engine.run(&mut history, "Tell me a very long story about a brave knight. Make it as detailed as possible with many paragraphs.").await;
Ok(out) => match out.result { match output.result {
EngineResult::Finished => println!("✅ Task completed normally"), EngineRunExit::Finished => println!("✅ Task completed normally"),
EngineResult::Paused => println!("⏸️ Task paused"), EngineRunExit::Paused => println!("⏸️ Task paused"),
EngineResult::LimitReached => println!("🔒 Turn limit reached"), EngineRunExit::Yielded => println!("↩️ Task yielded"),
EngineResult::Yielded => println!("↩️ Task yielded"), EngineRunExit::Interrupted(StopReason::LimitReached) => {
}, println!("🔒 Turn limit reached")
Err(e) => {
println!("❌ Task error: {}", e);
} }
EngineRunExit::Interrupted(reason) => println!("❌ Task interrupted: {reason:?}"),
} }
println!("\n✨ Demo complete!"); println!("\n✨ Demo complete!");
+11 -19
View File
@@ -39,7 +39,7 @@ use tracing::info;
use tracing_subscriber::EnvFilter; use tracing_subscriber::EnvFilter;
use agen::{ use agen::{
Engine, Engine, EngineRunExit, StopReason,
interceptor::{Interceptor, PostToolAction, ToolResultInfo}, interceptor::{Interceptor, PostToolAction, ToolResultInfo},
llm_client::{ llm_client::{
LlmClient, LlmClient,
@@ -451,6 +451,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// Create Engine // Create Engine
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let tool_call_names = Arc::new(Mutex::new(HashMap::new())); let tool_call_names = Arc::new(Mutex::new(HashMap::new()));
@@ -476,12 +477,9 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
// One-shot mode // One-shot mode
if let Some(prompt) = args.prompt { if let Some(prompt) = args.prompt {
match engine.run(&prompt).await { let output = engine.run(&mut history, &prompt).await;
Ok(_) => {} if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) = output.result {
Err(e) => { eprintln!("\n❌ Error: {error}");
eprintln!("\n❌ Error: {}", e);
std::process::exit(1);
}
} }
return Ok(()); return Ok(());
@@ -500,13 +498,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
return Ok(()); return Ok(());
} }
let mut locked = match engine.run(first_input).await { let output = engine.run(&mut history, first_input).await;
Ok(out) => out.engine, let mut locked = output.engine;
Err(e) => {
eprintln!("\n❌ Error: {}", e);
return Ok(());
}
};
loop { loop {
print!("\n👤 You: "); print!("\n👤 You: ");
@@ -525,11 +518,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
break; break;
} }
match locked.run(input).await { if let EngineRunExit::Interrupted(StopReason::Unexpected(error)) =
Ok(_) => {} locked.run(&mut history, input).await
Err(e) => { {
eprintln!("\n❌ Error: {}", e); eprintln!("\n❌ Error: {error}");
}
} }
} }
+795 -245
View File
File diff suppressed because it is too large Load Diff
+199
View File
@@ -0,0 +1,199 @@
//! Typed conversation history containers.
//!
//! Agen keeps provider-visible [`Item`](crate::Item) values separate from any
//! host-domain provenance. The host chooses the annotation type `A`, while Agen
//! preserves each item and annotation as one entry for clone/truncate/restore
//! style history operations.
use serde::{Deserialize, Serialize};
use crate::Item;
/// One conversation-history entry with host-owned annotation.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct HistoryEntry<A = ()> {
/// Provider/model-visible conversation item.
pub item: Item,
/// Host-domain metadata kept with the item and never projected to providers.
pub annotation: A,
}
impl<A> HistoryEntry<A> {
/// Build an entry from an item and its annotation.
pub fn new(item: Item, annotation: A) -> Self {
Self { item, annotation }
}
/// Split the entry into its item and annotation.
pub fn into_parts(self) -> (Item, A) {
(self.item, self.annotation)
}
}
impl HistoryEntry<()> {
/// Build a unit-annotated entry.
pub fn from_item(item: Item) -> Self {
Self {
item,
annotation: (),
}
}
}
/// Conversation history with one annotation per item.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
pub struct History<A = ()> {
entries: Vec<HistoryEntry<A>>,
}
impl<A> History<A> {
/// Create an empty history.
pub fn new() -> Self {
Self {
entries: Vec::new(),
}
}
/// Build history from already annotated entries, preserving order.
pub fn from_entries(entries: Vec<HistoryEntry<A>>) -> Self {
Self { entries }
}
/// Replace all entries as one restore/rebuild operation and return the old entries.
pub fn replace_entries(&mut self, entries: Vec<HistoryEntry<A>>) -> Vec<HistoryEntry<A>> {
std::mem::replace(&mut self.entries, entries)
}
/// Borrow annotated entries.
pub fn entries(&self) -> &[HistoryEntry<A>] {
&self.entries
}
/// Mutably borrow annotated entries for host-owned rebuild operations.
pub fn entries_mut(&mut self) -> &mut [HistoryEntry<A>] {
&mut self.entries
}
/// Consume the history into annotated entries.
pub fn into_entries(self) -> Vec<HistoryEntry<A>> {
self.entries
}
/// Number of entries.
pub fn len(&self) -> usize {
self.entries.len()
}
/// Whether the history is empty.
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
/// Iterate over annotated entries.
pub fn iter(&self) -> impl ExactSizeIterator<Item = &HistoryEntry<A>> {
self.entries.iter()
}
/// Iterate over provider-visible items only.
pub fn items(&self) -> impl ExactSizeIterator<Item = &Item> {
self.entries.iter().map(|entry| &entry.item)
}
/// Clone provider-visible items into a request-local projection.
pub fn items_cloned(&self) -> Vec<Item> {
self.items().cloned().collect()
}
/// Append an already annotated entry.
pub fn push_entry(&mut self, entry: HistoryEntry<A>) {
self.entries.push(entry);
}
/// Append many already annotated entries.
pub fn extend_entries(&mut self, entries: impl IntoIterator<Item = HistoryEntry<A>>) {
self.entries.extend(entries);
}
/// Commit one item through a trusted annotation callback before it becomes live.
///
/// The callback may durably persist the item and returns the annotation that
/// must be stored with it. If the callback fails, the history is left unchanged.
pub fn append_with(
&mut self,
item: Item,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<(), String> {
let annotation = annotate(&item)?;
self.entries.push(HistoryEntry { item, annotation });
Ok(())
}
/// Commit items through a trusted annotation callback before they become live.
///
/// Items before a failure remain appended; the failing item and later items do
/// not enter history. This mirrors append-only durable logs where each accepted
/// item is already committed before the next item is attempted.
pub fn extend_with(
&mut self,
items: impl IntoIterator<Item = Item>,
annotate: &mut impl FnMut(&Item) -> Result<A, String>,
) -> Result<(), String> {
for item in items {
self.append_with(item, annotate)?;
}
Ok(())
}
/// Truncate entries, preserving item+annotation pairing for retained entries.
pub fn truncate(&mut self, len: usize) {
self.entries.truncate(len);
}
/// Clear all entries.
pub fn clear(&mut self) {
self.entries.clear();
}
}
impl History<()> {
/// Build unit-annotated history from provider-visible items.
pub fn from_items(items: Vec<Item>) -> Self {
Self {
entries: items.into_iter().map(HistoryEntry::from_item).collect(),
}
}
/// Replace history from provider-visible items using unit annotations.
pub fn replace_items(&mut self, items: Vec<Item>) -> Vec<HistoryEntry<()>> {
self.replace_entries(items.into_iter().map(HistoryEntry::from_item).collect())
}
/// Append one item with unit annotation.
pub fn push(&mut self, item: Item) {
self.entries.push(HistoryEntry::from_item(item));
}
/// Append items with unit annotations.
pub fn extend_items(&mut self, items: impl IntoIterator<Item = Item>) {
self.entries
.extend(items.into_iter().map(HistoryEntry::from_item));
}
}
impl<A> IntoIterator for History<A> {
type Item = HistoryEntry<A>;
type IntoIter = std::vec::IntoIter<HistoryEntry<A>>;
fn into_iter(self) -> Self::IntoIter {
self.entries.into_iter()
}
}
impl<'a, A> IntoIterator for &'a History<A> {
type Item = &'a HistoryEntry<A>;
type IntoIter = std::slice::Iter<'a, HistoryEntry<A>>;
fn into_iter(self) -> Self::IntoIter {
self.entries.iter()
}
}
+9 -3
View File
@@ -2,6 +2,7 @@
mod engine; mod engine;
mod handler; mod handler;
mod history;
mod message; mod message;
pub(crate) mod callback; pub(crate) mod callback;
@@ -20,13 +21,18 @@ pub mod usage_record;
pub use agen_macros::{description, tool, tool_registry}; pub use agen_macros::{description, tool, tool_registry};
pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope}; pub use callback::{TextBlockScope, ThinkingBlockScope, ToolUseBlockScope};
pub use engine::{ pub use engine::{
Engine, EngineConfig, EngineError, EngineResult, EngineRunOutput, LlmRetryNotice, Engine, EngineConfig, EngineError, EngineResult, EngineRunExit, EngineRunOutput,
ToolRegistryError, LlmRetryNotice, StopReason, ToolRegistryError,
}; };
pub use handler::ToolUseBlockStart; pub use handler::ToolUseBlockStart;
pub use history::{History, HistoryEntry};
pub use interceptor::Interceptor; pub use interceptor::Interceptor;
pub use message::{ContentPart, Item, Message, Role}; pub use message::{ContentPart, Item, Message, Role};
pub use tool::{ToolCall, ToolExecutionContext, ToolOutputLimits, ToolResult}; pub use tool::{
ToolCall, ToolExecutionContext, ToolExecutionHandle, ToolExecutionPolicy,
ToolExecutionTerminal, ToolExecutionTerminalFuture, ToolOutputLimits, ToolResult,
ToolResultDisposition,
};
pub use usage_record::UsageRecord; pub use usage_record::UsageRecord;
/// Implementation dependencies used by code generated from `agen` macros. /// Implementation dependencies used by code generated from `agen` macros.
+8 -1
View File
@@ -18,6 +18,9 @@ pub enum ClientError {
message: String, message: String,
retry_after: Option<Duration>, retry_after: Option<Duration>,
}, },
/// The provider rejected the request because it exceeded the model context window.
/// Classified only from a structured provider error code, never message text.
ContextWindowExceeded,
/// A request lifecycle phase exceeded its hard timeout. /// A request lifecycle phase exceeded its hard timeout.
Timeout { Timeout {
phase: &'static str, phase: &'static str,
@@ -48,6 +51,7 @@ impl fmt::Display for ClientError {
} }
write!(f, ": {}", message) write!(f, ": {}", message)
} }
ClientError::ContextWindowExceeded => write!(f, "Model context window reached"),
ClientError::Timeout { phase, timeout } => { ClientError::Timeout { phase, timeout } => {
write!(f, "{phase} timed out after {}s", timeout.as_secs()) write!(f, "{phase} timed out after {}s", timeout.as_secs())
} }
@@ -112,7 +116,10 @@ pub fn is_retryable(error: &ClientError) -> bool {
ClientError::Api { status: None, .. } => false, ClientError::Api { status: None, .. } => false,
ClientError::Timeout { .. } => true, ClientError::Timeout { .. } => true,
ClientError::Http(e) => e.is_connect() || e.is_timeout(), ClientError::Http(e) => e.is_connect() || e.is_timeout(),
ClientError::Json(_) | ClientError::Sse(_) | ClientError::Config(_) => false, ClientError::ContextWindowExceeded
| ClientError::Json(_)
| ClientError::Sse(_)
| ClientError::Config(_) => false,
} }
} }
+4 -7
View File
@@ -431,13 +431,7 @@ fn api_error_code(error: &ClientError) -> Option<&str> {
} }
fn is_context_length_exceeded(error: &ClientError) -> bool { fn is_context_length_exceeded(error: &ClientError) -> bool {
match error { matches!(error, ClientError::ContextWindowExceeded)
ClientError::Api { code, message, .. } => {
code.as_deref() == Some("context_length_exceeded")
|| message.contains("context_length_exceeded")
}
_ => false,
}
} }
async fn response_with_timeout( async fn response_with_timeout(
@@ -487,6 +481,9 @@ async fn classify_error_response(resp: reqwest::Response) -> ClientError {
.and_then(|v| v.as_str()) .and_then(|v| v.as_str())
.unwrap_or(&text) .unwrap_or(&text)
.to_string(); .to_string();
if code.as_deref() == Some("context_length_exceeded") {
return ClientError::ContextWindowExceeded;
}
ClientError::Api { ClientError::Api {
status: Some(status), status: Some(status),
code, code,
+37 -2
View File
@@ -9,7 +9,7 @@
use std::{fmt, sync::Arc}; use std::{fmt, sync::Arc};
use crate::tool::Attachment; use crate::tool::{Attachment, ToolResultDisposition};
use base64::Engine as _; use base64::Engine as _;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
@@ -121,6 +121,9 @@ pub enum Item {
/// Detailed output (removed by pruning when old enough) /// Detailed output (removed by pruning when old enough)
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
content: Option<String>, content: Option<String>,
/// Typed terminal state used for replay and recovery.
#[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")]
disposition: ToolResultDisposition,
/// Whether the tool result represents an execution error. /// Whether the tool result represents an execution error.
#[serde(default, skip_serializing_if = "is_false")] #[serde(default, skip_serializing_if = "is_false")]
is_error: bool, is_error: bool,
@@ -261,7 +264,17 @@ impl Item {
content: Option<String>, content: Option<String>,
is_error: bool, is_error: bool,
) -> Self { ) -> Self {
Self::tool_result_item_with_attachments(call_id, summary, content, is_error, Vec::new()) Self::tool_result_item_with_disposition_and_attachments(
call_id,
summary,
content,
if is_error {
ToolResultDisposition::Error
} else {
ToolResultDisposition::Success
},
Vec::new(),
)
} }
/// Create a tool result item with durable, prunable structured attachments. /// Create a tool result item with durable, prunable structured attachments.
@@ -272,11 +285,33 @@ impl Item {
is_error: bool, is_error: bool,
attachments: Vec<Attachment>, attachments: Vec<Attachment>,
) -> Self { ) -> Self {
Self::tool_result_item_with_disposition_and_attachments(
call_id,
summary,
content,
if is_error {
ToolResultDisposition::Error
} else {
ToolResultDisposition::Success
},
attachments,
)
}
pub fn tool_result_item_with_disposition_and_attachments(
call_id: impl Into<String>,
summary: impl Into<String>,
content: Option<String>,
disposition: ToolResultDisposition,
attachments: Vec<Attachment>,
) -> Self {
let is_error = !disposition.is_success();
Self::ToolResult { Self::ToolResult {
id: None, id: None,
call_id: call_id.into(), call_id: call_id.into(),
summary: summary.into(), summary: summary.into(),
content, content,
disposition,
is_error, is_error,
attachments, attachments,
} }
+1 -1
View File
@@ -19,7 +19,7 @@ mod private {
/// - Editing message history (add, delete, clear) /// - Editing message history (add, delete, clear)
/// - Registering tools and hooks /// - Registering tools and hooks
/// ///
/// Can transition to [`Locked`] state via `Engine::lock()`. /// Can transition to [`Locked`] state via `Engine::lock(&history)`.
/// ///
/// # Examples /// # Examples
/// ///
+227 -2
View File
@@ -3,7 +3,14 @@
//! Traits for defining tools callable by LLM. //! Traits for defining tools callable by LLM.
//! Usually auto-implemented using the `#[tool]` macro. //! Usually auto-implemented using the `#[tool]` macro.
use std::{collections::HashMap, fmt, sync::Arc}; use std::{
collections::HashMap,
fmt,
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use async_trait::async_trait; use async_trait::async_trait;
use base64::{Engine as _, engine::general_purpose::STANDARD}; use base64::{Engine as _, engine::general_purpose::STANDARD};
@@ -23,6 +30,12 @@ pub enum ToolError {
/// Internal error /// Internal error
#[error("Internal error: {0}")] #[error("Internal error: {0}")]
Internal(String), Internal(String),
/// Cooperative cancellation completed with bounded terminal output.
#[error("Tool execution cancelled")]
Cancelled(ToolOutput),
/// Execution was interrupted with a confirmed bounded terminal output.
#[error("Tool execution interrupted")]
Interrupted(ToolOutput),
} }
// ============================================================================= // =============================================================================
@@ -158,6 +171,28 @@ pub enum Attachment {
Image(ImageAttachment), Image(ImageAttachment),
} }
/// Terminal disposition of one started tool call.
///
/// `Cancelled` means the tool confirmed cancellation. `OutcomeUnknown` means
/// execution stopped without confirmation, so neither completion nor side
/// effects may be inferred.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "snake_case")]
pub enum ToolResultDisposition {
#[default]
Success,
Error,
Interrupted,
Cancelled,
OutcomeUnknown,
}
impl ToolResultDisposition {
pub const fn is_success(&self) -> bool {
matches!(self, Self::Success)
}
}
/// Tool execution result. /// Tool execution result.
/// ///
/// Every output has a mandatory `summary` (1-2 lines) that persists in /// Every output has a mandatory `summary` (1-2 lines) that persists in
@@ -322,6 +357,12 @@ impl ToolExecutionContext {
} }
} }
/// Identifies one live execution attempt without making the batch id a durable
/// replay or idempotency authority.
pub fn execution_id(&self) -> String {
format!("{}:{}", self.batch_id, self.call_id)
}
/// Context for direct, non-engine calls in unit tests and low-level callers. /// Context for direct, non-engine calls in unit tests and low-level callers.
pub fn direct() -> Self { pub fn direct() -> Self {
Self::new("direct", "direct", 0) Self::new("direct", "direct", 0)
@@ -334,6 +375,142 @@ impl Default for ToolExecutionContext {
} }
} }
/// The provider-confirmed terminal result of one started tool execution.
///
/// `OutcomeUnknown` is reserved for an execution task that had to be force-closed
/// or failed before the provider could confirm its terminal result.
#[derive(Debug)]
pub enum ToolExecutionTerminal {
Confirmed(Result<ToolOutput, ToolError>),
OutcomeUnknown,
}
/// The completion future paired with a [`ToolExecutionHandle`]. Dropping this
/// future does not drop the provider execution: the spawned execution remains
/// owned by its handle until it completes or is explicitly force-closed.
pub struct ToolExecutionTerminalFuture {
task: tokio::task::JoinHandle<Result<ToolOutput, ToolError>>,
}
impl Future for ToolExecutionTerminalFuture {
type Output = ToolExecutionTerminal;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match Pin::new(&mut self.task).poll(cx) {
Poll::Ready(Ok(result)) => Poll::Ready(ToolExecutionTerminal::Confirmed(result)),
Poll::Ready(Err(_)) => Poll::Ready(ToolExecutionTerminal::OutcomeUnknown),
Poll::Pending => Poll::Pending,
}
}
}
/// Live ownership and control for one started tool execution.
///
/// Execution, cancellation, and terminal confirmation remain provider-owned:
/// this handle starts `Tool::execute`, delegates cooperative cancellation to
/// `Tool::cancel_execution`, and treats execution-future completion as the
/// provider's terminal confirmation. Agen may force-close only after its caller's
/// deadline expires, at which point the outcome is necessarily unknown.
#[derive(Clone)]
pub struct ToolExecutionHandle {
inner: Arc<ToolExecutionHandleInner>,
}
struct ToolExecutionHandleInner {
tool: Arc<dyn Tool>,
context: ToolExecutionContext,
abort: tokio::task::AbortHandle,
}
impl Drop for ToolExecutionHandleInner {
fn drop(&mut self) {
// Losing the final live owner is an explicit forced close, never a
// best-effort detached provider future.
self.abort.abort();
}
}
impl fmt::Debug for ToolExecutionHandle {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ToolExecutionHandle")
.field("call_id", &self.inner.context.call_id)
.field("batch_id", &self.inner.context.batch_id)
.finish_non_exhaustive()
}
}
impl ToolExecutionHandle {
pub fn start(
tool: Arc<dyn Tool>,
input_json: String,
context: ToolExecutionContext,
) -> (Self, ToolExecutionTerminalFuture) {
let execution_tool = Arc::clone(&tool);
let execution_context = context.clone();
let task =
tokio::spawn(
async move { execution_tool.execute(&input_json, execution_context).await },
);
let abort = task.abort_handle();
(
Self {
inner: Arc::new(ToolExecutionHandleInner {
tool,
context,
abort,
}),
},
ToolExecutionTerminalFuture { task },
)
}
pub fn context(&self) -> &ToolExecutionContext {
&self.inner.context
}
pub async fn cancel_before(&self, deadline: tokio::time::Instant) -> Result<(), ToolError> {
match tokio::time::timeout_at(
deadline,
self.inner.tool.cancel_execution(&self.inner.context),
)
.await
{
Ok(result) => result,
Err(_) => Err(ToolError::Internal(format!(
"tool cancellation request exceeded its deadline for call {}",
self.inner.context.call_id
))),
}
}
pub fn force_close(&self) {
self.inner.abort.abort();
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ToolExecutionPolicy {
/// Time a pause waits for already-started providers to reach a natural safe
/// boundary before escalating to explicit cooperative cancellation.
pub pause_safe_boundary_timeout: std::time::Duration,
/// Maximum time allowed for a provider to accept one cooperative
/// cancellation request.
pub cancellation_request_timeout: std::time::Duration,
/// Maximum time allowed for all providers to confirm terminal results after
/// cancellation has been requested.
pub terminal_confirmation_timeout: std::time::Duration,
}
impl Default for ToolExecutionPolicy {
fn default() -> Self {
Self {
pause_safe_boundary_timeout: std::time::Duration::from_millis(100),
cancellation_request_timeout: std::time::Duration::from_millis(100),
terminal_confirmation_timeout: std::time::Duration::from_millis(500),
}
}
}
// ============================================================================= // =============================================================================
// Tool trait // Tool trait
// ============================================================================= // =============================================================================
@@ -402,6 +579,26 @@ pub trait Tool: Send + Sync {
input_json: &str, input_json: &str,
ctx: ToolExecutionContext, ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError>; ) -> Result<ToolOutput, ToolError>;
/// Request cooperative cancellation for one started call.
///
/// Implementations that own cancellable provider operations should signal
/// every live execution identified by `call_id`, then let `execute` return
/// the confirmed bounded terminal output. Direct callers may use this
/// compatibility surface; Agen uses [`Tool::cancel_execution`] so providers
/// can bind cancellation to one exact live attempt.
async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> {
Ok(())
}
/// Request cooperative cancellation for one exact started execution.
///
/// The default preserves existing tools by delegating to `cancel(call_id)`.
/// Providers with their own execution registry should override this method
/// and key cancellation by [`ToolExecutionContext::execution_id`].
async fn cancel_execution(&self, ctx: &ToolExecutionContext) -> Result<(), ToolError> {
self.cancel(&ctx.call_id).await
}
} }
// ============================================================================= // =============================================================================
@@ -429,6 +626,9 @@ pub struct ToolCall {
pub struct ToolResult { pub struct ToolResult {
/// Corresponding tool call ID /// Corresponding tool call ID
pub tool_use_id: String, pub tool_use_id: String,
/// Typed terminal state.
#[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")]
pub disposition: ToolResultDisposition,
/// Short summary (always kept in history) /// Short summary (always kept in history)
pub summary: String, pub summary: String,
/// Detailed output (prunable) /// Detailed output (prunable)
@@ -445,11 +645,20 @@ pub struct ToolResult {
impl ToolResult { impl ToolResult {
/// Create a success result from a [`ToolOutput`]. /// Create a success result from a [`ToolOutput`].
pub fn from_output(tool_use_id: impl Into<String>, output: ToolOutput) -> Self { pub fn from_output(tool_use_id: impl Into<String>, output: ToolOutput) -> Self {
Self::from_output_with_disposition(tool_use_id, output, ToolResultDisposition::Success)
}
pub fn from_output_with_disposition(
tool_use_id: impl Into<String>,
output: ToolOutput,
disposition: ToolResultDisposition,
) -> Self {
Self { Self {
tool_use_id: tool_use_id.into(), tool_use_id: tool_use_id.into(),
disposition,
summary: output.summary, summary: output.summary,
content: output.content, content: output.content,
is_error: false, is_error: !disposition.is_success(),
attachments: output.attachments, attachments: output.attachments,
} }
} }
@@ -458,12 +667,28 @@ impl ToolResult {
pub fn error(tool_use_id: impl Into<String>, message: impl Into<String>) -> Self { pub fn error(tool_use_id: impl Into<String>, message: impl Into<String>) -> Self {
Self { Self {
tool_use_id: tool_use_id.into(), tool_use_id: tool_use_id.into(),
disposition: ToolResultDisposition::Error,
summary: message.into(), summary: message.into(),
content: None, content: None,
is_error: true, is_error: true,
attachments: Vec::new(), attachments: Vec::new(),
} }
} }
/// Close an execution whose completion and side effects cannot be confirmed.
pub fn outcome_unknown(tool_use_id: impl Into<String>) -> Self {
Self {
tool_use_id: tool_use_id.into(),
disposition: ToolResultDisposition::OutcomeUnknown,
summary: "Tool execution outcome unknown".to_string(),
content: Some(
"Execution was interrupted before completion could be confirmed. Completion and side effects are unknown."
.to_string(),
),
is_error: true,
attachments: Vec::new(),
}
}
} }
#[cfg(test)] #[cfg(test)]
@@ -0,0 +1,84 @@
mod common;
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::{Engine, EngineError, History, HistoryEntry, Item, Role};
use common::MockLlmClient;
fn completed_text_events(text: &str) -> Vec<Event> {
vec![
Event::text_block_start(0),
Event::text_delta(0, text),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]
}
#[tokio::test]
async fn run_preserves_item_annotations_without_projecting_them() {
let client = MockLlmClient::new(completed_text_events("assistant reply"));
let engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client);
let mut history = History::<String>::new();
let mut next = 0usize;
let mut annotate = |item: &Item| {
next += 1;
let kind = match item {
Item::Message { role, .. } => match role {
Role::User => "user",
Role::Assistant => "assistant",
Role::System => "system",
},
Item::ToolCall { .. } => "tool_call",
Item::ToolResult { .. } => "tool_result",
Item::Reasoning { .. } => "reasoning",
};
Ok(format!("{next}:{kind}"))
};
let output = engine
.run_with_annotation(&mut history, "hello", &mut annotate)
.await;
assert!(matches!(output.result, agen::EngineRunExit::Finished));
assert_eq!(history.len(), 2);
assert_eq!(history.entries()[0].annotation, "1:user");
assert_eq!(history.entries()[1].annotation, "2:assistant");
assert_eq!(history.items_cloned().len(), 2);
}
#[test]
fn append_failure_does_not_make_item_live() {
let client = MockLlmClient::new(vec![]);
let mut engine = Engine::<_, agen::state::Mutable, usize>::new_annotated(client);
let mut history = History::<usize>::new();
let mut fail = |_item: &Item| Err("commit failed".to_string());
let err = engine
.append_history_with(&mut history, [Item::user_message("uncommitted")], &mut fail)
.unwrap_err();
assert!(matches!(err, EngineError::HistoryAppend(message) if message == "commit failed"));
assert!(history.is_empty());
}
#[test]
fn replacement_keeps_items_and_annotations_together() {
let mut history = History::from_entries(vec![
HistoryEntry::new(Item::user_message("old"), "old-ann".to_string()),
HistoryEntry::new(Item::user_message("second"), "second-ann".to_string()),
]);
history.truncate(1);
assert_eq!(history.entries()[0].item.as_text(), Some("old"));
assert_eq!(history.entries()[0].annotation, "old-ann");
let previous = history.replace_entries(vec![HistoryEntry::new(
Item::user_message("restored"),
"restored-ann".to_string(),
)]);
assert_eq!(previous.len(), 1);
assert_eq!(history.entries()[0].item.as_text(), Some("restored"));
assert_eq!(history.entries()[0].annotation, "restored-ann");
}
+27 -14
View File
@@ -58,6 +58,7 @@ async fn test_callback_llm_retry_event() {
max_attempts: 2, max_attempts: 2,
total_timeout: Duration::from_secs(1), total_timeout: Duration::from_secs(1),
}); });
let mut history = agen::History::new();
let notices = Arc::new(Mutex::new(Vec::new())); let notices = Arc::new(Mutex::new(Vec::new()));
let sink = notices.clone(); let sink = notices.clone();
@@ -65,8 +66,11 @@ async fn test_callback_llm_retry_event() {
sink.lock().unwrap().push((llm_call, notice.clone())); sink.lock().unwrap().push((llm_call, notice.clone()));
}); });
let result = engine.run("retry once").await; let result = engine.run(&mut history, "retry once").await;
assert!(result.is_ok(), "engine should succeed after one retry"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"engine should succeed after one retry"
);
let notices = notices.lock().unwrap(); let notices = notices.lock().unwrap();
assert_eq!(notices.len(), 1); assert_eq!(notices.len(), 1);
@@ -91,6 +95,7 @@ async fn test_callback_text_block_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let text_deltas = Arc::new(Mutex::new(Vec::new())); let text_deltas = Arc::new(Mutex::new(Vec::new()));
let text_completes = Arc::new(Mutex::new(Vec::new())); let text_completes = Arc::new(Mutex::new(Vec::new()));
@@ -108,9 +113,12 @@ async fn test_callback_text_block_events() {
}); });
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run("Greet me").await; let result = engine.run(&mut history, "Greet me").await;
assert!(result.is_ok(), "Engine should complete"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete"
);
let deltas = text_deltas.lock().unwrap(); let deltas = text_deltas.lock().unwrap();
assert_eq!(deltas.len(), 2); assert_eq!(deltas.len(), 2);
@@ -137,6 +145,7 @@ async fn test_callback_tool_call_complete() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let tool_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new())); let tool_starts = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let tool_completes = Arc::new(Mutex::new(Vec::new())); let tool_completes = Arc::new(Mutex::new(Vec::new()));
@@ -154,8 +163,8 @@ async fn test_callback_tool_call_complete() {
}); });
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let _ = engine.run("Weather please").await; let _ = engine.run(&mut history, "Weather please").await;
let starts = tool_starts.lock().unwrap(); let starts = tool_starts.lock().unwrap();
assert_eq!(starts.len(), 1); assert_eq!(starts.len(), 1);
@@ -183,6 +192,7 @@ async fn test_callback_turn_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let turn_starts = Arc::new(Mutex::new(Vec::new())); let turn_starts = Arc::new(Mutex::new(Vec::new()));
let turn_ends = Arc::new(Mutex::new(Vec::new())); let turn_ends = Arc::new(Mutex::new(Vec::new()));
@@ -197,9 +207,9 @@ async fn test_callback_turn_events() {
ends.lock().unwrap().push(turn); ends.lock().unwrap().push(turn);
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let result = engine.run("Do something").await; let result = engine.run(&mut history, "Do something").await;
assert!(result.is_ok()); assert!(matches!(result.result, agen::EngineRunExit::Finished));
let starts = turn_starts.lock().unwrap(); let starts = turn_starts.lock().unwrap();
let ends = turn_ends.lock().unwrap(); let ends = turn_ends.lock().unwrap();
@@ -254,6 +264,7 @@ async fn test_callback_tool_result_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
engine.register_tool(fixed_tool( engine.register_tool(fixed_tool(
"fixed", "fixed",
@@ -276,7 +287,7 @@ async fn test_callback_tool_result_events() {
)); ));
}); });
let _ = engine.run("call it").await; let _ = engine.run(&mut history, "call it").await;
let observed = captured.lock().unwrap(); let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1); assert_eq!(observed.len(), 1);
@@ -330,6 +341,7 @@ async fn test_callback_tool_result_error_path() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
engine.register_tool(erroring_tool("erroring", "boom")); engine.register_tool(erroring_tool("erroring", "boom"));
@@ -345,7 +357,7 @@ async fn test_callback_tool_result_error_path() {
)); ));
}); });
let _ = engine.run("fail it").await; let _ = engine.run(&mut history, "fail it").await;
let observed = captured.lock().unwrap(); let observed = captured.lock().unwrap();
assert_eq!(observed.len(), 1); assert_eq!(observed.len(), 1);
@@ -374,6 +386,7 @@ async fn test_callback_usage_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
let usage_events = Arc::new(Mutex::new(Vec::new())); let usage_events = Arc::new(Mutex::new(Vec::new()));
@@ -382,8 +395,8 @@ async fn test_callback_usage_events() {
usages.lock().unwrap().push(event.clone()); usages.lock().unwrap().push(event.clone());
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineRunExit)
let _ = engine.run("Hello").await; let _ = engine.run(&mut history, "Hello").await;
let usages = usage_events.lock().unwrap(); let usages = usage_events.lock().unwrap();
assert_eq!(usages.len(), 1); assert_eq!(usages.len(), 1);
+8 -1
View File
@@ -19,6 +19,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
pub struct MockLlmClient { pub struct MockLlmClient {
responses: Arc<Vec<Vec<Event>>>, responses: Arc<Vec<Vec<Event>>>,
call_count: Arc<AtomicUsize>, call_count: Arc<AtomicUsize>,
requests: Arc<Mutex<Vec<Request>>>,
} }
impl MockLlmClient { impl MockLlmClient {
@@ -30,6 +31,7 @@ impl MockLlmClient {
Self { Self {
responses: Arc::new(responses), responses: Arc::new(responses),
call_count: Arc::new(AtomicUsize::new(0)), call_count: Arc::new(AtomicUsize::new(0)),
requests: Arc::new(Mutex::new(Vec::new())),
} }
} }
@@ -41,6 +43,10 @@ impl MockLlmClient {
pub fn event_count(&self) -> usize { pub fn event_count(&self) -> usize {
self.responses.iter().map(|v| v.len()).sum() self.responses.iter().map(|v| v.len()).sum()
} }
pub fn requests(&self) -> Vec<Request> {
self.requests.lock().unwrap().clone()
}
} }
#[async_trait] #[async_trait]
@@ -51,8 +57,9 @@ impl LlmClient for MockLlmClient {
async fn stream( async fn stream(
&self, &self,
_request: Request, request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> { ) -> Result<Pin<Box<dyn Stream<Item = Result<Event, ClientError>> + Send>>, ClientError> {
self.requests.lock().unwrap().push(request);
let count = self.call_count.fetch_add(1, Ordering::SeqCst); let count = self.call_count.fetch_add(1, Ordering::SeqCst);
if count >= self.responses.len() { if count >= self.responses.len() {
return Err(ClientError::Api { return Err(ClientError::Api {
+16 -5
View File
@@ -134,11 +134,15 @@ async fn test_engine_simple_text_response() {
let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history = agen::History::new();
// Send a simple message (Mutable::run consumes self, returns tuple) // Send a simple message (Mutable::run consumes self, returns tuple)
let result = engine.run("Hello").await; let result = engine.run(&mut history, "Hello").await;
assert!(result.is_ok(), "Engine should complete successfully"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete successfully"
);
} }
/// Verify that Engine can correctly process responses containing tool calls /// Verify that Engine can correctly process responses containing tool calls
@@ -156,6 +160,7 @@ async fn test_engine_tool_call() {
let client = MockLlmClient::from_fixture(&fixture_path).unwrap(); let client = MockLlmClient::from_fixture(&fixture_path).unwrap();
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history = agen::History::new();
// Register tool // Register tool
let weather_tool = MockWeatherTool::new(); let weather_tool = MockWeatherTool::new();
@@ -163,7 +168,9 @@ async fn test_engine_tool_call() {
engine.register_tool(weather_tool.definition()); engine.register_tool(weather_tool.definition());
// Send message (Mutable::run consumes self, returns tuple) // Send message (Mutable::run consumes self, returns tuple)
let _result = engine.run("What's the weather in Tokyo?").await; let _result = engine
.run(&mut history, "What's the weather in Tokyo?")
.await;
// Verify tool was called // Verify tool was called
// Note: max_turns=1 so no request is sent after tool result // Note: max_turns=1 so no request is sent after tool result
@@ -195,11 +202,15 @@ async fn test_engine_with_programmatic_events() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history = agen::History::new();
// Mutable::run consumes self, returns tuple // Mutable::run consumes self, returns tuple
let result = engine.run("Greet me").await; let result = engine.run(&mut history, "Greet me").await;
assert!(result.is_ok(), "Engine should complete successfully"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete successfully"
);
} }
/// Verify that ToolCallCollector correctly collects ToolCall from ToolUse block events /// Verify that ToolCallCollector correctly collects ToolCall from ToolUse block events
+154 -91
View File
@@ -14,7 +14,7 @@ use agen::interceptor::{
}; };
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, EngineError, EngineResult}; use agen::{Engine, EngineError, EngineRunExit, History, StopReason};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
@@ -42,36 +42,37 @@ fn test_mutable_set_system_prompt() {
fn test_mutable_history_manipulation() { fn test_mutable_history_manipulation() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// Initial state is empty // Initial state is empty
assert!(engine.history().is_empty()); assert!(history.is_empty());
// Add to history // Add to history
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Hi there!")]) .append_history(&mut history, vec![Item::assistant_message("Hi there!")])
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
// Append to history via the callback-aware API. // Append to history via the callback-aware API.
engine engine
.append_history(vec![Item::user_message("How are you?")]) .append_history(&mut history, vec![Item::user_message("How are you?")])
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 3); assert_eq!(history.len(), 3);
// Clear history // Clear history
engine.clear_history(); engine.clear_history(&mut history);
assert!(engine.history().is_empty()); assert!(history.is_empty());
// Set history // Set history
let items = vec![ let items = vec![
Item::user_message("Test"), Item::user_message("Test"),
Item::assistant_message("Response"), Item::assistant_message("Response"),
]; ];
engine.set_history(items); engine.set_history(&mut history, items);
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
} }
/// Verify that Engine can be constructed using builder pattern /// Verify that Engine can be constructed using builder pattern
@@ -79,9 +80,10 @@ fn test_mutable_history_manipulation() {
fn test_mutable_builder_pattern() { fn test_mutable_builder_pattern() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("System prompt"); let engine = Engine::new(client).system_prompt("System prompt");
let history: History = History::new();
assert_eq!(engine.get_system_prompt(), Some("System prompt")); assert_eq!(engine.get_system_prompt(), Some("System prompt"));
assert!(engine.history().is_empty()); assert!(history.is_empty());
} }
/// Verify that multiple items can be added with append_history and callbacks fire. /// Verify that multiple items can be added with append_history and callbacks fire.
@@ -91,6 +93,7 @@ fn test_mutable_append_history() {
let observed = Arc::new(Mutex::new(Vec::new())); let observed = Arc::new(Mutex::new(Vec::new()));
let observed_for_callback = Arc::clone(&observed); let observed_for_callback = Arc::clone(&observed);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.on_history_append(move |item| { engine.on_history_append(move |item| {
if let Some(text) = item.as_text() { if let Some(text) = item.as_text() {
observed_for_callback.lock().unwrap().push(text.to_string()); observed_for_callback.lock().unwrap().push(text.to_string());
@@ -99,18 +102,21 @@ fn test_mutable_append_history() {
}); });
engine engine
.append_history(vec![Item::user_message("First")]) .append_history(&mut history, vec![Item::user_message("First")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![ .append_history(
&mut history,
vec![
Item::assistant_message("Response 1"), Item::assistant_message("Response 1"),
Item::user_message("Second"), Item::user_message("Second"),
Item::assistant_message("Response 2"), Item::assistant_message("Response 2"),
]) ],
)
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 4); assert_eq!(history.len(), 4);
assert_eq!( assert_eq!(
observed.lock().unwrap().as_slice(), observed.lock().unwrap().as_slice(),
["First", "Response 1", "Second", "Response 2"] ["First", "Response 1", "Second", "Response 2"]
@@ -185,6 +191,7 @@ async fn history_append_failure_stops_before_tool_execution() {
]); ]);
let tool = CountingTool::new("count_tool"); let tool = CountingTool::new("count_tool");
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.register_tool(tool.definition()); engine.register_tool(tool.definition());
engine.on_history_append(|item| { engine.on_history_append(|item| {
if item.is_tool_call() { if item.is_tool_call() {
@@ -194,15 +201,15 @@ async fn history_append_failure_stops_before_tool_execution() {
} }
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
let error = engine.run("use the tool").await.unwrap_err(); let exit = engine.run(&mut history, "use the tool").await;
assert!( assert!(
matches!(error, EngineError::HistoryAppend(ref message) if message == "simulated ENOSPC") matches!(exit, EngineRunExit::Interrupted(StopReason::Unexpected(EngineError::HistoryAppend(ref message))) if message == "simulated ENOSPC")
); );
assert_eq!(tool.call_count(), 0); assert_eq!(tool.call_count(), 0);
assert_eq!(engine.history().len(), 1); assert_eq!(history.len(), 1);
assert_eq!(engine.history()[0].as_text(), Some("use the tool")); assert_eq!(history.entries()[0].item.as_text(), Some("use the tool"));
} }
// ============================================================================= // =============================================================================
@@ -214,21 +221,22 @@ async fn history_append_failure_stops_before_tool_execution() {
fn test_lock_transition() { fn test_lock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine.set_system_prompt("System"); engine.set_system_prompt("System");
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Hi")]) .append_history(&mut history, vec![Item::assistant_message("Hi")])
.unwrap(); .unwrap();
// Lock // Lock
let locked_engine = engine.lock(); let locked_engine = engine.lock(&history);
// History and system prompt are still accessible in Locked state // History and system prompt are still accessible in Locked state
assert_eq!(locked_engine.get_system_prompt(), Some("System")); assert_eq!(locked_engine.get_system_prompt(), Some("System"));
assert_eq!(locked_engine.history().len(), 2); assert_eq!(history.len(), 2);
assert_eq!(locked_engine.locked_prefix_len(), 2); assert_eq!(locked_engine.locked_prefix_len(), 2);
} }
@@ -237,21 +245,22 @@ fn test_lock_transition() {
fn test_unlock_transition() { fn test_unlock_transition() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine engine
.append_history(vec![Item::user_message("Hello")]) .append_history(&mut history, vec![Item::user_message("Hello")])
.unwrap(); .unwrap();
let locked_engine = engine.lock(); let locked_engine = engine.lock(&history);
// Unlock // Unlock
let mut engine = locked_engine.unlock(); let mut engine = locked_engine.unlock();
// History operations are available again in Mutable state // History operations are available again in Mutable state
engine engine
.append_history(vec![Item::assistant_message("Hi")]) .append_history(&mut history, vec![Item::assistant_message("Hi")])
.unwrap(); .unwrap();
engine.clear_history(); engine.clear_history(&mut history);
assert!(engine.history().is_empty()); assert!(history.is_empty());
} }
// ============================================================================= // =============================================================================
@@ -272,20 +281,20 @@ async fn test_mutable_run_updates_history() -> Result<(), EngineError> {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
// Execute (Mutable::run consumes self, returns EngineRunOutput) // Execute (Mutable::run consumes self, returns EngineRunOutput)
let out = engine.run("Hi there").await?; let _out = engine.run(&mut history, "Hi there").await;
let engine = out.engine;
// History is updated // History is updated
let history = engine.history(); let entries = history.entries();
assert_eq!(history.len(), 2); // user + assistant assert_eq!(history.len(), 2); // user + assistant
// User message // User message
assert_eq!(history[0].as_text(), Some("Hi there")); assert_eq!(entries[0].item.as_text(), Some("Hi there"));
// Assistant message // Assistant message
assert_eq!(history[1].as_text(), Some("Hello, I'm an assistant!")); assert_eq!(entries[1].item.as_text(), Some("Hello, I'm an assistant!"));
Ok(()) Ok(())
} }
@@ -316,35 +325,36 @@ async fn test_locked_multi_turn_history_accumulation() {
]); ]);
let engine = Engine::new(client).system_prompt("You are helpful."); let engine = Engine::new(client).system_prompt("You are helpful.");
let mut history: History = History::new();
// Lock (after setting system prompt) // Lock (after setting system prompt)
let mut locked_engine = engine.lock(); let mut locked_engine = engine.lock(&history);
assert_eq!(locked_engine.locked_prefix_len(), 0); // No items yet assert_eq!(locked_engine.locked_prefix_len(), 0); // No items yet
// Turn 1 // Turn 1
let result1 = locked_engine.run("Hello!").await; let result1 = locked_engine.run(&mut history, "Hello!").await;
assert!(result1.is_ok()); assert!(matches!(result1, EngineRunExit::Finished));
assert_eq!(locked_engine.history().len(), 2); // user + assistant assert_eq!(history.len(), 2); // user + assistant
// Turn 2 // Turn 2
let result2 = locked_engine.run("Can you help me?").await; let result2 = locked_engine.run(&mut history, "Can you help me?").await;
assert!(result2.is_ok()); assert!(matches!(result2, EngineRunExit::Finished));
assert_eq!(locked_engine.history().len(), 4); // 2 * (user + assistant) assert_eq!(history.len(), 4); // 2 * (user + assistant)
// Verify history contents // Verify history contents
let history = locked_engine.history(); let entries = history.entries();
// Turn 1 user message // Turn 1 user message
assert_eq!(history[0].as_text(), Some("Hello!")); assert_eq!(entries[0].item.as_text(), Some("Hello!"));
// Turn 1 assistant message // Turn 1 assistant message
assert_eq!(history[1].as_text(), Some("Nice to meet you!")); assert_eq!(entries[1].item.as_text(), Some("Nice to meet you!"));
// Turn 2 user message // Turn 2 user message
assert_eq!(history[2].as_text(), Some("Can you help me?")); assert_eq!(entries[2].item.as_text(), Some("Can you help me?"));
// Turn 2 assistant message // Turn 2 assistant message
assert_eq!(history[3].as_text(), Some("I can help with that.")); assert_eq!(entries[3].item.as_text(), Some("I can help with that."));
} }
/// Verify that locked_prefix_len correctly records history length at lock time /// Verify that locked_prefix_len correctly records history length at lock time
@@ -370,26 +380,33 @@ async fn test_locked_prefix_len_tracking() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// Add items beforehand // Add items beforehand
engine engine
.append_history(vec![Item::user_message("Pre-existing message 1")]) .append_history(
&mut history,
vec![Item::user_message("Pre-existing message 1")],
)
.unwrap(); .unwrap();
engine engine
.append_history(vec![Item::assistant_message("Pre-existing response 1")]) .append_history(
&mut history,
vec![Item::assistant_message("Pre-existing response 1")],
)
.unwrap(); .unwrap();
assert_eq!(engine.history().len(), 2); assert_eq!(history.len(), 2);
// Lock // Lock
let mut locked_engine = engine.lock(); let mut locked_engine = engine.lock(&history);
assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time assert_eq!(locked_engine.locked_prefix_len(), 2); // 2 items at lock time
// Execute turn // Execute turn
locked_engine.run("New message").await.unwrap(); locked_engine.run(&mut history, "New message").await;
// History grows but locked_prefix_len remains unchanged // History grows but locked_prefix_len remains unchanged
assert_eq!(locked_engine.history().len(), 4); // 2 + 2 assert_eq!(history.len(), 4); // 2 + 2
assert_eq!(locked_engine.locked_prefix_len(), 2); // Unchanged assert_eq!(locked_engine.locked_prefix_len(), 2); // Unchanged
} }
@@ -416,18 +433,22 @@ async fn test_turn_count_increment() -> Result<(), EngineError> {
]); ]);
let engine = Engine::new(client); let engine = Engine::new(client);
let mut history: History = History::new();
assert_eq!(engine.turn_count(), 0); assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.llm_call_count(), 0); assert_eq!(engine.llm_call_count(), 0);
// First run consumes Mutable, returns EngineRunOutput // First run consumes Mutable, returns EngineRunOutput
let mut engine = engine.run("First").await?.engine; let mut engine = engine.run(&mut history, "First").await.engine;
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
// Retry not yet implemented → AgentTurn:LlmCall is 1:1. // Retry not yet implemented → AgentTurn:LlmCall is 1:1.
assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.llm_call_count(), 1);
// Subsequent runs on Locked take &mut self // Subsequent runs on Locked take &mut self
engine.run("Second").await?; assert!(matches!(
engine.run(&mut history, "Second").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 2); assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.llm_call_count(), 2); assert_eq!(engine.llm_call_count(), 2);
@@ -447,28 +468,29 @@ async fn test_unlock_edit_relock() {
]]); ]]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
engine engine
.append_history(vec![ .append_history(
Item::user_message("Hello"), &mut history,
Item::assistant_message("Hi"), vec![Item::user_message("Hello"), Item::assistant_message("Hi")],
]) )
.unwrap(); .unwrap();
// Lock -> Unlock // Lock -> Unlock
let locked = engine.lock(); let locked = engine.lock(&history);
assert_eq!(locked.locked_prefix_len(), 2); assert_eq!(locked.locked_prefix_len(), 2);
let mut unlocked = locked.unlock(); let mut unlocked = locked.unlock();
// Edit history // Edit history
unlocked.clear_history(); unlocked.clear_history(&mut history);
unlocked unlocked
.append_history(vec![Item::user_message("Fresh start")]) .append_history(&mut history, vec![Item::user_message("Fresh start")])
.unwrap(); .unwrap();
// Re-lock // Re-lock
let relocked = unlocked.lock(); let relocked = unlocked.lock(&history);
assert_eq!(relocked.history().len(), 1); assert_eq!(history.len(), 1);
assert_eq!(relocked.locked_prefix_len(), 1); assert_eq!(relocked.locked_prefix_len(), 1);
} }
@@ -511,19 +533,26 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let tool_a = CountingTool::new("tool_a"); let tool_a = CountingTool::new("tool_a");
engine.register_tool(tool_a.definition()); engine.register_tool(tool_a.definition());
let mut locked = engine.lock(); let mut locked = engine.lock(&history);
locked.run("first").await.expect("first run"); assert!(matches!(
locked.run(&mut history, "first").await,
EngineRunExit::Finished
));
assert_eq!(tool_a.call_count(), 1, "tool_a should be called once"); assert_eq!(tool_a.call_count(), 1, "tool_a should be called once");
let mut unlocked = locked.unlock(); let mut unlocked = locked.unlock();
let tool_b = CountingTool::new("tool_b"); let tool_b = CountingTool::new("tool_b");
unlocked.register_tool(tool_b.definition()); unlocked.register_tool(tool_b.definition());
let mut relocked = unlocked.lock(); let mut relocked = unlocked.lock(&history);
relocked.run("second").await.expect("second run"); assert!(matches!(
relocked.run(&mut history, "second").await,
EngineRunExit::Finished
));
assert_eq!(tool_a.call_count(), 1, "tool_a should not be called again"); assert_eq!(tool_a.call_count(), 1, "tool_a should not be called again");
assert_eq!(tool_b.call_count(), 1, "tool_b should be called once"); assert_eq!(tool_b.call_count(), 1, "tool_b should be called once");
@@ -538,8 +567,9 @@ async fn test_lock_unlock_relock_tools_remain_effective() {
fn test_system_prompt_preserved_in_locked_state() { fn test_system_prompt_preserved_in_locked_state() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("Important system prompt"); let engine = Engine::new(client).system_prompt("Important system prompt");
let history: History = History::new();
let locked = engine.lock(); let locked = engine.lock(&history);
assert_eq!(locked.get_system_prompt(), Some("Important system prompt")); assert_eq!(locked.get_system_prompt(), Some("Important system prompt"));
let unlocked = locked.unlock(); let unlocked = locked.unlock();
@@ -554,14 +584,15 @@ fn test_system_prompt_preserved_in_locked_state() {
fn test_system_prompt_change_after_unlock() { fn test_system_prompt_change_after_unlock() {
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let engine = Engine::new(client).system_prompt("Original prompt"); let engine = Engine::new(client).system_prompt("Original prompt");
let history: History = History::new();
let locked = engine.lock(); let locked = engine.lock(&history);
let mut unlocked = locked.unlock(); let mut unlocked = locked.unlock();
unlocked.set_system_prompt("New prompt"); unlocked.set_system_prompt("New prompt");
assert_eq!(unlocked.get_system_prompt(), Some("New prompt")); assert_eq!(unlocked.get_system_prompt(), Some("New prompt"));
let relocked = unlocked.lock(); let relocked = unlocked.lock(&history);
assert_eq!(relocked.get_system_prompt(), Some("New prompt")); assert_eq!(relocked.get_system_prompt(), Some("New prompt"));
} }
@@ -623,40 +654,55 @@ impl Interceptor for ContinueTurnOnce {
#[tokio::test] #[tokio::test]
async fn max_turns_is_scoped_to_each_fresh_run() { async fn max_turns_is_scoped_to_each_fresh_run() {
let mut history: History = History::new();
let responses = vec![completed_text_events(), completed_text_events()]; let responses = vec![completed_text_events(), completed_text_events()];
let mut engine = Engine::new(MockLlmClient::with_responses(responses)); let mut engine = Engine::new(MockLlmClient::with_responses(responses));
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("first").await.unwrap(), EngineResult::Finished); assert!(matches!(
engine.run(&mut history, "first").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(engine.run("second").await.unwrap(), EngineResult::Finished); assert!(matches!(
engine.run(&mut history, "second").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 2); assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
} }
#[tokio::test] #[tokio::test]
async fn yielded_resume_keeps_the_same_unspent_turn_budget() { async fn yielded_resume_keeps_the_same_unspent_turn_budget() {
let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_interceptor(YieldOnce { engine.set_interceptor(YieldOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("start").await.unwrap(), EngineResult::Yielded); assert!(matches!(
engine.run(&mut history, "start").await,
EngineRunExit::Yielded
));
assert_eq!(engine.turn_count(), 0); assert_eq!(engine.turn_count(), 0);
assert_eq!(engine.active_run_turn_count(), Some(0)); assert_eq!(engine.active_run_turn_count(), Some(0));
assert_eq!(engine.resume().await.unwrap(), EngineResult::Finished); assert!(matches!(
engine.resume(&mut history).await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
} }
#[tokio::test] #[tokio::test]
async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() { async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
let mut history: History = History::new();
let events = vec![ let events = vec![
Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"), Event::tool_input_delta(0, "{}"),
@@ -672,14 +718,20 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
engine.set_interceptor(PauseToolOnce { engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("call it").await.unwrap(), EngineResult::Paused); assert!(matches!(
engine.run(&mut history, "call it").await,
EngineRunExit::Paused
));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(tool.call_count(), 0); assert_eq!(tool.call_count(), 0);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); assert!(matches!(
engine.resume(&mut history).await,
EngineRunExit::Interrupted(StopReason::LimitReached)
));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs"); assert_eq!(tool.call_count(), 1, "the consumed turn's tool still runs");
@@ -687,6 +739,7 @@ async fn paused_tool_resume_does_not_reset_the_consumed_turn_budget() {
#[tokio::test] #[tokio::test]
async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() { async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
let mut history: History = History::new();
let tool_events = vec![ let tool_events = vec![
Event::tool_use_start(0, "call_1", "count_tool"), Event::tool_use_start(0, "call_1", "count_tool"),
Event::tool_input_delta(0, "{}"), Event::tool_input_delta(0, "{}"),
@@ -703,12 +756,18 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
engine.set_interceptor(PauseToolOnce { engine.set_interceptor(PauseToolOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.run("pause").await.unwrap(), EngineResult::Paused); assert!(matches!(
engine.run(&mut history, "pause").await,
EngineRunExit::Paused
));
assert_eq!(engine.active_run_turn_count(), Some(1)); assert_eq!(engine.active_run_turn_count(), Some(1));
assert_eq!(engine.run("replace").await.unwrap(), EngineResult::Finished); assert!(matches!(
engine.run(&mut history, "replace").await,
EngineRunExit::Finished
));
assert_eq!(engine.turn_count(), 2); assert_eq!(engine.turn_count(), 2);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged"); assert_eq!(tool.call_count(), 1, "pending-tool semantics are unchanged");
@@ -716,17 +775,18 @@ async fn fresh_input_abandons_a_paused_run_and_starts_a_new_budget() {
#[tokio::test] #[tokio::test]
async fn interceptor_continuation_consumes_the_logical_run_budget() { async fn interceptor_continuation_consumes_the_logical_run_budget() {
let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_interceptor(ContinueTurnOnce { engine.set_interceptor(ContinueTurnOnce {
calls: AtomicUsize::new(0), calls: AtomicUsize::new(0),
}); });
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!( assert!(matches!(
engine.run("start").await.unwrap(), engine.run(&mut history, "start").await,
EngineResult::LimitReached EngineRunExit::Interrupted(StopReason::LimitReached)
); ));
assert_eq!(engine.turn_count(), 1); assert_eq!(engine.turn_count(), 1);
assert_eq!(engine.llm_call_count(), 1); assert_eq!(engine.llm_call_count(), 1);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
@@ -734,14 +794,17 @@ async fn interceptor_continuation_consumes_the_logical_run_budget() {
#[tokio::test] #[tokio::test]
async fn restored_active_run_budget_is_enforced_before_another_llm_call() { async fn restored_active_run_budget_is_enforced_before_another_llm_call() {
let mut history: History = History::new();
let mut engine = Engine::new(MockLlmClient::new(completed_text_events())); let mut engine = Engine::new(MockLlmClient::new(completed_text_events()));
engine.set_max_turns(Some(1)); engine.set_max_turns(Some(1));
engine.set_turn_count(7); engine.set_turn_count(7);
engine.set_last_run_interrupted(true);
engine.set_active_run_turn_count(Some(1)); engine.set_active_run_turn_count(Some(1));
let mut engine = engine.lock(); let mut engine = engine.lock(&history);
assert_eq!(engine.resume().await.unwrap(), EngineResult::LimitReached); assert!(matches!(
engine.resume(&mut history).await,
EngineRunExit::Interrupted(StopReason::LimitReached)
));
assert_eq!(engine.turn_count(), 7); assert_eq!(engine.turn_count(), 7);
assert_eq!(engine.llm_call_count(), 0); assert_eq!(engine.llm_call_count(), 0);
assert_eq!(engine.active_run_turn_count(), None); assert_eq!(engine.active_run_turn_count(), None);
+678 -10
View File
@@ -6,12 +6,13 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant}; use std::time::{Duration, Instant};
use agen::Engine;
use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo}; use agen::interceptor::{Interceptor, PostToolAction, PreToolAction, ToolCallInfo, ToolResultInfo};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::tool::{ use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult, Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolMeta, ToolOutput, ToolResult,
ToolResultDisposition,
}; };
use agen::{Engine, History, Item, ToolExecutionPolicy};
use async_trait::async_trait; use async_trait::async_trait;
mod common; mod common;
@@ -70,6 +71,144 @@ impl Tool for SlowTool {
} }
} }
#[derive(Clone)]
struct FirstAttemptHangsTool {
calls: Arc<AtomicUsize>,
}
impl FirstAttemptHangsTool {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new("hang_once")
.description("Hangs on the first execution attempt")
.input_schema(serde_json::json!({"type": "object"}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
fn call_count(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
}
#[async_trait]
impl Tool for FirstAttemptHangsTool {
async fn execute(
&self,
_input_json: &str,
_ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
let attempt = self.calls.fetch_add(1, Ordering::SeqCst);
if attempt == 0 {
std::future::pending::<()>().await;
}
Ok("completed on retry".to_string().into())
}
}
#[derive(Clone)]
struct CooperativeCancelTool {
calls: Arc<AtomicUsize>,
cancelled: Arc<tokio::sync::Notify>,
}
impl CooperativeCancelTool {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
cancelled: Arc::new(tokio::sync::Notify::new()),
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new("cooperative")
.description("Returns bounded progress after cancellation")
.input_schema(serde_json::json!({"type": "object"}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
}
#[async_trait]
impl Tool for CooperativeCancelTool {
async fn execute(
&self,
_input_json: &str,
_ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.cancelled.notified().await;
Err(ToolError::Cancelled(ToolOutput {
summary: "cooperative command cancelled".to_string(),
content: Some("stdout before cancellation\nstderr before cancellation".to_string()),
attachments: Vec::new(),
}))
}
async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> {
self.cancelled.notify_one();
Ok(())
}
}
#[derive(Clone)]
struct SafePauseTool {
calls: Arc<AtomicUsize>,
cancellations: Arc<AtomicUsize>,
release: Arc<tokio::sync::Notify>,
}
impl SafePauseTool {
fn new() -> Self {
Self {
calls: Arc::new(AtomicUsize::new(0)),
cancellations: Arc::new(AtomicUsize::new(0)),
release: Arc::new(tokio::sync::Notify::new()),
}
}
fn definition(&self) -> ToolDefinition {
let tool = self.clone();
Arc::new(move || {
let meta = ToolMeta::new("safe_pause")
.description("Waits for a safe-boundary release")
.input_schema(serde_json::json!({"type": "object"}));
(meta, Arc::new(tool.clone()) as Arc<dyn Tool>)
})
}
}
#[async_trait]
impl Tool for SafePauseTool {
async fn execute(
&self,
_input_json: &str,
_ctx: ToolExecutionContext,
) -> Result<ToolOutput, ToolError> {
self.calls.fetch_add(1, Ordering::SeqCst);
self.release.notified().await;
Ok(ToolOutput {
summary: "safe-boundary complete".to_string(),
content: Some("safe-boundary complete".to_string()),
attachments: Vec::new(),
})
}
async fn cancel(&self, _call_id: &str) -> Result<(), ToolError> {
self.cancellations.fetch_add(1, Ordering::SeqCst);
Ok(())
}
}
#[derive(Clone)] #[derive(Clone)]
struct ContextRecordingTool { struct ContextRecordingTool {
name: String, name: String,
@@ -145,6 +284,7 @@ async fn test_parallel_tool_execution() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let tool1 = SlowTool::new("slow_tool_1", 100); let tool1 = SlowTool::new("slow_tool_1", 100);
let tool2 = SlowTool::new("slow_tool_2", 100); let tool2 = SlowTool::new("slow_tool_2", 100);
let tool3 = SlowTool::new("slow_tool_3", 100); let tool3 = SlowTool::new("slow_tool_3", 100);
@@ -159,7 +299,7 @@ async fn test_parallel_tool_execution() {
let start = Instant::now(); let start = Instant::now();
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _result = engine.run("Run all tools").await; let _result = engine.run(&mut history, "Run all tools").await;
let elapsed = start.elapsed(); let elapsed = start.elapsed();
// Verify all tools were called // Verify all tools were called
@@ -178,6 +318,450 @@ async fn test_parallel_tool_execution() {
println!("Parallel execution completed in {:?}", elapsed); println!("Parallel execution completed in {:?}", elapsed);
} }
#[tokio::test]
async fn completed_results_commit_before_publish_without_waiting_for_siblings() {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_slow", "slow_first"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_fast", "fast_second"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::text_block_start(0),
Event::text_delta(0, "Done"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let client_probe = client.clone();
let mut engine = Engine::new(client);
engine.register_tool(SlowTool::new("slow_first", 100).definition());
engine.register_tool(SlowTool::new("fast_second", 5).definition());
let observed = Arc::new(Mutex::new(Vec::<String>::new()));
let published = observed.clone();
engine.on_tool_result(move |result| {
published
.lock()
.unwrap()
.push(format!("publish:{}", result.tool_use_id));
});
let committed = observed.clone();
let mut annotate = move |item: &Item| {
if let Item::ToolResult { call_id, .. } = item {
committed.lock().unwrap().push(format!("commit:{call_id}"));
}
Ok(())
};
let mut history = History::new();
let _ = engine
.run_with_annotation(&mut history, "run both", &mut annotate)
.await;
observed.lock().unwrap().push("run-returned".to_string());
assert_eq!(
observed.lock().unwrap().as_slice(),
[
"commit:call_fast",
"publish:call_fast",
"commit:call_slow",
"publish:call_slow",
"run-returned",
]
);
let committed_order: Vec<_> = history
.iter()
.filter_map(|entry| match &entry.item {
Item::ToolResult { call_id, .. } => Some(call_id.as_str()),
_ => None,
})
.collect();
assert_eq!(committed_order, ["call_fast", "call_slow"]);
let requests = client_probe.requests();
let projected_order: Vec<_> = requests[1]
.items
.iter()
.filter_map(|item| match item {
Item::ToolResult { call_id, .. } => Some(call_id.as_str()),
_ => None,
})
.collect();
assert_eq!(projected_order, ["call_slow", "call_fast"]);
}
#[tokio::test]
async fn cancellation_preserves_completed_results_and_resume_skips_them() {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_hang", "hang_once"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::tool_use_start(1, "call_fast_a", "fast_a"),
Event::tool_input_delta(1, r#"{}"#),
Event::tool_use_stop(1),
Event::tool_use_start(2, "call_fast_b", "fast_b"),
Event::tool_input_delta(2, r#"{}"#),
Event::tool_use_stop(2),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![
Event::text_block_start(0),
Event::text_delta(0, "Recovered"),
Event::text_block_stop(0, None),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
]);
let mut engine = Engine::new(client);
let hanging = FirstAttemptHangsTool::new();
let fast_a = SlowTool::new("fast_a", 1);
let fast_b = SlowTool::new("fast_b", 2);
engine.register_tool(hanging.definition());
engine.register_tool(fast_a.definition());
engine.register_tool(fast_b.definition());
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(30)).await;
cancel.send(()).await.unwrap();
});
let mut history = History::new();
let output = engine.run(&mut history, "start").await;
let mut engine = output.engine;
cancel_task.await.unwrap();
let completed_before_resume = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. }
if call_id == "call_fast_a" || call_id == "call_fast_b"
)
})
.count();
let unknown_before_resume = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_hang"
)
})
.count();
assert_eq!(completed_before_resume, 2);
assert_eq!(unknown_before_resume, 1);
assert_eq!(fast_a.call_count(), 1);
assert_eq!(fast_b.call_count(), 1);
assert_eq!(hanging.call_count(), 1);
let _ = engine.resume(&mut history).await;
assert_eq!(
fast_a.call_count(),
1,
"completed call must not be re-executed"
);
assert_eq!(
fast_b.call_count(),
1,
"completed call must not be re-executed"
);
assert_eq!(
hanging.call_count(),
1,
"OutcomeUnknown is terminal and must not be re-executed"
);
let completed_after_resume = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. }
if call_id == "call_fast_a" || call_id == "call_fast_b"
)
})
.count();
assert_eq!(completed_after_resume, 2);
assert_eq!(
history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_hang"
)
})
.count(),
1
);
}
#[tokio::test]
async fn cooperative_cancellation_commits_bounded_terminal_output() {
let client = MockLlmClient::with_responses(vec![vec![
Event::tool_use_start(0, "call_cooperative", "cooperative"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]]);
let mut engine = Engine::new(client);
let tool = CooperativeCancelTool::new();
engine.register_tool(tool.definition());
let observed = Arc::new(Mutex::new(Vec::<&'static str>::new()));
let published = observed.clone();
engine.on_tool_result(move |_| published.lock().unwrap().push("published"));
let committed = observed.clone();
let mut annotate = move |item: &Item| {
if matches!(item, Item::ToolResult { .. }) {
committed.lock().unwrap().push("committed");
}
Ok(())
};
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(30)).await;
cancel.send(()).await.unwrap();
});
let mut history = History::new();
let output = engine
.run_with_annotation(&mut history, "start", &mut annotate)
.await;
observed.lock().unwrap().push("run-returned");
cancel_task.await.unwrap();
assert_eq!(
observed.lock().unwrap().as_slice(),
["committed", "published", "run-returned"]
);
assert_eq!(tool.calls.load(Ordering::SeqCst), 1);
let terminal: Vec<_> = history
.iter()
.filter_map(|entry| match &entry.item {
Item::ToolResult {
call_id,
disposition,
content,
..
} if call_id == "call_cooperative" => Some((*disposition, content.as_deref())),
_ => None,
})
.collect();
assert_eq!(terminal.len(), 1);
assert_eq!(terminal[0].0, ToolResultDisposition::Cancelled);
assert_eq!(
terminal[0].1,
Some("stdout before cancellation\nstderr before cancellation")
);
assert!(matches!(
output.result,
agen::EngineRunExit::Interrupted(agen::StopReason::Cancelled)
));
}
#[tokio::test]
async fn pause_waits_for_started_tool_terminal_without_cancelling_provider() {
let client = MockLlmClient::with_responses(vec![vec![
Event::tool_use_start(0, "call_safe_pause", "safe_pause"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]]);
let mut engine = Engine::new(client);
let tool = SafePauseTool::new();
engine.register_tool(tool.definition());
let pause = engine.pause_sender();
let calls = Arc::clone(&tool.calls);
let release = Arc::clone(&tool.release);
let control = tokio::spawn(async move {
tokio::time::timeout(Duration::from_secs(1), async {
while calls.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("tool execution starts");
pause.send(()).await.unwrap();
tokio::time::sleep(Duration::from_millis(50)).await;
release.notify_one();
});
let started_at = std::time::Instant::now();
let mut history = History::new();
let output = engine.run(&mut history, "pause safely").await;
control.await.unwrap();
assert!(started_at.elapsed() >= Duration::from_millis(50));
assert_eq!(tool.calls.load(Ordering::SeqCst), 1);
assert_eq!(tool.cancellations.load(Ordering::SeqCst), 0);
assert!(matches!(output.result, agen::EngineRunExit::Paused));
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::Success,
..
} if call_id == "call_safe_pause"
)));
}
#[tokio::test]
async fn pause_escalates_to_explicit_cancel_and_confirm_after_safe_boundary_deadline() {
let client = MockLlmClient::with_responses(vec![vec![
Event::tool_use_start(0, "call_pause_cancel", "cooperative"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]]);
let mut engine = Engine::new(client);
engine.set_tool_execution_policy(ToolExecutionPolicy {
pause_safe_boundary_timeout: Duration::from_millis(20),
cancellation_request_timeout: Duration::from_millis(50),
terminal_confirmation_timeout: Duration::from_millis(100),
});
let tool = CooperativeCancelTool::new();
engine.register_tool(tool.definition());
let pause = engine.pause_sender();
let calls = Arc::clone(&tool.calls);
let control = tokio::spawn(async move {
tokio::time::timeout(Duration::from_secs(1), async {
while calls.load(Ordering::SeqCst) == 0 {
tokio::task::yield_now().await;
}
})
.await
.expect("tool execution starts");
pause.send(()).await.unwrap();
});
let mut history = History::new();
let output = engine.run(&mut history, "pause with escalation").await;
control.await.unwrap();
assert!(matches!(output.result, agen::EngineRunExit::Paused));
assert!(history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::Cancelled,
..
} if call_id == "call_pause_cancel"
)));
}
#[tokio::test]
async fn cancellation_completion_race_commits_one_terminal_output() {
for iteration in 0..24u64 {
let client = MockLlmClient::with_responses(vec![
vec![
Event::tool_use_start(0, "call_racy", "racy"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
],
vec![Event::Status(StatusEvent {
status: ResponseStatus::Completed,
})],
]);
let mut engine = Engine::new(client);
let delay = 2 + iteration % 3;
let tool = SlowTool::new("racy", delay);
engine.register_tool(tool.definition());
let cancel = engine.cancel_sender();
let cancel_task = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(delay)).await;
let _ = cancel.send(()).await;
});
let mut history = History::new();
let _ = engine.run(&mut history, "race").await;
cancel_task.await.unwrap();
let terminal_count = history
.iter()
.filter(|entry| {
matches!(
&entry.item,
Item::ToolResult { call_id, .. } if call_id == "call_racy"
)
})
.count();
assert_eq!(terminal_count, 1, "iteration {iteration}");
assert_eq!(tool.call_count(), 1, "iteration {iteration}");
}
}
#[tokio::test]
async fn tool_result_commit_failure_prevents_publication() {
let client = MockLlmClient::with_responses(vec![vec![
Event::tool_use_start(0, "call_fast", "fast"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]]);
let mut engine = Engine::new(client);
engine.register_tool(SlowTool::new("fast", 1).definition());
let published = Arc::new(AtomicUsize::new(0));
let published_probe = published.clone();
engine.on_tool_result(move |_| {
published_probe.fetch_add(1, Ordering::SeqCst);
});
let mut history = History::new();
let mut reject_tool_result = |item: &Item| {
if matches!(item, Item::ToolResult { .. }) {
Err("session log unavailable".to_string())
} else {
Ok(())
}
};
let _ = engine
.run_with_annotation(&mut history, "start", &mut reject_tool_result)
.await;
assert_eq!(published.load(Ordering::SeqCst), 0);
assert!(
history
.iter()
.all(|entry| !matches!(entry.item, Item::ToolResult { .. }))
);
}
#[tokio::test] #[tokio::test]
async fn test_tool_execution_context_order_and_batch_id() { async fn test_tool_execution_context_order_and_batch_id() {
let client = MockLlmClient::with_responses(vec![ let client = MockLlmClient::with_responses(vec![
@@ -205,13 +789,14 @@ async fn test_tool_execution_context_order_and_batch_id() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let contexts = Arc::new(Mutex::new(Vec::new())); let contexts = Arc::new(Mutex::new(Vec::new()));
engine.register_tool(ContextRecordingTool::new("record_a", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_a", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_b", contexts.clone()).definition());
engine.register_tool(ContextRecordingTool::new("record_c", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record_c", contexts.clone()).definition());
let _ = engine.run("record contexts").await; let _ = engine.run(&mut history, "record contexts").await;
let mut contexts = contexts.lock().unwrap().clone(); let mut contexts = contexts.lock().unwrap().clone();
contexts.sort_by_key(|ctx| ctx.call_index); contexts.sort_by_key(|ctx| ctx.call_index);
@@ -256,11 +841,12 @@ async fn test_tool_execution_context_batch_id_changes_between_batches() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let contexts = Arc::new(Mutex::new(Vec::new())); let contexts = Arc::new(Mutex::new(Vec::new()));
engine.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition()); engine.register_tool(ContextRecordingTool::new("record", contexts.clone()).definition());
let _ = engine.run("record batches").await; let _ = engine.run(&mut history, "record batches").await;
let contexts = contexts.lock().unwrap().clone(); let contexts = contexts.lock().unwrap().clone();
assert_eq!(contexts.len(), 2); assert_eq!(contexts.len(), 2);
@@ -298,6 +884,7 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let executed_contexts = Arc::new(Mutex::new(Vec::new())); let executed_contexts = Arc::new(Mutex::new(Vec::new()));
let pre_contexts = Arc::new(Mutex::new(Vec::new())); let pre_contexts = Arc::new(Mutex::new(Vec::new()));
let post_contexts = Arc::new(Mutex::new(Vec::new())); let post_contexts = Arc::new(Mutex::new(Vec::new()));
@@ -344,7 +931,9 @@ async fn test_tool_execution_context_for_skipped_and_synthetic_paths() {
post_contexts: post_contexts.clone(), post_contexts: post_contexts.clone(),
}); });
let _ = engine.run("record skipped and synthetic contexts").await; let _ = engine
.run(&mut history, "record skipped and synthetic contexts")
.await;
let mut pre_contexts = pre_contexts.lock().unwrap().clone(); let mut pre_contexts = pre_contexts.lock().unwrap().clone();
pre_contexts.sort_by_key(|ctx| ctx.call_index); pre_contexts.sort_by_key(|ctx| ctx.call_index);
@@ -389,6 +978,7 @@ async fn test_before_tool_call_skip() {
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let allowed_tool = SlowTool::new("allowed_tool", 10); let allowed_tool = SlowTool::new("allowed_tool", 10);
let blocked_tool = SlowTool::new("blocked_tool", 10); let blocked_tool = SlowTool::new("blocked_tool", 10);
@@ -416,7 +1006,7 @@ async fn test_before_tool_call_skip() {
engine.set_interceptor(BlockingPolicy); engine.set_interceptor(BlockingPolicy);
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let _result = engine.run("Test hook").await; let _result = engine.run(&mut history, "Test hook").await;
// allowed_tool is called, but blocked_tool is not // allowed_tool is called, but blocked_tool is not
assert_eq!( assert_eq!(
@@ -457,6 +1047,7 @@ async fn test_post_tool_call_modification() {
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
#[derive(Clone)] #[derive(Clone)]
struct SimpleTool; struct SimpleTool;
@@ -503,9 +1094,12 @@ async fn test_post_tool_call_modification() {
}); });
// Mutable::run consumes self, returns (Locked, EngineResult) // Mutable::run consumes self, returns (Locked, EngineResult)
let result = engine.run("Test modification").await; let result = engine.run(&mut history, "Test modification").await;
assert!(result.is_ok(), "Engine should complete"); assert!(
matches!(result.result, agen::EngineRunExit::Finished),
"Engine should complete"
);
// Verify hook was called and content was modified // Verify hook was called and content was modified
let content = modified_content.lock().unwrap().clone(); let content = modified_content.lock().unwrap().clone();
@@ -540,6 +1134,7 @@ async fn test_before_tool_call_synthetic_result_committed() {
], ],
]); ]);
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
let blocked_tool = SlowTool::new("blocked_tool", 10); let blocked_tool = SlowTool::new("blocked_tool", 10);
let blocked_clone = blocked_tool.clone(); let blocked_clone = blocked_tool.clone();
engine.register_tool(blocked_tool.definition()); engine.register_tool(blocked_tool.definition());
@@ -558,10 +1153,10 @@ async fn test_before_tool_call_synthetic_result_committed() {
engine.set_interceptor(SyntheticPolicy); engine.set_interceptor(SyntheticPolicy);
let result = engine.run("Test synthetic result").await.unwrap(); let _result = engine.run(&mut history, "Test synthetic result").await;
assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run"); assert_eq!(blocked_clone.call_count(), 0, "Blocked tool should not run");
assert!(result.engine.history().iter().any(|item| matches!( assert!(history.items().any(|item| matches!(
item, item,
agen::Item::ToolResult { agen::Item::ToolResult {
call_id, call_id,
@@ -571,3 +1166,76 @@ async fn test_before_tool_call_synthetic_result_committed() {
} if call_id == "call_1" && summary == "permission denied" } if call_id == "call_1" && summary == "permission denied"
))); )));
} }
#[tokio::test]
async fn post_tool_abort_commits_confirmed_result_before_stopping_run() {
let client = MockLlmClient::new(vec![
Event::tool_use_start(0, "call_confirmed", "confirmed"),
Event::tool_input_delta(0, r#"{}"#),
Event::tool_use_stop(0),
Event::Status(StatusEvent {
status: ResponseStatus::Completed,
}),
]);
let mut engine = Engine::new(client);
let tool = SlowTool::new("confirmed", 1);
engine.register_tool(tool.definition());
struct AbortAfterResult;
#[async_trait]
impl Interceptor for AbortAfterResult {
async fn post_tool_call(&self, _info: &mut ToolResultInfo) -> PostToolAction {
PostToolAction::Abort("policy stopped the run".to_string())
}
}
engine.set_interceptor(AbortAfterResult);
let observed = Arc::new(Mutex::new(Vec::<&'static str>::new()));
let published = observed.clone();
engine.on_tool_result(move |_| published.lock().unwrap().push("published"));
let committed = observed.clone();
let mut annotate = move |item: &Item| {
if matches!(item, Item::ToolResult { .. }) {
committed.lock().unwrap().push("committed");
}
Ok(())
};
let mut history = History::new();
let output = engine
.run_with_annotation(&mut history, "run confirmed tool", &mut annotate)
.await;
observed.lock().unwrap().push("run-returned");
assert_eq!(tool.call_count(), 1);
assert_eq!(
observed.lock().unwrap().as_slice(),
["committed", "published", "run-returned"]
);
assert!(matches!(
output.result,
agen::EngineRunExit::Interrupted(agen::StopReason::Unexpected(
agen::EngineError::Aborted(ref reason)
)) if reason == "policy stopped the run"
));
let terminal: Vec<_> = history
.iter()
.filter_map(|entry| match &entry.item {
Item::ToolResult {
call_id,
disposition,
..
} if call_id == "call_confirmed" => Some(*disposition),
_ => None,
})
.collect();
assert_eq!(terminal, [ToolResultDisposition::Success]);
assert!(!history.iter().any(|entry| matches!(
&entry.item,
Item::ToolResult {
call_id,
disposition: ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_confirmed"
)));
}
+23 -19
View File
@@ -13,12 +13,12 @@
mod common; mod common;
use agen::Engine;
use agen::Item; use agen::Item;
use agen::llm_client::event::{ use agen::llm_client::event::{
BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus, BlockMetadata, BlockStart, BlockStop, BlockType, Event, ReasoningBlockData, ResponseStatus,
StatusEvent, StatusEvent,
}; };
use agen::{Engine, History};
use common::MockLlmClient; use common::MockLlmClient;
fn reasoning_block(text: impl Into<String>, data: ReasoningBlockData) -> Vec<Event> { fn reasoning_block(text: impl Into<String>, data: ReasoningBlockData) -> Vec<Event> {
@@ -65,15 +65,15 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
]); ]);
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("question?").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "question?").await;
let history = engine.history(); let entries = history.entries();
// user / reasoning / assistant_message // user / reasoning / assistant_message
assert_eq!(history.len(), 3, "history: {history:?}"); assert_eq!(history.len(), 3, "history: {history:?}");
assert!(matches!(history[0], Item::Message { .. })); assert!(matches!(entries[0].item, Item::Message { .. }));
match &history[1] { match &entries[1].item {
Item::Reasoning { Item::Reasoning {
text, signature, .. text, signature, ..
} => { } => {
@@ -82,7 +82,7 @@ async fn anthropic_thinking_round_trips_signature_into_history() {
} }
other => panic!("expected Reasoning, got {other:?}"), other => panic!("expected Reasoning, got {other:?}"),
} }
assert_eq!(history[2].as_text(), Some("Here's the answer")); assert_eq!(entries[2].item.as_text(), Some("Here's the answer"));
} }
/// OpenAI Responses 風: encrypted_content + summary を持った reasoning が /// OpenAI Responses 風: encrypted_content + summary を持った reasoning が
@@ -109,11 +109,11 @@ async fn openai_reasoning_round_trips_encrypted_and_summary() {
]); ]);
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "q").await;
let history = engine.history(); let entries = history.entries();
match &history[1] { match &entries[1].item {
Item::Reasoning { Item::Reasoning {
text, text,
summary, summary,
@@ -155,13 +155,13 @@ async fn reasoning_precedes_text_in_assistant_burst() {
})); }));
let client = MockLlmClient::new(events); let client = MockLlmClient::new(events);
let engine = Engine::new(client); let engine = Engine::new(client);
let out = engine.run("q").await.expect("run ok"); let mut history: History = History::new();
let engine = out.engine; let _out = engine.run(&mut history, "q").await;
let history = engine.history(); let entries = history.entries();
// user / reasoning(先頭) / assistant_message // user / reasoning(先頭) / assistant_message
assert!(matches!(history[1], Item::Reasoning { .. })); assert!(matches!(entries[1].item, Item::Reasoning { .. }));
assert_eq!(history[2].as_text(), Some("intermediate")); assert_eq!(entries[2].item.as_text(), Some("intermediate"));
} }
/// resume シナリオ: history.json 由来の Item::Reasoning(signature) を Engine に /// resume シナリオ: history.json 由来の Item::Reasoning(signature) を Engine に
@@ -207,14 +207,18 @@ async fn injected_reasoning_survives_into_outgoing_request() {
}; };
let mut engine = Engine::new(client); let mut engine = Engine::new(client);
let mut history: History = History::new();
// resume: 既存 history を流し込む // resume: 既存 history を流し込む
engine.set_history(vec![ engine.set_history(
&mut history,
vec![
Item::user_message("prior question"), Item::user_message("prior question"),
Item::reasoning("prior thinking").with_signature("SIG-PRIOR"), Item::reasoning("prior thinking").with_signature("SIG-PRIOR"),
Item::assistant_message("prior answer"), Item::assistant_message("prior answer"),
]); ],
);
let _ = engine.run("follow up").await.expect("run ok"); let _ = engine.run(&mut history, "follow up").await;
let req = captured let req = captured
.lock() .lock()
+3 -2
View File
@@ -1,4 +1,4 @@
use agen::Engine; use agen::{Engine, History};
use agen::llm_client::capability::{ use agen::llm_client::capability::{
CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport, CacheStrategy, ModelCapability, StructuredOutput, ToolCallingSupport,
}; };
@@ -22,7 +22,8 @@ fn main() {
cap, cap,
); );
let engine = Engine::new(client); let engine = Engine::new(client);
let mut locked = engine.lock(); let history = History::new();
let mut locked = engine.lock(&history);
let def: agen::tool::ToolDefinition = Arc::new(|| panic!("unused")); let def: agen::tool::ToolDefinition = Arc::new(|| panic!("unused"));
let _ = locked.register_tool(def); let _ = locked.register_tool(def);
} }
@@ -1,8 +1,8 @@
error[E0599]: no method named `register_tool` found for struct `Engine<HttpTransport<AnthropicScheme>, Locked>` in the current scope error[E0599]: no method named `register_tool` found for struct `Engine<HttpTransport<AnthropicScheme>, Locked>` in the current scope
--> tests/ui/locked_register_tool.rs:27:20 --> tests/ui/locked_register_tool.rs:28:20
| |
27 | let _ = locked.register_tool(def); 28 | let _ = locked.register_tool(def);
| ^^^^^^^^^^^^^ method not found in `Engine<HttpTransport<AnthropicScheme>, Locked>` | ^^^^^^^^^^^^^ method not found in `Engine<HttpTransport<AnthropicScheme>, Locked>`
| |
= note: the method was found for = note: the method was found for
- `Engine<C>` - `Engine<C, Mutable, A>`
+4 -1
View File
@@ -9,7 +9,7 @@
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; use crate::schema::{EvidenceKind, EvidenceOrigin, SourceEvidenceRef, SourceRef};
/// Current flat staging schema version. /// Current flat staging schema version.
pub const STAGING_SCHEMA_VERSION: u32 = 2; pub const STAGING_SCHEMA_VERSION: u32 = 2;
@@ -80,6 +80,8 @@ pub struct StagingEvidence {
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub entry_range: Option<[u64; 2]>, pub entry_range: Option<[u64; 2]>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<EvidenceOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub excerpt: Option<String>, pub excerpt: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub summary: Option<String>, pub summary: Option<String>,
@@ -159,6 +161,7 @@ mod tests {
id: "E001".into(), id: "E001".into(),
kind: EvidenceKind::new(EvidenceKind::MESSAGE), kind: EvidenceKind::new(EvidenceKind::MESSAGE),
entry_range: Some([10, 12]), entry_range: Some([10, 12]),
origin: None,
excerpt: Some("extract candidate taxonomy".into()), excerpt: Some("extract candidate taxonomy".into()),
summary: Some("User and assistant discussed staging kinds".into()), summary: Some("User and assistant discussed staging kinds".into()),
}; };
+37
View File
@@ -67,6 +67,40 @@ impl EvidenceKind {
} }
} }
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "snake_case")]
pub enum EvidenceOriginKind {
HumanInput,
WorkerInput,
FlowInstruction,
BackendInstruction,
ModelOutput,
ToolOutput,
DerivedSummary,
LegacyUnknown,
}
/// Bounded origin snapshot attached to extraction evidence. This is audit
/// metadata only and cannot authorize Workspace operations.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
pub struct EvidenceOrigin {
pub kind: EvidenceOriginKind,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub account_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub worker_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_selector: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_definition_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub flow_definition_revision: Option<u64>,
}
/// Host-resolved source/evidence metadata for an individual staging claim. /// Host-resolved source/evidence metadata for an individual staging claim.
/// ///
/// This deliberately stores only bounded anchor metadata: stable ids, entry /// This deliberately stores only bounded anchor metadata: stable ids, entry
@@ -86,6 +120,9 @@ pub struct SourceEvidenceRef {
/// Host-assigned evidence id within the referenced evidence set. /// Host-assigned evidence id within the referenced evidence set.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub evidence_id: Option<String>, pub evidence_id: Option<String>,
/// Trusted typed origin snapshot for this logical evidence entry.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub origin: Option<EvidenceOrigin>,
/// Extensible evidence kind tag. /// Extensible evidence kind tag.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
pub evidence_kind: Option<EvidenceKind>, pub evidence_kind: Option<EvidenceKind>,
+4 -1
View File
@@ -10,7 +10,10 @@ mod decision;
mod request; mod request;
mod summary; mod summary;
pub use common::{EvidenceKind, Frontmatter, SourceEvidenceRef, SourceRef, split_frontmatter}; pub use common::{
EvidenceKind, EvidenceOrigin, EvidenceOriginKind, Frontmatter, SourceEvidenceRef, SourceRef,
split_frontmatter,
};
pub use decision::{DecisionFrontmatter, DecisionStatus}; pub use decision::{DecisionFrontmatter, DecisionStatus};
pub use request::RequestFrontmatter; pub use request::RequestFrontmatter;
pub use summary::SummaryFrontmatter; pub use summary::SummaryFrontmatter;
+20
View File
@@ -352,6 +352,18 @@ pub struct InternalWorkerSnapshot {
pub internal_workers: Vec<InternalWorkerSnapshot>, pub internal_workers: Vec<InternalWorkerSnapshot>,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(rename_all = "snake_case")]
pub enum ToolResultDisposition {
#[default]
Success,
Error,
Interrupted,
Cancelled,
OutcomeUnknown,
}
#[derive(Debug, Clone, Serialize, Deserialize)] #[derive(Debug, Clone, Serialize, Deserialize)]
#[cfg_attr(feature = "typescript", derive(ts_rs::TS))] #[cfg_attr(feature = "typescript", derive(ts_rs::TS))]
#[serde(tag = "event", content = "data", rename_all = "snake_case")] #[serde(tag = "event", content = "data", rename_all = "snake_case")]
@@ -501,6 +513,8 @@ pub enum Event {
/// summary-only, or when the result was pruned. /// summary-only, or when the result was pruned.
#[serde(default, skip_serializing_if = "Option::is_none")] #[serde(default, skip_serializing_if = "Option::is_none")]
output: Option<String>, output: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
disposition: Option<ToolResultDisposition>,
#[serde(default)] #[serde(default)]
is_error: bool, is_error: bool,
}, },
@@ -923,6 +937,7 @@ pub enum WorkerStatus {
Idle, Idle,
Running, Running,
Paused, Paused,
Stopped,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
@@ -1838,6 +1853,7 @@ mod tests {
id: "call_1".into(), id: "call_1".into(),
summary: "Read 128 bytes".into(), summary: "Read 128 bytes".into(),
output: Some("hello world".into()), output: Some("hello world".into()),
disposition: Some(ToolResultDisposition::Success),
is_error: false, is_error: false,
}; };
let json = serde_json::to_string(&event).unwrap(); let json = serde_json::to_string(&event).unwrap();
@@ -1854,11 +1870,13 @@ mod tests {
id, id,
summary, summary,
output, output,
disposition,
is_error, is_error,
} => { } => {
assert_eq!(id, "call_1"); assert_eq!(id, "call_1");
assert_eq!(summary, "Read 128 bytes"); assert_eq!(summary, "Read 128 bytes");
assert_eq!(output.as_deref(), Some("hello world")); assert_eq!(output.as_deref(), Some("hello world"));
assert_eq!(disposition, Some(ToolResultDisposition::Success));
assert!(!is_error); assert!(!is_error);
} }
other => panic!("expected ToolResult, got {other:?}"), other => panic!("expected ToolResult, got {other:?}"),
@@ -1871,6 +1889,7 @@ mod tests {
id: "call_2".into(), id: "call_2".into(),
summary: "ok".into(), summary: "ok".into(),
output: None, output: None,
disposition: Some(ToolResultDisposition::Success),
is_error: false, is_error: false,
}; };
let json = serde_json::to_string(&event).unwrap(); let json = serde_json::to_string(&event).unwrap();
@@ -1886,6 +1905,7 @@ mod tests {
id: "call_3".into(), id: "call_3".into(),
summary: "invalid argument".into(), summary: "invalid argument".into(),
output: None, output: None,
disposition: Some(ToolResultDisposition::Error),
is_error: true, is_error: true,
}; };
let json = serde_json::to_string(&event).unwrap(); let json = serde_json::to_string(&event).unwrap();
+2 -1
View File
@@ -8,7 +8,7 @@ use crate::{
CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot, CompletionKind, ErrorCode, Event, Greeting, InFlightBlock, InFlightSnapshot,
InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot, InFlightToolCallState, InternalWorkerKind, InternalWorkerRef, InternalWorkerSnapshot,
InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId, InvokeKind, MemoryWorkerEvent, Method, Permission, RewindSummary, RewindTarget, RewindTargetId,
RunResult, ScopeRule, Segment, TurnResult, WorkerEvent, WorkerStatus, RunResult, ScopeRule, Segment, ToolResultDisposition, TurnResult, WorkerEvent, WorkerStatus,
subscription::{ subscription::{
EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame, EventSubscriptionSelector, SubscriptionEvent, SubscriptionEventPayload, SubscriptionFrame,
SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest, SubscriptionFramePayload, SubscriptionId, SubscriptionRejectionCode, SubscriptionRequest,
@@ -45,6 +45,7 @@ pub fn generated_protocol_types() -> String {
push_decl::<TurnResult>(&cfg, &mut output); push_decl::<TurnResult>(&cfg, &mut output);
push_decl::<InvokeKind>(&cfg, &mut output); push_decl::<InvokeKind>(&cfg, &mut output);
push_decl::<RunResult>(&cfg, &mut output); push_decl::<RunResult>(&cfg, &mut output);
push_decl::<ToolResultDisposition>(&cfg, &mut output);
push_decl::<ErrorCode>(&cfg, &mut output); push_decl::<ErrorCode>(&cfg, &mut output);
push_decl::<Permission>(&cfg, &mut output); push_decl::<Permission>(&cfg, &mut output);
push_decl::<InFlightToolCallState>(&cfg, &mut output); push_decl::<InFlightToolCallState>(&cfg, &mut output);
+180
View File
@@ -0,0 +1,180 @@
//! Serializable history entries with restore-authoritative logical identity and origin.
use serde::{Deserialize, Serialize};
use crate::{LoggedItem, SessionId};
/// Stable logical identity of one model-visible history entry.
///
/// This value is generated at the trusted Worker session boundary and copied
/// unchanged across fork, rewind, compaction retention, restore, and reboot.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct LoggedSessionHistoryEntryId(pub String);
impl LoggedSessionHistoryEntryId {
pub fn new() -> Self {
Self(uuid::Uuid::now_v7().to_string())
}
}
impl Default for LoggedSessionHistoryEntryId {
fn default() -> Self {
Self::new()
}
}
/// Bounded subject snapshot. It is evidence, not a live authorization handle.
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LoggedWorkerSubject {
#[serde(default, skip_serializing_if = "Option::is_none")]
pub workspace_id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub runtime_id: Option<String>,
pub worker_id: String,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum LoggedSessionHistoryOrigin {
HumanInput {
account_id: String,
},
WorkerInput {
actor: LoggedWorkerSubject,
},
FlowInstruction {
selector: String,
definition_id: String,
definition_revision: u64,
instance_id: String,
state_id: String,
},
BackendInstruction {
#[serde(default, skip_serializing_if = "Option::is_none")]
operation_id: Option<String>,
},
ModelOutput {
worker: LoggedWorkerSubject,
},
ToolOutput {
worker: LoggedWorkerSubject,
},
DerivedSummary,
LegacyUnknown,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LoggedHistoryDerivation {
pub sources: Vec<LoggedSessionHistoryEntryId>,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct LoggedSessionHistoryMetadata {
pub entry_id: LoggedSessionHistoryEntryId,
pub origin: LoggedSessionHistoryOrigin,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub derivation: Option<LoggedHistoryDerivation>,
}
impl LoggedSessionHistoryMetadata {
pub fn legacy_unknown() -> Self {
Self {
entry_id: LoggedSessionHistoryEntryId::new(),
origin: LoggedSessionHistoryOrigin::LegacyUnknown,
derivation: None,
}
}
}
/// Persisted item and metadata are one value so transforms cannot reorder or
/// truncate one without the other.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct LoggedHistoryEntry {
pub item: LoggedItem,
pub metadata: LoggedSessionHistoryMetadata,
}
/// Typed system-item history record. The typed system event remains available
/// to client replay while its model-visible projection carries the same stable
/// metadata used by live history.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct LoggedSystemHistoryEntry {
pub item: crate::SystemItem,
pub metadata: LoggedSessionHistoryMetadata,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::LoggedRole;
use agen::llm_client::RequestConfig;
#[test]
fn logged_history_entry_round_trip_preserves_id_origin_and_derivation() {
let source_id = LoggedSessionHistoryEntryId::new();
let entry = LoggedHistoryEntry {
item: LoggedItem::Message {
role: LoggedRole::User,
content: vec![crate::LoggedContentPart::Text {
text: "preference".into(),
}],
},
metadata: LoggedSessionHistoryMetadata {
entry_id: LoggedSessionHistoryEntryId::new(),
origin: LoggedSessionHistoryOrigin::HumanInput {
account_id: "account-1".into(),
},
derivation: Some(LoggedHistoryDerivation {
sources: vec![source_id.clone()],
}),
},
};
let encoded = serde_json::to_vec(&entry).unwrap();
let decoded: LoggedHistoryEntry = serde_json::from_slice(&encoded).unwrap();
assert_eq!(decoded, entry);
assert_eq!(
decoded.metadata.derivation.unwrap().sources,
vec![source_id]
);
}
#[test]
fn annotated_segment_start_is_restore_visible_without_projecting_metadata() {
let session_id = uuid::Uuid::now_v7();
let history_entry = legacy_logged_history(LoggedItem::Message {
role: LoggedRole::Assistant,
content: vec![crate::LoggedContentPart::Text {
text: "answer".into(),
}],
});
let state = crate::collect_state(&[crate::LogEntry::AnnotatedSegmentStart {
ts: 1,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: vec![history_entry],
forked_from: None,
compacted_from: None,
}]);
assert_eq!(state.history[0].as_text(), Some("answer"));
}
}
/// Legacy Session Logs did not persist annotations. Decode helpers explicitly
/// create `LegacyUnknown`; they never infer Human/System authority from role or
/// plaintext.
pub fn legacy_logged_history(item: LoggedItem) -> LoggedHistoryEntry {
LoggedHistoryEntry {
item,
metadata: LoggedSessionHistoryMetadata::legacy_unknown(),
}
}
pub fn legacy_segment_history(
session_id: SessionId,
items: impl IntoIterator<Item = LoggedItem>,
) -> Vec<LoggedHistoryEntry> {
let _ = session_id;
items.into_iter().map(legacy_logged_history).collect()
}
+6
View File
@@ -32,6 +32,7 @@
pub mod event_trace; pub mod event_trace;
pub mod fs_store; pub mod fs_store;
pub mod history;
pub mod logged_item; pub mod logged_item;
pub mod segment; pub mod segment;
pub mod segment_log; pub mod segment_log;
@@ -44,6 +45,11 @@ pub use agen::UsageRecord;
pub use agen::llm_client::types::{ContentPart, Item, Role}; pub use agen::llm_client::types::{ContentPart, Item, Role};
pub use event_trace::{TraceEntry, TracePayload}; pub use event_trace::{TraceEntry, TracePayload};
pub use fs_store::FsStore; pub use fs_store::FsStore;
pub use history::{
LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId,
LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedSystemHistoryEntry,
LoggedWorkerSubject, legacy_logged_history, legacy_segment_history,
};
pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged}; pub use logged_item::{LoggedContentPart, LoggedItem, LoggedRole, from_logged, to_logged};
pub use segment::{ pub use segment::{
SegmentStartState, append_entry, append_system_item, classify_history_item, SegmentStartState, append_entry, append_system_item, classify_history_item,
+52 -3
View File
@@ -14,7 +14,7 @@
use agen::{ use agen::{
llm_client::types::{ContentPart, Item, Role}, llm_client::types::{ContentPart, Item, Role},
tool::{Attachment, ImageAttachment}, tool::{Attachment, ImageAttachment, ToolResultDisposition},
}; };
use base64::{Engine as _, engine::general_purpose::STANDARD}; use base64::{Engine as _, engine::general_purpose::STANDARD};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _}; use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error as _};
@@ -61,6 +61,8 @@ pub enum LoggedItem {
content: Option<String>, content: Option<String>,
#[serde(default, skip_serializing_if = "Vec::is_empty")] #[serde(default, skip_serializing_if = "Vec::is_empty")]
attachments: Vec<LoggedAttachment>, attachments: Vec<LoggedAttachment>,
#[serde(default, skip_serializing_if = "ToolResultDisposition::is_success")]
disposition: ToolResultDisposition,
#[serde(default, skip_serializing_if = "is_false")] #[serde(default, skip_serializing_if = "is_false")]
is_error: bool, is_error: bool,
}, },
@@ -128,6 +130,7 @@ impl From<&Item> for LoggedItem {
summary, summary,
content, content,
attachments, attachments,
disposition,
is_error, is_error,
.. ..
} => Self::ToolResult { } => Self::ToolResult {
@@ -135,6 +138,7 @@ impl From<&Item> for LoggedItem {
summary: summary.clone(), summary: summary.clone(),
content: content.clone(), content: content.clone(),
attachments: attachments.iter().map(LoggedAttachment::from).collect(), attachments: attachments.iter().map(LoggedAttachment::from).collect(),
disposition: *disposition,
is_error: *is_error, is_error: *is_error,
}, },
Item::Reasoning { Item::Reasoning {
@@ -184,15 +188,24 @@ impl From<LoggedItem> for Item {
summary, summary,
content, content,
attachments, attachments,
disposition,
is_error, is_error,
} => Item::ToolResult { } => {
let disposition = if is_error && disposition.is_success() {
ToolResultDisposition::Error
} else {
disposition
};
Item::ToolResult {
id: None, id: None,
call_id, call_id,
summary, summary,
content, content,
disposition,
is_error, is_error,
attachments: attachments.into_iter().map(Attachment::from).collect(), attachments: attachments.into_iter().map(Attachment::from).collect(),
}, }
}
LoggedItem::Reasoning { LoggedItem::Reasoning {
text, text,
summary, summary,
@@ -430,6 +443,42 @@ mod tests {
} }
} }
#[test]
fn outcome_unknown_tool_result_round_trips_as_terminal() {
let original = Item::tool_result_item_with_disposition_and_attachments(
"call_unknown",
"outcome unknown",
Some("bounded progress".to_string()),
ToolResultDisposition::OutcomeUnknown,
Vec::new(),
);
let logged: LoggedItem = (&original).into();
let json = serde_json::to_string(&logged).unwrap();
assert!(json.contains(r#""disposition":"outcome_unknown""#));
match Item::from(serde_json::from_str::<LoggedItem>(&json).unwrap()) {
Item::ToolResult {
disposition,
is_error,
..
} => {
assert_eq!(disposition, ToolResultDisposition::OutcomeUnknown);
assert!(is_error);
}
other => panic!("unexpected variant: {other:?}"),
}
}
#[test]
fn legacy_error_tool_result_infers_error_disposition() {
let legacy = r#"{"kind":"tool_result","call_id":"call_old","summary":"failed","content":null,"is_error":true}"#;
match Item::from(serde_json::from_str::<LoggedItem>(legacy).unwrap()) {
Item::ToolResult { disposition, .. } => {
assert_eq!(disposition, ToolResultDisposition::Error)
}
other => panic!("unexpected variant: {other:?}"),
}
}
#[test] #[test]
fn tool_result_persistence_round_trips_binary_attachments() { fn tool_result_persistence_round_trips_binary_attachments() {
let original = Item::tool_result_item_with_attachments( let original = Item::tool_result_item_with_attachments(
+77
View File
@@ -14,6 +14,7 @@ use agen::{EngineResult, UsageRecord};
use protocol::{InvokeKind, Segment}; use protocol::{InvokeKind, Segment};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use crate::history::{LoggedHistoryEntry, LoggedSystemHistoryEntry};
use crate::logged_item::LoggedItem; use crate::logged_item::LoggedItem;
use crate::system_item::SystemItem; use crate::system_item::SystemItem;
@@ -70,6 +71,20 @@ pub enum LogEntry {
compacted_from: Option<SegmentOrigin>, compacted_from: Option<SegmentOrigin>,
}, },
/// Schema-v2 segment seed. Retained entries keep their stable logical
/// identity and origin across fork/compaction/restore.
AnnotatedSegmentStart {
ts: u64,
session_id: crate::SessionId,
system_prompt: Option<String>,
config: RequestConfig,
history: Vec<LoggedHistoryEntry>,
#[serde(default, skip_serializing_if = "Option::is_none")]
forked_from: Option<SegmentOrigin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
compacted_from: Option<SegmentOrigin>,
},
/// IDLE → active marker. Records the start of a new self-driving /// IDLE → active marker. Records the start of a new self-driving
/// cycle (Invoke range). The range extends implicitly until the /// cycle (Invoke range). The range extends implicitly until the
/// next `Invoke` entry; this entry carries the trigger only — the /// next `Invoke` entry; this entry carries the trigger only — the
@@ -105,14 +120,37 @@ pub enum LogEntry {
extensions: Vec<SessionExtension>, extensions: Vec<SessionExtension>,
}, },
/// Schema-v2 user submission with its exact model-visible entries. Typed
/// Flow instructions and caller-attributed input remain separate entries.
AnnotatedUserInput {
ts: u64,
segments: Vec<Segment>,
#[serde(default, skip_serializing_if = "Vec::is_empty")]
extensions: Vec<SessionExtension>,
history: Vec<LoggedHistoryEntry>,
},
/// Schema-v2 model output and metadata committed as one journal record.
AnnotatedAssistantItem { ts: u64, entry: LoggedHistoryEntry },
/// One assistant-side item appended to history — assistant message, /// One assistant-side item appended to history — assistant message,
/// reasoning, or tool call. Singular: one entry per history item so /// reasoning, or tool call. Singular: one entry per history item so
/// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1. /// the wire-side `Event::*` lane and on-disk LogEntry stay 1:1.
AssistantItem { ts: u64, item: LoggedItem }, AssistantItem { ts: u64, item: LoggedItem },
/// Schema-v2 tool output and metadata committed as one journal record.
AnnotatedToolResult { ts: u64, entry: LoggedHistoryEntry },
/// One tool-execution result appended to history. /// One tool-execution result appended to history.
ToolResult { ts: u64, item: LoggedItem }, ToolResult { ts: u64, item: LoggedItem },
/// Schema-v2 typed system event and model-visible metadata committed
/// together.
AnnotatedSystemItem {
ts: u64,
entry: LoggedSystemHistoryEntry,
},
/// One typed agent-injected system item: notification, child-Worker /// One typed agent-injected system item: notification, child-Worker
/// lifecycle event, `@<path>` / `/<slug>` resolution payload. Each /// lifecycle event, `@<path>` / `/<slug>` resolution payload. Each
/// `SystemItem` carries kind metadata that the LLM /// `SystemItem` carries kind metadata that the LLM
@@ -278,6 +316,22 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
state.config = config.clone(); state.config = config.clone();
state.history = history.iter().cloned().map(Item::from).collect(); state.history = history.iter().cloned().map(Item::from).collect();
} }
LogEntry::AnnotatedSegmentStart {
session_id,
system_prompt,
config,
history,
..
} => {
state.session_id = Some(*session_id);
state.system_prompt = system_prompt.clone();
state.config = config.clone();
state.history = history
.iter()
.cloned()
.map(|entry| Item::from(entry.item))
.collect();
}
LogEntry::Invoke { .. } => { LogEntry::Invoke { .. } => {
// A terminal run record below clears or refines this. If the // A terminal run record below clears or refines this. If the
// log ends first, restore must treat the turn as interrupted. // log ends first, restore must treat the turn as interrupted.
@@ -298,6 +352,29 @@ pub fn collect_state(entries: &[LogEntry]) -> RestoredState {
.map(|extension| (extension.domain.clone(), extension.payload.clone())), .map(|extension| (extension.domain.clone(), extension.payload.clone())),
); );
} }
LogEntry::AnnotatedUserInput {
segments,
extensions,
history,
..
} => {
state
.history
.extend(history.iter().cloned().map(|entry| Item::from(entry.item)));
state.user_segments.push(segments.clone());
state.extensions.extend(
extensions
.iter()
.map(|extension| (extension.domain.clone(), extension.payload.clone())),
);
}
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => {
state.history.push(Item::from(entry.item.clone()));
}
LogEntry::AnnotatedSystemItem { entry, .. } => {
state.history.push(entry.item.to_history_item());
}
LogEntry::AssistantItem { item, .. } => { LogEntry::AssistantItem { item, .. } => {
state.history.push(Item::from(item.clone())); state.history.push(Item::from(item.clone()));
} }
@@ -20,7 +20,8 @@ use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use std::time::SystemTime; use std::time::SystemTime;
const SESSION_SCHEMA_VERSION: u32 = 1; const SESSION_SCHEMA_VERSION: u32 = 2;
const LEGACY_SESSION_SCHEMA_VERSION: u32 = 1;
const SESSION_FILE: &str = "session.json"; const SESSION_FILE: &str = "session.json";
const SEGMENTS_DIR: &str = "segments"; const SEGMENTS_DIR: &str = "segments";
@@ -44,16 +45,23 @@ impl WorkerSessionStore {
fs::create_dir_all(root.join(SEGMENTS_DIR))?; fs::create_dir_all(root.join(SEGMENTS_DIR))?;
let session_id = match fs::read(root.join(SESSION_FILE)) { let session_id = match fs::read(root.join(SESSION_FILE)) {
Ok(bytes) => { Ok(bytes) => {
let manifest: SessionManifest = serde_json::from_slice(&bytes)?; let mut manifest: SessionManifest = serde_json::from_slice(&bytes)?;
if manifest.schema_version != SESSION_SCHEMA_VERSION { match manifest.schema_version {
SESSION_SCHEMA_VERSION => {}
LEGACY_SESSION_SCHEMA_VERSION => {
validate_legacy_segment_logs(&root)?;
manifest.schema_version = SESSION_SCHEMA_VERSION;
atomic_write_json(&root.join(SESSION_FILE), &manifest)?;
}
version => {
return Err(StoreError::Corrupt { return Err(StoreError::Corrupt {
line: 0, line: 0,
message: format!( message: format!(
"unsupported Worker Session schema version {}, expected {}", "unsupported Worker Session schema version {version}, expected {SESSION_SCHEMA_VERSION}"
manifest.schema_version, SESSION_SCHEMA_VERSION
), ),
}); });
} }
}
Some(manifest.session_id) Some(manifest.session_id)
} }
Err(error) if error.kind() == std::io::ErrorKind::NotFound => None, Err(error) if error.kind() == std::io::ErrorKind::NotFound => None,
@@ -278,6 +286,37 @@ impl Store for WorkerSessionStore {
} }
} }
fn validate_legacy_segment_logs(root: &Path) -> Result<(), StoreError> {
let segments = root.join(SEGMENTS_DIR);
if !segments.exists() {
return Ok(());
}
for entry in fs::read_dir(&segments)? {
let entry = entry?;
let path = entry.path();
let Some(name) = path.file_name().and_then(|name| name.to_str()) else {
continue;
};
if !name.ends_with(".jsonl") || name.ends_with(".trace.jsonl") {
continue;
}
let contents = fs::read_to_string(&path)?;
for (line_index, line) in contents.lines().enumerate() {
if line.trim().is_empty() {
continue;
}
serde_json::from_str::<LogEntry>(line).map_err(|error| StoreError::Corrupt {
line: line_index + 1,
message: format!(
"cannot migrate legacy Worker Session log {}: {error}",
path.display()
),
})?;
}
}
Ok(())
}
fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> Result<(), StoreError> { fn atomic_write_json<T: Serialize>(path: &Path, value: &T) -> Result<(), StoreError> {
let mut bytes = serde_json::to_vec_pretty(value)?; let mut bytes = serde_json::to_vec_pretty(value)?;
bytes.push(b'\n'); bytes.push(b'\n');
@@ -405,6 +444,54 @@ mod tests {
assert_eq!(store.list_sessions().unwrap(), vec![session_id]); assert_eq!(store.list_sessions().unwrap(), vec![session_id]);
} }
#[test]
fn schema_v1_logs_are_validated_and_promoted_to_v2() {
let root = tempfile::tempdir().unwrap();
let session_id = new_session_id();
let segment_id = new_segment_id();
WorkerSessionStore::new(root.path())
.unwrap()
.create_segment(session_id, segment_id, &[])
.unwrap();
let manifest_path = root.path().join(SESSION_FILE);
let mut manifest: SessionManifest =
serde_json::from_slice(&fs::read(&manifest_path).unwrap()).unwrap();
manifest.schema_version = LEGACY_SESSION_SCHEMA_VERSION;
atomic_write_json(&manifest_path, &manifest).unwrap();
let reopened = WorkerSessionStore::new(root.path()).unwrap();
assert_eq!(reopened.session_id().unwrap(), Some(session_id));
let migrated: SessionManifest =
serde_json::from_slice(&fs::read(&manifest_path).unwrap()).unwrap();
assert_eq!(migrated.schema_version, SESSION_SCHEMA_VERSION);
}
#[test]
fn schema_v1_migration_rejects_corrupt_log_before_manifest_update() {
let root = tempfile::tempdir().unwrap();
let session_id = new_session_id();
let manifest = SessionManifest {
schema_version: LEGACY_SESSION_SCHEMA_VERSION,
session_id,
};
atomic_write_json(&root.path().join(SESSION_FILE), &manifest).unwrap();
fs::create_dir_all(root.path().join(SEGMENTS_DIR)).unwrap();
fs::write(
root.path().join(SEGMENTS_DIR).join("broken.jsonl"),
"{not-json}\n",
)
.unwrap();
let error = match WorkerSessionStore::new(root.path()) {
Ok(_) => panic!("corrupt legacy Session log must reject migration"),
Err(error) => error,
};
assert!(matches!(error, StoreError::Corrupt { .. }));
let persisted: SessionManifest =
serde_json::from_slice(&fs::read(root.path().join(SESSION_FILE)).unwrap()).unwrap();
assert_eq!(persisted.schema_version, LEGACY_SESSION_SCHEMA_VERSION);
}
#[test] #[test]
fn reopen_preserves_session_and_segment_ids() { fn reopen_preserves_session_and_segment_ids() {
let root = tempfile::tempdir().unwrap(); let root = tempfile::tempdir().unwrap();
+89 -37
View File
@@ -1,12 +1,13 @@
mod common; mod common;
use std::ops::{Deref, DerefMut};
use std::sync::Arc; use std::sync::Arc;
use agen::Engine;
use agen::interceptor::{Interceptor, TurnEndAction}; use agen::interceptor::{Interceptor, TurnEndAction};
use agen::llm_client::event::{Event, ResponseStatus, StatusEvent}; use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
use agen::llm_client::types::{Item, RequestConfig}; use agen::llm_client::types::{Item, RequestConfig};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use agen::{Engine, History};
use async_trait::async_trait; use async_trait::async_trait;
use common::MockLlmClient; use common::MockLlmClient;
use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state}; use session_store::{FsStore, LogEntry, SegmentStartState, Store, collect_state};
@@ -94,15 +95,47 @@ fn make_store() -> (tempfile::TempDir, FsStore) {
(dir, store) (dir, store)
} }
struct TestWorker {
engine: Engine<MockLlmClient>,
history: History,
}
impl TestWorker {
fn new(engine: Engine<MockLlmClient>) -> Self {
Self {
engine,
history: History::new(),
}
}
fn history(&self) -> Vec<Item> {
self.history.items_cloned()
}
}
impl Deref for TestWorker {
type Target = Engine<MockLlmClient>;
fn deref(&self) -> &Self::Target {
&self.engine
}
}
impl DerefMut for TestWorker {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.engine
}
}
/// Run a worker turn and persist via session-store functions. /// Run a worker turn and persist via session-store functions.
/// Takes ownership of the worker (needed for lock/unlock) and returns it. /// Takes ownership of the worker (needed for lock/unlock) and returns it.
async fn run_and_persist( async fn run_and_persist(
worker: Engine<MockLlmClient>, mut worker: TestWorker,
store: &FsStore, store: &FsStore,
session_id: session_store::SessionId, session_id: session_store::SessionId,
segment_id: session_store::SegmentId, segment_id: session_store::SegmentId,
input: &str, input: &str,
) -> (Engine<MockLlmClient>, agen::EngineResult) { ) -> (TestWorker, agen::EngineRunExit) {
// Mirror Worker's run-entry contract: log the user input as segments // Mirror Worker's run-entry contract: log the user input as segments
// before the worker pushes its flattened user_message; save_delta // before the worker pushes its flattened user_message; save_delta
// skips the resulting user_message item to avoid double-write. // skips the resulting user_message item to avoid double-write.
@@ -114,42 +147,61 @@ async fn run_and_persist(
) )
.unwrap(); .unwrap();
let history_before = worker.history().len(); let history_before = worker.history.len();
let mut locked = worker.lock(); let mut locked = worker.engine.lock(&worker.history);
let result = locked.run(input).await; let result = locked.run(&mut worker.history, input).await;
let worker = locked.unlock(); worker.engine = locked.unlock();
let new_items = &worker.history()[history_before..]; let projected = worker.history();
let new_items = &projected[history_before..];
session_store::save_delta(store, session_id, segment_id, new_items).unwrap(); session_store::save_delta(store, session_id, segment_id, new_items).unwrap();
session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap(); session_store::save_turn_end(store, session_id, segment_id, worker.turn_count()).unwrap();
match &result { match &result {
Ok(r) => { agen::EngineRunExit::Finished
| agen::EngineRunExit::Paused
| agen::EngineRunExit::Yielded => {
let (legacy_result, interrupted) = match &result {
agen::EngineRunExit::Finished => (agen::EngineResult::Finished, false),
agen::EngineRunExit::Paused => (agen::EngineResult::Paused, true),
agen::EngineRunExit::Yielded => (agen::EngineResult::Yielded, true),
agen::EngineRunExit::Interrupted(_) => unreachable!(),
};
session_store::save_run_completed( session_store::save_run_completed(
store, store,
session_id, session_id,
segment_id, segment_id,
r.clone(), legacy_result,
worker.last_run_interrupted(), interrupted,
worker.active_run_turn_count(), worker.active_run_turn_count(),
) )
.unwrap(); .unwrap();
} }
Err(e) => { agen::EngineRunExit::Interrupted(agen::StopReason::LimitReached) => {
session_store::save_run_completed(
store,
session_id,
segment_id,
agen::EngineResult::LimitReached,
false,
worker.active_run_turn_count(),
)
.unwrap();
}
agen::EngineRunExit::Interrupted(reason) => {
session_store::save_run_errored( session_store::save_run_errored(
store, store,
session_id, session_id,
segment_id, segment_id,
e.to_string(), format!("{reason:?}"),
worker.last_run_interrupted(), true,
) )
.unwrap(); .unwrap();
} }
} }
let r = result.unwrap(); (worker, result)
(worker, r)
} }
// ============================================================================= // =============================================================================
@@ -160,14 +212,14 @@ async fn run_and_persist(
async fn session_run_logs_entries() { async fn session_run_logs_entries() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); let client = MockLlmClient::new(simple_text_events());
let worker = Engine::new(client); let worker = TestWorker::new(Engine::new(client));
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -204,7 +256,7 @@ async fn session_run_logs_entries() {
async fn session_restore_round_trip() { async fn session_restore_round_trip() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); let client = MockLlmClient::new(simple_text_events());
let mut worker = Engine::new(client); let mut worker = TestWorker::new(Engine::new(client));
worker.set_system_prompt("You are helpful."); worker.set_system_prompt("You are helpful.");
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -212,7 +264,7 @@ async fn session_restore_round_trip() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -243,7 +295,7 @@ async fn session_restore_round_trip() {
async fn session_run_with_tool_call() { async fn session_run_with_tool_call() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::with_responses(tool_call_events()); let client = MockLlmClient::with_responses(tool_call_events());
let mut worker = Engine::new(client); let mut worker = TestWorker::new(Engine::new(client));
worker.register_tool(weather_tool_definition()); worker.register_tool(weather_tool_definition());
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -251,7 +303,7 @@ async fn session_run_with_tool_call() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -277,7 +329,7 @@ async fn session_resume_after_pause() {
// First run: tool call with pause policy → Paused // First run: tool call with pause policy → Paused
let client = MockLlmClient::with_responses(tool_call_events()); let client = MockLlmClient::with_responses(tool_call_events());
let mut worker = Engine::new(client); let mut worker = TestWorker::new(Engine::new(client));
worker.register_tool(weather_tool_definition()); worker.register_tool(weather_tool_definition());
worker.set_interceptor(PausePolicy); worker.set_interceptor(PausePolicy);
@@ -286,13 +338,13 @@ async fn session_resume_after_pause() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
let (_worker, result) = run_and_persist(worker, &store, sid, segid, "Weather?").await; let (_worker, result) = run_and_persist(worker, &store, sid, segid, "Weather?").await;
assert!(matches!(result, agen::EngineResult::Paused)); assert!(matches!(result, agen::EngineRunExit::Paused));
// Check RunCompleted is Paused // Check RunCompleted is Paused
let entries = store.read_all(sid, segid).unwrap(); let entries = store.read_all(sid, segid).unwrap();
@@ -317,7 +369,7 @@ async fn session_resume_after_pause() {
async fn session_fork_creates_new_session() { async fn session_fork_creates_new_session() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); let client = MockLlmClient::new(simple_text_events());
let mut worker = Engine::new(client); let mut worker = TestWorker::new(Engine::new(client));
worker.set_system_prompt("System prompt"); worker.set_system_prompt("System prompt");
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
@@ -325,7 +377,7 @@ async fn session_fork_creates_new_session() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -338,7 +390,7 @@ async fn session_fork_creates_new_session() {
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -359,14 +411,14 @@ async fn session_fork_creates_new_session() {
async fn session_fork_at_truncates_within_session() { async fn session_fork_at_truncates_within_session() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); let client = MockLlmClient::new(simple_text_events());
let worker = Engine::new(client); let worker = TestWorker::new(Engine::new(client));
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -404,14 +456,14 @@ async fn session_fork_at_truncates_within_session() {
async fn session_config_changed_logged() { async fn session_config_changed_logged() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(vec![]); let client = MockLlmClient::new(vec![]);
let mut worker = Engine::new(client); let mut worker = TestWorker::new(Engine::new(client));
let (sid, segid) = session_store::create_segment( let (sid, segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -437,14 +489,14 @@ async fn session_auto_forks_on_conflict() {
// Create a segment // Create a segment
let client_a = MockLlmClient::new(simple_text_events()); let client_a = MockLlmClient::new(simple_text_events());
let worker_a = Engine::new(client_a); let worker_a = TestWorker::new(Engine::new(client_a));
let (sid, original_segid) = session_store::create_segment( let (sid, original_segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker_a.get_system_prompt(), system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(), config: worker_a.request_config(),
history: worker_a.history(), history: &worker_a.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -470,7 +522,7 @@ async fn session_auto_forks_on_conflict() {
SegmentStartState { SegmentStartState {
system_prompt: worker_a.get_system_prompt(), system_prompt: worker_a.get_system_prompt(),
config: worker_a.request_config(), config: worker_a.request_config(),
history: worker_a.history(), history: &worker_a.history(),
}, },
) )
.unwrap(); .unwrap();
@@ -522,14 +574,14 @@ async fn session_auto_forks_on_conflict() {
async fn nested_past_fork_leaves_ancestors_immutable() { async fn nested_past_fork_leaves_ancestors_immutable() {
let (_dir, store) = make_store(); let (_dir, store) = make_store();
let client = MockLlmClient::new(simple_text_events()); let client = MockLlmClient::new(simple_text_events());
let worker = Engine::new(client); let worker = TestWorker::new(Engine::new(client));
let (sid, root_segid) = session_store::create_segment( let (sid, root_segid) = session_store::create_segment(
&store, &store,
SegmentStartState { SegmentStartState {
system_prompt: worker.get_system_prompt(), system_prompt: worker.get_system_prompt(),
config: worker.request_config(), config: worker.request_config(),
history: worker.history(), history: &worker.history(),
}, },
) )
.unwrap(); .unwrap();
+139 -5
View File
@@ -1223,10 +1223,17 @@ impl Tool for TicketQueueTool {
) -> Result<ToolOutput, ToolError> { ) -> Result<ToolOutput, ToolError> {
let params: TicketQueueParams = parse_input("TicketQueue", input_json)?; let params: TicketQueueParams = parse_input("TicketQueue", input_json)?;
let queued_by = default_author(); let queued_by = default_author();
let outcome = self let mut outcome = self
.backend .backend
.queue_ready(TicketIdOrSlug::Query(params.ticket.clone()), &queued_by) .queue_ready(TicketIdOrSlug::Query(params.ticket.clone()), &queued_by)
.map_err(|error| backend_error("TicketQueue", error))?; .map_err(|error| backend_error("TicketQueue", error))?;
outcome.requested_ticket =
model_ticket_reference(&self.backend, &outcome.requested_ticket, "TicketQueue")?;
outcome.queued_tickets = outcome
.queued_tickets
.into_iter()
.map(|ticket| model_ticket_reference(&self.backend, &ticket, "TicketQueue"))
.collect::<Result<Vec<_>, _>>()?;
Ok(json_output( Ok(json_output(
format!( format!(
"Queued {} ticket(s) for Orchestrator", "Queued {} ticket(s) for Orchestrator",
@@ -1264,15 +1271,17 @@ impl Tool for TicketWorkflowStateTool {
self.backend self.backend
.set_workflow_state(TicketIdOrSlug::Query(params.ticket.clone()), change) .set_workflow_state(TicketIdOrSlug::Query(params.ticket.clone()), change)
.map_err(|error| backend_error("TicketWorkflowState", error))?; .map_err(|error| backend_error("TicketWorkflowState", error))?;
let ticket_ref =
model_ticket_reference(&self.backend, &params.ticket, "TicketWorkflowState")?;
Ok(json_output( Ok(json_output(
format!( format!(
"Transitioned ticket {} state {} -> {}", "Transitioned ticket {} state {} -> {}",
params.ticket, ticket_ref,
from.as_str(), from.as_str(),
to.as_str() to.as_str()
), ),
json!({ json!({
"ticket": params.ticket, "ticket": ticket_ref,
"from": from.as_str(), "from": from.as_str(),
"to": to.as_str(), "to": to.as_str(),
"state": to.as_str(), "state": to.as_str(),
@@ -1296,9 +1305,10 @@ impl Tool for TicketCloseTool {
MarkdownText::new(params.resolution), MarkdownText::new(params.resolution),
) )
.map_err(|error| backend_error("TicketClose", error))?; .map_err(|error| backend_error("TicketClose", error))?;
let ticket_ref = model_ticket_reference(&self.backend, &params.ticket, "TicketClose")?;
Ok(json_output( Ok(json_output(
format!("Closed ticket {}", params.ticket), format!("Closed ticket {ticket_ref}"),
json!({ "ticket": params.ticket, "state": "closed", "ok": true }), json!({ "ticket": ticket_ref, "state": "closed", "ok": true }),
)) ))
} }
} }
@@ -1525,6 +1535,29 @@ impl Tool for TicketDependencyCheckTool {
} }
} }
fn model_ticket_reference(
backend: &TicketToolBackend,
reference: &str,
tool_name: &str,
) -> Result<String, ToolError> {
let ticket = backend
.show(TicketIdOrSlug::Id(reference.to_string()))
.map_err(|error| backend_error(tool_name, error))?;
match ticket.meta.resource_key {
Some(resource_key) if is_canonical_ticket_resource_key(&resource_key) => Ok(resource_key),
Some(_) => Err(ToolError::ExecutionFailed(format!(
"{tool_name} failed: required Ticket human key is unavailable"
))),
None => Ok(ticket.meta.id),
}
}
fn is_canonical_ticket_resource_key(resource_key: &str) -> bool {
resource_key.strip_prefix("T-").is_some_and(|sequence| {
!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit())
})
}
fn parse_input<T: for<'de> Deserialize<'de>>(tool: &str, input_json: &str) -> Result<T, ToolError> { fn parse_input<T: for<'de> Deserialize<'de>>(tool: &str, input_json: &str) -> Result<T, ToolError> {
serde_json::from_str(input_json) serde_json::from_str(input_json)
.map_err(|error| ToolError::InvalidArgument(format!("invalid {tool} input: {error}"))) .map_err(|error| ToolError::InvalidArgument(format!("invalid {tool} input: {error}")))
@@ -1922,6 +1955,12 @@ mod tests {
.with_target_authority(Arc::new(TestTargetAuthority)) .with_target_authority(Arc::new(TestTargetAuthority))
} }
fn sqlite_backend(temp: &TempDir) -> crate::SqliteTicketBackend {
crate::SqliteTicketBackend::open(temp.path().join("tickets.db"), "workspace")
.unwrap()
.with_target_authority(Arc::new(TestTargetAuthority))
}
fn tool(definition: ToolDefinition) -> Arc<dyn Tool> { fn tool(definition: ToolDefinition) -> Arc<dyn Tool> {
let (_, tool) = definition(); let (_, tool) = definition();
tool tool
@@ -2549,6 +2588,101 @@ mod tests {
); );
} }
#[tokio::test]
async fn queue_workflow_and_close_project_internal_inputs_to_ticket_keys() {
let temp = TempDir::new().unwrap();
let inner = sqlite_backend(&temp);
let mut dependency_input = NewTicket::new("Dependency");
dependency_input.repository_id = Some("main".to_string());
let dependency = inner.create(dependency_input).unwrap();
let mut target_input = NewTicket::new("Target");
target_input.repository_id = Some("main".to_string());
let target = inner.create(target_input).unwrap();
inner
.add_ticket_relation(
TicketIdOrSlug::Id(target.id.clone()),
NewTicketRelation {
kind: TicketRelationKind::DependsOn,
target: dependency.id.clone(),
note: None,
author: None,
},
)
.unwrap();
for id in [&dependency.id, &target.id] {
inner
.mark_ready(
TicketIdOrSlug::Id(id.clone()),
TicketMarkReady {
operation_key: format!("ready-{id}"),
reason: None,
author: None,
intake_summary: None,
},
)
.unwrap();
}
let target_key = target.resource_key.clone().unwrap();
let dependency_key = dependency.resource_key.clone().unwrap();
let backend = inner;
let queue = tool_by_name(TicketToolBackend::new(backend.clone()), "TicketQueue");
let workflow = tool_by_name(
TicketToolBackend::new(backend.clone()),
"TicketWorkflowState",
);
let close = tool_by_name(TicketToolBackend::new(backend), "TicketClose");
let queued = queue
.execute(
&json!({"ticket": target.id.clone()}).to_string(),
Default::default(),
)
.await
.unwrap();
assert!(queued.summary.contains("2 ticket(s)"));
let queued_content = queued.content.unwrap();
assert!(queued_content.contains(&target_key));
assert!(queued_content.contains(&dependency_key));
assert!(!queued_content.contains(&target.id));
assert!(!queued_content.contains(&dependency.id));
for (from, to) in [("queued", "inprogress"), ("inprogress", "done")] {
let transitioned = workflow
.execute(
&json!({
"ticket": target.id.clone(),
"from": from,
"to": to,
"reason": "test_transition",
"body": "transitioned",
"author": "tester"
})
.to_string(),
Default::default(),
)
.await
.unwrap();
assert!(transitioned.summary.contains(&target_key));
assert!(!transitioned.summary.contains(&target.id));
let content = transitioned.content.unwrap();
assert!(content.contains(&target_key));
assert!(!content.contains(&target.id));
}
let closed = close
.execute(
&json!({"ticket": target.id.clone(), "resolution": "Done"}).to_string(),
Default::default(),
)
.await
.unwrap();
assert!(closed.summary.contains(&target_key));
assert!(!closed.summary.contains(&target.id));
let content = closed.content.unwrap();
assert!(content.contains(&target_key));
assert!(!content.contains(&target.id));
}
#[tokio::test] #[tokio::test]
async fn ticket_workflow_tools_mark_ready_and_transition_state() { async fn ticket_workflow_tools_mark_ready_and_transition_state() {
let temp = TempDir::new().unwrap(); let temp = TempDir::new().unwrap();
+156 -11
View File
@@ -1,5 +1,6 @@
use std::collections::{HashMap, HashSet};
use std::path::PathBuf; use std::path::PathBuf;
use std::sync::Arc; use std::sync::{Arc, Mutex};
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait; use async_trait::async_trait;
@@ -20,21 +21,65 @@ struct BashParams {
pub(crate) struct BashTool { pub(crate) struct BashTool {
session: WorkdirSessionHandle, session: WorkdirSessionHandle,
state: Arc<Mutex<BashExecutionState>>,
}
#[derive(Clone)]
struct ActiveCommand {
call_id: String,
execution_nonce: u64,
handle: CommandHandle,
}
#[derive(Default)]
struct BashExecutionState {
active: HashMap<String, ActiveCommand>,
cancellation_requested: HashSet<String>,
legacy_cancellation_requested: HashSet<String>,
next_execution_nonce: u64,
} }
struct CommandGuard { struct CommandGuard {
session: WorkdirSessionHandle, session: WorkdirSessionHandle,
state: Arc<Mutex<BashExecutionState>>,
execution_id: String,
execution_nonce: u64,
handle: Option<CommandHandle>, handle: Option<CommandHandle>,
} }
impl Drop for CommandGuard { impl Drop for CommandGuard {
fn drop(&mut self) { fn drop(&mut self) {
if let Some(handle) = self.handle.take() { let Some(handle) = self.handle.take() else {
return;
};
let workdir = self.session.clone(); let workdir = self.session.clone();
let state = Arc::clone(&self.state);
let execution_id = self.execution_id.clone();
let execution_nonce = self.execution_nonce;
// A dropped provider future is not terminal confirmation. Keep the live
// execution registered until cleanup has both requested cancellation and
// observed terminal command output, so cancellation/session teardown
// cannot race with an apparently empty registry.
tokio::spawn(async move { tokio::spawn(async move {
let _ = workdir.cancel_command(handle).await; let _ = workdir.cancel_command(handle.clone()).await;
}); let _ = workdir
.command_output(CommandOutputRequest {
handle,
cursor: 0,
limit: INLINE_BYTE_BUDGET,
wait: true,
})
.await;
let mut state = state.lock().unwrap();
if state
.active
.get(&execution_id)
.is_some_and(|active| active.execution_nonce == execution_nonce)
{
state.active.remove(&execution_id);
state.cancellation_requested.remove(&execution_id);
} }
});
} }
} }
@@ -52,20 +97,50 @@ impl Tool for BashTool {
.unwrap_or(DEFAULT_TIMEOUT_SECS) .unwrap_or(DEFAULT_TIMEOUT_SECS)
.clamp(1, MAX_TIMEOUT_SECS); .clamp(1, MAX_TIMEOUT_SECS);
let cmd_summary = truncate_for_summary(&params.command); let cmd_summary = truncate_for_summary(&params.command);
let execution_id = ctx.execution_id();
let call_id = ctx.call_id;
let execution_nonce = {
let mut state = self.state.lock().unwrap();
state.next_execution_nonce = state.next_execution_nonce.wrapping_add(1);
state.next_execution_nonce
};
let mut guard = CommandGuard {
session: self.session.clone(),
state: self.state.clone(),
execution_id: execution_id.clone(),
execution_nonce,
handle: None,
};
let handle = self let handle = self
.session .session
.start_command(CommandRequest { .start_command(CommandRequest {
command: params.command, command: params.command,
timeout_secs, timeout_secs,
output_limit: INLINE_BYTE_BUDGET, output_limit: INLINE_BYTE_BUDGET,
tool_call_id: Some(ctx.call_id), tool_call_id: Some(call_id.clone()),
}) })
.await .await
.map_err(crate::ToolsError::from)?; .map_err(crate::ToolsError::from)?;
let mut guard = CommandGuard { let cancel_after_start = {
session: self.session.clone(), let mut state = self.state.lock().unwrap();
handle: Some(handle.clone()), state.active.insert(
execution_id.clone(),
ActiveCommand {
call_id: call_id.clone(),
execution_nonce,
handle: handle.clone(),
},
);
state.cancellation_requested.contains(&execution_id)
|| state.legacy_cancellation_requested.contains(&call_id)
}; };
guard.handle = Some(handle.clone());
if cancel_after_start {
self.session
.cancel_command(handle.clone())
.await
.map_err(crate::ToolsError::from)?;
}
let output = self let output = self
.session .session
.command_output(CommandOutputRequest { .command_output(CommandOutputRequest {
@@ -76,9 +151,27 @@ impl Tool for BashTool {
}) })
.await .await
.map_err(crate::ToolsError::from)?; .map_err(crate::ToolsError::from)?;
let cancellation_requested = {
let mut state = self.state.lock().unwrap();
let owns_registration = state
.active
.get(&execution_id)
.is_some_and(|active| active.execution_nonce == execution_nonce);
let exact = if owns_registration {
state.active.remove(&execution_id);
state.cancellation_requested.remove(&execution_id)
} else {
false
};
let legacy = state.legacy_cancellation_requested.remove(&call_id);
exact || legacy
};
guard.handle = None; guard.handle = None;
let summary = if output.timed_out { let timed_out = output.timed_out;
let summary = if cancellation_requested {
format!("$ {cmd_summary} (cancelled)")
} else if output.timed_out {
format!("$ {cmd_summary} (timed out after {timeout_secs}s)") format!("$ {cmd_summary} (timed out after {timeout_secs}s)")
} else { } else {
match output.exit_code { match output.exit_code {
@@ -97,11 +190,62 @@ impl Tool for BashTool {
} else { } else {
Some(output.content) Some(output.content)
}; };
Ok(ToolOutput { let output = ToolOutput {
summary, summary,
content, content,
attachments: Vec::new(), attachments: Vec::new(),
}) };
if cancellation_requested {
Err(ToolError::Cancelled(output))
} else if timed_out {
Err(ToolError::Interrupted(output))
} else {
Ok(output)
}
}
async fn cancel(&self, call_id: &str) -> Result<(), ToolError> {
let handles = {
let mut state = self.state.lock().unwrap();
state
.legacy_cancellation_requested
.insert(call_id.to_string());
state
.active
.values()
.filter(|active| active.call_id == call_id)
.map(|active| active.handle.clone())
.collect::<Vec<_>>()
};
for handle in handles {
self.session
.cancel_command(handle)
.await
.map_err(crate::ToolsError::from)?;
}
Ok(())
}
async fn cancel_execution(
&self,
ctx: &agen::tool::ToolExecutionContext,
) -> Result<(), ToolError> {
let execution_id = ctx.execution_id();
let handle = {
let mut state = self.state.lock().unwrap();
state.cancellation_requested.insert(execution_id.clone());
state
.active
.get(&execution_id)
.map(|active| active.handle.clone())
};
if let Some(handle) = handle {
self.session
.cancel_command(handle)
.await
.map_err(crate::ToolsError::from)?;
}
Ok(())
} }
} }
@@ -123,6 +267,7 @@ pub fn bash_tool(session: WorkdirSessionHandle, _output_dir: PathBuf) -> ToolDef
.input_schema(serde_json::to_value(schema).expect("Bash schema serialization")); .input_schema(serde_json::to_value(schema).expect("Bash schema serialization"));
let tool: Arc<dyn Tool> = Arc::new(BashTool { let tool: Arc<dyn Tool> = Arc::new(BashTool {
session: session.clone(), session: session.clone(),
state: Arc::new(Mutex::new(BashExecutionState::default())),
}); });
(meta, tool) (meta, tool)
}) })
+83 -1
View File
@@ -7,7 +7,10 @@
use std::path::Path; use std::path::Path;
use std::sync::Arc; use std::sync::Arc;
use agen::tool::{Tool, ToolDefinition, ToolMeta}; use agen::tool::{
Tool, ToolDefinition, ToolError, ToolExecutionContext, ToolExecutionHandle,
ToolExecutionTerminal, ToolMeta,
};
use manifest::{Permission, Scope, ScopeConfig, ScopeRule}; use manifest::{Permission, Scope, ScopeConfig, ScopeRule};
use serde_json::json; use serde_json::json;
use tempfile::TempDir; use tempfile::TempDir;
@@ -401,5 +404,84 @@ async fn bash_provider_output_does_not_expose_internal_paths() {
assert_eq!(std::fs::read_dir(spill.path()).unwrap().count(), 0); assert_eq!(std::fs::read_dir(spill.path()).unwrap().count(), 0);
} }
#[tokio::test]
async fn bash_cancellation_returns_bounded_progress_as_terminal_output() {
let (dir, _spill, reg) = setup();
let marker = dir.path().join("must-not-run-after-cancel");
let command = format!(
"printf 'before\\n'; printf 'err-before\\n' >&2; sleep 1; touch {}; printf 'after\\n'",
marker.display()
);
let input = serde_json::to_string(&json!({ "command": command })).unwrap();
let context = ToolExecutionContext::new("call-heavy", "attempt-heavy", 0);
let bash = reg.get("Bash");
let executing = bash.clone();
let execution_context = context.clone();
let execution = tokio::spawn(async move { executing.execute(&input, execution_context).await });
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
bash.cancel_execution(&context)
.await
.expect("signal exact execution cancellation");
let error = tokio::time::timeout(std::time::Duration::from_secs(2), execution)
.await
.expect("cancelled Bash should terminate inside the Engine grace budget")
.expect("Bash task join");
let ToolError::Cancelled(output) = error.expect_err("cancelled command is non-success") else {
panic!("expected typed cancellation result");
};
let content = output.content.expect("bounded progress output");
assert!(
content.contains("before"),
"missing pre-cancel stdout: {content}"
);
assert!(
content.contains("err-before"),
"missing pre-cancel stderr: {content}"
);
assert!(
!content.contains("after"),
"post-cancel output leaked: {content}"
);
assert!(content.len() <= 16 * 1024, "output must remain bounded");
tokio::time::sleep(std::time::Duration::from_millis(1_100)).await;
assert!(
!marker.exists(),
"the cancelled command continued executing after terminal confirmation"
);
}
#[tokio::test]
async fn bash_force_close_cleanup_stops_command_and_keeps_session_reusable() {
let (dir, _spill, reg) = setup();
let marker = dir.path().join("must-not-survive-force-close");
let command = format!("sleep 1; touch {}", marker.display());
let input = serde_json::to_string(&json!({ "command": command })).unwrap();
let bash = reg.get("Bash");
let context = ToolExecutionContext::new("call-force", "attempt-force", 0);
let (handle, terminal) = ToolExecutionHandle::start(bash.clone(), input, context);
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
handle.force_close();
assert!(matches!(
terminal.await,
ToolExecutionTerminal::OutcomeUnknown
));
tokio::time::sleep(std::time::Duration::from_millis(1_100)).await;
assert!(
!marker.exists(),
"CommandGuard cleanup allowed a force-closed command to continue"
);
let output = bash
.execute(r#"{"command":"printf 'reused'"}"#, Default::default())
.await
.expect("workdir session remains reusable after cleanup");
assert_eq!(output.content.as_deref(), Some("reused"));
}
// Sanity: unused Path import guard // Sanity: unused Path import guard
const _: fn() -> &'static Path = || Path::new("/"); const _: fn() -> &'static Path = || Path::new("/");
+1
View File
@@ -1244,6 +1244,7 @@ impl App {
id, id,
summary, summary,
output, output,
disposition: _,
is_error, is_error,
} => { } => {
self.latest_llm_wait_event = None; self.latest_llm_wait_event = None;
+1 -1
View File
@@ -1016,7 +1016,7 @@ fn handle_key(app: &mut App, key: KeyEvent) -> Option<Method> {
app.clear_queued_inputs(); app.clear_queued_inputs();
Some(Method::Cancel) Some(Method::Cancel)
} }
WorkerStatus::Idle => Some(Method::Shutdown), WorkerStatus::Idle | WorkerStatus::Stopped => Some(Method::Shutdown),
}), }),
KeyCode::Char('d') if ctrl => { KeyCode::Char('d') if ctrl => {
app.quit = true; app.quit = true;
+87 -90
View File
@@ -29,7 +29,6 @@ use ratatui::layout::{Constraint, Layout, Position, Rect};
use ratatui::style::{Color, Modifier, Style}; use ratatui::style::{Color, Modifier, Style};
use ratatui::text::{Line, Span}; use ratatui::text::{Line, Span};
use ratatui::widgets::{Block, Borders, Clear, Paragraph, Widget, Wrap}; use ratatui::widgets::{Block, Borders, Clear, Paragraph, Widget, Wrap};
use serde::Serialize;
use session_store::FsStore; use session_store::FsStore;
use session_store::FsWorkerStore; use session_store::FsWorkerStore;
use ticket::config::{GitBranchName, TicketConfig, TicketOrchestrationConfig}; use ticket::config::{GitBranchName, TicketConfig, TicketOrchestrationConfig};
@@ -70,10 +69,6 @@ use render::{PanelListRow, row_hit_boxes};
const MAX_ENTRIES: usize = 50; const MAX_ENTRIES: usize = 50;
const CLOSED_VISIBLE_ROWS: usize = 3; const CLOSED_VISIBLE_ROWS: usize = 3;
const ORCHESTRATOR_IDLE_QUEUE_NOTICE_PROMPT: &str = "panel.orchestrator_idle_queue_notice";
const ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS: usize = 6;
const ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS: usize = 120;
const ORCHESTRATOR_QUEUE_ATTENTION_MAX_MESSAGE_CHARS: usize = 2_400;
const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(3); const SOCKET_OP_TIMEOUT: Duration = Duration::from_secs(3);
const DASHBOARD_POLL_INTERVAL: Duration = Duration::from_millis(1_500); const DASHBOARD_POLL_INTERVAL: Duration = Duration::from_millis(1_500);
const TERMINAL_EVENT_POLL_INTERVAL: Duration = Duration::from_millis(100); const TERMINAL_EVENT_POLL_INTERVAL: Duration = Duration::from_millis(100);
@@ -581,6 +576,7 @@ pub(crate) enum IntakeRegistryUpdate {
pub(crate) struct ReadyTicketPlanningReturnRequest { pub(crate) struct ReadyTicketPlanningReturnRequest {
workspace_root: PathBuf, workspace_root: PathBuf,
ticket_id: String, ticket_id: String,
ticket_key: String,
user_instruction: String, user_instruction: String,
followup: ReadyTicketPlanningReturnFollowup, followup: ReadyTicketPlanningReturnFollowup,
} }
@@ -910,6 +906,7 @@ struct OrchestratorActiveWorkItem {
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
struct OrchestratorQueuedWorkItem { struct OrchestratorQueuedWorkItem {
id: String, id: String,
resource_key: Option<String>,
title: String, title: String,
classification: OrchestratorQueuedClassification, classification: OrchestratorQueuedClassification,
waiting_reason: Option<String>, waiting_reason: Option<String>,
@@ -974,22 +971,6 @@ impl OrchestratorQueueAttentionNoticeResult {
} }
} }
#[derive(Debug, Serialize)]
struct OrchestratorQueueTemplateContext {
workspace: String,
actionable_tickets: Vec<OrchestratorQueueTemplateTicket>,
waiting_tickets: Vec<OrchestratorQueueTemplateTicket>,
omitted_ticket_count: usize,
}
#[derive(Debug, Serialize)]
struct OrchestratorQueueTemplateTicket {
id: String,
title: String,
classification: &'static str,
waiting_reason: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq)] #[derive(Debug, Clone, PartialEq, Eq)]
struct PanelRowHitBox { struct PanelRowHitBox {
rect: Rect, rect: Rect,
@@ -1325,7 +1306,16 @@ impl DashboardApp {
if self.orchestrator_work_set.is_empty() { if self.orchestrator_work_set.is_empty() {
self.refresh_orchestrator_work_set(); self.refresh_orchestrator_work_set();
} }
let notice = orchestrator_queue_attention_notice(&self.panel, &self.orchestrator_work_set)?; let notice = match orchestrator_queue_attention_notice(&self.orchestrator_work_set) {
Ok(Some(notice)) => notice,
Ok(None) => return None,
Err(error) => {
self.notice = Some(format!(
"Orchestrator queued-work attention not delivered: {error}"
));
return None;
}
};
if self if self
.orchestrator_queue_attention .orchestrator_queue_attention
.as_ref() .as_ref()
@@ -2042,11 +2032,18 @@ impl DashboardApp {
return None; return None;
}; };
let ticket_id = ticket.id.clone(); let ticket_id = ticket.id.clone();
let ticket_key = match required_ticket_handoff_key(ticket.resource_key.as_deref()) {
Ok(ticket_key) => ticket_key.to_string(),
Err(error) => {
self.notice = Some(error);
return None;
}
};
let mut context = let mut context =
TicketRoleLaunchContext::new(current_workspace_root(), TicketRole::Intake); TicketRoleLaunchContext::new(current_workspace_root(), TicketRole::Intake);
context.ticket = Some(TicketRef::id(ticket_id.clone())); context.ticket = Some(TicketRef::id(ticket_id.clone()));
context.user_instruction = Some(format!( context.user_instruction = Some(format!(
"Continue Intake for existing Ticket {ticket_id}. Do not create a duplicate Ticket unless the user explicitly requests one. Read ShowTicket body/thread/artifacts before making routing or requirements decisions." "Continue Intake for existing Ticket {ticket_key}. Do not create a duplicate Ticket unless the user explicitly requests one. Read ShowTicket body/thread/artifacts before making routing or requirements decisions."
)); ));
let store = match PanelRegistryStore::default_for_workspace(&context.workspace_root) { let store = match PanelRegistryStore::default_for_workspace(&context.workspace_root) {
Ok(store) => store, Ok(store) => store,
@@ -2059,7 +2056,7 @@ impl DashboardApp {
Ok(Some(claim)) => { Ok(Some(claim)) => {
let status = local_claim_status_for_pod(&claim.worker_name, &self.list); let status = local_claim_status_for_pod(&claim.worker_name, &self.list);
self.notice = Some(existing_ticket_claim_notice( self.notice = Some(existing_ticket_claim_notice(
&ticket_id, &ticket_key,
&claim.worker_name, &claim.worker_name,
status, status,
)); ));
@@ -2087,7 +2084,7 @@ impl DashboardApp {
self.sending = true; self.sending = true;
self.notice = Some(format!( self.notice = Some(format!(
"Launching Ticket Intake for {} as {}…", "Launching Ticket Intake for {} as {}…",
ticket_id, planned.worker_name ticket_key, planned.worker_name
)); ));
Some(IntakeLaunchRequest { Some(IntakeLaunchRequest {
context, context,
@@ -2158,10 +2155,17 @@ impl DashboardApp {
return None; return None;
}; };
let ticket_id = ticket.id.clone(); let ticket_id = ticket.id.clone();
let ticket_key = match required_ticket_handoff_key(ticket.resource_key.as_deref()) {
Ok(ticket_key) => ticket_key.to_string(),
Err(error) => {
self.notice = Some(error);
return None;
}
};
if ticket.workflow_state != TicketWorkflowState::Ready { if ticket.workflow_state != TicketWorkflowState::Ready {
self.notice = Some(format!( self.notice = Some(format!(
"Ticket {} is {}; expected ready before returning to planning.", "Ticket {} is {}; expected ready before returning to planning.",
ticket_id, ticket_key,
ticket.workflow_state.as_str() ticket.workflow_state.as_str()
)); ));
return None; return None;
@@ -2213,7 +2217,7 @@ impl DashboardApp {
TicketRoleLaunchContext::new(workspace_root.clone(), TicketRole::Intake); TicketRoleLaunchContext::new(workspace_root.clone(), TicketRole::Intake);
context.ticket = Some(TicketRef::id(ticket_id.clone())); context.ticket = Some(TicketRef::id(ticket_id.clone()));
context.user_instruction = Some(build_ready_ticket_refinement_launch_instruction( context.user_instruction = Some(build_ready_ticket_refinement_launch_instruction(
&ticket_id, &ticket_key,
&user_instruction, &user_instruction,
)); ));
let peer_registration = self.prepare_intake_peer_registration(&mut context); let peer_registration = self.prepare_intake_peer_registration(&mut context);
@@ -2237,11 +2241,12 @@ impl DashboardApp {
self.sending = true; self.sending = true;
self.notice = Some(format!( self.notice = Some(format!(
"Returning ready Ticket {} to planning for refinement…", "Returning ready Ticket {} to planning for refinement…",
ticket_id ticket_key
)); ));
Some(ReadyTicketPlanningReturnRequest { Some(ReadyTicketPlanningReturnRequest {
workspace_root, workspace_root,
ticket_id, ticket_id,
ticket_key,
user_instruction, user_instruction,
followup, followup,
}) })
@@ -3645,6 +3650,7 @@ fn derive_orchestrator_work_set(
}; };
Some(OrchestratorQueuedWorkItem { Some(OrchestratorQueuedWorkItem {
id: ticket.id.clone(), id: ticket.id.clone(),
resource_key: ticket.resource_key.clone(),
title: ticket.title.clone(), title: ticket.title.clone(),
classification, classification,
waiting_reason, waiting_reason,
@@ -3728,72 +3734,46 @@ fn orchestrator_work_set_fingerprint(
} }
fn orchestrator_queue_attention_notice( fn orchestrator_queue_attention_notice(
panel: &WorkspacePanelViewModel,
work_set: &OrchestratorWorkSet, work_set: &OrchestratorWorkSet,
) -> Option<OrchestratorQueueAttentionNotice> { ) -> Result<Option<OrchestratorQueueAttentionNotice>, &'static str> {
if work_set.has_active_inprogress() { if work_set.has_active_inprogress() {
return None; return Ok(None);
} }
let actionable = work_set.actionable_queued(); let actionable = work_set.actionable_queued();
if actionable.is_empty() { if actionable.is_empty() {
return None; return Ok(None);
} }
let waiting = work_set let waiting = work_set
.queued .queued
.iter() .iter()
.filter(|item| item.waiting_reason.is_some()) .filter(|item| item.waiting_reason.is_some());
.collect::<Vec<_>>(); let tickets = actionable
let ticket_count = actionable.len() + waiting.len(); .into_iter()
let actionable_tickets = actionable .chain(waiting)
.iter() .map(|item| {
.take(ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS) let resource_key = item
.map(|item| orchestrator_queue_template_ticket(item)) .resource_key
.collect::<Vec<_>>(); .clone()
let remaining_capacity = .ok_or("queued Ticket is missing its required resource key")?;
ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS.saturating_sub(actionable_tickets.len()); worker::OrchestratorQueueAttentionTicket::new(resource_key, item.title.clone())
let waiting_tickets = waiting .map_err(|_| "queued Ticket has an invalid resource key")
.iter()
.take(remaining_capacity)
.map(|item| orchestrator_queue_template_ticket(item))
.collect::<Vec<_>>();
let rendered =
render_orchestrator_queue_attention_template(&OrchestratorQueueTemplateContext {
workspace: bounded_progress_text(
&panel.header.workspace_label,
ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS,
),
actionable_tickets,
waiting_tickets,
omitted_ticket_count: ticket_count
.saturating_sub(ORCHESTRATOR_QUEUE_ATTENTION_MAX_TICKETS),
}) })
.ok()?; .collect::<Result<Vec<_>, _>>()?;
let message = bounded_progress_text(&rendered, ORCHESTRATOR_QUEUE_ATTENTION_MAX_MESSAGE_CHARS); let context = worker::OrchestratorQueueAttentionContext::new(tickets);
let message = render_orchestrator_queue_attention_template(&context)
.map_err(|_| "queued-work attention prompt rendering failed")?;
let fingerprint = format!("idle-queue:{}", work_set.fingerprint); let fingerprint = format!("idle-queue:{}", work_set.fingerprint);
Some(OrchestratorQueueAttentionNotice { Ok(Some(OrchestratorQueueAttentionNotice {
message, message,
fingerprint, fingerprint,
}) }))
}
fn orchestrator_queue_template_ticket(
item: &&OrchestratorQueuedWorkItem,
) -> OrchestratorQueueTemplateTicket {
OrchestratorQueueTemplateTicket {
id: bounded_progress_text(&item.id, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS),
title: bounded_progress_text(&item.title, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS),
classification: item.classification.as_str(),
waiting_reason: item.waiting_reason.as_ref().map(|reason| {
bounded_progress_text(reason, ORCHESTRATOR_QUEUE_ATTENTION_MAX_TEXT_CHARS)
}),
}
} }
fn render_orchestrator_queue_attention_template( fn render_orchestrator_queue_attention_template(
context: &OrchestratorQueueTemplateContext, context: &worker::OrchestratorQueueAttentionContext,
) -> Result<String, worker::CatalogError> { ) -> Result<String, worker::CatalogError> {
worker::PromptCatalog::builtins_only()? worker::PromptCatalog::builtins_only()?
.render_serializable(ORCHESTRATOR_IDLE_QUEUE_NOTICE_PROMPT, context) .orchestrator_queue_attention(worker::OrchestratorQueueAttentionPrompt::Tui, context)
} }
fn orchestrator_work_set_detail( fn orchestrator_work_set_detail(
@@ -3918,21 +3898,35 @@ fn bounded_refinement_instruction(input: &str) -> String {
.to_string() .to_string()
} }
fn build_ready_ticket_refinement_thread_body(ticket_id: &str, instruction: &str) -> String { fn required_ticket_handoff_key(resource_key: Option<&str>) -> Result<&str, String> {
let resource_key = resource_key.ok_or_else(|| {
"Ticket handoff is unavailable because the canonical T-* resource key is missing. Refresh the panel and retry."
.to_string()
})?;
let sequence = resource_key.strip_prefix("T-").filter(|sequence| {
!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit())
});
sequence.map(|_| resource_key).ok_or_else(|| {
"Ticket handoff is unavailable because the canonical T-* resource key is invalid. Refresh the panel and retry."
.to_string()
})
}
fn build_ready_ticket_refinement_thread_body(ticket_key: &str, instruction: &str) -> String {
format!( format!(
"Panel returned ready Ticket {ticket_id} to planning for requirements sync. This is not Queue routing and must not start implementation.\n\n## User refinement instruction\n\n{instruction}\n" "Panel returned ready Ticket {ticket_key} to planning for requirements sync. This is not Queue routing and must not start implementation.\n\n## User refinement instruction\n\n{instruction}\n"
) )
} }
fn build_ready_ticket_refinement_launch_instruction(ticket_id: &str, instruction: &str) -> String { fn build_ready_ticket_refinement_launch_instruction(ticket_key: &str, instruction: &str) -> String {
format!( format!(
"Continue Ticket Intake / requirements sync for existing Ticket {ticket_id}. The Panel has returned the Ticket from ready to planning; do not queue the Ticket, do not route implementation, and do not create a duplicate unless the user explicitly asks for one. Read ShowTicket body/thread/artifacts before making requirements or readiness decisions.\n\nUser refinement instruction:\n\n{instruction}" "Continue Ticket Intake / requirements sync for existing Ticket {ticket_key}. The Panel has returned the Ticket from ready to planning; do not queue the Ticket, do not route implementation, and do not create a duplicate unless the user explicitly asks for one. Read ShowTicket body/thread/artifacts before making requirements or readiness decisions.\n\nUser refinement instruction:\n\n{instruction}"
) )
} }
fn build_ready_ticket_refinement_notify(ticket_id: &str, instruction: &str) -> String { fn build_ready_ticket_refinement_notify(ticket_key: &str, instruction: &str) -> String {
format!( format!(
"Ticket {ticket_id} was returned from ready to planning from the Panel for requirements sync. Continue Intake/refinement only; do not Queue or route implementation. Read the Ticket thread for the recorded state change and user instruction.\n\nUser refinement instruction:\n\n{instruction}" "Ticket {ticket_key} was returned from ready to planning from the Panel for requirements sync. Continue Intake/refinement only; do not Queue or route implementation. Read the Ticket thread for the recorded state change and user instruction.\n\nUser refinement instruction:\n\n{instruction}"
) )
} }
@@ -3961,10 +3955,12 @@ async fn dispatch_ready_ticket_planning_return(
let ticket = backend let ticket = backend
.show(id.clone()) .show(id.clone())
.map_err(|error| TicketActionError::Ticket(error.to_string()))?; .map_err(|error| TicketActionError::Ticket(error.to_string()))?;
let ticket_key =
required_ticket_handoff_key(Some(&request.ticket_key)).map_err(TicketActionError::Stale)?;
if ticket.meta.workflow_state != TicketWorkflowState::Ready { if ticket.meta.workflow_state != TicketWorkflowState::Ready {
return Err(TicketActionError::Stale(format!( return Err(TicketActionError::Stale(format!(
"Ticket {} is {}; expected ready before returning it to planning. Refresh the panel and retry if appropriate.", "Ticket {} is {}; expected ready before returning it to planning. Refresh the panel and retry if appropriate.",
ticket.meta.id, ticket_key,
ticket.meta.workflow_state.as_str() ticket.meta.workflow_state.as_str()
))); )));
} }
@@ -3973,7 +3969,7 @@ async fn dispatch_ready_ticket_planning_return(
TicketWorkflowState::Planning.as_str(), TicketWorkflowState::Planning.as_str(),
"panel_return_to_planning", "panel_return_to_planning",
MarkdownText::from(build_ready_ticket_refinement_thread_body( MarkdownText::from(build_ready_ticket_refinement_thread_body(
&ticket.meta.id, ticket_key,
&request.user_instruction, &request.user_instruction,
)), )),
); );
@@ -3987,7 +3983,7 @@ async fn dispatch_ready_ticket_planning_return(
ReadyTicketPlanningReturnOutcome { ReadyTicketPlanningReturnOutcome {
notice: format!( notice: format!(
"Ticket {} returned to planning for refinement; launching Ticket Intake…", "Ticket {} returned to planning for refinement; launching Ticket Intake…",
ticket.meta.id ticket_key
), ),
followup: ReadyTicketPlanningReturnAfterMutation::LaunchIntake(request), followup: ReadyTicketPlanningReturnAfterMutation::LaunchIntake(request),
} }
@@ -3997,19 +3993,19 @@ async fn dispatch_ready_ticket_planning_return(
socket_path, socket_path,
} => { } => {
let message = let message =
build_ready_ticket_refinement_notify(&ticket.meta.id, &request.user_instruction); build_ready_ticket_refinement_notify(ticket_key, &request.user_instruction);
match send_notify_only(&socket_path, message, true).await { match send_notify_only(&socket_path, message, true).await {
Ok(()) => ReadyTicketPlanningReturnOutcome { Ok(()) => ReadyTicketPlanningReturnOutcome {
notice: format!( notice: format!(
"Ticket {} returned to planning for refinement; notified live Intake Worker {}.", "Ticket {} returned to planning for refinement; notified live Intake Worker {}.",
ticket.meta.id, worker_name ticket_key, worker_name
), ),
followup: ReadyTicketPlanningReturnAfterMutation::None, followup: ReadyTicketPlanningReturnAfterMutation::None,
}, },
Err(error) => ReadyTicketPlanningReturnOutcome { Err(error) => ReadyTicketPlanningReturnOutcome {
notice: bounded_panel_diagnostic(format!( notice: bounded_panel_diagnostic(format!(
"Ticket {} returned to planning and instruction was recorded, but notifying Intake Worker {} failed: {}", "Ticket {} returned to planning and instruction was recorded, but notifying Intake Worker {} failed: {}",
ticket.meta.id, worker_name, error ticket_key, worker_name, error
)), )),
followup: ReadyTicketPlanningReturnAfterMutation::None, followup: ReadyTicketPlanningReturnAfterMutation::None,
}, },
@@ -4020,7 +4016,7 @@ async fn dispatch_ready_ticket_planning_return(
ReadyTicketPlanningReturnOutcome { ReadyTicketPlanningReturnOutcome {
notice: format!( notice: format!(
"Ticket {} returned to planning for refinement; opening/restoring claimed Intake Worker {}…", "Ticket {} returned to planning for refinement; opening/restoring claimed Intake Worker {}…",
ticket.meta.id, worker_name ticket_key, worker_name
), ),
followup: ReadyTicketPlanningReturnAfterMutation::OpenClaim(request), followup: ReadyTicketPlanningReturnAfterMutation::OpenClaim(request),
} }
@@ -4029,7 +4025,7 @@ async fn dispatch_ready_ticket_planning_return(
ReadyTicketPlanningReturnOutcome { ReadyTicketPlanningReturnOutcome {
notice: bounded_panel_diagnostic(format!( notice: bounded_panel_diagnostic(format!(
"Ticket {} returned to planning and instruction was recorded, but Intake launch was not attempted because existing Intake claim {} is stale; inspect or clear the local claim before launching another Intake Worker.", "Ticket {} returned to planning and instruction was recorded, but Intake launch was not attempted because existing Intake claim {} is stale; inspect or clear the local claim before launching another Intake Worker.",
ticket.meta.id, worker_name ticket_key, worker_name
)), )),
followup: ReadyTicketPlanningReturnAfterMutation::None, followup: ReadyTicketPlanningReturnAfterMutation::None,
} }
@@ -5204,6 +5200,7 @@ fn row_status_label(entry: &WorkerListEntry) -> (&'static str, Style) {
.fg(Color::Cyan) .fg(Color::Cyan)
.add_modifier(Modifier::BOLD), .add_modifier(Modifier::BOLD),
), ),
Some(WorkerStatus::Stopped) => ("live stopped", Style::default().fg(Color::DarkGray)),
None => ("live", Style::default().fg(Color::DarkGray)), None => ("live", Style::default().fg(Color::DarkGray)),
}; };
} }
+111 -8
View File
@@ -390,6 +390,7 @@ fn planning_return_request(
ReadyTicketPlanningReturnRequest { ReadyTicketPlanningReturnRequest {
workspace_root: temp.path().to_path_buf(), workspace_root: temp.path().to_path_buf(),
ticket_id, ticket_id,
ticket_key: "T-482".to_string(),
user_instruction: instruction.to_string(), user_instruction: instruction.to_string(),
followup: ReadyTicketPlanningReturnFollowup::BlockedByStaleClaim { followup: ReadyTicketPlanningReturnFollowup::BlockedByStaleClaim {
worker_name: "stale-intake".to_string(), worker_name: "stale-intake".to_string(),
@@ -494,6 +495,7 @@ fn ready_ticket_intake_enter_prepares_planning_return_not_queue_or_generic_launc
}; };
assert_eq!(request.ticket_id, "20260608-000123-ready"); assert_eq!(request.ticket_id, "20260608-000123-ready");
assert_eq!(request.ticket_key, "T-1");
assert_eq!(request.user_instruction, "clarify expected behavior"); assert_eq!(request.user_instruction, "clarify expected behavior");
assert!(matches!( assert!(matches!(
request.followup, request.followup,
@@ -515,6 +517,7 @@ async fn planning_return_with_launch_followup_changes_state_before_launch_follow
let request = ReadyTicketPlanningReturnRequest { let request = ReadyTicketPlanningReturnRequest {
workspace_root: temp.path().to_path_buf(), workspace_root: temp.path().to_path_buf(),
ticket_id: ticket_id.clone(), ticket_id: ticket_id.clone(),
ticket_key: "T-482".to_string(),
user_instruction: "launch intake after state change".to_string(), user_instruction: "launch intake after state change".to_string(),
followup: ReadyTicketPlanningReturnFollowup::LaunchIntake(IntakeLaunchRequest { followup: ReadyTicketPlanningReturnFollowup::LaunchIntake(IntakeLaunchRequest {
context: TicketRoleLaunchContext::new(temp.path().to_path_buf(), TicketRole::Intake), context: TicketRoleLaunchContext::new(temp.path().to_path_buf(), TicketRole::Intake),
@@ -2969,7 +2972,7 @@ fn dashboard_empty_enter_on_non_openable_row_reports_open_diagnostic() {
} }
#[test] #[test]
fn idle_orchestrator_gets_bounded_attention_for_new_queued_work() { fn idle_orchestrator_gets_sanitized_attention_for_new_queued_work() {
let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]); let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]);
app.panel.rows = vec![panel_test_ticket_row( app.panel.rows = vec![panel_test_ticket_row(
"00001QUEUE", "00001QUEUE",
@@ -2989,11 +2992,87 @@ fn idle_orchestrator_gets_bounded_attention_for_new_queued_work() {
request request
.notice .notice
.message .message
.starts_with("Workspace Dashboard observed") .starts_with("Queued Tickets require attention:")
); );
assert!(request.notice.message.contains("00001QUEUE")); assert!(request.notice.message.contains("- T-1 — Queued work"));
assert!(request.notice.message.contains("new_queued")); assert!(
assert!(request.notice.message.contains("queued -> inprogress")); request
.notice
.message
.contains("Reread the current Ticket state before acting")
);
assert!(
!request
.notice
.message
.contains(&app.panel.header.workspace_label)
);
for hidden in [
"00001QUEUE",
"Workspace:",
"workspace_id",
"new_queued",
"bounded",
"queued -> inprogress",
] {
assert!(!request.notice.message.contains(hidden), "leaked {hidden}");
}
}
#[test]
fn queued_attention_missing_resource_key_fails_closed_with_panel_notice() {
let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]);
let mut row = panel_test_ticket_row(
"00001QUEUE",
"Queued work",
ActionPriority::Background,
NextUserAction::Wait,
"queued",
);
row.ticket.as_mut().unwrap().resource_key = None;
app.panel.rows = vec![row];
app.refresh_orchestrator_work_set();
assert!(app.prepare_orchestrator_queue_attention_notice().is_none());
assert_eq!(
app.notice.as_deref(),
Some(
"Orchestrator queued-work attention not delivered: queued Ticket is missing its required resource key"
)
);
}
#[test]
fn queued_attention_truncates_only_when_tickets_are_omitted() {
let mut app = ticket_enabled_app(vec![live_info("test-orchestrator", WorkerStatus::Idle)]);
app.panel.rows = (1..=worker::OrchestratorQueueAttentionContext::MAX_TICKETS + 1)
.map(|index| {
let mut row = panel_test_ticket_row(
&format!("opaque-{index}"),
&format!("Queued work {index}"),
ActionPriority::Background,
NextUserAction::Wait,
"queued",
);
row.ticket.as_mut().unwrap().resource_key = Some(format!("T-{index}"));
row
})
.collect();
app.refresh_orchestrator_work_set();
let request = app
.prepare_orchestrator_queue_attention_notice()
.expect("bounded queued-work attention");
assert!(request.notice.message.contains("- T-20 — Queued work 20"));
assert!(!request.notice.message.contains("T-21"));
assert!(
request
.notice
.message
.contains("were omitted from this notice: 1")
);
assert!(!request.notice.message.contains("opaque-"));
} }
#[test] #[test]
@@ -3083,7 +3162,9 @@ fn planned_queued_prompts_when_active_work_clears() {
.prepare_orchestrator_queue_attention_notice() .prepare_orchestrator_queue_attention_notice()
.expect("planned queued work should prompt after active work clears"); .expect("planned queued work should prompt after active work clears");
assert!(request.notice.message.contains("planned_queued")); assert!(request.notice.message.contains("- T-1 — Queued work"));
assert!(!request.notice.message.contains("planned_queued"));
assert!(!request.notice.message.contains("00001QUEUE"));
assert!( assert!(
!request !request
.notice .notice
@@ -3138,8 +3219,9 @@ fn rediscovered_queued_work_is_actionable_when_session_work_set_is_empty() {
.prepare_orchestrator_queue_attention_notice() .prepare_orchestrator_queue_attention_notice()
.expect("queued ticket state should be rediscovered safely"); .expect("queued ticket state should be rediscovered safely");
assert!(request.notice.message.contains("new_queued")); assert!(request.notice.message.contains("- T-1 — Queued work"));
assert!(request.notice.message.contains("00001QUEUE")); assert!(!request.notice.message.contains("new_queued"));
assert!(!request.notice.message.contains("00001QUEUE"));
} }
#[test] #[test]
@@ -3425,6 +3507,27 @@ fn ticket_action_error_records_f2_diagnostic_details() {
assert!(!app.panel_diagnostic_open); assert!(!app.panel_diagnostic_open);
} }
#[test]
fn ready_ticket_refinement_projection_uses_only_canonical_resource_key() {
const INTERNAL_ID: &str = "00001KZVNXFNK";
let thread = build_ready_ticket_refinement_thread_body("T-482", "Clarify rollback.");
let launch = build_ready_ticket_refinement_launch_instruction("T-482", "Clarify rollback.");
let notify = build_ready_ticket_refinement_notify("T-482", "Clarify rollback.");
for projection in [&thread, &launch, &notify] {
assert!(projection.contains("T-482"));
assert!(!projection.contains(INTERNAL_ID));
}
}
#[test]
fn ticket_handoff_fails_closed_without_canonical_resource_key() {
assert_eq!(required_ticket_handoff_key(Some("T-482")), Ok("T-482"));
for invalid in [None, Some(""), Some("00001KZVNXFNK"), Some("T-key")] {
assert!(required_ticket_handoff_key(invalid).is_err());
}
}
fn plain_line(line: &Line<'_>) -> String { fn plain_line(line: &Line<'_>) -> String {
line.spans line.spans
.iter() .iter()
+1
View File
@@ -1530,6 +1530,7 @@ fn worker_status_label(entry: &WorkerListEntry) -> &'static str {
Some(WorkerStatus::Idle) => "live idle", Some(WorkerStatus::Idle) => "live idle",
Some(WorkerStatus::Running) => "live running", Some(WorkerStatus::Running) => "live running",
Some(WorkerStatus::Paused) => "live paused", Some(WorkerStatus::Paused) => "live paused",
Some(WorkerStatus::Stopped) => "live stopped",
None => "live", None => "live",
}; };
} }
+2 -1
View File
@@ -2742,6 +2742,7 @@ impl RuntimeState {
protocol::WorkerStatus::Running => Some(WorkerStatus::Running), protocol::WorkerStatus::Running => Some(WorkerStatus::Running),
protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle), protocol::WorkerStatus::Idle => Some(WorkerStatus::Idle),
protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused), protocol::WorkerStatus::Paused => Some(WorkerStatus::Paused),
protocol::WorkerStatus::Stopped => Some(WorkerStatus::Stopped),
}, },
protocol::Event::RunEnd { result } => match result { protocol::Event::RunEnd { result } => match result {
protocol::RunResult::Finished | protocol::RunResult::RolledBack => { protocol::RunResult::Finished | protocol::RunResult::RolledBack => {
@@ -3104,7 +3105,7 @@ mod tests {
&mut activity, &mut activity,
&internal_worker_status_event( &internal_worker_status_event(
internal_worker_ref("child-b", None), internal_worker_ref("child-b", None),
protocol::WorkerStatus::Idle, protocol::WorkerStatus::Stopped,
), ),
)); ));
} }
+21 -14
View File
@@ -38,9 +38,7 @@ use crate::working_directory::{
}; };
use async_trait::async_trait; use async_trait::async_trait;
use protocol::{Event, Method, Segment, WorkerStatus}; use protocol::{Event, Method, Segment, WorkerStatus};
use session_store::{ use session_store::{CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore};
CombinedStore, LogEntry, WorkerAggregateStore, WorkerSessionStore, collect_state,
};
#[cfg(test)] #[cfg(test)]
use session_store::{FsStore, FsWorkerStore}; use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime; use tokio::runtime::Runtime;
@@ -68,8 +66,10 @@ const RUNTIME_TASK_TIMEOUT: Duration = Duration::from_secs(10);
const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9); const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool { fn user_input_has_submission(entry: &LogEntry, submission_id: &str) -> bool {
let LogEntry::UserInput { extensions, .. } = entry else { let extensions = match entry {
return false; LogEntry::UserInput { extensions, .. }
| LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
_ => return false,
}; };
extensions.iter().any(|extension| { extensions.iter().any(|extension| {
extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN extension.domain == WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN
@@ -212,11 +212,11 @@ impl WorkerObservationProvider for RuntimeGrantedWorkerObservationProvider {
return Err(WorkerObservationError::NotFound); return Err(WorkerObservationError::NotFound);
} }
let entries = sink.subscribe_with_snapshot().0; let entries = sink.subscribe_with_snapshot().0;
let state = collect_state(&entries); WorkerSessionCapture::from_log_entries(
Ok(WorkerSessionCapture { format!("runtime:{runtime_id}:worker:{worker_id}"),
segment_id: format!("runtime:{runtime_id}:worker:{worker_id}"), &entries,
items: state.history, )
}) .map_err(WorkerObservationError::Unavailable)
} }
} }
@@ -1546,7 +1546,9 @@ fn accepted_notify_run_state(status: WorkerStatus, auto_run: bool) -> WorkerExec
match status { match status {
WorkerStatus::Running => WorkerExecutionRunState::Busy, WorkerStatus::Running => WorkerExecutionRunState::Busy,
WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy, WorkerStatus::Idle if auto_run => WorkerExecutionRunState::Busy,
WorkerStatus::Idle | WorkerStatus::Paused => WorkerExecutionRunState::Idle, WorkerStatus::Idle | WorkerStatus::Paused | WorkerStatus::Stopped => {
WorkerExecutionRunState::Idle
}
} }
} }
@@ -2505,7 +2507,9 @@ mod tests {
let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?; let scope = Scope::writable(&scope_root).map_err(|err| err.to_string())?;
let worker = Worker::new( let worker = Worker::new(
manifest, manifest,
Engine::new(self.client.clone()), Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(
self.client.clone(),
),
store, store,
workspace_context, workspace_context,
filesystem_authority, filesystem_authority,
@@ -3241,14 +3245,17 @@ mod tests {
matches!( matches!(
entry, entry,
LogEntry::UserInput { segments, .. } LogEntry::UserInput { segments, .. }
| LogEntry::AnnotatedUserInput { segments, .. }
if segments == &vec![Segment::text("start the ticket")] if segments == &vec![Segment::text("start the ticket")]
) )
})); }));
let submission_id = entries let submission_id = entries
.iter() .iter()
.find_map(|entry| { .find_map(|entry| {
let LogEntry::UserInput { extensions, .. } = entry else { let extensions = match entry {
return None; LogEntry::UserInput { extensions, .. }
| LogEntry::AnnotatedUserInput { extensions, .. } => extensions,
_ => return None,
}; };
extensions extensions
.iter() .iter()
+2 -1
View File
@@ -66,11 +66,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
WorkerRunResult::Finished => println!("(finished)"), WorkerRunResult::Finished => println!("(finished)"),
WorkerRunResult::Paused => println!("(paused)"), WorkerRunResult::Paused => println!("(paused)"),
WorkerRunResult::LimitReached => println!("(turn limit reached)"), WorkerRunResult::LimitReached => println!("(turn limit reached)"),
WorkerRunResult::Interrupted { message, .. } => println!("(interrupted: {message})"),
WorkerRunResult::RolledBack => println!("(empty turn rolled back)"), WorkerRunResult::RolledBack => println!("(empty turn rolled back)"),
} }
// 5. Extract the assistant's reply from history // 5. Extract the assistant's reply from history
let history = worker.engine().history(); let history = worker.history();
if let Some(text) = history if let Some(text) = history
.iter() .iter()
.rev() .rev()
+1 -1
View File
@@ -22,7 +22,7 @@ use crate::compact::token_counter::{
EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl, EstimateSource, savings_for_prune_impl, token_estimates_for_prune_impl,
}; };
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// Enable prune projection on the underlying Engine. /// Enable prune projection on the underlying Engine.
/// ///
/// Registers the config and token/savings-estimator closures on the Engine. /// Registers the config and token/savings-estimator closures on the Engine.
+4 -4
View File
@@ -242,13 +242,13 @@ pub(crate) fn savings_for_prune_impl(
// ── Worker に生やす公開 API ─────────────────────────────────────────────── // ── Worker に生やす公開 API ───────────────────────────────────────────────
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
/// 現在の history 全体の推定トークン数。 /// 現在の history 全体の推定トークン数。
/// ///
/// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。 /// 最後の measurement と、その後に追加された未測定分の byte/4 外挿。
pub fn total_tokens(&self) -> TokenEstimate { pub fn total_tokens(&self) -> TokenEstimate {
let usage = self.usage_history(); let usage = self.usage_history();
agen::token_counter::total_tokens(self.history(), &usage) agen::token_counter::total_tokens(&self.history(), &usage)
} }
/// 任意の history index 時点でのプロンプト全長推定。 /// 任意の history index 時点でのプロンプト全長推定。
@@ -259,7 +259,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// pointer 以降に増えたプロンプト長を測るのに使う。 /// pointer 以降に増えたプロンプト長を測るのに使う。
pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate { pub fn total_tokens_at(&self, history_len: usize) -> TokenEstimate {
let usage = self.usage_history(); let usage = self.usage_history();
agen::token_counter::total_tokens_at(self.history(), &usage, history_len) agen::token_counter::total_tokens_at(&self.history(), &usage, history_len)
} }
/// 末尾から `retained` トークン以上を残すための分割位置。 /// 末尾から `retained` トークン以上を残すための分割位置。
@@ -267,7 +267,7 @@ impl<C: LlmClient, St: Store> Worker<C, St> {
/// `history[..cut.index]` が要約/破棄される側、`history[cut.index..]` が残る側。 /// `history[..cut.index]` が要約/破棄される側、`history[cut.index..]` が残る側。
pub fn split_for_retained(&self, retained: u64) -> SplitPoint { pub fn split_for_retained(&self, retained: u64) -> SplitPoint {
let usage = self.usage_history(); let usage = self.usage_history();
split_for_retained_impl(self.history(), &usage, retained) split_for_retained_impl(&self.history(), &usage, retained)
} }
} }
+172 -13
View File
@@ -485,6 +485,7 @@ impl WorkerController {
// into the controller task so the in-flight turn can be reached // into the controller task so the in-flight turn can be reached
// via these handles while worker itself is borrowed by drive_turn. // via these handles while worker itself is borrowed by drive_turn.
let cancel_tx = worker.engine_mut().cancel_sender(); let cancel_tx = worker.engine_mut().cancel_sender();
let pause_tx = worker.engine_mut().pause_sender();
let notify_buffer = worker.notify_buffer_handle(); let notify_buffer = worker.notify_buffer_handle();
tokio::spawn(controller_loop( tokio::spawn(controller_loop(
@@ -494,6 +495,7 @@ impl WorkerController {
shared_state, shared_state,
runtime_dir, runtime_dir,
cancel_tx, cancel_tx,
pause_tx,
notify_buffer, notify_buffer,
self_parent_socket, self_parent_socket,
spawner_name, spawner_name,
@@ -763,6 +765,19 @@ pub(crate) fn wire_event_bridges_on_engine<C, St>(
id: result.tool_use_id.clone(), id: result.tool_use_id.clone(),
summary: result.summary.clone(), summary: result.summary.clone(),
output: result.content.clone(), output: result.content.clone(),
disposition: Some(match result.disposition {
agen::ToolResultDisposition::Success => protocol::ToolResultDisposition::Success,
agen::ToolResultDisposition::Error => protocol::ToolResultDisposition::Error,
agen::ToolResultDisposition::Interrupted => {
protocol::ToolResultDisposition::Interrupted
}
agen::ToolResultDisposition::Cancelled => {
protocol::ToolResultDisposition::Cancelled
}
agen::ToolResultDisposition::OutcomeUnknown => {
protocol::ToolResultDisposition::OutcomeUnknown
}
}),
is_error: result.is_error, is_error: result.is_error,
}); });
}); });
@@ -1123,6 +1138,7 @@ async fn controller_loop<C, St>(
shared_state: Arc<WorkerSharedState>, shared_state: Arc<WorkerSharedState>,
runtime_dir: Arc<RuntimeDir>, runtime_dir: Arc<RuntimeDir>,
cancel_tx: mpsc::Sender<()>, cancel_tx: mpsc::Sender<()>,
pause_tx: mpsc::Sender<()>,
notify_buffer: NotifyBuffer, notify_buffer: NotifyBuffer,
self_parent_socket: Option<PathBuf>, self_parent_socket: Option<PathBuf>,
spawner_name: String, spawner_name: String,
@@ -1169,6 +1185,9 @@ async fn controller_loop<C, St>(
// clear at run start prevents stale partial output left by an older // clear at run start prevents stale partial output left by an older
// interrupted/error turn from being carried into the next snapshot. // interrupted/error turn from being carried into the next snapshot.
worker.clear_in_flight_events(); worker.clear_in_flight_events();
let parent_originated = run.is_parent_originated();
let user_input_run = matches!(&run, PendingRun::Run(_) | PendingRun::RunTracked { .. });
if !user_input_run {
set_controller_status( set_controller_status(
&shared_state, &shared_state,
&runtime_dir, &runtime_dir,
@@ -1176,15 +1195,25 @@ async fn controller_loop<C, St>(
WorkerStatus::Running, WorkerStatus::Running,
) )
.await; .await;
let parent_originated = run.is_parent_originated(); }
let (mut new_status, shutdown) = match run { let (mut new_status, shutdown) = match run {
PendingRun::Run(input) => { PendingRun::Run(input) => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
drive_turn( drive_turn(
worker.run(input), worker.run_with_input_extensions_and_commit_hook(
input,
Vec::new(),
move || {
let _ = input_commit_tx.send(());
},
),
&mut method_rx, &mut method_rx,
&event_tx, &event_tx,
&cancel_tx, &cancel_tx,
&pause_tx,
&shared_state, &shared_state,
&runtime_dir,
Some(input_commit_rx),
&notify_buffer, &notify_buffer,
self_parent_socket.as_ref(), self_parent_socket.as_ref(),
&spawner_name, &spawner_name,
@@ -1194,12 +1223,22 @@ async fn controller_loop<C, St>(
.await .await
} }
PendingRun::RunTracked { input, extension } => { PendingRun::RunTracked { input, extension } => {
let (input_commit_tx, input_commit_rx) = oneshot::channel();
drive_turn( drive_turn(
worker.run_with_input_extensions(input, vec![extension]), worker.run_with_input_extensions_and_commit_hook(
input,
vec![extension],
move || {
let _ = input_commit_tx.send(());
},
),
&mut method_rx, &mut method_rx,
&event_tx, &event_tx,
&cancel_tx, &cancel_tx,
&pause_tx,
&shared_state, &shared_state,
&runtime_dir,
Some(input_commit_rx),
&notify_buffer, &notify_buffer,
self_parent_socket.as_ref(), self_parent_socket.as_ref(),
&spawner_name, &spawner_name,
@@ -1214,7 +1253,10 @@ async fn controller_loop<C, St>(
&mut method_rx, &mut method_rx,
&event_tx, &event_tx,
&cancel_tx, &cancel_tx,
&pause_tx,
&shared_state, &shared_state,
&runtime_dir,
None,
&notify_buffer, &notify_buffer,
self_parent_socket.as_ref(), self_parent_socket.as_ref(),
&spawner_name, &spawner_name,
@@ -1229,7 +1271,10 @@ async fn controller_loop<C, St>(
&mut method_rx, &mut method_rx,
&event_tx, &event_tx,
&cancel_tx, &cancel_tx,
&pause_tx,
&shared_state, &shared_state,
&runtime_dir,
None,
&notify_buffer, &notify_buffer,
self_parent_socket.as_ref(), self_parent_socket.as_ref(),
&spawner_name, &spawner_name,
@@ -1354,7 +1399,7 @@ async fn controller_loop<C, St>(
}); });
} }
}, },
WorkerStatus::Idle => { WorkerStatus::Idle | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error { let _ = event_tx.send(Event::Error {
code: ErrorCode::NotRunning, code: ErrorCode::NotRunning,
message: "Worker is not running".into(), message: "Worker is not running".into(),
@@ -1395,7 +1440,7 @@ async fn controller_loop<C, St>(
.into(), .into(),
}); });
} }
WorkerStatus::Running => { WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error { let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning, code: ErrorCode::AlreadyRunning,
message: message:
@@ -1409,7 +1454,7 @@ async fn controller_loop<C, St>(
WorkerStatus::Idle | WorkerStatus::Paused => { WorkerStatus::Idle | WorkerStatus::Paused => {
emit_rewind_targets(&worker, &event_tx) emit_rewind_targets(&worker, &event_tx)
} }
WorkerStatus::Running => { WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error { let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning, code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused" message: "Worker is already executing a turn; rewind can only run while idle or paused"
@@ -1438,7 +1483,7 @@ async fn controller_loop<C, St>(
.into(), .into(),
}); });
} }
WorkerStatus::Running => { WorkerStatus::Running | WorkerStatus::Stopped => {
let _ = event_tx.send(Event::Error { let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning, code: ErrorCode::AlreadyRunning,
message: "Worker is already executing a turn; rewind can only run while idle or paused" message: "Worker is already executing a turn; rewind can only run while idle or paused"
@@ -1626,7 +1671,10 @@ async fn drive_turn<F>(
method_rx: &mut mpsc::Receiver<Method>, method_rx: &mut mpsc::Receiver<Method>,
event_tx: &broadcast::Sender<Event>, event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>, cancel_tx: &mpsc::Sender<()>,
pause_tx: &mpsc::Sender<()>,
shared_state: &Arc<WorkerSharedState>, shared_state: &Arc<WorkerSharedState>,
runtime_dir: &RuntimeDir,
mut input_commit_rx: Option<oneshot::Receiver<()>>,
notify_buffer: &NotifyBuffer, notify_buffer: &NotifyBuffer,
parent_socket: Option<&PathBuf>, parent_socket: Option<&PathBuf>,
self_name: &str, self_name: &str,
@@ -1642,14 +1690,58 @@ where
loop { loop {
tokio::select! { tokio::select! {
// If input commit and provider completion become ready together, expose
// Running only after processing the commit fence. This makes the
// Running snapshot contract deterministic even for immediate clients.
biased;
committed = async {
input_commit_rx
.as_mut()
.expect("input commit receiver guarded by select condition")
.await
}, if input_commit_rx.is_some() => {
input_commit_rx = None;
if committed.is_ok() {
set_controller_status(
shared_state,
runtime_dir,
event_tx,
WorkerStatus::Running,
)
.await;
}
}
result = &mut worker_future => { result = &mut worker_future => {
return match result { return match result {
Ok(r) => { Ok(r) => {
let (status, run_result) = match r { let (status, run_result) = match r {
WorkerRunResult::Finished if pause_requested => {
(WorkerStatus::Paused, RunResult::Paused)
}
WorkerRunResult::Finished => (WorkerStatus::Idle, RunResult::Finished), WorkerRunResult::Finished => (WorkerStatus::Idle, RunResult::Finished),
WorkerRunResult::Paused => (WorkerStatus::Paused, RunResult::Paused), WorkerRunResult::Paused => (WorkerStatus::Paused, RunResult::Paused),
WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached), WorkerRunResult::LimitReached => (WorkerStatus::Idle, RunResult::LimitReached),
WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack), WorkerRunResult::RolledBack => (WorkerStatus::Idle, RunResult::RolledBack),
WorkerRunResult::Interrupted { .. } if pause_requested => {
let _ = event_tx.send(Event::RunEnd { result: RunResult::Paused });
return (WorkerStatus::Paused, shutdown_requested);
}
WorkerRunResult::Interrupted { code, message } => {
let _ = event_tx.send(Event::Error {
code,
message: message.clone(),
});
if parent_originated {
crate::ipc::event::fire_and_forget(
parent_socket.cloned(),
protocol::WorkerEvent::Errored {
worker_name: self_name.to_string(),
message,
},
);
}
return (WorkerStatus::Idle, shutdown_requested);
}
}; };
let _ = event_tx.send(Event::RunEnd { result: run_result }); let _ = event_tx.send(Event::RunEnd { result: run_result });
if parent_originated && matches!(run_result, RunResult::Finished) { if parent_originated && matches!(run_result, RunResult::Finished) {
@@ -1698,7 +1790,7 @@ where
} }
Some(Method::Pause) => { Some(Method::Pause) => {
pause_requested = true; pause_requested = true;
let _ = cancel_tx.try_send(()); let _ = pause_tx.try_send(());
} }
Some(Method::Shutdown) => { Some(Method::Shutdown) => {
shutdown_requested = true; shutdown_requested = true;
@@ -1760,7 +1852,7 @@ where
fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, event_tx: &broadcast::Sender<Event>) fn emit_rewind_targets<C, St>(worker: &Worker<C, St>, event_tx: &broadcast::Sender<Event>)
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
match worker.list_rewind_targets() { match worker.list_rewind_targets() {
@@ -1786,7 +1878,7 @@ fn apply_rewind<C, St>(
expected_head_entries: usize, expected_head_entries: usize,
) -> bool ) -> bool
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
match worker.rewind_to(target, expected_head_entries) { match worker.rewind_to(target, expected_head_entries) {
@@ -1834,7 +1926,7 @@ fn model_supports_image_attachments(model: &manifest::ModelManifest) -> bool {
fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting fn build_greeting<C, St>(worker: &Worker<C, St>) -> protocol::Greeting
where where
C: LlmClient, C: LlmClient + 'static,
St: Store, St: Store,
{ {
let manifest = worker.manifest(); let manifest = worker.manifest();
@@ -1950,11 +2042,13 @@ mod tests {
event_tx: broadcast::Sender<Event>, event_tx: broadcast::Sender<Event>,
cancel_tx: mpsc::Sender<()>, cancel_tx: mpsc::Sender<()>,
_cancel_rx: mpsc::Receiver<()>, _cancel_rx: mpsc::Receiver<()>,
pause_tx: mpsc::Sender<()>,
_pause_rx: mpsc::Receiver<()>,
shared_state: Arc<WorkerSharedState>, shared_state: Arc<WorkerSharedState>,
notify_buffer: NotifyBuffer, notify_buffer: NotifyBuffer,
spawned_registry: Arc<SpawnedWorkerRegistry>, spawned_registry: Arc<SpawnedWorkerRegistry>,
parent_socket_path: PathBuf, parent_socket_path: PathBuf,
_runtime_dir: Arc<RuntimeDir>, runtime_dir: Arc<RuntimeDir>,
_temp: TempDir, _temp: TempDir,
} }
@@ -1968,6 +2062,7 @@ mod tests {
let (method_tx, method_rx) = mpsc::channel::<Method>(16); let (method_tx, method_rx) = mpsc::channel::<Method>(16);
let (event_tx, _) = broadcast::channel::<Event>(16); let (event_tx, _) = broadcast::channel::<Event>(16);
let (cancel_tx, cancel_rx) = mpsc::channel::<()>(1); let (cancel_tx, cancel_rx) = mpsc::channel::<()>(1);
let (pause_tx, pause_rx) = mpsc::channel::<()>(1);
let shared_state = Arc::new(WorkerSharedState::new( let shared_state = Arc::new(WorkerSharedState::new(
"child-worker".to_string(), "child-worker".to_string(),
session_store::new_segment_id(), session_store::new_segment_id(),
@@ -1993,11 +2088,13 @@ mod tests {
event_tx, event_tx,
cancel_tx, cancel_tx,
_cancel_rx: cancel_rx, _cancel_rx: cancel_rx,
pause_tx,
_pause_rx: pause_rx,
shared_state, shared_state,
notify_buffer, notify_buffer,
spawned_registry, spawned_registry,
parent_socket_path, parent_socket_path,
_runtime_dir: runtime_dir, runtime_dir,
_temp: temp, _temp: temp,
} }
} }
@@ -2050,7 +2147,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"child-worker", "child-worker",
@@ -2071,6 +2171,44 @@ mod tests {
} }
} }
#[tokio::test]
async fn pause_waits_for_run_boundary_and_uses_safe_pause_channel() {
let mut env = make_env().await;
let method_tx = env._method_tx.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
method_tx.send(Method::Pause).await.expect("send pause");
});
let worker_future = async {
tokio::time::sleep(Duration::from_millis(100)).await;
Ok::<_, WorkerError>(WorkerRunResult::Finished)
};
let started_at = std::time::Instant::now();
let (status, shutdown) = drive_turn(
worker_future,
&mut env.method_rx,
&env.event_tx,
&env.cancel_tx,
&env.pause_tx,
&env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer,
None,
"child-worker",
&env.spawned_registry,
true,
)
.await;
assert_eq!(status, WorkerStatus::Paused);
assert!(!shutdown);
assert!(started_at.elapsed() >= Duration::from_millis(100));
assert!(env._pause_rx.try_recv().is_ok());
assert!(env._cancel_rx.try_recv().is_err());
}
#[tokio::test] #[tokio::test]
async fn non_parent_originated_finished_stays_silent() { async fn non_parent_originated_finished_stays_silent() {
let mut env = make_env().await; let mut env = make_env().await;
@@ -2082,7 +2220,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"child-worker", "child-worker",
@@ -2117,7 +2258,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"child-worker", "child-worker",
@@ -2158,7 +2302,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"child-worker", "child-worker",
@@ -2197,7 +2344,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"parent", "parent",
@@ -2233,7 +2383,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"parent", "parent",
@@ -2267,7 +2420,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"parent", "parent",
@@ -2300,7 +2456,10 @@ mod tests {
&mut env.method_rx, &mut env.method_rx,
&env.event_tx, &env.event_tx,
&env.cancel_tx, &env.cancel_tx,
&env.pause_tx,
&env.shared_state, &env.shared_state,
&env.runtime_dir,
None,
&env.notify_buffer, &env.notify_buffer,
Some(&env.parent_socket_path), Some(&env.parent_socket_path),
"child-worker", "child-worker",
+2 -2
View File
@@ -1795,9 +1795,9 @@ impl FeatureRegistryBuilder {
} }
/// Install modules into the existing Engine tool path and hook builder. /// Install modules into the existing Engine tool path and hook builder.
pub(crate) fn install_into_engine<C: LlmClient>( pub(crate) fn install_into_engine<C: LlmClient, A>(
self, self,
worker: &mut Engine<C, Mutable>, worker: &mut Engine<C, Mutable, A>,
hook_builder: &mut HookRegistryBuilder, hook_builder: &mut HookRegistryBuilder,
) -> FeatureRegistryInstallReport { ) -> FeatureRegistryInstallReport {
let mut pending_tools = Vec::new(); let mut pending_tools = Vec::new();
@@ -6,7 +6,9 @@ use memory::backend::{
MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation, MemoryBackendOperation, MemoryBackendOperationResult, MemoryStageCandidateOperation,
}; };
use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence}; use memory::extract::{CandidateKind, ExtractedCandidate, StagingEvidence};
use memory::schema::{EvidenceKind, SourceEvidenceRef, SourceRef}; use memory::schema::{
EvidenceKind, EvidenceOrigin, EvidenceOriginKind, SourceEvidenceRef, SourceRef,
};
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::Deserialize; use serde::Deserialize;
@@ -174,17 +176,29 @@ impl Tool for StageMemoryCandidateTool {
"StageMemoryCandidate requires at least one entry_ref".to_string(), "StageMemoryCandidate requires at least one entry_ref".to_string(),
)); ));
} }
let mut evidence = Vec::with_capacity(params.entry_refs.len()); let mut entries = Vec::with_capacity(params.entry_refs.len());
let mut source_refs = Vec::with_capacity(params.entry_refs.len());
for entry_ref in &params.entry_refs { for entry_ref in &params.entry_refs {
let projection = self.state.view.evidence_for(entry_ref).ok_or_else(|| { entries.push(self.state.view.evidence_for(entry_ref).ok_or_else(|| {
ToolError::InvalidArgument(format!( ToolError::InvalidArgument(format!(
"unknown SessionEntryRef {entry_ref:?} for this extraction capture" "unknown SessionEntryRef {entry_ref:?} for this extraction capture"
)) ))
})?; })?);
evidence.push(staging_evidence(&projection));
source_refs.push(source_evidence_ref(&projection));
} }
if matches!(params.kind, CandidateKind::Preference)
&& entries.iter().any(|entry| {
!matches!(
entry.origin,
crate::WorkerHistoryProvenance::HumanInput { .. }
)
})
{
return Err(ToolError::InvalidArgument(
"preference candidates require exclusively HumanInput evidence; model, Worker, Flow, backend, derived, and legacy-unknown origins are not preference authority"
.to_string(),
));
}
let evidence = entries.iter().map(staging_evidence).collect();
let source_refs = entries.iter().map(source_evidence_ref).collect();
let candidate = ExtractedCandidate { let candidate = ExtractedCandidate {
kind: params.kind, kind: params.kind,
claim: params.claim, claim: params.claim,
@@ -310,11 +324,65 @@ fn evidence_kind(entry: &SessionEntryEvidence) -> EvidenceKind {
} }
} }
fn evidence_origin(origin: &crate::WorkerHistoryProvenance) -> EvidenceOrigin {
use crate::WorkerHistoryProvenance as Origin;
let mut evidence = EvidenceOrigin {
kind: EvidenceOriginKind::LegacyUnknown,
account_id: None,
workspace_id: None,
runtime_id: None,
worker_id: None,
flow_selector: None,
flow_definition_id: None,
flow_definition_revision: None,
};
match origin {
Origin::HumanInput { account_id } => {
evidence.kind = EvidenceOriginKind::HumanInput;
evidence.account_id = Some(account_id.clone());
}
Origin::WorkerInput { actor } => {
evidence.kind = EvidenceOriginKind::WorkerInput;
evidence.workspace_id = actor.workspace_id.clone();
evidence.runtime_id = actor.runtime_id.clone();
evidence.worker_id = Some(actor.worker_id.clone());
}
Origin::FlowInstruction {
selector,
definition_id,
definition_revision,
..
} => {
evidence.kind = EvidenceOriginKind::FlowInstruction;
evidence.flow_selector = Some(selector.clone());
evidence.flow_definition_id = Some(definition_id.clone());
evidence.flow_definition_revision = Some(*definition_revision);
}
Origin::BackendInstruction { .. } => evidence.kind = EvidenceOriginKind::BackendInstruction,
Origin::ModelOutput { worker } => {
evidence.kind = EvidenceOriginKind::ModelOutput;
evidence.workspace_id = worker.workspace_id.clone();
evidence.runtime_id = worker.runtime_id.clone();
evidence.worker_id = Some(worker.worker_id.clone());
}
Origin::ToolOutput { worker } => {
evidence.kind = EvidenceOriginKind::ToolOutput;
evidence.workspace_id = worker.workspace_id.clone();
evidence.runtime_id = worker.runtime_id.clone();
evidence.worker_id = Some(worker.worker_id.clone());
}
Origin::DerivedSummary => evidence.kind = EvidenceOriginKind::DerivedSummary,
Origin::LegacyUnknown => evidence.kind = EvidenceOriginKind::LegacyUnknown,
}
evidence
}
fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence { fn staging_evidence(entry: &SessionEntryEvidence) -> StagingEvidence {
StagingEvidence { StagingEvidence {
id: entry.entry_ref.to_string(), id: entry.entry_ref.to_string(),
kind: evidence_kind(entry), kind: evidence_kind(entry),
entry_range: Some(entry.entry_range), entry_range: Some(entry.entry_range),
origin: Some(evidence_origin(&entry.origin)),
excerpt: Some(entry.excerpt.clone()), excerpt: Some(entry.excerpt.clone()),
summary: Some(entry.summary.clone()), summary: Some(entry.summary.clone()),
} }
@@ -325,6 +393,7 @@ fn source_evidence_ref(entry: &SessionEntryEvidence) -> SourceEvidenceRef {
segment_id: Some(entry.segment_id.clone()), segment_id: Some(entry.segment_id.clone()),
entry_range: Some(entry.entry_range), entry_range: Some(entry.entry_range),
evidence_id: Some(entry.entry_ref.to_string()), evidence_id: Some(entry.entry_ref.to_string()),
origin: Some(evidence_origin(&entry.origin)),
evidence_kind: Some(evidence_kind(entry)), evidence_kind: Some(evidence_kind(entry)),
label: Some(entry.label.clone()), label: Some(entry.label.clone()),
summary: Some(entry.summary.clone()), summary: Some(entry.summary.clone()),
@@ -432,6 +501,15 @@ mod tests {
assert!(input.contains("StageMemoryCandidate.entry_refs")); assert!(input.contains("StageMemoryCandidate.entry_refs"));
} }
#[test]
fn human_origin_projects_account_authority_into_evidence() {
let origin = evidence_origin(&crate::WorkerHistoryProvenance::HumanInput {
account_id: "account-1".into(),
});
assert_eq!(origin.kind, EvidenceOriginKind::HumanInput);
assert_eq!(origin.account_id.as_deref(), Some("account-1"));
}
#[test] #[test]
fn backend_input_failures_remain_invalid_argument_tool_errors() { fn backend_input_failures_remain_invalid_argument_tool_errors() {
let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend( let backend = map_memory_stage_error(WorkspaceMemoryBackendError::Backend(
@@ -445,6 +523,19 @@ mod tests {
assert!(matches!(http, ToolError::InvalidArgument(_))); assert!(matches!(http, ToolError::InvalidArgument(_)));
} }
#[tokio::test]
async fn preference_rejects_legacy_unknown_before_backend_mutation() {
let tool = StageMemoryCandidateTool { state: state() };
let error = tool
.execute(
r#"{"kind":"preference","claim":"claim","why_useful":"useful","entry_refs":["E00000000"]}"#,
agen::tool::ToolExecutionContext::direct(),
)
.await
.unwrap_err();
assert!(format!("{error:?}").contains("exclusively HumanInput evidence"));
}
#[tokio::test] #[tokio::test]
async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() { async fn stage_rejects_entry_ref_outside_capture_before_backend_mutation() {
let tool = StageMemoryCandidateTool { state: state() }; let tool = StageMemoryCandidateTool { state: state() };
+175 -4
View File
@@ -62,8 +62,9 @@ impl WorkspaceHttpObjectiveBackend {
.await .await
.map_err(backend_error)?; .map_err(backend_error)?;
let response = project_objective_detail(response).map_err(ToolError::ExecutionFailed)?; let response = project_objective_detail(response).map_err(ToolError::ExecutionFailed)?;
let objective_ref = response.objective_ref().to_string();
Ok(ToolOutput { Ok(ToolOutput {
summary: format!("Read objective {id}"), summary: format!("Read objective {objective_ref}"),
content: Some(serde_json::to_string_pretty(&response).map_err(decode_error)?), content: Some(serde_json::to_string_pretty(&response).map_err(decode_error)?),
attachments: Vec::new(), attachments: Vec::new(),
}) })
@@ -146,6 +147,7 @@ impl WorkspaceHttpObjectiveBackend {
async fn link_ticket(&self, input: ObjectiveLinkTicketInput) -> Result<ToolOutput, ToolError> { async fn link_ticket(&self, input: ObjectiveLinkTicketInput) -> Result<ToolOutput, ToolError> {
let id = validate_id(&input.id, "ObjectiveLinkTicket")?; let id = validate_id(&input.id, "ObjectiveLinkTicket")?;
let ticket_id = validate_id(&input.ticket_id, "ObjectiveLinkTicket")?; let ticket_id = validate_id(&input.ticket_id, "ObjectiveLinkTicket")?;
let ticket_resource_key = self.ticket_resource_key(ticket_id).await?;
let url = format!("{}/ticket-links", self.objective_url(id)); let url = format!("{}/ticket-links", self.objective_url(id));
let response = send_json::<ObjectiveLinkTicketRequest, ObjectiveDetail>( let response = send_json::<ObjectiveLinkTicketRequest, ObjectiveDetail>(
self.client.as_ref(), self.client.as_ref(),
@@ -159,7 +161,7 @@ impl WorkspaceHttpObjectiveBackend {
.map_err(backend_error)?; .map_err(backend_error)?;
Ok(objective_output( Ok(objective_output(
format!( format!(
"Linked ticket {ticket_id} to objective {}", "Linked ticket {ticket_resource_key} to objective {}",
&response.resource_key &response.resource_key
), ),
response, response,
@@ -172,19 +174,46 @@ impl WorkspaceHttpObjectiveBackend {
) -> Result<ToolOutput, ToolError> { ) -> Result<ToolOutput, ToolError> {
let id = validate_id(&input.id, "ObjectiveUnlinkTicket")?; let id = validate_id(&input.id, "ObjectiveUnlinkTicket")?;
let ticket_id = validate_id(&input.ticket_id, "ObjectiveUnlinkTicket")?; let ticket_id = validate_id(&input.ticket_id, "ObjectiveUnlinkTicket")?;
let ticket_resource_key = self.ticket_resource_key(ticket_id).await?;
let url = format!("{}/ticket-links/{}", self.objective_url(id), ticket_id); let url = format!("{}/ticket-links/{}", self.objective_url(id), ticket_id);
let response = delete_json::<ObjectiveDetail>(self.client.as_ref(), &url) let response = delete_json::<ObjectiveDetail>(self.client.as_ref(), &url)
.await .await
.map_err(backend_error)?; .map_err(backend_error)?;
Ok(objective_output( Ok(objective_output(
format!( format!(
"Unlinked ticket {ticket_id} from objective {}", "Unlinked ticket {ticket_resource_key} from objective {}",
&response.resource_key &response.resource_key
), ),
response, response,
)?) )?)
} }
async fn ticket_resource_key(&self, ticket_reference: &str) -> Result<String, ToolError> {
let workspace_id = self.client.workspace_id().unwrap_or_default();
let response: serde_json::Value = decode_response(
self.client
.execute(WorkspaceRequest::get(format!(
"/api/w/{workspace_id}/tickets/{ticket_reference}"
)))
.map_err(WorkspaceObjectiveBackendError::from)
.map_err(backend_error)?,
)
.map_err(backend_error)?;
response
.get("resource_key")
.or_else(|| {
response
.get("meta")
.and_then(|meta| meta.get("resource_key"))
})
.and_then(serde_json::Value::as_str)
.filter(|key| is_canonical_resource_key(key, "T-"))
.map(ToOwned::to_owned)
.ok_or_else(|| {
ToolError::ExecutionFailed("required T- human key is unavailable".to_string())
})
}
fn objective_url(&self, id: &str) -> String { fn objective_url(&self, id: &str) -> String {
let workspace_id = self.client.workspace_id().unwrap_or_default(); let workspace_id = self.client.workspace_id().unwrap_or_default();
format!("/api/w/{workspace_id}/objectives/{id}") format!("/api/w/{workspace_id}/objectives/{id}")
@@ -257,8 +286,14 @@ fn decode_response<T: for<'de> Deserialize<'de>>(
serde_json::from_str(&response.body).map_err(Into::into) serde_json::from_str(&response.body).map_err(Into::into)
} }
fn is_canonical_resource_key(resource_key: &str, prefix: &str) -> bool {
resource_key.strip_prefix(prefix).is_some_and(|sequence| {
!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit())
})
}
fn objective_output(summary: String, response: ObjectiveDetail) -> Result<ToolOutput, ToolError> { fn objective_output(summary: String, response: ObjectiveDetail) -> Result<ToolOutput, ToolError> {
if !response.resource_key.starts_with("O-") { if !is_canonical_resource_key(&response.resource_key, "O-") {
return Err(ToolError::ExecutionFailed( return Err(ToolError::ExecutionFailed(
"required O- human key is unavailable".to_string(), "required O- human key is unavailable".to_string(),
)); ));
@@ -624,6 +659,11 @@ struct ObjectiveDetail {
mod tests { mod tests {
use super::*; use super::*;
use agen::tool::ToolDefinition; use agen::tool::ToolDefinition;
use std::{
io::{Read, Write},
net::TcpListener,
thread,
};
fn tool_names(definitions: Vec<ToolDefinition>) -> Vec<String> { fn tool_names(definitions: Vec<ToolDefinition>) -> Vec<String> {
let mut names = definitions let mut names = definitions
@@ -670,4 +710,135 @@ mod tests {
let link = link_ticket_schema(); let link = link_ticket_schema();
assert_eq!(link["required"], json!(["id", "ticket_id"])); assert_eq!(link["required"], json!(["id", "ticket_id"]));
} }
#[tokio::test(flavor = "multi_thread")]
async fn objective_show_summary_uses_projected_human_key() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let server = thread::spawn(move || {
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(
request.starts_with("POST /api/w/workspace/objectives/00001INTERNAL/show HTTP/1.1")
);
let body = serde_json::json!({
"id": "00001INTERNAL",
"resource_key": "O-3",
"title": "Objective",
"body": "Body",
"state": "active",
"created_at": null,
"updated_at": null,
"linked_ticket_summaries": [],
"events": [],
"event_page": {"next_cursor": null, "has_more": false}
})
.to_string();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
body.len(),
body
)
.unwrap();
});
let backend = WorkspaceHttpObjectiveBackend::new(Arc::new(
crate::worker::TestWorkspaceHttpClient::new("workspace", base_url),
));
let output = backend
.show(ShowObjectiveInput {
id: "00001INTERNAL".to_string(),
event_limit: None,
event_cursor: None,
})
.await
.unwrap();
server.join().unwrap();
assert_eq!(output.summary, "Read objective O-3");
assert!(!output.summary.contains("00001INTERNAL"));
assert!(!output.content.unwrap().contains("00001INTERNAL"));
}
#[tokio::test(flavor = "multi_thread")]
async fn objective_link_summaries_resolve_internal_ticket_ids_to_human_keys() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap());
let server = thread::spawn(move || {
for mutation in ["POST", "DELETE"] {
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(request.starts_with("GET /api/w/workspace/tickets/00001INTERNAL HTTP/1.1"));
let response_body = serde_json::json!({"resource_key": "T-7"}).to_string();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
response_body.len(),
response_body
)
.unwrap();
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(request.starts_with(&format!(
"{mutation} /api/w/workspace/objectives/O-3/ticket-links"
)));
let response_body = serde_json::json!({
"resource_key": "O-3",
"title": "Objective",
"state": "active"
})
.to_string();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
response_body.len(),
response_body
)
.unwrap();
}
});
let backend = WorkspaceHttpObjectiveBackend::new(Arc::new(
crate::worker::TestWorkspaceHttpClient::new("workspace", base_url),
));
let linked = backend
.link_ticket(ObjectiveLinkTicketInput {
id: "O-3".to_string(),
ticket_id: "00001INTERNAL".to_string(),
})
.await
.unwrap();
let unlinked = backend
.unlink_ticket(ObjectiveUnlinkTicketInput {
id: "O-3".to_string(),
ticket_id: "00001INTERNAL".to_string(),
})
.await
.unwrap();
server.join().unwrap();
for output in [linked, unlinked] {
assert!(output.summary.contains("T-7"));
assert!(!output.summary.contains("00001INTERNAL"));
assert!(!output.content.unwrap().contains("00001INTERNAL"));
}
}
#[test]
fn objective_output_rejects_noncanonical_human_keys() {
let response = ObjectiveDetail {
resource_key: "O-internal".to_string(),
title: "Objective".to_string(),
state: "active".to_string(),
};
assert!(objective_output("created".to_string(), response).is_err());
}
} }
@@ -89,18 +89,19 @@ impl Tool for SpawnTicketCoderTool {
let input: SpawnTicketCoderInput = serde_json::from_str(input_json).map_err(|error| { let input: SpawnTicketCoderInput = serde_json::from_str(input_json).map_err(|error| {
ToolError::InvalidArgument(format!("invalid {TOOL_NAME} input: {error}")) ToolError::InvalidArgument(format!("invalid {TOOL_NAME} input: {error}"))
})?; })?;
let ticket_id = authority_id(input.ticket_id, "ticket_id")?; let ticket_ref = authority_id(input.ticket_id, "ticket_id")?;
let workflow_state = self let ticket = self
.ticket_service .ticket_service
.workflow_state(&ticket_id) .ticket_handoff(&ticket_ref)
.map_err(|error| ToolError::ExecutionFailed(error.to_string()))?; .map_err(|error| ToolError::ExecutionFailed(error.to_string()))?;
if !matches!( if !matches!(
workflow_state, ticket.workflow_state,
ticket::TicketWorkflowState::Queued | ticket::TicketWorkflowState::InProgress ticket::TicketWorkflowState::Queued | ticket::TicketWorkflowState::InProgress
) { ) {
return Err(ToolError::ExecutionFailed(format!( return Err(ToolError::ExecutionFailed(format!(
"Ticket {ticket_id} must be queued or inprogress before spawning its Coder; current state is {}", "Ticket {} must be queued or inprogress before spawning its Coder; current state is {}",
workflow_state.as_str() ticket.resource_key,
ticket.workflow_state.as_str()
))); )));
} }
let call_id = non_empty(ctx.call_id, "tool call_id")?; let call_id = non_empty(ctx.call_id, "tool call_id")?;
@@ -115,14 +116,14 @@ impl Tool for SpawnTicketCoderTool {
)?, )?,
relative_cwd, relative_cwd,
profile: CODER_PROFILE.to_string(), profile: CODER_PROFILE.to_string(),
ticket_id: Some(ticket_id.clone()), ticket_id: Some(ticket.id.clone()),
operation_id: Some(format!("spawn-ticket-coder:{ticket_id}:{call_id}")), operation_id: Some(format!("spawn-ticket-coder:{}:{call_id}", ticket.id)),
display_name: format!("Coder · {ticket_id}"), display_name: format!("Coder · {}", ticket.resource_key),
initial_submit: vec![ initial_submit: vec![
Segment::Flow { Segment::Flow {
selector: CODER_FLOW.to_string(), selector: CODER_FLOW.to_string(),
}, },
Segment::text(format!("Implement Ticket {ticket_id}.")), Segment::text(format!("Implement Ticket {}.", ticket.resource_key)),
], ],
}) })
.await .await
@@ -134,7 +135,7 @@ impl Tool for SpawnTicketCoderTool {
))); )));
} }
Ok(ToolOutput { Ok(ToolOutput {
summary: format!("Spawned Coder for Ticket {ticket_id}"), summary: format!("Spawned Coder for Ticket {}", ticket.resource_key),
content: Some(response.body), content: Some(response.body),
attachments: Vec::new(), attachments: Vec::new(),
}) })
@@ -201,21 +202,31 @@ mod tests {
use crate::worker::{WorkspaceClientError, WorkspaceResponse}; use crate::worker::{WorkspaceClientError, WorkspaceResponse};
use super::*; use super::*;
use crate::feature::builtin::ticket::TicketHandoff;
#[derive(Default)] #[derive(Default)]
struct RecordingTicketService; struct RecordingTicketService;
impl TicketService for RecordingTicketService { impl TicketService for RecordingTicketService {
fn workflow_state(&self, _ticket_id: &str) -> Result<TicketWorkflowState, TicketError> { fn ticket_handoff(&self, ticket_ref: &str) -> Result<TicketHandoff, TicketError> {
Ok(TicketWorkflowState::Queued) assert_eq!(ticket_ref, "T-482");
Ok(TicketHandoff {
id: "00001KZXN51C7".to_string(),
resource_key: "T-482".to_string(),
workflow_state: TicketWorkflowState::Queued,
})
} }
} }
struct FixedTicketService(TicketWorkflowState); struct FixedTicketService(TicketWorkflowState);
impl TicketService for FixedTicketService { impl TicketService for FixedTicketService {
fn workflow_state(&self, _ticket_id: &str) -> Result<TicketWorkflowState, TicketError> { fn ticket_handoff(&self, _ticket_ref: &str) -> Result<TicketHandoff, TicketError> {
Ok(self.0) Ok(TicketHandoff {
id: "00001KZXN51C7".to_string(),
resource_key: "T-482".to_string(),
workflow_state: self.0,
})
} }
} }
@@ -247,7 +258,7 @@ mod tests {
}; };
tool.execute( tool.execute(
&serde_json::json!({ &serde_json::json!({
"ticket_id": "00001KZXN51C7", "ticket_id": "T-482",
"runtime_id": "runtime-1", "runtime_id": "runtime-1",
"working_directory_id": "workdir-1" "working_directory_id": "workdir-1"
}) })
@@ -265,16 +276,20 @@ mod tests {
request.operation_id.as_deref(), request.operation_id.as_deref(),
Some("spawn-ticket-coder:00001KZXN51C7:call-7") Some("spawn-ticket-coder:00001KZXN51C7:call-7")
); );
assert_eq!(request.display_name, "Coder · 00001KZXN51C7"); assert_eq!(request.display_name, "Coder · T-482");
assert_eq!( assert_eq!(
request.initial_submit, request.initial_submit,
vec![ vec![
Segment::Flow { Segment::Flow {
selector: CODER_FLOW.to_string() selector: CODER_FLOW.to_string()
}, },
Segment::text("Implement Ticket 00001KZXN51C7.") Segment::text("Implement Ticket T-482.")
] ]
); );
assert!(!request.display_name.contains("00001KZXN51C7"));
assert!(request.initial_submit.iter().all(|segment| {
!Segment::flatten_to_text(std::slice::from_ref(segment)).contains("00001KZXN51C7")
}));
} }
#[tokio::test] #[tokio::test]
@@ -63,7 +63,7 @@ pub(super) struct ModelObjectiveQueryResponse {
struct ModelObjectiveQueryItem { struct ModelObjectiveQueryItem {
objective: String, objective: String,
title: String, title: String,
summary: String, summary: Option<String>,
state: String, state: String,
created_at: Option<String>, created_at: Option<String>,
updated_at: Option<String>, updated_at: Option<String>,
@@ -84,6 +84,12 @@ pub(super) struct ModelObjectiveDetail {
event_page: ModelObjectiveEventPage, event_page: ModelObjectiveEventPage,
} }
impl ModelObjectiveDetail {
pub(super) fn objective_ref(&self) -> &str {
&self.objective
}
}
#[derive(Debug, Serialize)] #[derive(Debug, Serialize)]
struct ModelWorkerSummary { struct ModelWorkerSummary {
worker: String, worker: String,
@@ -331,7 +337,7 @@ fn project_objective_query_item(value: &Value) -> Result<ModelObjectiveQueryItem
Ok(ModelObjectiveQueryItem { Ok(ModelObjectiveQueryItem {
objective: human_ref(item, "resource_key", "O-")?, objective: human_ref(item, "resource_key", "O-")?,
title: string_field(item, "title")?, title: string_field(item, "title")?,
summary: string_field(item, "snippet")?, summary: optional_string(item, "snippet")?,
state: string_field(item, "state")?, state: string_field(item, "state")?,
created_at: optional_string(item, "created_at")?, created_at: optional_string(item, "created_at")?,
updated_at: optional_string(item, "updated_at")?, updated_at: optional_string(item, "updated_at")?,
@@ -649,7 +655,10 @@ fn human_ref(object: &Map<String, Value>, key: &str, prefix: &str) -> Result<Str
} }
fn validate_human_ref(value: String, prefix: &str) -> Result<String, String> { fn validate_human_ref(value: String, prefix: &str) -> Result<String, String> {
if value.starts_with(prefix) && value.len() > prefix.len() { let valid = value.strip_prefix(prefix).is_some_and(|sequence| {
!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit())
});
if valid {
Ok(value) Ok(value)
} else { } else {
Err(format!("required {prefix} human key is unavailable")) Err(format!("required {prefix} human key is unavailable"))
@@ -754,7 +763,7 @@ mod tests {
"created_at": null, "created_at": null,
"updated_at": null, "updated_at": null,
"matched_fields": [], "matched_fields": [],
"snippet": "Objective summary", "snippet": null,
"linked_ticket_count": 1, "linked_ticket_count": 1,
"linked_tickets": ["00001TICKETINTERNAL"], "linked_tickets": ["00001TICKETINTERNAL"],
"linked_ticket_keys": ["T-543"] "linked_ticket_keys": ["T-543"]
@@ -764,10 +773,18 @@ mod tests {
let objective_json = serde_json::to_string(&objective).expect("serialize Objective query"); let objective_json = serde_json::to_string(&objective).expect("serialize Objective query");
assert!(objective_json.contains("O-6")); assert!(objective_json.contains("O-6"));
assert!(objective_json.contains("T-543")); assert!(objective_json.contains("T-543"));
assert!(objective_json.contains("\"summary\":null"));
assert!(!objective_json.contains("00001OBJECTIVEINTERNAL")); assert!(!objective_json.contains("00001OBJECTIVEINTERNAL"));
assert!(!objective_json.contains("00001TICKETINTERNAL")); assert!(!objective_json.contains("00001TICKETINTERNAL"));
} }
#[test]
fn human_resource_projection_rejects_noncanonical_keys() {
for (key, prefix) in [("T-key", "T-"), ("O-", "O-"), ("W-1x", "W-")] {
assert!(validate_human_ref(key.to_string(), prefix).is_err());
}
}
#[test] #[test]
fn ticket_projection_fails_closed_without_worker_resource_key() { fn ticket_projection_fails_closed_without_worker_resource_key() {
let error = project_worker(&json!({"worker_resource_key": null})) let error = project_worker(&json!({"worker_resource_key": null}))
@@ -193,6 +193,7 @@ impl Tool for ShowOverviewTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"label": entry.label, "label": entry.label,
@@ -234,15 +235,16 @@ impl Tool for SearchEntriesTool {
.transpose()?; .transpose()?;
let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.as_deref().map(parse_entry_ref).transpose()?; let through = params.through.as_deref().map(parse_entry_ref).transpose()?;
let view = self.state.view();
if let (Some(from), Some(through)) = (&from, &through) { if let (Some(from), Some(through)) = (&from, &through) {
if from.source_index() > through.source_index() { if view.source_index_for_ref(from) > view.source_index_for_ref(through) {
return Err(ToolError::InvalidArgument( return Err(ToolError::InvalidArgument(
"SearchEntries from must not be after through".to_string(), "SearchEntries from must not be after through".to_string(),
)); ));
} }
} }
let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT); let limit = bounded_limit(params.limit, DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT);
let hits = self.state.view().search(&SearchOptions { let hits = view.search(&SearchOptions {
query: params.query, query: params.query,
kind, kind,
tool_part, tool_part,
@@ -318,6 +320,7 @@ impl Tool for ReadEntryTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
+180 -25
View File
@@ -286,7 +286,20 @@ pub const TICKET_SERVICE_ID: &str = "ticket.authority";
const TICKET_SERVICE_VERSION: &str = "1"; const TICKET_SERVICE_VERSION: &str = "1";
pub trait TicketService: Send + Sync { pub trait TicketService: Send + Sync {
fn workflow_state(&self, ticket_id: &str) -> Result<TicketWorkflowState, TicketError>; fn ticket_handoff(&self, ticket_ref: &str) -> Result<TicketHandoff, TicketError>;
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TicketHandoff {
pub id: String,
pub resource_key: String,
pub workflow_state: TicketWorkflowState,
}
fn is_canonical_ticket_resource_key(resource_key: &str) -> bool {
resource_key.strip_prefix("T-").is_some_and(|sequence| {
!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit())
})
} }
struct BackendTicketService { struct BackendTicketService {
@@ -294,10 +307,18 @@ struct BackendTicketService {
} }
impl TicketService for BackendTicketService { impl TicketService for BackendTicketService {
fn workflow_state(&self, ticket_id: &str) -> Result<TicketWorkflowState, TicketError> { fn ticket_handoff(&self, ticket_ref: &str) -> Result<TicketHandoff, TicketError> {
self.backend let ticket = self.backend.show(ticket_ref.into())?;
.show(ticket_id.into()) let resource_key = ticket
.map(|ticket| ticket.meta.workflow_state) .meta
.resource_key
.filter(|key| is_canonical_ticket_resource_key(key))
.ok_or_else(|| TicketError::Conflict("ticket resource key is unavailable".into()))?;
Ok(TicketHandoff {
id: ticket.meta.id,
resource_key,
workflow_state: ticket.meta.workflow_state,
})
} }
} }
@@ -757,7 +778,7 @@ impl WorkspaceHttpTicketBackend {
if let Some(resource_key) = object if let Some(resource_key) = object
.get("resource_key") .get("resource_key")
.and_then(Value::as_str) .and_then(Value::as_str)
.filter(|key| key.starts_with("T-")) .filter(|key| is_canonical_ticket_resource_key(key))
.map(ToOwned::to_owned) .map(ToOwned::to_owned)
&& object.contains_key("id") && object.contains_key("id")
{ {
@@ -768,6 +789,32 @@ impl WorkspaceHttpTicketBackend {
} }
} }
fn resolve_ticket_resource_key(
client: Arc<dyn WorkspaceClient>,
base: &str,
reference: &TicketIdOrSlug,
) -> TicketResult<String> {
let response: Value = Self::request(
client,
WorkspaceRequestMethod::Get,
format!("{base}/{}", Self::ticket_path(reference)),
None,
)?;
response
.get("resource_key")
.or_else(|| {
response
.get("meta")
.and_then(|meta| meta.get("resource_key"))
})
.and_then(Value::as_str)
.filter(|key| is_canonical_ticket_resource_key(key))
.map(ToOwned::to_owned)
.ok_or_else(|| {
TicketError::Conflict("required Ticket human key is unavailable".to_string())
})
}
fn request_unit( fn request_unit(
client: Arc<dyn WorkspaceClient>, client: Arc<dyn WorkspaceClient>,
method: WorkspaceRequestMethod, method: WorkspaceRequestMethod,
@@ -829,12 +876,22 @@ impl WorkspaceHttpTicketBackend {
Ok(TicketBackendOperationResult::Tickets(tickets)) Ok(TicketBackendOperationResult::Tickets(tickets))
} }
TicketBackendOperation::Show { id } => { TicketBackendOperation::Show { id } => {
let ticket = Self::request( let ticket: Ticket = Self::request(
client, client,
WorkspaceRequestMethod::Get, WorkspaceRequestMethod::Get,
format!("{base}/{}/record", Self::ticket_path(&id)), format!("{base}/{}/record", Self::ticket_path(&id)),
None, None,
)?; )?;
if !ticket
.meta
.resource_key
.as_deref()
.is_some_and(is_canonical_ticket_resource_key)
{
return Err(TicketError::Conflict(
"required Ticket human key is unavailable".to_string(),
));
}
Ok(TicketBackendOperationResult::Ticket(ticket)) Ok(TicketBackendOperationResult::Ticket(ticket))
} }
TicketBackendOperation::Create { input } => { TicketBackendOperation::Create { input } => {
@@ -937,12 +994,13 @@ impl WorkspaceHttpTicketBackend {
})?), })?),
), ),
TicketBackendOperation::AddTicketRelation { id, relation } => { TicketBackendOperation::AddTicketRelation { id, relation } => {
let source_reference = match &id { let source_resource_key =
TicketIdOrSlug::Id(value) Self::resolve_ticket_resource_key(client.clone(), &base, &id)?;
| TicketIdOrSlug::Slug(value) let target_resource_key = Self::resolve_ticket_resource_key(
| TicketIdOrSlug::Query(value) => value.clone(), client.clone(),
}; &base,
let target_reference = relation.target.clone(); &TicketIdOrSlug::Id(relation.target.clone()),
)?;
let mut relation: TicketRelation = Self::request( let mut relation: TicketRelation = Self::request(
client, client,
WorkspaceRequestMethod::Post, WorkspaceRequestMethod::Post,
@@ -951,31 +1009,29 @@ impl WorkspaceHttpTicketBackend {
TicketError::Conflict(format!("serialize Ticket relation: {error}")) TicketError::Conflict(format!("serialize Ticket relation: {error}"))
})?), })?),
)?; )?;
relation.ticket_id = source_reference; relation.ticket_id = source_resource_key;
relation.target = target_reference; relation.target = target_resource_key;
relation.author = "workspace".to_string(); relation.author = "workspace".to_string();
Ok(TicketBackendOperationResult::Relation(relation)) Ok(TicketBackendOperationResult::Relation(relation))
} }
TicketBackendOperation::RemoveTicketRelation { id, kind, target } => { TicketBackendOperation::RemoveTicketRelation { id, kind, target } => {
let source_reference = match &id { let source_resource_key =
TicketIdOrSlug::Id(value) Self::resolve_ticket_resource_key(client.clone(), &base, &id)?;
| TicketIdOrSlug::Slug(value) let target_resource_key =
| TicketIdOrSlug::Query(value) => value.clone(), Self::resolve_ticket_resource_key(client.clone(), &base, &target)?;
};
let target = match target { let target = match target {
TicketIdOrSlug::Id(value) TicketIdOrSlug::Id(value)
| TicketIdOrSlug::Slug(value) | TicketIdOrSlug::Slug(value)
| TicketIdOrSlug::Query(value) => value, | TicketIdOrSlug::Query(value) => value,
}; };
let target_reference = target.clone();
let mut relation: TicketRelation = Self::request( let mut relation: TicketRelation = Self::request(
client, client,
WorkspaceRequestMethod::Delete, WorkspaceRequestMethod::Delete,
format!("{base}/{}/relations", Self::ticket_path(&id)), format!("{base}/{}/relations", Self::ticket_path(&id)),
Some(serde_json::json!({ "kind": kind, "target": target })), Some(serde_json::json!({ "kind": kind, "target": target })),
)?; )?;
relation.ticket_id = source_reference; relation.ticket_id = source_resource_key;
relation.target = target_reference; relation.target = target_resource_key;
relation.author = "workspace".to_string(); relation.author = "workspace".to_string();
Ok(TicketBackendOperationResult::Relation(relation)) Ok(TicketBackendOperationResult::Relation(relation))
} }
@@ -1804,11 +1860,102 @@ provider = "github"
server.join().unwrap(); server.join().unwrap();
} }
#[test]
fn workspace_http_backend_records_relation_with_authoritative_human_keys() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let server = thread::spawn(move || {
for (expected_path, resource_key) in [
("GET /api/w/workspace-a/tickets/01SOURCE HTTP/1.1", "T-1"),
("GET /api/w/workspace-a/tickets/01TARGET HTTP/1.1", "T-2"),
] {
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(request.starts_with(expected_path));
let body = serde_json::json!({"meta": {"resource_key": resource_key}}).to_string();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
body.len(), body
)
.unwrap();
}
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(
request.starts_with("POST /api/w/workspace-a/tickets/01SOURCE/relations HTTP/1.1")
);
let body = serde_json::to_string(&TicketRelation {
ticket_id: "01SOURCE".to_string(),
kind: TicketRelationKind::DependsOn,
target: "01TARGET".to_string(),
note: None,
author: "worker-internal".to_string(),
at: "2026-08-06T00:00:00Z".to_string(),
})
.unwrap();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
body.len(),
body
)
.unwrap();
});
let backend = WorkspaceHttpTicketBackend::new(Arc::new(
crate::worker::TestWorkspaceHttpClient::new("workspace-a", format!("http://{addr}")),
));
let relation = backend
.add_ticket_relation(
TicketIdOrSlug::Id("01SOURCE".to_string()),
NewTicketRelation {
kind: TicketRelationKind::DependsOn,
target: "01TARGET".to_string(),
note: None,
author: None,
},
)
.unwrap();
server.join().unwrap();
assert_eq!(relation.ticket_id, "T-1");
assert_eq!(relation.target, "T-2");
assert_eq!(relation.author, "workspace");
}
#[test] #[test]
fn workspace_http_backend_deletes_exact_ticket_relation() { fn workspace_http_backend_deletes_exact_ticket_relation() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let base_url = format!("http://{}", listener.local_addr().unwrap()); let base_url = format!("http://{}", listener.local_addr().unwrap());
let server = thread::spawn(move || { let server = thread::spawn(move || {
for (expected_path, resource_key) in [
("GET /api/w/workspace-a/tickets/01SOURCE HTTP/1.1", "T-1"),
("GET /api/w/workspace-a/tickets/01TARGET HTTP/1.1", "T-2"),
] {
let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap();
let request = String::from_utf8_lossy(&buffer[..len]);
assert!(request.starts_with(expected_path));
let response_body = serde_json::json!({
"meta": {"resource_key": resource_key}
})
.to_string();
write!(
stream,
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
response_body.len(),
response_body
)
.unwrap();
}
let (mut stream, _) = listener.accept().unwrap(); let (mut stream, _) = listener.accept().unwrap();
let mut buffer = [0_u8; 8192]; let mut buffer = [0_u8; 8192];
let len = stream.read(&mut buffer).unwrap(); let len = stream.read(&mut buffer).unwrap();
@@ -1849,8 +1996,16 @@ provider = "github"
.unwrap(); .unwrap();
server.join().unwrap(); server.join().unwrap();
assert_eq!(removed.ticket_id, "01SOURCE"); assert_eq!(removed.ticket_id, "T-1");
assert_eq!(removed.target, "01TARGET"); assert_eq!(removed.target, "T-2");
}
#[test]
fn ticket_handoff_accepts_only_canonical_ticket_resource_keys() {
assert!(is_canonical_ticket_resource_key("T-482"));
for invalid in ["", "00001KZVNXFNK", "T-", "T-key", "O-482"] {
assert!(!is_canonical_ticket_resource_key(invalid));
}
} }
#[test] #[test]
@@ -1,11 +1,12 @@
use std::sync::Arc; use std::sync::Arc;
#[cfg(test)]
use agen::Item; use agen::Item;
use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput}; use agen::tool::{Tool, ToolDefinition, ToolError, ToolMeta, ToolOutput};
use async_trait::async_trait; use async_trait::async_trait;
use schemars::JsonSchema; use schemars::JsonSchema;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
use session_store::collect_state; use session_store::{LogEntry, collect_state};
use super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService}; use super::manage_worker::{WORKER_CONTROL_SERVICE_ID, WorkerControlService};
use crate::feature::{ use crate::feature::{
@@ -60,7 +61,27 @@ pub struct WorkerObservationSubject {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub struct WorkerSessionCapture { pub struct WorkerSessionCapture {
pub segment_id: String, pub segment_id: String,
pub items: Vec<Item>, pub entries: Vec<agen::HistoryEntry<crate::SessionHistoryMetadata>>,
}
impl WorkerSessionCapture {
pub fn from_log_entries(
segment_id: impl Into<String>,
log_entries: &[LogEntry],
) -> Result<Self, String> {
let segment_id = segment_id.into();
let state = collect_state(log_entries);
let parsed_segment_id = segment_id.parse().unwrap_or_default();
let entries = crate::session_history::restore_history_entries(
state.session_id.unwrap_or_default(),
parsed_segment_id,
log_entries,
)?;
Ok(Self {
segment_id,
entries,
})
}
} }
#[derive(Debug, thiserror::Error)] #[derive(Debug, thiserror::Error)]
@@ -161,9 +182,17 @@ impl WorkerObservationProvider for WorkspaceClientWorkerObservationProvider {
}) })
.collect::<Result<Vec<session_store::LogEntry>, _>>()?; .collect::<Result<Vec<session_store::LogEntry>, _>>()?;
let state = collect_state(&entries); let state = collect_state(&entries);
let segment_id = response.segment_id;
let parsed_segment_id = segment_id.parse().unwrap_or_default();
let typed_entries = crate::session_history::restore_history_entries(
state.session_id.unwrap_or_default(),
parsed_segment_id,
&entries,
)
.map_err(WorkerObservationError::Unavailable)?;
Ok(WorkerSessionCapture { Ok(WorkerSessionCapture {
segment_id: response.segment_id, segment_id,
items: state.history, entries: typed_entries,
}) })
} }
} }
@@ -392,9 +421,15 @@ impl WorkerObservationProvider for SpawnedSubWorkerObservationProvider {
.ok_or(WorkerObservationError::NotFound)?; .ok_or(WorkerObservationError::NotFound)?;
let entries = record.session.entries(); let entries = record.session.entries();
let state = collect_state(&entries); let state = collect_state(&entries);
let typed_entries = crate::session_history::restore_history_entries(
state.session_id.unwrap_or_default(),
Default::default(),
&entries,
)
.map_err(WorkerObservationError::Unavailable)?;
Ok(WorkerSessionCapture { Ok(WorkerSessionCapture {
segment_id: format!("subworker:{name}"), segment_id: format!("subworker:{name}"),
items: state.history, entries: typed_entries,
}) })
} }
} }
@@ -508,6 +543,7 @@ impl Tool for ViewSessionOverviewTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"label": entry.label, "label": entry.label,
@@ -547,7 +583,7 @@ impl Tool for SearchSessionEntriesTool {
let from = params.from.as_deref().map(parse_entry_ref).transpose()?; let from = params.from.as_deref().map(parse_entry_ref).transpose()?;
let through = params.through.as_deref().map(parse_entry_ref).transpose()?; let through = params.through.as_deref().map(parse_entry_ref).transpose()?;
if let (Some(from), Some(through)) = (&from, &through) { if let (Some(from), Some(through)) = (&from, &through) {
if from.source_index() > through.source_index() { if view.source_index_for_ref(from) > view.source_index_for_ref(through) {
return Err(ToolError::InvalidArgument( return Err(ToolError::InvalidArgument(
"SearchSessionEntries from must not be after through".to_string(), "SearchSessionEntries from must not be after through".to_string(),
)); ));
@@ -573,6 +609,7 @@ impl Tool for SearchSessionEntriesTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
@@ -628,6 +665,7 @@ impl Tool for ReadSessionEntryTool {
.map(|entry| { .map(|entry| {
serde_json::json!({ serde_json::json!({
"entry_ref": entry.id, "entry_ref": entry.id,
"origin": entry.origin,
"entry_range": entry.entry_range, "entry_range": entry.entry_range,
"kind": entry.kind.as_str(), "kind": entry.kind.as_str(),
"tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()), "tool_part": entry.tool_part.map(|part| format!("{part:?}").to_lowercase()),
@@ -661,7 +699,10 @@ async fn latest_view(
.capture_worker_session(subject) .capture_worker_session(subject)
.await .await
.map_err(tool_error)?; .map_err(tool_error)?;
Ok(SessionCapture::new(capture.segment_id, capture.items)) Ok(SessionCapture::from_history_entries(
capture.segment_id,
capture.entries,
))
} }
fn parse_input<T: serde::de::DeserializeOwned>( fn parse_input<T: serde::de::DeserializeOwned>(
@@ -751,9 +792,23 @@ mod tests {
if subject != &granted_subject() { if subject != &granted_subject() {
return Err(WorkerObservationError::NotFound); return Err(WorkerObservationError::NotFound);
} }
let entries = self
.captures
.lock()
.unwrap()
.clone()
.into_iter()
.enumerate()
.map(|(index, item)| {
let mut metadata = crate::SessionHistoryMetadata::legacy_unknown();
metadata.entry_id =
session_store::LoggedSessionHistoryEntryId(format!("fake-{index:08}"));
agen::HistoryEntry::new(item, metadata)
})
.collect();
Ok(WorkerSessionCapture { Ok(WorkerSessionCapture {
segment_id: "segment".to_string(), segment_id: "segment".to_string(),
items: self.captures.lock().unwrap().clone(), entries,
}) })
} }
} }
@@ -796,7 +851,7 @@ mod tests {
let read = read_definition(provider.clone())().1; let read = read_definition(provider.clone())().1;
let hidden = read let hidden = read
.execute( .execute(
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"unauthorized"},"entry_ref":"E00000000"}"#, r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"unauthorized"},"entry_ref":"Efake-00000000"}"#,
agen::tool::ToolExecutionContext::direct(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
@@ -810,7 +865,7 @@ mod tests {
.push(message("a1", Role::Assistant, "second")); .push(message("a1", Role::Assistant, "second"));
let output = read let output = read
.execute( .execute(
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000000"}"#, r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"Efake-00000000"}"#,
agen::tool::ToolExecutionContext::direct(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
@@ -819,7 +874,7 @@ mod tests {
let output = read let output = read
.execute( .execute(
r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"E00000001"}"#, r#"{"subject":{"kind":"runtime_worker","runtime_id":"runtime-1","worker_id":"granted"},"entry_ref":"Efake-00000001"}"#,
agen::tool::ToolExecutionContext::direct(), agen::tool::ToolExecutionContext::direct(),
) )
.await .await
+177 -25
View File
@@ -10,7 +10,7 @@ use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use agen::timeline::event::UsageEvent; use agen::timeline::event::UsageEvent;
use agen::{Engine, llm_client::LlmClient}; use agen::{Engine, EngineError, llm_client::LlmClient};
use manifest::{Scope, WorkerManifest}; use manifest::{Scope, WorkerManifest};
use protocol::{Event, InFlightSnapshot, WorkerStatus}; use protocol::{Event, InFlightSnapshot, WorkerStatus};
use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry}; use session_store::{LogEntry, SegmentId, SessionId, Store, StoreError, TraceEntry};
@@ -55,7 +55,17 @@ pub(crate) struct InternalWorkerSpec {
pub input: String, pub input: String,
pub cache_key: Option<String>, pub cache_key: Option<String>,
pub max_turns: Option<u32>, pub max_turns: Option<u32>,
pub engine_configurator: Option<Box<dyn FnOnce(&mut Engine<Box<dyn LlmClient>>) + Send>>, pub engine_configurator: Option<
Box<
dyn FnOnce(
&mut Engine<
Box<dyn LlmClient>,
agen::state::Mutable,
crate::SessionHistoryMetadata,
>,
) + Send,
>,
>,
pub features: FeatureRegistryBuilder, pub features: FeatureRegistryBuilder,
pub required_tools: &'static [&'static str], pub required_tools: &'static [&'static str],
pub authority: InternalWorkerAuthority, pub authority: InternalWorkerAuthority,
@@ -124,7 +134,9 @@ where
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>)); let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
let usage_slot = last_usage.clone(); let usage_slot = last_usage.clone();
let mut engine = Engine::new(client).system_prompt(system_prompt); let mut engine =
Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client)
.system_prompt(system_prompt);
engine.on_usage(move |usage| { engine.on_usage(move |usage| {
if let Ok(mut slot) = usage_slot.lock() { if let Ok(mut slot) = usage_slot.lock() {
*slot = Some(usage.clone()); *slot = Some(usage.clone());
@@ -199,12 +211,28 @@ where
on_cancel_sender(worker.engine_mut().cancel_sender()); on_cancel_sender(worker.engine_mut().cancel_sender());
match worker.run_text(&input).await { match worker.run_text(&input).await {
Ok(lifecycle) => Ok(InternalWorkerResult { Ok(lifecycle @ WorkerRunResult::Finished)
| Ok(lifecycle @ WorkerRunResult::Paused)
| Ok(lifecycle @ WorkerRunResult::RolledBack) => Ok(InternalWorkerResult {
usage: last_usage.lock().ok().and_then(|slot| slot.clone()), usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity, identity,
lifecycle, lifecycle,
history_entries: store.entries_count(session_id, segment_id), history_entries: store.entries_count(session_id, segment_id),
}), }),
Ok(WorkerRunResult::LimitReached) => Err(InternalWorkerError {
source: WorkerError::Engine(EngineError::Aborted(
"internal Worker reached its turn limit".to_string(),
)),
usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity,
history_entries: store.entries_count(session_id, segment_id),
}),
Ok(WorkerRunResult::Interrupted { message, .. }) => Err(InternalWorkerError {
source: WorkerError::Engine(EngineError::Aborted(message)),
usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
identity,
history_entries: store.entries_count(session_id, segment_id),
}),
Err(source) => Err(InternalWorkerError { Err(source) => Err(InternalWorkerError {
source, source,
usage: last_usage.lock().ok().and_then(|slot| slot.clone()), usage: last_usage.lock().ok().and_then(|slot| slot.clone()),
@@ -232,6 +260,7 @@ impl Default for InternalWorkerVisibility {
pub(crate) enum InternalWorkerSessionStatus { pub(crate) enum InternalWorkerSessionStatus {
Idle, Idle,
Running, Running,
Paused,
Stopping, Stopping,
Stopped, Stopped,
Failed, Failed,
@@ -242,9 +271,10 @@ impl InternalWorkerSessionStatus {
match self { match self {
Self::Idle => 0, Self::Idle => 0,
Self::Running => 1, Self::Running => 1,
Self::Stopping => 2, Self::Paused => 2,
Self::Stopped => 3, Self::Stopping => 3,
Self::Failed => 4, Self::Stopped => 4,
Self::Failed => 5,
} }
} }
@@ -252,13 +282,35 @@ impl InternalWorkerSessionStatus {
match value { match value {
0 => Self::Idle, 0 => Self::Idle,
1 => Self::Running, 1 => Self::Running,
2 => Self::Stopping, 2 => Self::Paused,
3 => Self::Stopped, 3 => Self::Stopping,
4 => Self::Stopped,
_ => Self::Failed, _ => Self::Failed,
} }
} }
} }
fn classify_internal_turn_result(
result: Result<WorkerRunResult, WorkerError>,
) -> (InternalWorkerSessionStatus, Option<String>) {
match result {
Ok(WorkerRunResult::Finished) => (InternalWorkerSessionStatus::Idle, None),
Ok(WorkerRunResult::Paused) => (InternalWorkerSessionStatus::Paused, None),
Ok(WorkerRunResult::LimitReached) => (
InternalWorkerSessionStatus::Stopped,
Some("internal Worker reached its turn limit".to_string()),
),
Ok(WorkerRunResult::Interrupted { message, .. }) => {
(InternalWorkerSessionStatus::Stopped, Some(message))
}
Ok(WorkerRunResult::RolledBack) => (
InternalWorkerSessionStatus::Stopped,
Some("internal Worker run was cancelled before AI output".to_string()),
),
Err(error) => (InternalWorkerSessionStatus::Failed, Some(error.to_string())),
}
}
#[derive(Debug, thiserror::Error)] #[derive(Debug, thiserror::Error)]
pub(crate) enum InternalWorkerSessionError { pub(crate) enum InternalWorkerSessionError {
#[error("failed to build internal Worker session: {message}")] #[error("failed to build internal Worker session: {message}")]
@@ -353,10 +405,11 @@ impl InternalWorkerSessionHandle {
entries, entries,
status: match self.status() { status: match self.status() {
InternalWorkerSessionStatus::Running => WorkerStatus::Running, InternalWorkerSessionStatus::Running => WorkerStatus::Running,
InternalWorkerSessionStatus::Paused => WorkerStatus::Paused,
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle, InternalWorkerSessionStatus::Idle => WorkerStatus::Idle,
InternalWorkerSessionStatus::Stopping InternalWorkerSessionStatus::Stopping
| InternalWorkerSessionStatus::Stopped | InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => WorkerStatus::Paused, | InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped,
}, },
error: self.last_error.lock().unwrap().clone(), error: self.last_error.lock().unwrap().clone(),
in_flight, in_flight,
@@ -388,6 +441,7 @@ impl InternalWorkerSessionHandle {
.map_err( .map_err(
|current| match InternalWorkerSessionStatus::decode(current) { |current| match InternalWorkerSessionStatus::decode(current) {
InternalWorkerSessionStatus::Running InternalWorkerSessionStatus::Running
| InternalWorkerSessionStatus::Paused
| InternalWorkerSessionStatus::Stopping => InternalWorkerSessionError::Busy, | InternalWorkerSessionStatus::Stopping => InternalWorkerSessionError::Busy,
InternalWorkerSessionStatus::Stopped | InternalWorkerSessionStatus::Failed => { InternalWorkerSessionStatus::Stopped | InternalWorkerSessionStatus::Failed => {
InternalWorkerSessionError::Stopped InternalWorkerSessionError::Stopped
@@ -494,7 +548,9 @@ pub(crate) async fn spawn_internal_worker_session(
let last_usage = Arc::new(Mutex::new(None::<UsageEvent>)); let last_usage = Arc::new(Mutex::new(None::<UsageEvent>));
let usage_slot = last_usage.clone(); let usage_slot = last_usage.clone();
let mut engine = Engine::new(client).system_prompt(system_prompt); let mut engine =
Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client)
.system_prompt(system_prompt);
engine.on_usage(move |usage| { engine.on_usage(move |usage| {
if let Ok(mut slot) = usage_slot.lock() { if let Ok(mut slot) = usage_slot.lock() {
*slot = Some(usage.clone()); *slot = Some(usage.clone());
@@ -591,7 +647,9 @@ pub(crate) fn prepare_internal_worker_from_spec(
manifest.compaction = None; manifest.compaction = None;
manifest.memory = None; manifest.memory = None;
let mut engine = Engine::new(client).system_prompt(system_prompt); let mut engine =
Engine::<_, agen::state::Mutable, crate::SessionHistoryMetadata>::new_annotated(client)
.system_prompt(system_prompt);
engine.set_cache_key(cache_key); engine.set_cache_key(cache_key);
engine.set_max_turns(max_turns); engine.set_max_turns(max_turns);
if let Some(configure) = engine_configurator { if let Some(configure) = engine_configurator {
@@ -733,13 +791,7 @@ pub(crate) async fn prepare_internal_worker_session(
loop { loop {
tokio::select! { tokio::select! {
result = &mut run => { result = &mut run => {
let (turn_status, error) = match result { let (turn_status, error) = classify_internal_turn_result(result);
Ok(_) => (InternalWorkerSessionStatus::Idle, None),
Err(error) => (
InternalWorkerSessionStatus::Failed,
Some(error.to_string()),
),
};
actor_in_flight.clear(); actor_in_flight.clear();
status.store(turn_status.encode(), std::sync::atomic::Ordering::Release); status.store(turn_status.encode(), std::sync::atomic::Ordering::Release);
if let Some(message) = error { if let Some(message) = error {
@@ -748,11 +800,20 @@ pub(crate) async fn prepare_internal_worker_session(
code: protocol::ErrorCode::Internal, code: protocol::ErrorCode::Internal,
message, message,
}); });
} else {
let _ = event_tx.send(Event::Status {
status: WorkerStatus::Idle,
});
} }
let protocol_status = match turn_status {
InternalWorkerSessionStatus::Idle => WorkerStatus::Idle,
InternalWorkerSessionStatus::Paused => WorkerStatus::Paused,
InternalWorkerSessionStatus::Stopped
| InternalWorkerSessionStatus::Failed => WorkerStatus::Stopped,
InternalWorkerSessionStatus::Running
| InternalWorkerSessionStatus::Stopping => {
unreachable!("run completion cannot remain active")
}
};
let _ = event_tx.send(Event::Status {
status: protocol_status,
});
if let Some(callback) = &on_turn_end { if let Some(callback) = &on_turn_end {
callback(turn_status); callback(turn_status);
} }
@@ -766,7 +827,7 @@ pub(crate) async fn prepare_internal_worker_session(
let _ = (&mut run).await; let _ = (&mut run).await;
actor_in_flight.clear(); actor_in_flight.clear();
status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release); status.store(InternalWorkerSessionStatus::Stopped.encode(), std::sync::atomic::Ordering::Release);
let _ = event_tx.send(Event::Status { status: WorkerStatus::Paused }); let _ = event_tx.send(Event::Status { status: WorkerStatus::Stopped });
let _ = event_tx.send(Event::Shutdown); let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters(); state_changed.notify_waiters();
let _ = done.send(()); let _ = done.send(());
@@ -792,7 +853,7 @@ pub(crate) async fn prepare_internal_worker_session(
std::sync::atomic::Ordering::Release, std::sync::atomic::Ordering::Release,
); );
let _ = event_tx.send(Event::Status { let _ = event_tx.send(Event::Status {
status: WorkerStatus::Paused, status: WorkerStatus::Stopped,
}); });
let _ = event_tx.send(Event::Shutdown); let _ = event_tx.send(Event::Shutdown);
state_changed.notify_waiters(); state_changed.notify_waiters();
@@ -1102,6 +1163,26 @@ mod tests {
} }
} }
#[derive(Clone)]
struct FailingClient;
#[async_trait]
impl LlmClient for FailingClient {
fn clone_boxed(&self) -> Box<dyn LlmClient> {
Box::new(self.clone())
}
async fn stream(
&self,
_request: Request,
) -> Result<Pin<Box<dyn Stream<Item = Result<LlmEvent, ClientError>> + Send>>, ClientError>
{
Err(ClientError::Config(
"intentional internal failure".to_string(),
))
}
}
#[derive(Clone)] #[derive(Clone)]
struct CancelBeforeAiClient { struct CancelBeforeAiClient {
calls: Arc<AtomicUsize>, calls: Arc<AtomicUsize>,
@@ -1215,6 +1296,77 @@ permission = "write"
assert_eq!(result.identity.kind, "test"); assert_eq!(result.identity.kind, "test");
} }
#[test]
fn internal_turn_result_mapping_is_exhaustive() {
let cases = [
(
WorkerRunResult::Finished,
InternalWorkerSessionStatus::Idle,
false,
),
(
WorkerRunResult::Paused,
InternalWorkerSessionStatus::Paused,
false,
),
(
WorkerRunResult::LimitReached,
InternalWorkerSessionStatus::Stopped,
true,
),
(
WorkerRunResult::Interrupted {
code: protocol::ErrorCode::Internal,
message: "cancelled".to_string(),
},
InternalWorkerSessionStatus::Stopped,
true,
),
(
WorkerRunResult::RolledBack,
InternalWorkerSessionStatus::Stopped,
true,
),
];
for (result, expected_status, expects_error) in cases {
let (status, error) = classify_internal_turn_result(Ok(result));
assert_eq!(status, expected_status);
assert_eq!(error.is_some(), expects_error);
}
let (status, error) = classify_internal_turn_result(Err(WorkerError::Engine(
EngineError::Aborted("fatal".to_string()),
)));
assert_eq!(status, InternalWorkerSessionStatus::Failed);
assert!(error.is_some_and(|message| message.contains("fatal")));
}
#[tokio::test]
async fn fatal_internal_run_transitions_to_stopped_protocol_status() {
let calls = Arc::new(AtomicUsize::new(0));
let mut internal_spec = spec(calls, &[]);
internal_spec.client = Box::new(FailingClient);
let handle = spawn_internal_worker_session(internal_spec)
.await
.expect("spawn failing Internal Worker session");
assert_eq!(
handle.wait_until_idle().await,
InternalWorkerSessionStatus::Stopped
);
assert_eq!(handle.status(), InternalWorkerSessionStatus::Stopped);
assert_eq!(handle.protocol_snapshot().status, WorkerStatus::Stopped);
assert!(
handle
.last_error
.lock()
.unwrap()
.as_ref()
.is_some_and(|message| message.contains("intentional internal failure"))
);
}
#[tokio::test] #[tokio::test]
async fn session_accepts_follow_up_turns_and_stops_without_runtime_registration() { async fn session_accepts_follow_up_turns_and_stops_without_runtime_registration() {
let calls = Arc::new(AtomicUsize::new(0)); let calls = Arc::new(AtomicUsize::new(0));
+13 -2
View File
@@ -13,7 +13,7 @@
#[cfg(test)] #[cfg(test)]
use crate::prompt::catalog::PromptCatalog; use crate::prompt::catalog::PromptCatalog;
use agen::Item; use agen::{Item, ToolResultDisposition};
/// Build synthetic `Item::ToolResult` items for every unanswered /// Build synthetic `Item::ToolResult` items for every unanswered
/// `Item::ToolCall` in `history`, preserving order. /// `Item::ToolCall` in `history`, preserving order.
@@ -28,7 +28,16 @@ pub(crate) fn orphan_tool_result_closures(history: &[Item], summary: &str) -> Ve
for item in history { for item in history {
if let Item::ToolCall { call_id, .. } = item { if let Item::ToolCall { call_id, .. } = item {
if !answered.contains(call_id.as_str()) { if !answered.contains(call_id.as_str()) {
out.push(Item::tool_result(call_id.clone(), summary)); out.push(Item::tool_result_item_with_disposition_and_attachments(
call_id.clone(),
summary,
Some(
"Execution ended before completion could be confirmed. Completion and side effects are unknown."
.to_string(),
),
ToolResultDisposition::OutcomeUnknown,
Vec::new(),
));
} }
} }
} }
@@ -77,10 +86,12 @@ mod tests {
Item::ToolResult { Item::ToolResult {
call_id, call_id,
summary: got, summary: got,
disposition,
.. ..
} => { } => {
assert_eq!(call_id, "c1"); assert_eq!(call_id, "c1");
assert_eq!(got, &summary); assert_eq!(got, &summary);
assert_eq!(*disposition, ToolResultDisposition::OutcomeUnknown);
} }
other => panic!("expected ToolResult, got {other:?}"), other => panic!("expected ToolResult, got {other:?}"),
} }
+39 -2
View File
@@ -8,6 +8,7 @@
//! decisions (continue / skip / abort / pause). //! decisions (continue / skip / abort / pause).
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::VecDeque;
use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
@@ -33,7 +34,9 @@ use crate::hook::{
}; };
use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance}; use crate::ipc::notify_buffer::{NotifyBuffer, build_system_item_with_provenance};
use crate::prompt::catalog::PromptCatalog; use crate::prompt::catalog::PromptCatalog;
use crate::session_history::SessionHistoryMetadata;
use crate::worker::SystemItemCommitter; use crate::worker::SystemItemCommitter;
use agen::HistoryEntry;
use agen::token_counter::total_tokens; use agen::token_counter::total_tokens;
/// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`. /// Maximum number of bytes copied into `TurnEndInfo::final_text_preview`.
@@ -73,6 +76,7 @@ pub(crate) struct WorkerInterceptor {
/// worker. `None` in tests / `Worker::new` paths where no writer is /// worker. `None` in tests / `Worker::new` paths where no writer is
/// attached. /// attached.
log_writer: Option<Arc<dyn SystemItemCommitter>>, log_writer: Option<Arc<dyn SystemItemCommitter>>,
pending_committed_history: Arc<Mutex<VecDeque<HistoryEntry<SessionHistoryMetadata>>>>,
/// Next turn index assigned by `on_prompt_submit`. /// Next turn index assigned by `on_prompt_submit`.
next_turn_index: AtomicUsize, next_turn_index: AtomicUsize,
/// Tool calls observed in the current turn (reset on each new prompt). /// Tool calls observed in the current turn (reset on each new prompt).
@@ -80,6 +84,7 @@ pub(crate) struct WorkerInterceptor {
} }
impl WorkerInterceptor { impl WorkerInterceptor {
#[cfg(test)]
pub(crate) fn new( pub(crate) fn new(
registry: Arc<HookRegistry>, registry: Arc<HookRegistry>,
compact_state: Option<Arc<CompactState>>, compact_state: Option<Arc<CompactState>>,
@@ -88,6 +93,28 @@ impl WorkerInterceptor {
pending_attachments: Arc<Mutex<Vec<SystemItem>>>, pending_attachments: Arc<Mutex<Vec<SystemItem>>>,
prompts: Arc<ArcSwap<PromptCatalog>>, prompts: Arc<ArcSwap<PromptCatalog>>,
log_writer: Option<Arc<dyn SystemItemCommitter>>, log_writer: Option<Arc<dyn SystemItemCommitter>>,
) -> Self {
Self::new_with_history_queue(
registry,
compact_state,
usage_history,
pending_notifies,
pending_attachments,
prompts,
log_writer,
Arc::new(Mutex::new(VecDeque::new())),
)
}
pub(crate) fn new_with_history_queue(
registry: Arc<HookRegistry>,
compact_state: Option<Arc<CompactState>>,
usage_history: Option<Arc<Mutex<Vec<UsageRecord>>>>,
pending_notifies: NotifyBuffer,
pending_attachments: Arc<Mutex<Vec<SystemItem>>>,
prompts: Arc<ArcSwap<PromptCatalog>>,
log_writer: Option<Arc<dyn SystemItemCommitter>>,
pending_committed_history: Arc<Mutex<VecDeque<HistoryEntry<SessionHistoryMetadata>>>>,
) -> Self { ) -> Self {
Self { Self {
registry, registry,
@@ -99,6 +126,7 @@ impl WorkerInterceptor {
prompts, prompts,
prompt_workspace_id: None, prompt_workspace_id: None,
log_writer, log_writer,
pending_committed_history,
next_turn_index: AtomicUsize::new(0), next_turn_index: AtomicUsize::new(0),
tool_calls_this_turn: AtomicUsize::new(0), tool_calls_this_turn: AtomicUsize::new(0),
} }
@@ -125,7 +153,11 @@ impl WorkerInterceptor {
return Ok(()); return Ok(());
}; };
for item in items { for item in items {
writer.commit_system_item(item.clone())?; let entry = writer.commit_system_item(item.clone())?;
self.pending_committed_history
.lock()
.expect("pending committed history poisoned")
.push_back(entry);
} }
Ok(()) Ok(())
} }
@@ -507,7 +539,12 @@ mod tests {
&self, &self,
entry: session_store::LogEntry, entry: session_store::LogEntry,
) -> Result<(), session_store::StoreError> { ) -> Result<(), session_store::StoreError> {
if let session_store::LogEntry::SystemItem { item, .. } = entry { let item = match entry {
session_store::LogEntry::SystemItem { item, .. } => Some(item),
session_store::LogEntry::AnnotatedSystemItem { entry, .. } => Some(entry.item),
_ => None,
};
if let Some(item) = item {
self.committed self.committed
.lock() .lock()
.expect("committed system-item list poisoned") .expect("committed system-item list poisoned")
+8 -2
View File
@@ -29,15 +29,21 @@ pub fn subscribe_worker_protocol_session(handle: &WorkerHandle) -> WorkerProtoco
pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> { pub fn live_log_entry_event(entry: LogEntry) -> Option<Event> {
match entry { match entry {
LogEntry::SegmentStart { .. } => { entry @ (LogEntry::SegmentStart { .. } | LogEntry::AnnotatedSegmentStart { .. }) => {
let value = serde_json::to_value(&entry).expect("LogEntry is Serialize"); let value = serde_json::to_value(&entry).expect("LogEntry is Serialize");
Some(Event::SegmentRotated { entry: value }) Some(Event::SegmentRotated { entry: value })
} }
LogEntry::UserInput { segments, .. } => Some(Event::UserMessage { segments }), LogEntry::UserInput { segments, .. } | LogEntry::AnnotatedUserInput { segments, .. } => {
Some(Event::UserMessage { segments })
}
LogEntry::SystemItem { item, .. } => { LogEntry::SystemItem { item, .. } => {
let value = serde_json::to_value(&item).expect("SystemItem is Serialize"); let value = serde_json::to_value(&item).expect("SystemItem is Serialize");
Some(Event::SystemItem { item: value }) Some(Event::SystemItem { item: value })
} }
LogEntry::AnnotatedSystemItem { entry, .. } => {
let value = serde_json::to_value(&entry.item).expect("SystemItem is Serialize");
Some(Event::SystemItem { item: value })
}
LogEntry::Invoke { trigger, .. } => Some(Event::InvokeStart { kind: trigger }), LogEntry::Invoke { trigger, .. } => Some(Event::InvokeStart { kind: trigger }),
other => { other => {
// `SegmentLogSink::is_live_relevant` keeps non-live-relevant // `SegmentLogSink::is_live_relevant` keeps non-live-relevant
+8 -2
View File
@@ -12,6 +12,7 @@ pub mod prompt;
pub mod runtime; pub mod runtime;
pub mod segment_log_sink; pub mod segment_log_sink;
mod session_capture; mod session_capture;
mod session_history;
pub mod shared_state; pub mod shared_state;
mod shutdown_after_idle; mod shutdown_after_idle;
pub mod skill; pub mod skill;
@@ -33,14 +34,19 @@ pub use manifest::{
}; };
pub use model_client::{ProviderError, build_client}; pub use model_client::{ProviderError, build_client};
pub use prompt::catalog::{ pub use prompt::catalog::{
CatalogError, EffectivePromptCatalog, PromptCatalog, WorkerPrompt, WorkspacePromptProjection, CatalogError, EffectivePromptCatalog, OrchestratorQueueAttentionContext,
prompt_schema_source, OrchestratorQueueAttentionPrompt, OrchestratorQueueAttentionTicket, PromptCatalog,
WorkerPrompt, WorkspacePromptProjection, prompt_schema_source,
}; };
pub use prompt::source::PromptCatalogSource; pub use prompt::source::PromptCatalogSource;
pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate}; pub use prompt::system::{SystemPromptContext, SystemPromptError, SystemPromptTemplate};
pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus}; pub use protocol::{ErrorCode, Event, Method, TurnResult, WorkerStatus};
pub use runtime::dir::RuntimeDir; pub use runtime::dir::RuntimeDir;
pub use segment_log_sink::SegmentLogSink; pub use segment_log_sink::SegmentLogSink;
pub use session_history::{
SessionHistoryDerivation, SessionHistoryEntryId, SessionHistoryMetadata,
WorkerHistoryProvenance, WorkerSubjectSnapshot,
};
pub use shared_state::WorkerSharedState; pub use shared_state::WorkerSharedState;
pub use worker::{ pub use worker::{
LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError, LocalWorkingDirectory, WORKER_INPUT_SUBMISSION_EXTENSION_DOMAIN, Worker, WorkerError,
+1 -1
View File
@@ -34,7 +34,7 @@ impl PermissionHook {
} }
} }
impl<C: LlmClient, St: Store> Worker<C, St> { impl<C: LlmClient + 'static, St: Store> Worker<C, St> {
pub(crate) fn apply_permissions_from_manifest(&mut self) { pub(crate) fn apply_permissions_from_manifest(&mut self) {
let Some(permissions) = self.manifest().permissions.clone() else { let Some(permissions) = self.manifest().permissions.clone() else {
return; return;
+155 -3
View File
@@ -102,7 +102,6 @@ pub enum WorkerPrompt {
AgentsMdSection, AgentsMdSection,
ResidentMemorySummarySection, ResidentMemorySummarySection,
WorkerOrchestrationGuidanceSection, WorkerOrchestrationGuidanceSection,
TicketEventCompanionNotice,
SubWorkerSpawnToolDescription, SubWorkerSpawnToolDescription,
} }
@@ -122,7 +121,6 @@ impl WorkerPrompt {
Self::WorkerOrchestrationGuidanceSection => { Self::WorkerOrchestrationGuidanceSection => {
"internal.worker_orchestration_guidance_section" "internal.worker_orchestration_guidance_section"
} }
Self::TicketEventCompanionNotice => "worker.ticket_event_companion_notice",
Self::SubWorkerSpawnToolDescription => "internal.sub_worker_spawn_tool_description", Self::SubWorkerSpawnToolDescription => "internal.sub_worker_spawn_tool_description",
} }
} }
@@ -139,13 +137,97 @@ impl WorkerPrompt {
WorkerPrompt::AgentsMdSection, WorkerPrompt::AgentsMdSection,
WorkerPrompt::ResidentMemorySummarySection, WorkerPrompt::ResidentMemorySummarySection,
WorkerPrompt::WorkerOrchestrationGuidanceSection, WorkerPrompt::WorkerOrchestrationGuidanceSection,
WorkerPrompt::TicketEventCompanionNotice,
WorkerPrompt::SubWorkerSpawnToolDescription, WorkerPrompt::SubWorkerSpawnToolDescription,
]; ];
} }
/// Model-visible queued Ticket projection shared by Server and TUI backlog attention paths.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct OrchestratorQueueAttentionTicket {
resource_key: String,
title: String,
}
impl OrchestratorQueueAttentionTicket {
pub fn new(
resource_key: impl Into<String>,
title: impl Into<String>,
) -> Result<Self, CatalogError> {
let resource_key = resource_key.into();
if !is_ticket_resource_key(&resource_key) {
return Err(CatalogError::InvalidQueueAttentionResourceKey);
}
Ok(Self {
resource_key,
title: bounded_queue_attention_text(&title.into(), 240),
})
}
}
/// Shared model-visible context for every Orchestrator backlog attention renderer.
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct OrchestratorQueueAttentionContext {
tickets: Vec<OrchestratorQueueAttentionTicket>,
separator: &'static str,
omitted_ticket_count: usize,
}
impl OrchestratorQueueAttentionContext {
pub const MAX_TICKETS: usize = 20;
pub fn new(tickets: Vec<OrchestratorQueueAttentionTicket>) -> Self {
let omitted_ticket_count = tickets.len().saturating_sub(Self::MAX_TICKETS);
Self {
tickets: tickets.into_iter().take(Self::MAX_TICKETS).collect(),
separator: "",
omitted_ticket_count,
}
}
}
/// Prompt-catalog entries that must share the same backlog-attention body contract.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OrchestratorQueueAttentionPrompt {
Server,
Tui,
}
impl OrchestratorQueueAttentionPrompt {
fn key(self) -> &'static str {
match self {
Self::Server => "internal.workspace_orchestrator_queue_attention",
Self::Tui => "panel.orchestrator_idle_queue_notice",
}
}
}
fn is_ticket_resource_key(input: &str) -> bool {
input.len() <= 32
&& input.strip_prefix("T-").is_some_and(|suffix| {
!suffix.is_empty() && suffix.bytes().all(|byte| byte.is_ascii_digit())
})
}
fn bounded_queue_attention_text(input: &str, max_chars: usize) -> String {
let mut output = String::new();
for (index, character) in input.chars().enumerate() {
if index == max_chars {
output.push('…');
break;
}
output.push(if character.is_control() {
' '
} else {
character
});
}
output
}
#[derive(Debug, Error)] #[derive(Debug, Error)]
pub enum CatalogError { pub enum CatalogError {
#[error("queued Ticket resource key is missing or invalid")]
InvalidQueueAttentionResourceKey,
#[error("failed to build builtin Prompt source tree: {0}")] #[error("failed to build builtin Prompt source tree: {0}")]
BuiltinTree(String), BuiltinTree(String),
#[error("failed to evaluate builtin Prompt source tree: {0}")] #[error("failed to evaluate builtin Prompt source tree: {0}")]
@@ -322,6 +404,14 @@ impl PromptCatalog {
self.render_name(key, Value::from_serialize(context)) self.render_name(key, Value::from_serialize(context))
} }
pub fn orchestrator_queue_attention(
&self,
prompt: OrchestratorQueueAttentionPrompt,
context: &OrchestratorQueueAttentionContext,
) -> Result<String, CatalogError> {
self.render_serializable(prompt.key(), context)
}
pub fn render_name(&self, key: &str, ctx: Value) -> Result<String, CatalogError> { pub fn render_name(&self, key: &str, ctx: Value) -> Result<String, CatalogError> {
let template = self let template = self
.env .env
@@ -593,6 +683,12 @@ mod tests {
fn builtin_dcdl_catalog_loads() { fn builtin_dcdl_catalog_loads() {
let catalog = PromptCatalog::builtins_only().unwrap(); let catalog = PromptCatalog::builtins_only().unwrap();
assert!(!catalog.projection.templates.is_empty()); assert!(!catalog.projection.templates.is_empty());
assert!(
!catalog
.projection
.templates
.contains_key("worker.ticket_event_companion_notice")
);
} }
#[test] #[test]
@@ -650,6 +746,62 @@ mod tests {
assert!(reviewer.contains("target-only movement does not invalidate approval")); assert!(reviewer.contains("target-only movement does not invalidate approval"));
} }
#[test]
fn queue_attention_prompts_share_sanitized_contract_and_true_truncation() {
let catalog = PromptCatalog::builtins_only().unwrap();
let tickets = (1..=OrchestratorQueueAttentionContext::MAX_TICKETS + 1)
.map(|index| {
OrchestratorQueueAttentionTicket::new(
format!("T-{index}"),
format!("Ticket {index}\nwith control\u{7}"),
)
.unwrap()
})
.collect();
let context = OrchestratorQueueAttentionContext::new(tickets);
let server = catalog
.orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Server, &context)
.unwrap();
let tui = catalog
.orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Tui, &context)
.unwrap();
assert_eq!(server, tui);
assert!(server.starts_with("Queued Tickets require attention:"));
assert!(server.contains("- T-1 — Ticket 1 with control "));
assert!(!server.contains("T-21"));
assert!(server.contains("were omitted from this notice: 1"));
assert!(server.contains("Re-query current Ticket authority"));
assert!(server.contains("Reread the current Ticket state before acting"));
for secret in [
"workspace_id",
"Workspace:",
"runtime_id",
"worker_id",
"bounded",
] {
assert!(!server.contains(secret), "leaked {secret}: {server}");
}
}
#[test]
fn queue_attention_prompt_omits_truncation_text_for_complete_list() {
let catalog = PromptCatalog::builtins_only().unwrap();
let context = OrchestratorQueueAttentionContext::new(vec![
OrchestratorQueueAttentionTicket::new("T-541", "Attention contract").unwrap(),
]);
let rendered = catalog
.orchestrator_queue_attention(OrchestratorQueueAttentionPrompt::Server, &context)
.unwrap();
assert!(rendered.contains("- T-541 — Attention contract"));
assert!(!rendered.contains("omitted"));
assert!(matches!(
OrchestratorQueueAttentionTicket::new("opaque-id", "must fail"),
Err(CatalogError::InvalidQueueAttentionResourceKey)
));
}
#[test] #[test]
fn graph_rejects_dynamic_legacy_missing_and_cycles() { fn graph_rejects_dynamic_legacy_missing_and_cycles() {
let invalid = BTreeMap::from([ let invalid = BTreeMap::from([
+3
View File
@@ -121,8 +121,11 @@ impl SegmentLogSink {
matches!( matches!(
entry, entry,
LogEntry::SegmentStart { .. } LogEntry::SegmentStart { .. }
| LogEntry::AnnotatedSegmentStart { .. }
| LogEntry::UserInput { .. } | LogEntry::UserInput { .. }
| LogEntry::AnnotatedUserInput { .. }
| LogEntry::SystemItem { .. } | LogEntry::SystemItem { .. }
| LogEntry::AnnotatedSystemItem { .. }
| LogEntry::Invoke { .. } | LogEntry::Invoke { .. }
) )
} }
+144 -20
View File
@@ -6,7 +6,8 @@
use std::sync::Arc; use std::sync::Arc;
use agen::{Item, Role}; use crate::session_history::{SessionHistoryMetadata, WorkerHistoryProvenance};
use agen::{HistoryEntry, Item, Role};
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Serialize};
const DEFAULT_SEARCH_LIMIT: usize = 20; const DEFAULT_SEARCH_LIMIT: usize = 20;
@@ -21,14 +22,21 @@ const OVERVIEW_ANCHOR_STRIDE: usize = 8;
pub(crate) struct SessionEntryRef(String); pub(crate) struct SessionEntryRef(String);
impl SessionEntryRef { impl SessionEntryRef {
pub(crate) fn new(source_index: usize) -> Self { pub(crate) fn from_history_entry_id(entry_id: &crate::SessionHistoryEntryId) -> Self {
Self(format!("E{source_index:08}")) Self(format!("E{}", entry_id.0))
} }
pub(crate) fn parse(value: &str) -> Option<Self> { pub(crate) fn parse(value: &str) -> Option<Self> {
let reference = Self(value.to_string()); let suffix = value.strip_prefix('E')?;
reference.source_index()?; if suffix.is_empty()
Some(reference) || suffix.len() > 64
|| !suffix
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || byte == b'-')
{
return None;
}
Some(Self(value.to_string()))
} }
pub(crate) fn as_str(&self) -> &str { pub(crate) fn as_str(&self) -> &str {
@@ -97,6 +105,7 @@ impl ToolPart {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct OverviewItem { pub(crate) struct OverviewItem {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub label: String, pub label: String,
@@ -107,6 +116,7 @@ pub(crate) struct OverviewItem {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReferenceEntry { pub(crate) struct ReferenceEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -132,6 +142,7 @@ pub(crate) struct SearchOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SearchHit { pub(crate) struct SearchHit {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -177,6 +188,7 @@ impl Default for ReadOptions {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct ReadEntry { pub(crate) struct ReadEntry {
pub id: SessionEntryRef, pub id: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
pub tool_name: Option<String>, pub tool_name: Option<String>,
@@ -195,6 +207,7 @@ pub(crate) struct ReadResult {
pub(crate) struct SessionEntryEvidence { pub(crate) struct SessionEntryEvidence {
pub segment_id: String, pub segment_id: String,
pub entry_ref: SessionEntryRef, pub entry_ref: SessionEntryRef,
pub origin: WorkerHistoryProvenance,
pub entry_range: [u64; 2], pub entry_range: [u64; 2],
pub kind: ReferenceKind, pub kind: ReferenceKind,
pub tool_part: Option<ToolPart>, pub tool_part: Option<ToolPart>,
@@ -206,26 +219,42 @@ pub(crate) struct SessionEntryEvidence {
#[derive(Debug, Clone)] #[derive(Debug, Clone)]
pub(crate) struct SessionCapture { pub(crate) struct SessionCapture {
segment_id: String, segment_id: String,
items: Arc<Vec<Item>>, entries: Arc<Vec<HistoryEntry<SessionHistoryMetadata>>>,
overview: Vec<OverviewItem>, overview: Vec<OverviewItem>,
index: Vec<ReferenceEntry>, index: Vec<ReferenceEntry>,
} }
impl SessionCapture { impl SessionCapture {
pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self { pub(crate) fn new(segment_id: impl Into<String>, items: Vec<Item>) -> Self {
let entries = items
.into_iter()
.enumerate()
.map(|(index, item)| {
let mut metadata = SessionHistoryMetadata::legacy_unknown();
metadata.entry_id =
session_store::LoggedSessionHistoryEntryId(format!("{index:08}"));
HistoryEntry::new(item, metadata)
})
.collect();
Self::from_history_entries(segment_id, entries)
}
pub(crate) fn from_history_entries(
segment_id: impl Into<String>,
entries: Vec<HistoryEntry<SessionHistoryMetadata>>,
) -> Self {
let segment_id = segment_id.into(); let segment_id = segment_id.into();
let items = Arc::new(items); let entries = Arc::new(entries);
let mut overview = Vec::new(); let mut overview = Vec::new();
let mut index = Vec::new(); let mut index = Vec::new();
for (idx, item) in items.iter().enumerate() { for (idx, entry) in entries.iter().enumerate() {
let item = &entry.item;
let entry_range = [idx as u64, idx as u64]; let entry_range = [idx as u64, idx as u64];
match item { match item {
Item::Message { role, content, .. } => { Item::Message { role, content, .. } => {
let kind = match role { let Some(kind) = message_reference_kind(&entry.annotation.origin, role) else {
Role::User => ReferenceKind::User, continue;
Role::Assistant => ReferenceKind::Assistant,
Role::System => continue,
}; };
let text = content let text = content
.iter() .iter()
@@ -234,9 +263,10 @@ impl SessionCapture {
.join(""); .join("");
let label = format!("{} message", kind.as_str()); let label = format!("{} message", kind.as_str());
let summary = truncate_chars(&text, 240); let summary = truncate_chars(&text, 240);
let id = SessionEntryRef::new(idx); let id = SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id);
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind, kind,
tool_part: None, tool_part: None,
@@ -248,6 +278,7 @@ impl SessionCapture {
if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) { if matches!(kind, ReferenceKind::User | ReferenceKind::Assistant) {
overview.push(OverviewItem { overview.push(OverviewItem {
id: id.clone(), id: id.clone(),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind, kind,
label, label,
@@ -261,7 +292,8 @@ impl SessionCapture {
} => { } => {
let text = format!("{name}\n{arguments}"); let text = format!("{name}\n{arguments}");
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::new(idx), id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Input), tool_part: Some(ToolPart::Input),
@@ -287,7 +319,8 @@ impl SessionCapture {
content.as_deref().unwrap_or_default(), content.as_deref().unwrap_or_default(),
); );
index.push(ReferenceEntry { index.push(ReferenceEntry {
id: SessionEntryRef::new(idx), id: SessionEntryRef::from_history_entry_id(&entry.annotation.entry_id),
origin: entry.annotation.origin.clone(),
entry_range, entry_range,
kind: ReferenceKind::Tool, kind: ReferenceKind::Tool,
tool_part: Some(ToolPart::Output), tool_part: Some(ToolPart::Output),
@@ -327,7 +360,7 @@ impl SessionCapture {
Self { Self {
segment_id, segment_id,
items, entries,
overview, overview,
index, index,
} }
@@ -337,6 +370,14 @@ impl SessionCapture {
&self.overview &self.overview
} }
pub(crate) fn source_index_for_ref(&self, reference: &SessionEntryRef) -> Option<u64> {
self.index
.iter()
.find(|entry| entry.id == *reference)
.map(|entry| entry.entry_range[0])
.or_else(|| reference.source_index())
}
pub(crate) fn search(&self, options: &SearchOptions) -> Vec<SearchHit> { pub(crate) fn search(&self, options: &SearchOptions) -> Vec<SearchHit> {
let query = options.query.trim().to_lowercase(); let query = options.query.trim().to_lowercase();
let limit = options let limit = options
@@ -347,12 +388,12 @@ impl SessionCapture {
let min_entry_index = options let min_entry_index = options
.from .from
.as_ref() .as_ref()
.and_then(SessionEntryRef::source_index) .and_then(|reference| self.source_index_for_ref(reference))
.unwrap_or_else(|| options.min_entry_index.unwrap_or(0)); .unwrap_or_else(|| options.min_entry_index.unwrap_or(0));
let max_entry_index = options let max_entry_index = options
.through .through
.as_ref() .as_ref()
.and_then(SessionEntryRef::source_index) .and_then(|reference| self.source_index_for_ref(reference))
.unwrap_or(u64::MAX); .unwrap_or(u64::MAX);
let mut skipped = 0usize; let mut skipped = 0usize;
let mut hits = Vec::new(); let mut hits = Vec::new();
@@ -391,6 +432,7 @@ impl SessionCapture {
} }
hits.push(SearchHit { hits.push(SearchHit {
id: entry.id.clone(), id: entry.id.clone(),
origin: entry.origin.clone(),
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
tool_name: entry.tool_name.clone(), tool_name: entry.tool_name.clone(),
@@ -442,13 +484,18 @@ impl SessionCapture {
} }
} }
} }
let Some(item) = self.items.get(entry.entry_range[0] as usize) else { let Some(item) = self
.entries
.get(entry.entry_range[0] as usize)
.map(|entry| &entry.item)
else {
continue; continue;
}; };
let text = render_item(item, entry, options.detail, max_bytes.saturating_sub(bytes)); let text = render_item(item, entry, options.detail, max_bytes.saturating_sub(bytes));
bytes = bytes.saturating_add(text.len()); bytes = bytes.saturating_add(text.len());
entries.push(ReadEntry { entries.push(ReadEntry {
id: entry.id.clone(), id: entry.id.clone(),
origin: entry.origin.clone(),
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
tool_name: entry.tool_name.clone(), tool_name: entry.tool_name.clone(),
@@ -485,6 +532,7 @@ impl SessionCapture {
Some(SessionEntryEvidence { Some(SessionEntryEvidence {
segment_id: self.segment_id.clone(), segment_id: self.segment_id.clone(),
entry_ref: entry.id.clone(), entry_ref: entry.id.clone(),
origin: entry.origin.clone(),
entry_range: entry.entry_range, entry_range: entry.entry_range,
kind: entry.kind, kind: entry.kind,
tool_part: entry.tool_part, tool_part: entry.tool_part,
@@ -495,6 +543,28 @@ impl SessionCapture {
} }
} }
fn message_reference_kind(
origin: &WorkerHistoryProvenance,
provider_role: &Role,
) -> Option<ReferenceKind> {
match origin {
WorkerHistoryProvenance::HumanInput { .. }
| WorkerHistoryProvenance::WorkerInput { .. } => Some(ReferenceKind::User),
WorkerHistoryProvenance::ModelOutput { .. } => Some(ReferenceKind::Assistant),
WorkerHistoryProvenance::ToolOutput { .. } => Some(ReferenceKind::Tool),
WorkerHistoryProvenance::LegacyUnknown => match provider_role {
Role::User => Some(ReferenceKind::User),
Role::Assistant => Some(ReferenceKind::Assistant),
Role::System => None,
},
// Flow/backend/system content remains out of the observation surface
// even when represented with a provider user/system role.
WorkerHistoryProvenance::FlowInstruction { .. }
| WorkerHistoryProvenance::BackendInstruction { .. }
| WorkerHistoryProvenance::DerivedSummary => None,
}
}
fn render_item( fn render_item(
item: &Item, item: &Item,
entry: &ReferenceEntry, entry: &ReferenceEntry,
@@ -563,6 +633,60 @@ fn truncate_chars(text: &str, max_chars: usize) -> String {
mod tests { mod tests {
use super::*; use super::*;
#[test]
fn flow_user_role_is_excluded_while_explicit_human_origin_remains_evidence() {
let entries = vec![
crate::session_history::history_entry(
Item::user_message("trusted flow instruction"),
WorkerHistoryProvenance::FlowInstruction {
selector: "builtin:coder-review".into(),
definition_id: "coder-review".into(),
definition_revision: 3,
instance_id: "instance".into(),
state_id: "implement".into(),
},
),
crate::session_history::history_entry(
Item::user_message("remember my preference"),
WorkerHistoryProvenance::HumanInput {
account_id: "account-1".into(),
},
),
];
let capture = SessionCapture::from_history_entries("segment", entries);
let overview = capture.overview();
assert_eq!(overview.len(), 1);
assert!(matches!(
overview[0].origin,
WorkerHistoryProvenance::HumanInput { .. }
));
let evidence = capture.evidence_for(overview[0].id.as_str()).unwrap();
assert!(evidence.excerpt.ends_with("remember my preference"));
assert!(matches!(
evidence.origin,
WorkerHistoryProvenance::HumanInput { .. }
));
}
#[test]
fn stable_logical_ref_survives_retention_and_restore_projection() {
let retained = crate::session_history::history_entry(
Item::assistant_message("retained"),
WorkerHistoryProvenance::ModelOutput {
worker: crate::session_history::worker_subject(Default::default()),
},
);
let expected_ref = SessionEntryRef::from_history_entry_id(&retained.annotation.entry_id);
let before = SessionCapture::from_history_entries("old", vec![retained.clone()]);
let after = SessionCapture::from_history_entries("new", vec![retained]);
assert_eq!(before.overview()[0].id, expected_ref);
assert_eq!(after.overview()[0].id, expected_ref);
assert_eq!(
after.evidence_for(expected_ref.as_str()).unwrap().entry_ref,
expected_ref
);
}
#[test] #[test]
fn overview_contains_user_and_assistant_only() { fn overview_contains_user_and_assistant_only() {
let view = SessionCapture::new( let view = SessionCapture::new(
+219
View File
@@ -0,0 +1,219 @@
//! Restore-authoritative metadata for model-visible Worker history.
//!
//! Agen transports this annotation without interpreting it. Session Log v2
//! stores each item and metadata in one typed record; legacy records are
//! retained only as explicit `LegacyUnknown` entries.
use agen::{HistoryEntry, Item};
use protocol::Segment;
use session_store::{
LogEntry, LoggedHistoryDerivation, LoggedHistoryEntry, LoggedSessionHistoryEntryId,
LoggedSessionHistoryMetadata, LoggedSessionHistoryOrigin, LoggedWorkerSubject, SegmentId,
SessionId,
};
pub type SessionHistoryEntryId = LoggedSessionHistoryEntryId;
pub type SessionHistoryMetadata = LoggedSessionHistoryMetadata;
pub type WorkerHistoryProvenance = LoggedSessionHistoryOrigin;
pub type SessionHistoryDerivation = LoggedHistoryDerivation;
pub type WorkerSubjectSnapshot = LoggedWorkerSubject;
pub(crate) fn worker_subject(session_id: SessionId) -> WorkerSubjectSnapshot {
WorkerSubjectSnapshot {
workspace_id: None,
runtime_id: None,
worker_id: session_id.to_string(),
}
}
pub(crate) fn metadata(
origin: WorkerHistoryProvenance,
derivation: Option<SessionHistoryDerivation>,
) -> SessionHistoryMetadata {
SessionHistoryMetadata {
entry_id: SessionHistoryEntryId::new(),
origin,
derivation,
}
}
pub(crate) fn history_entry(
item: Item,
origin: WorkerHistoryProvenance,
) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(item, metadata(origin, None))
}
pub(crate) fn to_logged_history_entry(
entry: &HistoryEntry<SessionHistoryMetadata>,
) -> LoggedHistoryEntry {
LoggedHistoryEntry {
item: entry.item.clone().into(),
metadata: entry.annotation.clone(),
}
}
fn legacy_entry(item: Item) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(item, SessionHistoryMetadata::legacy_unknown())
}
fn from_logged(entry: &LoggedHistoryEntry) -> HistoryEntry<SessionHistoryMetadata> {
HistoryEntry::new(Item::from(entry.item.clone()), entry.metadata.clone())
}
/// Rebuild typed Worker history directly from the append-only Session Log.
/// Missing legacy metadata is never inferred from role or plaintext.
pub(crate) fn restore_history_entries(
_session_id: SessionId,
_segment_id: SegmentId,
entries: &[LogEntry],
) -> Result<Vec<HistoryEntry<SessionHistoryMetadata>>, String> {
let mut history = Vec::new();
for entry in entries {
match entry {
LogEntry::AnnotatedSegmentStart { history: seed, .. } => {
history = seed.iter().map(from_logged).collect();
}
LogEntry::SegmentStart { history: seed, .. } => {
history = seed
.iter()
.cloned()
.map(Item::from)
.map(legacy_entry)
.collect();
}
LogEntry::AnnotatedUserInput { history: input, .. } => {
history.extend(input.iter().map(from_logged))
}
LogEntry::UserInput { segments, .. } => history.push(legacy_entry(Item::user_message(
Segment::flatten_to_text(segments),
))),
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => history.push(from_logged(entry)),
LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => {
history.push(legacy_entry(Item::from(item.clone())));
}
LogEntry::AnnotatedSystemItem { entry, .. } => history.push(HistoryEntry::new(
entry.item.to_history_item(),
entry.metadata.clone(),
)),
LogEntry::SystemItem { item, .. } => {
history.push(legacy_entry(item.to_history_item()));
}
_ => {}
}
}
Ok(history)
}
#[cfg(test)]
mod tests {
use super::*;
use agen::llm_client::RequestConfig;
use session_store::LogEntry;
#[test]
fn legacy_user_role_is_not_inferred_as_human_authority() {
let entries = vec![LogEntry::UserInput {
ts: 1,
segments: vec![Segment::text("legacy")],
extensions: Vec::new(),
}];
let restored =
restore_history_entries(SessionId::now_v7(), SegmentId::now_v7(), &entries).unwrap();
assert!(matches!(
restored[0].annotation.origin,
WorkerHistoryProvenance::LegacyUnknown
));
}
#[test]
fn typed_flow_and_unknown_caller_input_round_trip_without_role_inference() {
let session_id = SessionId::now_v7();
let projected = vec![
history_entry(
Item::user_message("flow instructions"),
WorkerHistoryProvenance::FlowInstruction {
selector: "builtin:coder-review".to_string(),
definition_id: "coder-review".to_string(),
definition_revision: 7,
instance_id: "flow-instance".to_string(),
state_id: "implement".to_string(),
},
),
history_entry(
Item::user_message("implement"),
WorkerHistoryProvenance::LegacyUnknown,
),
];
let entries = vec![
LogEntry::AnnotatedSegmentStart {
ts: 0,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: Vec::new(),
forked_from: None,
compacted_from: None,
},
LogEntry::AnnotatedUserInput {
ts: 1,
segments: vec![
Segment::Flow {
selector: "builtin:coder-review".to_string(),
},
Segment::text("implement"),
],
extensions: Vec::new(),
history: projected.iter().map(to_logged_history_entry).collect(),
},
];
let restored = restore_history_entries(session_id, SegmentId::now_v7(), &entries).unwrap();
assert_eq!(restored, projected);
}
#[test]
fn annotated_restore_preserves_logical_ids_across_reboot() {
let session_id = SessionId::now_v7();
let entry = history_entry(
Item::assistant_message("persisted"),
WorkerHistoryProvenance::ModelOutput {
worker: worker_subject(session_id),
},
);
let log = vec![LogEntry::AnnotatedSegmentStart {
ts: 0,
session_id,
system_prompt: None,
config: RequestConfig::default(),
history: vec![to_logged_history_entry(&entry)],
forked_from: None,
compacted_from: None,
}];
let first = restore_history_entries(session_id, SegmentId::now_v7(), &log).unwrap();
let second = restore_history_entries(session_id, SegmentId::now_v7(), &log).unwrap();
assert_eq!(first[0].annotation.entry_id, entry.annotation.entry_id);
assert_eq!(second[0].annotation.entry_id, entry.annotation.entry_id);
}
#[test]
fn compacted_derivation_uses_stable_logical_entry_ids() {
let source = history_entry(
Item::user_message("source"),
WorkerHistoryProvenance::LegacyUnknown,
);
let summary = HistoryEntry::new(
Item::system_message("summary"),
metadata(
WorkerHistoryProvenance::DerivedSummary,
Some(SessionHistoryDerivation {
sources: vec![source.annotation.entry_id.clone()],
}),
),
);
assert_eq!(
summary.annotation.derivation.unwrap().sources,
vec![source.annotation.entry_id]
);
}
}
+9 -6
View File
@@ -499,7 +499,10 @@ impl Tool for SubWorkerSpawnTool {
InternalWorkerVisibility::ParentClient, InternalWorkerVisibility::ParentClient,
Some(child_registry.clone()), Some(child_registry.clone()),
Some(Arc::new(move |status| { Some(Arc::new(move |status| {
if status == InternalWorkerSessionStatus::Failed { if matches!(
status,
InternalWorkerSessionStatus::Failed | InternalWorkerSessionStatus::Stopped
) {
if let Some(registry) = registry.upgrade() { if let Some(registry) = registry.upgrade() {
if let Err(error) = registry.reclaim_internal_scope(&child_name) { if let Err(error) = registry.reclaim_internal_scope(&child_name) {
tracing::warn!( tracing::warn!(
@@ -1249,7 +1252,7 @@ extract_threshold = 4000
) )
.await .await
.unwrap(); .unwrap();
assert!(first_capture.items.iter().any(|item| { assert!(first_capture.entries.iter().map(|entry| &entry.item).any(|item| {
matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed")))) matches!(item, Item::Message { role: Role::Assistant, content, .. } if content.iter().any(|part| matches!(part, ContentPart::Text { text } if text.contains("reviewed"))))
})); }));
@@ -1271,7 +1274,7 @@ extract_threshold = 4000
) )
.await .await
.unwrap(); .unwrap();
assert!(latest_capture.items.len() > first_capture.items.len()); assert!(latest_capture.entries.len() > first_capture.entries.len());
fail_requests.store(true, Ordering::SeqCst); fail_requests.store(true, Ordering::SeqCst);
send.execute( send.execute(
@@ -1282,16 +1285,16 @@ extract_threshold = 4000
.unwrap(); .unwrap();
assert_eq!( assert_eq!(
record.session.wait_until_idle().await, record.session.wait_until_idle().await,
InternalWorkerSessionStatus::Failed InternalWorkerSessionStatus::Stopped
); );
assert_eq!(calls.load(Ordering::SeqCst), 3); assert_eq!(calls.load(Ordering::SeqCst), 3);
assert!( assert!(
spawner_scope.snapshot().is_writable(&workspace_root), spawner_scope.snapshot().is_writable(&workspace_root),
"Failed terminal child must release its delegated Workdir session" "Stopped terminal child must release its delegated Workdir session"
); );
assert!( assert!(
!record.workdir_delegation.is_active(), !record.workdir_delegation.is_active(),
"failed child must revoke cloned scoped sessions" "stopped child must revoke cloned scoped sessions"
); );
assert!(registry.get_internal("reviewer-child").is_some()); assert!(registry.get_internal("reviewer-child").is_some());
+1186 -240
View File
File diff suppressed because it is too large Load Diff
+16 -4
View File
@@ -163,7 +163,8 @@ async fn make_worker_with_manifest(
let scope = worker::Scope::writable(&pwd).unwrap(); let scope = worker::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp); std::mem::forget(pwd_tmp);
let worker = Engine::new(client); let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let mut worker = Worker::new( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
@@ -204,7 +205,14 @@ fn system_texts_in_sink_session_start(
) -> Vec<String> { ) -> Vec<String> {
let (entries, _rx) = worker.sink().subscribe_with_snapshot(); let (entries, _rx) = worker.sink().subscribe_with_snapshot();
for entry in entries.into_iter().rev() { for entry in entries.into_iter().rev() {
if let session_store::LogEntry::SegmentStart { history, .. } = entry { let history = match entry {
session_store::LogEntry::AnnotatedSegmentStart { history, .. } => history
.into_iter()
.map(|entry| entry.item)
.collect::<Vec<_>>(),
session_store::LogEntry::SegmentStart { history, .. } => history,
_ => continue,
};
return history return history
.into_iter() .into_iter()
.filter_map(|logged| { .filter_map(|logged| {
@@ -226,7 +234,6 @@ fn system_texts_in_sink_session_start(
}) })
.collect(); .collect();
} }
}
Vec::new() Vec::new()
} }
@@ -337,7 +344,12 @@ permission = "write"
// New segment records forked_from pointing at the source. // New segment records forked_from pointing at the source.
let new_entries = store.read_all(session_id, new_segment_id).unwrap(); let new_entries = store.read_all(session_id, new_segment_id).unwrap();
match &new_entries[0] { match &new_entries[0] {
LogEntry::SegmentStart { LogEntry::AnnotatedSegmentStart {
session_id: seg_session,
forked_from: Some(origin),
..
}
| LogEntry::SegmentStart {
session_id: seg_session, session_id: seg_session,
forked_from: Some(origin), forked_from: Some(origin),
.. ..
+72 -29
View File
@@ -32,16 +32,29 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec<Item> {
let mut items = Vec::new(); let mut items = Vec::new();
for entry in entries { for entry in entries {
match entry { match entry {
LogEntry::AnnotatedSegmentStart { history, .. } => {
items.extend(history.into_iter().map(|entry| Item::from(entry.item)));
}
LogEntry::SegmentStart { history, .. } => { LogEntry::SegmentStart { history, .. } => {
items.extend(history.into_iter().map(Item::from)); items.extend(history.into_iter().map(Item::from));
} }
LogEntry::AnnotatedUserInput { history, .. } => {
items.extend(history.into_iter().map(|entry| Item::from(entry.item)));
}
LogEntry::UserInput { segments, .. } => { LogEntry::UserInput { segments, .. } => {
let text = protocol::Segment::flatten_to_text(&segments); let text = protocol::Segment::flatten_to_text(&segments);
items.push(Item::user_message(text)); items.push(Item::user_message(text));
} }
LogEntry::AnnotatedAssistantItem { entry, .. }
| LogEntry::AnnotatedToolResult { entry, .. } => {
items.push(Item::from(entry.item));
}
LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => { LogEntry::AssistantItem { item, .. } | LogEntry::ToolResult { item, .. } => {
items.push(Item::from(item)); items.push(Item::from(item));
} }
LogEntry::AnnotatedSystemItem { entry, .. } => {
items.push(entry.item.to_history_item());
}
LogEntry::SystemItem { item, .. } => { LogEntry::SystemItem { item, .. } => {
items.push(item.to_history_item()); items.push(item.to_history_item());
} }
@@ -51,6 +64,14 @@ fn history_from_sink(handle: &WorkerHandle) -> Vec<Item> {
items items
} }
fn system_item(entry: &LogEntry) -> Option<&session_store::SystemItem> {
match entry {
LogEntry::AnnotatedSystemItem { entry, .. } => Some(&entry.item),
LogEntry::SystemItem { item, .. } => Some(item),
_ => None,
}
}
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
// Mock LLM Client // Mock LLM Client
// --------------------------------------------------------------------------- // ---------------------------------------------------------------------------
@@ -192,7 +213,8 @@ async fn make_worker_with_pwd_and_manifest(
let scope = manifest::Scope::writable(&pwd).unwrap(); let scope = manifest::Scope::writable(&pwd).unwrap();
std::mem::forget(pwd_tmp); std::mem::forget(pwd_tmp);
let worker = Engine::new(client); let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone()); let authority = WorkerFilesystemAuthority::local(pwd.clone(), pwd.clone());
let worker = Worker::new( let worker = Worker::new(
manifest, manifest,
@@ -784,13 +806,30 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]); let client = MockClient::sequential(vec![MockResponse::Hang(simple_text_events())]);
let worker = make_worker(client).await; let worker = make_worker(client).await;
let handle = spawn_controller(worker).await; let handle = spawn_controller(worker).await;
let mut events = handle.subscribe();
handle handle
.send(Method::run_text("hello in-flight")) .send(Method::run_text("hello in-flight"))
.await .await
.unwrap(); .unwrap();
wait_for_status(&handle, WorkerStatus::Running).await; tokio::time::timeout(std::time::Duration::from_secs(2), async {
loop {
if matches!(
events.recv().await,
Ok(Event::Status {
status: WorkerStatus::Running,
})
) {
break;
}
}
})
.await
.expect("running status event");
// The Running event is the in-flight visibility fence: the committed
// annotated input must already be available to an immediately attaching
// subscriber rather than racing behind this status transition.
let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path()) let stream = tokio::net::UnixStream::connect(handle.runtime_dir.socket_path())
.await .await
.unwrap(); .unwrap();
@@ -804,10 +843,12 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
// Walk the entries, find a `LogEntry::UserInput` and // Walk the entries, find a `LogEntry::UserInput` and
// confirm its segments flatten to our submitted text. // confirm its segments flatten to our submitted text.
let mut found = false; let mut found = false;
for value in entries { for value in &entries {
let entry: session_store::LogEntry = let entry: session_store::LogEntry =
serde_json::from_value(value).expect("LogEntry deserialise"); serde_json::from_value(value.clone()).expect("LogEntry deserialise");
if let session_store::LogEntry::UserInput { segments, .. } = entry { if let session_store::LogEntry::UserInput { segments, .. }
| session_store::LogEntry::AnnotatedUserInput { segments, .. } = entry
{
let text = protocol::Segment::flatten_to_text(&segments); let text = protocol::Segment::flatten_to_text(&segments);
if text == "hello in-flight" { if text == "hello in-flight" {
found = true; found = true;
@@ -815,7 +856,10 @@ async fn snapshot_includes_user_input_for_in_flight_turn() {
} }
} }
} }
assert!(found, "snapshot must carry the in-flight UserInput entry"); assert!(
found,
"snapshot must carry the in-flight UserInput entry: {entries:?}"
);
return; return;
} }
Event::Alert(_) => continue, Event::Alert(_) => continue,
@@ -1086,7 +1130,7 @@ async fn run_with_paste_segment_inlines_content_and_emits_typed_user_message() {
_ => {} _ => {}
}, },
entry = entry_rx.recv() => match entry { entry = entry_rx.recv() => match entry {
Ok(session_store::LogEntry::UserInput { segments, .. }) => { Ok(session_store::LogEntry::UserInput { segments, .. } | session_store::LogEntry::AnnotatedUserInput { segments, .. }) => {
user_input_segments = Some(segments); user_input_segments = Some(segments);
if saw_turn_end { if saw_turn_end {
break; break;
@@ -1317,11 +1361,8 @@ async fn notify_while_idle_auto_starts_turn_and_injects_system_message() {
let (entries, _) = handle.sink.subscribe_with_snapshot(); let (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_notify_in_mirror = entries.iter().any(|e| { let saw_notify_in_mirror = entries.iter().any(|e| {
matches!( matches!(
e, system_item(e),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::Notification { message, .. }) if message == "turn finished"
item: session_store::SystemItem::Notification { message, .. },
..
} if message == "turn finished"
) )
}); });
assert!( assert!(
@@ -1463,14 +1504,11 @@ async fn worker_event_turn_ended_while_idle_auto_starts_turn_and_injects_system_
let (entries, _) = handle.sink.subscribe_with_snapshot(); let (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_worker_event_in_mirror = entries.iter().any(|e| { let saw_worker_event_in_mirror = entries.iter().any(|e| {
matches!( matches!(
e, system_item(e),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::WorkerEvent {
item: session_store::SystemItem::WorkerEvent {
event: protocol::WorkerEvent::TurnEnded { worker_name }, event: protocol::WorkerEvent::TurnEnded { worker_name },
.. ..
}, }) if worker_name == "child"
..
} if worker_name == "child"
) )
}); });
assert!( assert!(
@@ -1552,14 +1590,11 @@ async fn worker_event_scope_sub_delegated_while_idle_stays_control_plane_only()
let (entries, _) = handle.sink.subscribe_with_snapshot(); let (entries, _) = handle.sink.subscribe_with_snapshot();
let saw_scope_event_in_mirror = entries.iter().any(|entry| { let saw_scope_event_in_mirror = entries.iter().any(|entry| {
matches!( matches!(
entry, system_item(entry),
session_store::LogEntry::SystemItem { Some(session_store::SystemItem::WorkerEvent {
item: session_store::SystemItem::WorkerEvent {
event: protocol::WorkerEvent::ScopeSubDelegated { .. }, event: protocol::WorkerEvent::ScopeSubDelegated { .. },
.. ..
}, })
..
}
) )
}); });
assert!( assert!(
@@ -2134,9 +2169,13 @@ async fn paused_then_run_closes_orphan_tool_use_for_next_request() {
for item in items { for item in items {
match item { match item {
agen::Item::ToolResult { agen::Item::ToolResult {
call_id, summary, .. call_id,
summary,
disposition,
..
} if call_id == "call_orphan" => { } if call_id == "call_orphan" => {
assert_eq!(summary, "[Interrupted by user]"); assert_eq!(summary, "Tool execution outcome unknown");
assert_eq!(*disposition, agen::ToolResultDisposition::OutcomeUnknown);
saw_synthetic_tool_result = true; saw_synthetic_tool_result = true;
} }
agen::Item::Message { role, content, .. } if *role == agen::Role::System => { agen::Item::Message { role, content, .. } if *role == agen::Role::System => {
@@ -2327,8 +2366,11 @@ async fn paused_cancel_abandons_resume_and_next_input_is_fresh_run() {
assert!( assert!(
items.iter().any(|item| matches!( items.iter().any(|item| matches!(
item, item,
agen::Item::ToolResult { call_id, summary, .. } agen::Item::ToolResult {
if call_id == "call_cancelled" && summary == "[Interrupted by user]" call_id,
disposition: agen::ToolResultDisposition::OutcomeUnknown,
..
} if call_id == "call_cancelled"
)), )),
"paused cancel should close orphan tool_use before future requests: {items:?}" "paused cancel should close orphan tool_use before future requests: {items:?}"
); );
@@ -2373,7 +2415,8 @@ async fn snapshot_contains_user_input(handle: &WorkerHandle, needle: &str) -> bo
let entry: session_store::LogEntry = let entry: session_store::LogEntry =
serde_json::from_value(value).expect("LogEntry deserialise"); serde_json::from_value(value).expect("LogEntry deserialise");
match entry { match entry {
session_store::LogEntry::UserInput { segments, .. } => { session_store::LogEntry::UserInput { segments, .. }
| session_store::LogEntry::AnnotatedUserInput { segments, .. } => {
protocol::Segment::flatten_to_text(&segments).contains(needle) protocol::Segment::flatten_to_text(&segments).contains(needle)
} }
_ => false, _ => false,
+6 -3
View File
@@ -188,7 +188,8 @@ async fn make_worker(
let pwd = pwd_tmp.path().to_path_buf(); let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap(); let scope = worker::Scope::writable(&pwd).unwrap();
let mut worker = Engine::new(client); let mut worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
worker.register_tool(big_content_tool_definition(tool_name)); worker.register_tool(big_content_tool_definition(tool_name));
let worker = Worker::new( let worker = Worker::new(
@@ -460,7 +461,8 @@ async fn metric_write_failure_emits_warn_alert_and_does_not_abort_run() {
// protected token budget covers the only user message). That is enough to drive // protected token budget covers the only user message). That is enough to drive
// the failure path: at least one metric attempts to write. // the failure path: at least one metric attempts to write.
let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]); let client = MockClient::new(vec![text_response_with_cache("hi", 0, 0)]);
let worker = Engine::new(client); let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let mut worker = Worker::new( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
@@ -536,7 +538,8 @@ permission = "write"
let pwd_tmp = tempfile::tempdir().unwrap(); let pwd_tmp = tempfile::tempdir().unwrap();
let pwd = pwd_tmp.path().to_path_buf(); let pwd = pwd_tmp.path().to_path_buf();
let scope = worker::Scope::writable(&pwd).unwrap(); let scope = worker::Scope::writable(&pwd).unwrap();
let worker = Engine::new(client); let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let mut worker = Worker::new( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
@@ -130,7 +130,8 @@ async fn make_worker_with_body(
EffectivePromptCatalog::new(templates, 1, "test-schema", "test-toolchain").unwrap(); EffectivePromptCatalog::new(templates, 1, "test-schema", "test-toolchain").unwrap();
let loader = PromptCatalogSource::builtins_only().with_effective_catalog(projection); let loader = PromptCatalogSource::builtins_only().with_effective_catalog(projection);
let worker = Engine::new(client); let worker =
Engine::<_, agen::state::Mutable, worker::SessionHistoryMetadata>::new_annotated(client);
let mut worker = Worker::new( let mut worker = Worker::new(
manifest, manifest,
worker, worker,
+157 -51
View File
@@ -327,8 +327,6 @@ fn repository_local_path(source: &workspace_api::RepositorySource) -> Option<Pat
} }
} }
const ORCHESTRATOR_ATTENTION_TICKET_LIMIT: usize = 20;
const ORCHESTRATOR_ATTENTION_PROMPT_NAME: &str = "internal.workspace_orchestrator_queue_attention";
static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock< static EMBEDDED_RUNTIME_REQUEST_IDENTITY: std::sync::LazyLock<
worker_runtime::auth::RuntimeIdentityMaterial, worker_runtime::auth::RuntimeIdentityMaterial,
> = std::sync::LazyLock::new(|| { > = std::sync::LazyLock::new(|| {
@@ -6414,9 +6412,15 @@ fn worker_ticket_source_context(
} }
} }
fn ticket_notification_content(ticket_id: &str, current_state: &str) -> String { fn canonical_ticket_resource_key(resource_key: &str) -> Option<&str> {
let sequence = resource_key.strip_prefix("T-")?;
(!sequence.is_empty() && sequence.bytes().all(|byte| byte.is_ascii_digit()))
.then_some(resource_key)
}
fn ticket_notification_content(resource_key: &str, current_state: &str) -> String {
format!( format!(
"Ticket notification: ticket_id={ticket_id} current_state={current_state}. Reread the Ticket before acting." "Ticket {resource_key} changed to {current_state}. Reread the current Ticket before acting."
) )
} }
@@ -6498,6 +6502,16 @@ fn notify_ticket_recipients(
current_state: &str, current_state: &str,
source: Option<RuntimeWorkerRef>, source: Option<RuntimeWorkerRef>,
) { ) {
let Ok(Some(resource_key)) =
api.store
.resource_key(workspace_id, WorkspaceResourceKind::Ticket, ticket_id)
else {
return;
};
let Some(resource_key) = canonical_ticket_resource_key(&resource_key) else {
return;
};
let mut recipients = Vec::new(); let mut recipients = Vec::new();
if let Some(assignment) = api if let Some(assignment) = api
.store .store
@@ -6518,7 +6532,7 @@ fn notify_ticket_recipients(
recipients.sort(); recipients.sort();
recipients.dedup(); recipients.dedup();
let content = ticket_notification_content(ticket_id, current_state); let content = ticket_notification_content(resource_key, current_state);
for recipient in recipients { for recipient in recipients {
if source.as_ref().is_some_and(|source| source == &recipient) { if source.as_ref().is_some_and(|source| source == &recipient) {
continue; continue;
@@ -7241,25 +7255,21 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) {
return; return;
} }
let shown = queued let attention_context = match orchestrator_queue_attention_context(
.iter() &api.config.workspace_id,
.take(ORCHESTRATOR_ATTENTION_TICKET_LIMIT) &api.config.workspace_id,
.map(|ticket| { &queued,
format!( ) {
"- {} — {}", Ok(context) => context,
bounded_orchestrator_attention_text(&ticket.id, 80), Err(error) => {
bounded_orchestrator_attention_text(&ticket.title, 240) tracing::warn!(
) workspace_id = %api.config.workspace_id,
}) candidate_count = queued.len(),
.collect::<Vec<_>>() diagnostic = error,
.join("\n"); "orchestrator backlog attention projection rejected"
let omitted = queued );
.len() return;
.saturating_sub(ORCHESTRATOR_ATTENTION_TICKET_LIMIT); }
let omitted_line = if omitted == 0 {
String::new()
} else {
format!("Additional queued Tickets omitted from this notice: {omitted}\n")
}; };
let Ok(Some(config_state)) = api let Ok(Some(config_state)) = api
.config_store .config_store
@@ -7276,16 +7286,19 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) {
let Ok(catalog) = worker::PromptCatalog::from_projection(projection.catalog().clone()) else { let Ok(catalog) = worker::PromptCatalog::from_projection(projection.catalog().clone()) else {
return; return;
}; };
let content = match catalog.render_serializable( let content = match catalog.orchestrator_queue_attention(
ORCHESTRATOR_ATTENTION_PROMPT_NAME, worker::OrchestratorQueueAttentionPrompt::Server,
&BTreeMap::from([ &attention_context,
("omitted_line", omitted_line.as_str()),
("workspace_id", api.config.workspace_id.as_str()),
("ticket_lines", shown.as_str()),
]),
) { ) {
Ok(content) => content, Ok(content) => content,
Err(_) => return, Err(error) => {
tracing::warn!(
workspace_id = %api.config.workspace_id,
diagnostic = %error,
"orchestrator backlog attention rendering failed"
);
return;
}
}; };
let accepted = api let accepted = api
.runtime .runtime
@@ -7305,20 +7318,26 @@ fn dispatch_orchestrator_queue_attention(api: &WorkspaceApi) {
} }
} }
fn bounded_orchestrator_attention_text(input: &str, max_chars: usize) -> String { fn orchestrator_queue_attention_context(
let mut output = String::new(); expected_workspace_id: &str,
for (index, character) in input.chars().enumerate() { candidate_workspace_id: &str,
if index == max_chars { tickets: &[ticket::TicketSummary],
output.push('…'); ) -> std::result::Result<worker::OrchestratorQueueAttentionContext, &'static str> {
break; if candidate_workspace_id != expected_workspace_id {
return Err("foreign_workspace_ticket_projection");
} }
output.push(if character.is_control() { let tickets = tickets
' ' .iter()
} else { .map(|ticket| {
character let resource_key = ticket
}); .resource_key
} .clone()
output .ok_or("missing_ticket_resource_key")?;
worker::OrchestratorQueueAttentionTicket::new(resource_key, ticket.title.clone())
.map_err(|_| "invalid_ticket_resource_key")
})
.collect::<std::result::Result<Vec<_>, _>>()?;
Ok(worker::OrchestratorQueueAttentionContext::new(tickets))
} }
fn require_online_workspace_orchestrator_source( fn require_online_workspace_orchestrator_source(
@@ -18216,15 +18235,25 @@ mod tests {
} }
#[test] #[test]
fn ticket_notification_projection_exposes_only_ticket_and_current_state() { fn ticket_notification_requires_canonical_ticket_resource_key() {
assert_eq!(canonical_ticket_resource_key("T-429"), Some("T-429"));
for invalid in ["", "00001KZ9SR97B", "T-", "T-key", "O-429"] {
assert_eq!(canonical_ticket_resource_key(invalid), None);
}
}
#[test]
fn ticket_notification_projection_exposes_only_resource_key_and_current_state() {
const INTERNAL_ID: &str = "00001KZ9SR97B";
for current_state in ["queued", "inprogress"] { for current_state in ["queued", "inprogress"] {
let content = ticket_notification_content("00001KZ9SR97B", current_state); let content = ticket_notification_content("T-429", current_state);
assert_eq!( assert_eq!(
content, content,
format!( format!(
"Ticket notification: ticket_id=00001KZ9SR97B current_state={current_state}. Reread the Ticket before acting." "Ticket T-429 changed to {current_state}. Reread the current Ticket before acting."
) )
); );
assert!(!content.contains(INTERNAL_ID));
for forbidden in [ for forbidden in [
"workspace_id", "workspace_id",
"event_sequence", "event_sequence",
@@ -18402,9 +18431,13 @@ mod tests {
assert_eq!(inputs.len(), expected_states.len()); assert_eq!(inputs.len(), expected_states.len());
for ((recipient, content), current_state) in inputs.iter().zip(expected_states) { for ((recipient, content), current_state) in inputs.iter().zip(expected_states) {
assert_eq!(recipient.worker_id.to_string(), orchestrator.worker_id); assert_eq!(recipient.worker_id.to_string(), orchestrator.worker_id);
assert!(!content.contains(&ticket.id));
assert_eq!( assert_eq!(
content, content,
&ticket_notification_content(&ticket.id, current_state) &ticket_notification_content(
ticket.resource_key.as_deref().unwrap(),
current_state,
)
); );
} }
} }
@@ -19520,7 +19553,7 @@ mod tests {
assert_eq!( assert_eq!(
notifications[0].1, notifications[0].1,
ticket_notification_content( ticket_notification_content(
ticket_ref.id.as_str(), ticket_ref.resource_key.as_deref().unwrap(),
TicketWorkflowState::Queued.as_str() TicketWorkflowState::Queued.as_str()
) )
); );
@@ -19585,7 +19618,8 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn orchestrator_running_to_idle_recovers_queued_ticket_without_notification_memory() { async fn orchestrator_running_to_idle_recovers_queued_ticket_without_notification_memory() {
let dir = tempfile::tempdir().unwrap(); let dir = tempfile::tempdir().unwrap();
let api = test_api(dir.path()).await; init_clean_git_workspace(dir.path());
let (api, execution) = test_api_with_recording_backend(dir.path()).await;
let backend = browser_ticket_backend(&api).unwrap(); let backend = browser_ticket_backend(&api).unwrap();
let mut input = ticket::NewTicket::new("Recover queued work"); let mut input = ticket::NewTicket::new("Recover queued work");
input.workflow_state = Some(TicketWorkflowState::Queued); input.workflow_state = Some(TicketWorkflowState::Queued);
@@ -19604,6 +19638,8 @@ mod tests {
.await .await
.unwrap(); .unwrap();
assert!(started.online); assert!(started.online);
let startup_inputs = execution.take_inputs();
assert_eq!(startup_inputs.len(), 1);
assert_eq!( assert_eq!(
api.orchestrator_attention_fingerprint api.orchestrator_attention_fingerprint
.lock() .lock()
@@ -19639,6 +19675,76 @@ mod tests {
.as_deref(), .as_deref(),
Some(ticket_ref.id.as_str()) Some(ticket_ref.id.as_str())
); );
let notifications = execution.take_inputs();
assert_eq!(notifications.len(), 1);
assert_eq!(notifications[0].0.worker_id.to_string(), worker_id);
let content = &notifications[0].1;
assert!(content.starts_with("Queued Tickets require attention:"));
assert!(
content.contains(&format!(
"- {} — Recover queued work",
ticket_ref.resource_key.as_deref().unwrap()
)),
"unexpected notification body: {content:?}"
);
assert!(content.contains("Reread the current Ticket state before acting"));
assert!(!content.contains(ticket_ref.id.as_str()));
assert!(!content.contains(TEST_WORKSPACE_ID));
assert!(!content.contains("bounded"));
assert!(!content.contains("omitted"));
let candidates = backend
.list(ticket::TicketListQuery::states([
ticket::TicketListState::Queued,
]))
.unwrap();
let mut truncated_candidates = (1..=worker::OrchestratorQueueAttentionContext::MAX_TICKETS
+ 1)
.map(|index| {
let mut candidate = candidates[0].clone();
candidate.id = format!("opaque-{index}");
candidate.resource_key = Some(format!("T-{index}"));
candidate.title = format!("Queued {index}");
candidate
})
.collect::<Vec<_>>();
let truncated = orchestrator_queue_attention_context(
TEST_WORKSPACE_ID,
TEST_WORKSPACE_ID,
&truncated_candidates,
)
.unwrap();
let rendered = worker::PromptCatalog::builtins_only()
.unwrap()
.orchestrator_queue_attention(
worker::OrchestratorQueueAttentionPrompt::Server,
&truncated,
)
.unwrap();
assert!(rendered.contains("- T-20 — Queued 20"));
assert!(!rendered.contains("T-21"));
assert!(rendered.contains("were omitted from this notice: 1"));
assert!(!rendered.contains("opaque-"));
truncated_candidates[0].resource_key = None;
assert_eq!(
orchestrator_queue_attention_context(
TEST_WORKSPACE_ID,
TEST_WORKSPACE_ID,
&truncated_candidates
)
.unwrap_err(),
"missing_ticket_resource_key"
);
assert_eq!(
orchestrator_queue_attention_context(
TEST_WORKSPACE_ID,
"foreign-workspace",
&truncated_candidates
)
.unwrap_err(),
"foreign_workspace_ticket_projection"
);
} }
#[tokio::test] #[tokio::test]
-4
View File
@@ -30,7 +30,6 @@ internalAgentsMdSection = import "./internal/agents_md_section.md";
internalResidentMemorySummarySection = import "./internal/resident_memory_summary_section.md"; internalResidentMemorySummarySection = import "./internal/resident_memory_summary_section.md";
internalSubWorkerSpawnToolDescription = import "./internal/sub_worker_spawn_tool_description.md"; internalSubWorkerSpawnToolDescription = import "./internal/sub_worker_spawn_tool_description.md";
panelOrchestratorIdleQueueNotice = import "./panel/orchestrator_idle_queue_notice.md"; panelOrchestratorIdleQueueNotice = import "./panel/orchestrator_idle_queue_notice.md";
workerTicketEventCompanionNotice = import "./worker/ticket_event_companion_notice.md";
in in
{ {
default_prompt = defaultDocument.content; default_prompt = defaultDocument.content;
@@ -69,7 +68,4 @@ in
panel = { panel = {
orchestrator_idle_queue_notice = panelOrchestratorIdleQueueNotice.content; orchestrator_idle_queue_notice = panelOrchestratorIdleQueueNotice.content;
}; };
worker = {
ticket_event_companion_notice = workerTicketEventCompanionNotice.content;
};
} }
@@ -1,7 +1,8 @@
Workspace Orchestrator attention: authoritative Ticket state still contains queued work after the previous turn or after Server recovery. Queued Tickets require attention:
{% for ticket in tickets -%}
Workspace: {{workspace_id}} - {{ ticket.resource_key }} {{ separator }} {{ ticket.title }}
Remaining queued Tickets (bounded): {% endfor -%}
{{ticket_lines}} {% if omitted_ticket_count > 0 -%}
{{omitted_line}} Additional queued Tickets were omitted from this notice: {{ omitted_ticket_count }}. Re-query current Ticket authority for the complete set.
Reread the listed Tickets, their relations, orchestration plans, current assignments, Workers, and Workdirs before acting. Continue only work already authorized by the human `ready -> queued` transition. Do not drain the queue automatically and do not create duplicate assignments, Workers, Workdirs, or merges. If no Ticket is currently actionable, record the durable waiting reason on the authoritative Ticket or orchestration plan and stop. For an actionable queued Ticket, call the guarded `SpawnTicketCoder` operation without first changing Ticket state; it records `queued -> inprogress` only after the Coder, initial input, current assignment, and Workdir finalization are durably accepted. {% endif -%}
Reread the current Ticket state before acting. Preserve the human queue gate and current assignment, dependency, Worker, and Workdir authority; do not create duplicate work.
@@ -1,22 +1,8 @@
Workspace Dashboard observed that this Orchestrator Worker is idle while queued Ticket work is present. Queued Tickets require attention:
{% for ticket in tickets -%}
This is bounded attention only, not scheduler authority. Do not drain the queue automatically. Verify the Ticket is still `queued`, then use the guarded `SpawnTicketCoder` operation without a separate state transition; that operation records `queued -> inprogress` only after Worker creation, initial input, assignment, and Workdir finalization are durably accepted. - {{ ticket.resource_key }} {{ separator }} {{ ticket.title }}
Workspace: {{ workspace }}
Actionable queued Tickets:
{% for ticket in actionable_tickets -%}
- {{ ticket.id }} — {{ ticket.title }} [{{ ticket.classification }}]
{% endfor -%} {% endfor -%}
{% if waiting_tickets | length > 0 -%}
Queued Tickets retained in the session work set but currently waiting:
{% for ticket in waiting_tickets -%}
- {{ ticket.id }} — {{ ticket.title }} [{{ ticket.classification }}]: {{ ticket.waiting_reason }}
{% endfor -%}
{% endif -%}
{% if omitted_ticket_count > 0 -%} {% if omitted_ticket_count > 0 -%}
Additional queued Tickets omitted from this bounded notice: {{ omitted_ticket_count }} Additional queued Tickets were omitted from this notice: {{ omitted_ticket_count }}. Re-query current Ticket authority for the complete set.
{% endif -%} {% endif -%}
Reread the current Ticket state before acting. Preserve the human queue gate and current assignment, dependency, Worker, and Workdir authority; do not create duplicate work.
Preserve the existing human gate, dependency/conflict/capacity/dirty-workspace checks, and duplicate-start checks using actual Ticket state, role/session claims, visible Workers, and worktrees.
@@ -1,7 +0,0 @@
Ticket event notice (weak; auto_run=false)
ticket: {{ ticket_id }}
title: {{ title }}
state: {{ state }}
event: {{ event_kind }}
summary: {{ summary }}
ref: {{ ref_path }}
+4 -2
View File
@@ -8,7 +8,7 @@ export type AlertSource = "worker" | "engine" | "compactor" | "agents_md";
export type CompletionKind = "file"; export type CompletionKind = "file";
export type WorkerStatus = "idle" | "running" | "paused"; export type WorkerStatus = "idle" | "running" | "paused" | "stopped";
export type TurnResult = "finished" | "paused"; export type TurnResult = "finished" | "paused";
@@ -16,6 +16,8 @@ export type InvokeKind = "user_send" | "notify" | "worker_event" | "system_remin
export type RunResult = "finished" | "paused" | "limit_reached" | "rolled_back"; export type RunResult = "finished" | "paused" | "limit_reached" | "rolled_back";
export type ToolResultDisposition = "success" | "error" | "interrupted" | "cancelled" | "outcome_unknown";
export type ErrorCode = "already_running" | "not_running" | "not_paused" | "provider_error" | "tool_error" | "invalid_request" | "internal"; export type ErrorCode = "already_running" | "not_running" | "not_paused" | "provider_error" | "tool_error" | "invalid_request" | "internal";
export type Permission = "read" | "write"; export type Permission = "read" | "write";
@@ -191,7 +193,7 @@ summary: string,
* Full tool output. Absent when the tool chose to return * Full tool output. Absent when the tool chose to return
* summary-only, or when the result was pruned. * summary-only, or when the result was pruned.
*/ */
output?: string | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { entries: Array<unknown>, greeting: Greeting, status: WorkerStatus, output?: string | null, disposition?: ToolResultDisposition | null, is_error: boolean, } } | { "event": "usage", "data": { input_tokens: number | null, output_tokens: number | null, cache_read_input_tokens?: number | null, } } | { "event": "run_end", "data": { result: RunResult, } } | { "event": "error", "data": { code: ErrorCode, message: string, } } | { "event": "snapshot", "data": { entries: Array<unknown>, greeting: Greeting, status: WorkerStatus,
/** /**
* Unfinished model output that has already streamed in the current * Unfinished model output that has already streamed in the current
* run but is not yet represented by committed snapshot entries. * run but is not yet represented by committed snapshot entries.
@@ -5,7 +5,7 @@
workspaceWorkersStore, workspaceWorkersStore,
type SidebarWorker, type SidebarWorker,
} from './worker-subscription'; } from './worker-subscription';
import { canShowWorkerInSidebar } from './workers'; import { canShowWorkerInSidebar, sidebarWorkerActivity } from './workers';
const COLLAPSED_WORKER_COUNT = 6; const COLLAPSED_WORKER_COUNT = 6;
@@ -69,6 +69,7 @@
<ul class="nav-list" aria-label="Workers"> <ul class="nav-list" aria-label="Workers">
{#each visibleWorkers as worker (`${worker.runtime_id}:${worker.worker_id}`)} {#each visibleWorkers as worker (`${worker.runtime_id}:${worker.worker_id}`)}
{@const href = workerConsoleHref(worker, workspaceId)} {@const href = workerConsoleHref(worker, workspaceId)}
{@const activity = sidebarWorkerActivity(worker)}
<li> <li>
<a <a
href={href} href={href}
@@ -77,11 +78,11 @@
aria-current={currentPath === href ? 'page' : undefined} aria-current={currentPath === href ? 'page' : undefined}
> >
<span class="worker-status-indicator"> <span class="worker-status-indicator">
{#if worker.state === 'running'} {#if activity === 'worker-running'}
<span class="worker-status-spinner"><Spinner label="Running" /></span> <span class="worker-status-spinner"><Spinner label="Running" /></span>
{:else if worker.has_running_internal_workers} {:else if activity === 'subworker-running'}
<span class="worker-status-spinner is-subworker"><Spinner label="SubWorker running" /></span> <span class="worker-status-spinner is-subworker"><Spinner label="SubWorker running" /></span>
{:else if worker.state === 'idle'} {:else if activity === 'idle'}
<span class="worker-status-dot" aria-label="Idle"></span> <span class="worker-status-dot" aria-label="Idle"></span>
{/if} {/if}
</span> </span>
@@ -14,13 +14,18 @@ declare const Deno: {
test(name: string, fn: () => void | Promise<void>): void; test(name: string, fn: () => void | Promise<void>): void;
}; };
function worker(runtimeId: string, workerId: string, revision: number): SubscriptionWorker { function worker(
runtimeId: string,
workerId: string,
revision: number,
hasRunningInternalWorkers = false,
): SubscriptionWorker {
return { return {
worker_id: workerId, worker_id: workerId,
runtime_id: runtimeId, runtime_id: runtimeId,
subject_revision: revision, subject_revision: revision,
state: 'idle', state: 'idle',
has_running_internal_workers: false, has_running_internal_workers: hasRunningInternalWorkers,
workspace_id: 'workspace-test', workspace_id: 'workspace-test',
display_name: null, display_name: null,
profile: null, profile: null,
@@ -88,3 +93,28 @@ Deno.test('workspace Worker reducer ignores stale events and removes composite s
assertEquals(projection.workers.size, 0); assertEquals(projection.workers.size, 0);
assertEquals(projection.revisions.get('runtime-a:1'), 4); assertEquals(projection.revisions.get('runtime-a:1'), 4);
}); });
Deno.test('fatal child stop replaces the running-child sidebar projection', () => {
const projection = createWorkspaceWorkersProjection();
projection.workers.set('runtime-a:1', worker('runtime-a', '1', 1, true));
projection.revisions.set('runtime-a:1', 1);
applyWorkspaceWorkersFrame(projection, {
protocol_version: 1,
frame: 'event',
message: {
event: 'event',
data: {
subscription_id: 'subscription-1',
subject_revision: 2,
payload: {
event: 'worker_upserted',
data: { worker: worker('runtime-a', '1', 2, false) },
},
},
},
});
assertEquals(projection.workers.get('runtime-a:1')?.has_running_internal_workers, false);
assertEquals(projection.revisions.get('runtime-a:1'), 2);
});
@@ -2,6 +2,7 @@ import {
canOpenWorkerConsole, canOpenWorkerConsole,
canShowWorkerInSidebar, canShowWorkerInSidebar,
compareWorkersForSidebar, compareWorkersForSidebar,
sidebarWorkerActivity,
} from "./workers.ts"; } from "./workers.ts";
import type { Worker } from "./types.ts"; import type { Worker } from "./types.ts";
@@ -77,3 +78,21 @@ Deno.test("sidebar workers sort running then idle then stopped", () => {
workers.sort(compareWorkersForSidebar); workers.sort(compareWorkersForSidebar);
assertEquals(workers.map((candidate) => candidate.worker_id).join(","), "2,1,4,3"); assertEquals(workers.map((candidate) => candidate.worker_id).join(","), "2,1,4,3");
}); });
Deno.test("fatal child stop clears the sidebar SubWorker spinner activity", () => {
const parent = { state: "idle", has_running_internal_workers: true };
assertEquals(sidebarWorkerActivity(parent), "subworker-running");
parent.has_running_internal_workers = false;
assertEquals(sidebarWorkerActivity(parent), "idle");
});
Deno.test("stopped parents do not fall back to the idle indicator", () => {
assertEquals(
sidebarWorkerActivity({
state: "stopped",
has_running_internal_workers: false,
}),
"none",
);
});
@@ -1,5 +1,24 @@
import type { Worker } from './types'; import type { Worker } from './types';
export type SidebarWorkerActivity =
| 'worker-running'
| 'subworker-running'
| 'idle'
| 'none';
type WorkerActivitySource = Pick<Worker, 'state'> & {
has_running_internal_workers: boolean;
};
export function sidebarWorkerActivity(
worker: WorkerActivitySource,
): SidebarWorkerActivity {
if (worker.state === 'running') return 'worker-running';
if (worker.has_running_internal_workers) return 'subworker-running';
if (worker.state === 'idle') return 'idle';
return 'none';
}
export function canShowWorkerInSidebar(worker: Worker): boolean { export function canShowWorkerInSidebar(worker: Worker): boolean {
return worker.implementation.kind !== 'backend_worker_registry'; return worker.implementation.kind !== 'backend_worker_registry';
} }