diff --git a/crates/tinyagents-graph/src/checkpoint/file.rs b/crates/tinyagents-graph/src/checkpoint/file.rs index 9c052f56b..2854a884e 100644 --- a/crates/tinyagents-graph/src/checkpoint/file.rs +++ b/crates/tinyagents-graph/src/checkpoint/file.rs @@ -340,19 +340,8 @@ fn io_err(context: &str, err: impl std::fmt::Display) -> TinyAgentsError { /// addressing/parent/pending-writes wiring of the default /// [`Checkpointer::get_tuple`]. fn tuple_from_checkpoint(checkpoint: Checkpoint) -> CheckpointTuple { - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let parent_config = checkpoint - .parent_checkpoint_id - .as_ref() - .map(|parent| CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(parent.clone()), - namespace: checkpoint.namespace.clone(), - }); + let config = checkpoint.config(); + let parent_config = checkpoint.parent_config(); let pending_writes = checkpoint.pending_writes.clone(); CheckpointTuple { config, diff --git a/crates/tinyagents-graph/src/checkpoint/mod.rs b/crates/tinyagents-graph/src/checkpoint/mod.rs index 05dcf2bae..ddd201b9d 100644 --- a/crates/tinyagents-graph/src/checkpoint/mod.rs +++ b/crates/tinyagents-graph/src/checkpoint/mod.rs @@ -209,11 +209,7 @@ where checkpoint: Checkpoint, writes: &[PendingWrite], ) -> Result { - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let config = checkpoint.config(); let id = self.put(checkpoint).await?; self.put_writes(&config, writes).await?; Ok(id) @@ -342,20 +338,8 @@ where else { return Ok(None); }; - let resolved = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let parent_config = - checkpoint - .parent_checkpoint_id - .as_ref() - .map(|parent| CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(parent.clone()), - namespace: checkpoint.namespace.clone(), - }); + let resolved = checkpoint.config(); + let parent_config = checkpoint.parent_config(); let pending_writes = self.resolved_writes(&resolved, &checkpoint).await?; Ok(Some(CheckpointTuple { config: resolved, diff --git a/crates/tinyagents-graph/src/checkpoint/sqlite.rs b/crates/tinyagents-graph/src/checkpoint/sqlite.rs index 810e6a8d3..90ce2a32d 100644 --- a/crates/tinyagents-graph/src/checkpoint/sqlite.rs +++ b/crates/tinyagents-graph/src/checkpoint/sqlite.rs @@ -524,11 +524,7 @@ where self.put(checkpoint).await?; return Ok(id); } - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let config = checkpoint.config(); let checkpoint_id = checkpoint.checkpoint_id.clone(); let writes = writes.to_vec(); let conn = self.conn.clone(); @@ -718,20 +714,8 @@ where // needed here. let mut out = Vec::with_capacity(records.len()); for checkpoint in records { - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; - let parent_config = - checkpoint - .parent_checkpoint_id - .as_ref() - .map(|parent| CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(parent.clone()), - namespace: checkpoint.namespace.clone(), - }); + let config = checkpoint.config(); + let parent_config = checkpoint.parent_config(); let pending_writes = writes .get(&checkpoint.checkpoint_id) .cloned() diff --git a/crates/tinyagents-graph/src/checkpoint/types.rs b/crates/tinyagents-graph/src/checkpoint/types.rs index 76fe7be98..ff3e74549 100644 --- a/crates/tinyagents-graph/src/checkpoint/types.rs +++ b/crates/tinyagents-graph/src/checkpoint/types.rs @@ -410,6 +410,26 @@ pub struct BarrierArrivals { } impl Checkpoint { + /// The address of this checkpoint: its thread, its id, and its namespace. + pub fn config(&self) -> CheckpointConfig { + CheckpointConfig { + thread_id: self.thread_id.clone(), + checkpoint_id: Some(self.checkpoint_id.clone()), + namespace: self.namespace.clone(), + } + } + + /// The address of the checkpoint this one was written after, if any. + pub fn parent_config(&self) -> Option { + self.parent_checkpoint_id + .as_ref() + .map(|parent| CheckpointConfig { + thread_id: self.thread_id.clone(), + checkpoint_id: Some(parent.clone()), + namespace: self.namespace.clone(), + }) + } + /// Builds a fresh checkpoint format v2 record. /// /// Sensible defaults for everything except `state` and `tasks`: a diff --git a/crates/tinyagents-graph/src/compiled/boundary.rs b/crates/tinyagents-graph/src/compiled/boundary.rs index b7ebd7194..5136b1ce1 100644 --- a/crates/tinyagents-graph/src/compiled/boundary.rs +++ b/crates/tinyagents-graph/src/compiled/boundary.rs @@ -604,11 +604,7 @@ where .with_versions_seen(versions_seen) .with_metadata(metadata); let writes = checkpoint.pending_writes.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let config = checkpoint.config(); let id = checkpointer.put(checkpoint).await?; if !writes.is_empty() { checkpointer.put_writes(&config, &writes).await?; @@ -748,11 +744,7 @@ where boundary.pending, )); let writes = checkpoint.pending_writes.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let config = checkpoint.config(); let id = checkpointer.put(checkpoint).await?; // Also record the ledger through the write protocol, so backends // that implement it can answer "did this task run?" without loading @@ -784,11 +776,7 @@ where let checkpoint = self.build_loop_checkpoint(ctx, thread, boundary, step, interrupts, interrupted); let writes = checkpoint.pending_writes.clone(); - let config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let config = checkpoint.config(); let id = checkpointer.put(checkpoint).await?; checkpointer.put_writes(&config, &writes).await?; self.emit( @@ -838,11 +826,7 @@ where // the ledger tooling sees no completion markers for any checkpoint // written under `DurabilityMode::Async`. let writes = checkpoint.pending_writes.clone(); - let write_config = CheckpointConfig { - thread_id: checkpoint.thread_id.clone(), - checkpoint_id: Some(checkpoint.checkpoint_id.clone()), - namespace: checkpoint.namespace.clone(), - }; + let write_config = checkpoint.config(); match tokio::runtime::Handle::try_current() { Ok(handle) => { diff --git a/crates/tinyagents-harness/src/agent_loop/dialect.rs b/crates/tinyagents-harness/src/agent_loop/dialect.rs index 91fec0699..d0e1e5b54 100644 --- a/crates/tinyagents-harness/src/agent_loop/dialect.rs +++ b/crates/tinyagents-harness/src/agent_loop/dialect.rs @@ -14,10 +14,11 @@ use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use sha2::{Digest, Sha256}; -use tinyinference_llm::message::{ContentBlock, Message}; +use tinyinference_llm::message::Message; use tinyinference_llm::model::{ ModelRequest, ModelResponse, PromptSegment, SegmentRole, ToolChoice, }; +use tinyinference_llm::prompt_tools::replace_text_blocks; use tinyinference_llm::tool::{ToolCall, ToolSchema}; use tinytools_agent::dialect::{CodeDialect, CodeStyle, PFormatDialect}; use tinytools_agent::types::{ParseDiagnostic, ParseOptions, ParsedToolCall}; @@ -617,30 +618,6 @@ pub(super) fn recover_text_calls( replace_text_blocks(std::mem::take(&mut response.message.content), outcome.text); } -/// Keeps every non-text block in place and substitutes one cleaned text at -/// the first text block's position; an empty `cleaned` emits no text block. -fn replace_text_blocks(content: Vec, cleaned: String) -> Vec { - let mut out = Vec::with_capacity(content.len()); - let mut inserted = false; - for block in content { - match block { - ContentBlock::Text(_) => { - if !inserted { - if !cleaned.is_empty() { - out.push(ContentBlock::Text(cleaned.clone())); - } - inserted = true; - } - } - other => out.push(other), - } - } - if !inserted && !cleaned.is_empty() { - out.push(ContentBlock::Text(cleaned)); - } - out -} - /// Scrubs tool-call markup from streamed visible text and collects the /// calls it completes, minting harness ids for them. /// diff --git a/crates/tinyagents-harness/src/agent_loop/mod_tests.rs b/crates/tinyagents-harness/src/agent_loop/mod_tests.rs index 743fd6bae..79d2e2206 100644 --- a/crates/tinyagents-harness/src/agent_loop/mod_tests.rs +++ b/crates/tinyagents-harness/src/agent_loop/mod_tests.rs @@ -363,23 +363,8 @@ impl ChatModel<()> for ProfiledModel { /// Builds a tool-call assistant response (no text, one tool call). fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(7, 3)), - origin: None, - }, - usage: Some(Usage::new(7, 3)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(7, 3)) } /// Builds a tool-call response whose arguments the provider could not parse, @@ -391,44 +376,13 @@ fn invalid_tool_call_response(id: &str, name: &str, raw: &str) -> ModelResponse "openai response contained invalid JSON arguments for tool call `{id}` (`{name}`): \ EOF while parsing a value; raw arguments: {raw:?}" ); - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::invalid(id, name, raw, reason)], - usage: Some(Usage::new(7, 3)), - origin: None, - }, - usage: Some(Usage::new(7, 3)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::tool_call_response(ToolCall::invalid(id, name, raw, reason)) + .with_usage(Usage::new(7, 3)) } /// Builds a plain-text assistant response with explicit usage. fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(input, output)), - origin: None, - }, - usage: Some(Usage::new(input, output)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::text_response(text).with_usage(Usage::new(input, output)) } /// Builds a *truncated empty* completion: `finish_reason == "length"` with no @@ -437,49 +391,23 @@ fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { /// channel. `reasoning_tokens` is folded into `output_tokens` so the usage /// mirrors a real length-truncated response. fn truncated_empty_response(reasoning_tokens: u64) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: Vec::new(), - tool_calls: Vec::new(), - usage: Some(Usage::new(4, reasoning_tokens)), - origin: None, - }, - usage: Some(Usage::new(4, reasoning_tokens)), - finish_reason: Some("length".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + let mut response = crate::testkit::text_response(String::new()) + .with_usage(Usage::new(4, reasoning_tokens)) + .with_finish_reason("length"); + response.message.content = Vec::new(); + response } /// A provider may finish normally after emitting only the hidden reasoning /// channel. Unlike a length-truncated response, another call should keep the /// same output-token limit rather than doubling it. fn reasoning_only_stop_response() -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Thinking { - text: "A greeting that never reached visible content".into(), - signature: None, - }], - tool_calls: Vec::new(), - usage: Some(Usage::new(4, 20)), - origin: None, - }, - usage: Some(Usage::new(4, 20)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + let mut response = crate::testkit::text_response(String::new()).with_usage(Usage::new(4, 20)); + response.message.content = vec![ContentBlock::Thinking { + text: "A greeting that never reached visible content".into(), + signature: None, + }]; + response } /// Middleware that appends a user message to every model request. diff --git a/crates/tinyagents-harness/src/cache/hash.rs b/crates/tinyagents-harness/src/cache/hash.rs index 518853ce3..c49c06477 100644 --- a/crates/tinyagents-harness/src/cache/hash.rs +++ b/crates/tinyagents-harness/src/cache/hash.rs @@ -7,6 +7,7 @@ use serde_json::Value; use sha2::{Digest, Sha256}; +use tinyinference_llm::cache::canonical_value; /// Renders a finalized SHA-256 digest as a 64-character lowercase hex string. pub(super) fn hex_digest(digest: impl AsRef<[u8]>) -> String { @@ -57,22 +58,3 @@ pub(super) fn fnv1a_hex(data: &[u8]) -> String { } format!("{hash:016x}") } - -/// Recursively sorts the keys of every JSON object so that the serialized form -/// is canonical regardless of insertion order. -pub(super) fn canonical_value(v: Value) -> Value { - match v { - Value::Object(map) => { - let mut pairs: Vec<(String, Value)> = map.into_iter().collect(); - pairs.sort_by(|a, b| a.0.cmp(&b.0)); - Value::Object( - pairs - .into_iter() - .map(|(k, val)| (k, canonical_value(val))) - .collect(), - ) - } - Value::Array(arr) => Value::Array(arr.into_iter().map(canonical_value).collect()), - other => other, - } -} diff --git a/crates/tinyagents-harness/src/media/mod_tests.rs b/crates/tinyagents-harness/src/media/mod_tests.rs index b33235ab3..b99ebc99b 100644 --- a/crates/tinyagents-harness/src/media/mod_tests.rs +++ b/crates/tinyagents-harness/src/media/mod_tests.rs @@ -5,7 +5,7 @@ use std::time::Duration; use serde_json::json; use tinyinference_image::{MediaReference, MockImageGenerator}; -use tinyinference_llm::message::{AssistantMessage, ContentBlock, Message}; +use tinyinference_llm::message::Message; use tinyinference_llm::model::ModelResponse; use tinyinference_llm::providers::MockModel; use tinyinference_llm::tool::ToolCall; @@ -374,43 +374,12 @@ async fn video_tool_waits_for_delivery_and_saves_the_clip() { } fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(1, 1)), - origin: None, - }, - usage: Some(Usage::new(1, 1)), - finish_reason: Some("tool_calls".to_owned()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(1, 1)) } fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_owned())], - tool_calls: Vec::new(), - usage: Some(Usage::new(1, 1)), - origin: None, - }, - usage: Some(Usage::new(1, 1)), - finish_reason: Some("stop".to_owned()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::text_response(text).with_usage(Usage::new(1, 1)) } #[tokio::test] diff --git a/crates/tinyagents-harness/src/steering/mod_tests.rs b/crates/tinyagents-harness/src/steering/mod_tests.rs index b2a4dac1e..6b9d88598 100644 --- a/crates/tinyagents-harness/src/steering/mod_tests.rs +++ b/crates/tinyagents-harness/src/steering/mod_tests.rs @@ -31,25 +31,7 @@ use tinytools::{Tool, ToolResult}; /// Builds a plain-text assistant response. fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: tinyinference_llm::message::AssistantMessage { - id: None, - content: vec![tinyinference_llm::message::ContentBlock::Text( - text.to_string(), - )], - tool_calls: Vec::new(), - usage: Some(Usage::new(1, 1)), - origin: None, - }, - usage: Some(Usage::new(1, 1)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + crate::testkit::text_response(text).with_usage(Usage::new(1, 1)) } /// A model that records every request it receives. On its first call it pushes diff --git a/crates/tinyagents-harness/src/structured/validate.rs b/crates/tinyagents-harness/src/structured/validate.rs index 393ef9b1e..52d18770f 100644 --- a/crates/tinyagents-harness/src/structured/validate.rs +++ b/crates/tinyagents-harness/src/structured/validate.rs @@ -29,6 +29,7 @@ //! them here would produce confident wrong answers. use serde_json::Value; +use tinyinference_llm::tool::validate_json_value; use crate::error::{Result, TinyAgentsError}; @@ -37,149 +38,22 @@ use crate::error::{Result, TinyAgentsError}; /// `root` names the value in error messages — the caller passes something like /// `schema 'review'` so the message reads `schema 'review'.items[2].id must be /// integer, got string`. +/// +/// The checking itself is [`tinyinference_llm::tool::validate_json_value`], +/// the same validator the tool-call boundary uses, so a structured output and +/// a tool call are held to one definition of the supported subset. +/// +/// # Errors +/// +/// Returns [`TinyAgentsError::StructuredOutput`] naming the first failing +/// instance path. pub fn validate_value(schema: &Value, value: &Value, root: &str) -> Result<()> { - validate_at(schema, value, root) -} - -/// Recursive validation worker behind [`validate_value`]: checks `enum`, -/// `type`, `required`, `properties` (+ `additionalProperties: false`), and -/// `items` at one schema node, then recurses into matched properties/items -/// with `path` extended to name them. -fn validate_at(schema: &Value, value: &Value, path: &str) -> Result<()> { - if schema.is_null() || schema.as_object().is_some_and(|map| map.is_empty()) { - return Ok(()); - } - - if let Some(allowed) = schema.get("enum").and_then(Value::as_array) - && !allowed.iter().any(|candidate| candidate == value) - { - return Err(invalid(format!( - "{path} must be one of the declared enum values" - ))); - } - - if let Some(type_spec) = schema.get("type") { - validate_type(type_spec, value, path)?; - } - - // `required` is enforced independently of `properties`: a schema may name - // required fields without describing them, and nesting the check under - // `properties` would let such a schema fail open. - if let Some(required) = schema.get("required").and_then(Value::as_array) { - if let Some(object) = value.as_object() { - for field in required.iter().filter_map(Value::as_str) { - if !object.contains_key(field) { - return Err(invalid(format!("{path}.{field} is required"))); - } - } - } else if schema.get("type").is_none() { - return Err(invalid(format!( - "{path} must be an object with the declared fields, got {}", - kind_of(value) - ))); - } - } - - if let Some(properties) = schema.get("properties").and_then(Value::as_object) { - if let Some(object) = value.as_object() { - if schema.get("additionalProperties").and_then(Value::as_bool) == Some(false) { - for field in object.keys() { - if !properties.contains_key(field) { - return Err(invalid(format!("{path}.{field} is not allowed"))); - } - } - } - for (field, field_schema) in properties { - if let Some(field_value) = object.get(field) { - validate_at(field_schema, field_value, &format!("{path}.{field}"))?; - } - } - } else if schema.get("type").is_none() { - return Err(invalid(format!( - "{path} must be an object with the declared fields, got {}", - kind_of(value) - ))); - } - } - - if let Some(items_schema) = schema.get("items") - && let Some(items) = value.as_array() - { - for (index, item) in items.iter().enumerate() { - validate_at(items_schema, item, &format!("{path}[{index}]"))?; - } - } - - Ok(()) -} - -/// Checks `value` against a `type` keyword, which may be a single type name -/// or an array of alternatives (a union), matching any one. -fn validate_type(type_spec: &Value, value: &Value, path: &str) -> Result<()> { - if let Some(kind) = type_spec.as_str() { - if matches_type(value, kind) { - return Ok(()); - } - return Err(invalid(format!( - "{path} must be {kind}, got {}", - kind_of(value) - ))); - } - - if let Some(kinds) = type_spec.as_array() { - let allowed: Vec<&str> = kinds.iter().filter_map(Value::as_str).collect(); - if allowed.iter().any(|kind| matches_type(value, kind)) { - return Ok(()); - } - return Err(invalid(format!( - "{path} must be one of {}, got {}", - allowed.join(", "), - kind_of(value) - ))); - } - - Ok(()) -} - -/// Whether `value`'s runtime JSON kind satisfies the named schema `kind`. -/// An unrecognised `kind` always matches (see the module doc's "supported -/// subset" note). -fn matches_type(value: &Value, kind: &str) -> bool { - match kind { - "null" => value.is_null(), - "boolean" => value.is_boolean(), - "object" => value.is_object(), - "array" => value.is_array(), - "number" => value.is_number(), - "integer" => value.as_i64().is_some() || value.as_u64().is_some(), - "string" => value.is_string(), - // An unknown type keyword must not fail closed: providers accept richer - // vocabularies than this subset understands. - _ => true, - } -} - -/// Names a JSON value's kind for error messages, distinguishing `integer` -/// from `number` (a schema `type` vocabulary distinction JSON itself does not -/// make). -fn kind_of(value: &Value) -> &'static str { - match value { - Value::Null => "null", - Value::Bool(_) => "boolean", - Value::Number(number) if number.as_i64().is_some() || number.as_u64().is_some() => { - "integer" - } - Value::Number(_) => "number", - Value::String(_) => "string", - Value::Array(_) => "array", - Value::Object(_) => "object", - } -} - -/// Wraps a validation failure message as the error variant used throughout -/// this module. -fn invalid(message: String) -> TinyAgentsError { - TinyAgentsError::StructuredOutput(message) + validate_json_value(schema, value, root).map_err(|error| { + TinyAgentsError::StructuredOutput(match error { + tinyinference_llm::Error::Validation(message) => message, + other => other.to_string(), + }) + }) } #[cfg(test)] diff --git a/crates/tinyagents-harness/src/testkit/mod.rs b/crates/tinyagents-harness/src/testkit/mod.rs index d2398b2df..37044491c 100644 --- a/crates/tinyagents-harness/src/testkit/mod.rs +++ b/crates/tinyagents-harness/src/testkit/mod.rs @@ -23,6 +23,7 @@ //! | [`DeterministicIds`] | Monotonic `"{prefix}-N"` id generator | //! | [`EventRecorder`] | Captures `AgentEvent`s from an `EventSink` | //! | [`Trajectory`] | Structural assertions over a sequence of events | +//! | [`text_response`], [`tool_call_response`] | Canned model responses for scripted doubles | mod types; @@ -45,6 +46,30 @@ use tinytools::{Tool, ToolResult}; pub use types::*; +// --------------------------------------------------------------------------- +// Canned responses +// --------------------------------------------------------------------------- + +/// A plain-text assistant response that finished with `stop`. +/// +/// Chain [`ModelResponse::with_usage`] for token counts. +pub fn text_response(text: impl Into) -> ModelResponse { + ModelResponse::assistant(text).with_finish_reason("stop") +} + +/// An assistant response that makes exactly `call` and says nothing, finishing +/// with `tool_calls`. The message id is `msg-{call id}`, so each scripted turn +/// is distinguishable in a transcript. +/// +/// Chain [`ModelResponse::with_usage`] for token counts. +pub fn tool_call_response(call: ToolCall) -> ModelResponse { + let mut response = ModelResponse::assistant(String::new()).with_finish_reason("tool_calls"); + response.message.id = Some(format!("msg-{}", call.id)); + response.message.content = Vec::new(); + response.message.tool_calls = vec![call]; + response +} + // --------------------------------------------------------------------------- // StreamingMock // --------------------------------------------------------------------------- diff --git a/crates/tinyagents-harness/src/testkit/mod_tests.rs b/crates/tinyagents-harness/src/testkit/mod_tests.rs index b4f40ab01..7fa2d0c43 100644 --- a/crates/tinyagents-harness/src/testkit/mod_tests.rs +++ b/crates/tinyagents-harness/src/testkit/mod_tests.rs @@ -587,3 +587,24 @@ async fn schema_driven_model_with_no_tools_returns_final_response_immediately() assert!(response.tool_calls().is_empty()); assert_eq!(response.text(), "no tools needed"); } + +#[test] +fn canned_responses_carry_their_finish_reason_and_content() { + use super::{text_response, tool_call_response}; + use serde_json::json; + use tinyinference_llm::tool::ToolCall; + use tinyinference_llm::usage::Usage; + + let text = text_response("done").with_usage(Usage::new(3, 2)); + assert_eq!(text.text(), "done"); + assert_eq!(text.finish_reason.as_deref(), Some("stop")); + assert_eq!(text.usage, Some(Usage::new(3, 2))); + assert_eq!(text.message.usage, Some(Usage::new(3, 2))); + + let call = tool_call_response(ToolCall::new("c1", "lookup", json!({"q": 1}))); + assert_eq!(call.finish_reason.as_deref(), Some("tool_calls")); + assert_eq!(call.message.id.as_deref(), Some("msg-c1")); + assert!(call.message.content.is_empty()); + assert_eq!(call.tool_calls().len(), 1); + assert_eq!(call.tool_calls()[0].name, "lookup"); +} diff --git a/crates/tinyagents-harness/src/tool/prompt.rs b/crates/tinyagents-harness/src/tool/prompt.rs index d9c0b96ef..9bb67e409 100644 --- a/crates/tinyagents-harness/src/tool/prompt.rs +++ b/crates/tinyagents-harness/src/tool/prompt.rs @@ -16,6 +16,7 @@ use serde_json::{Map, Value}; use tinyinference_llm::message::{ContentBlock, Message}; use tinyinference_llm::model::ModelResponse; +use tinyinference_llm::prompt_tools::replace_text_blocks; use tinyinference_llm::tool::{ToolCall, ToolSchema}; /// Opening / closing delimiters for a text-mode tool call. @@ -519,33 +520,6 @@ pub fn apply_prompt_tool_calls(mut response: ModelResponse) -> ModelResponse { response } -/// Rebuild a content vector, keeping every non-[`ContentBlock::Text`] block (e.g. -/// `Thinking`) in place and substituting the single cleaned text at the position -/// of the first original `Text` block. If the original content had no `Text` -/// block, the cleaned text (when non-empty) is appended; if `cleaned` is empty, -/// no text block is emitted at all. -fn replace_text_blocks(content: Vec, cleaned: String) -> Vec { - let mut out = Vec::with_capacity(content.len()); - let mut inserted = false; - for block in content { - match block { - ContentBlock::Text(_) => { - if !inserted { - if !cleaned.is_empty() { - out.push(ContentBlock::Text(cleaned.clone())); - } - inserted = true; - } - } - other => out.push(other), - } - } - if !inserted && !cleaned.is_empty() { - out.push(ContentBlock::Text(cleaned)); - } - out -} - /// Keys a model may put its arguments under inside a tool-call object. /// /// `arguments` is the OpenAI spelling; `parameters` is what a model copying the diff --git a/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs b/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs index 8f079f1a8..6c193eee2 100644 --- a/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs +++ b/crates/tinyagents-integration-tests/examples/agent_loop_tools.rs @@ -58,44 +58,13 @@ impl Tool for CalculatorTool { /// Builds an assistant response that requests a single tool call. fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(12, 4)), - origin: None, - }, - usage: Some(Usage::new(12, 4)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(12, 4)) } /// Builds a final, plain-text assistant response. fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(20, 8)), - origin: None, - }, - usage: Some(Usage::new(20, 8)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(20, 8)) } #[tokio::main] diff --git a/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs b/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs index c5a8dcc71..34ab1aead 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_agent_graph.rs @@ -41,44 +41,13 @@ struct AgentGraphState { /// Builds an assistant response that asks for a single tool call. fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(6, 2)), - origin: None, - }, - usage: Some(Usage::new(6, 2)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(6, 2)) } /// Builds a plain-text final assistant response. fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(input, output)), - origin: None, - }, - usage: Some(Usage::new(input, output)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(input, output)) } #[tokio::test] diff --git a/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs b/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs index 5dd0775f7..c9fd21b15 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_graph_todos.rs @@ -20,43 +20,12 @@ use tinyinference_llm::tool::ToolCall; use tinyinference_llm::usage::Usage; fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(7, 3)), - origin: None, - }, - usage: Some(Usage::new(7, 3)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(7, 3)) } fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(4, 2)), - origin: None, - }, - usage: Some(Usage::new(4, 2)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(4, 2)) } #[tokio::test] diff --git a/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs b/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs index e09198147..8505d37a1 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_harness_mock_capabilities.rs @@ -43,23 +43,7 @@ fn tool_turn(calls: Vec) -> ModelResponse { } fn text_turn(text: impl Into) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.into())], - tool_calls: Vec::new(), - usage: Some(Usage::new(5, 2)), - origin: None, - }, - usage: Some(Usage::new(5, 2)), - finish_reason: Some("stop".into()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(5, 2)) } /// A probe that observes the request after all earlier `before_model` hooks diff --git a/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs b/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs index a7545a13b..6f75740f1 100644 --- a/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs +++ b/crates/tinyagents-integration-tests/tests/e2e_unknown_tool_policy.rs @@ -27,43 +27,12 @@ use tinyinference_llm::usage::Usage; // ── Scripted response helpers ───────────────────────────────────────────────── fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(6, 2)), - origin: None, - }, - usage: Some(Usage::new(6, 2)), - finish_reason: Some("tool_calls".into()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(6, 2)) } fn text_response(text: &str) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.into())], - tool_calls: Vec::new(), - usage: Some(Usage::new(3, 1)), - origin: None, - }, - usage: Some(Usage::new(3, 1)), - finish_reason: Some("stop".into()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(3, 1)) } /// Finds the single recorded [`AgentEvent::UnknownToolCall`], asserting the diff --git a/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs b/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs index e85ba115a..21e970ed6 100644 --- a/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs +++ b/crates/tinyagents-integration-tests/tests/harness_agent_loop.rs @@ -59,43 +59,12 @@ fn trajectory(listener: &Arc) -> Trajectory { } fn tool_call_response(id: &str, name: &str, arguments: serde_json::Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(7, 3)), - origin: None, - }, - usage: Some(Usage::new(7, 3)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(7, 3)) } fn text_response(text: &str, input: u64, output: u64) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: None, - content: vec![ContentBlock::Text(text.to_string())], - tool_calls: Vec::new(), - usage: Some(Usage::new(input, output)), - origin: None, - }, - usage: Some(Usage::new(input, output)), - finish_reason: Some("stop".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::text_response(text).with_usage(Usage::new(input, output)) } struct ProfiledIntegrationModel { diff --git a/crates/tinyagents-integration-tests/tests/tool_deferral.rs b/crates/tinyagents-integration-tests/tests/tool_deferral.rs index e387a6428..c21d3c123 100644 --- a/crates/tinyagents-integration-tests/tests/tool_deferral.rs +++ b/crates/tinyagents-integration-tests/tests/tool_deferral.rs @@ -136,23 +136,8 @@ impl ChatModel<()> for RecordingModel { } fn tool_call(id: &str, name: &str, arguments: Value) -> ModelResponse { - ModelResponse { - message: AssistantMessage { - id: Some(format!("msg-{id}")), - content: Vec::new(), - tool_calls: vec![ToolCall::new(id, name, arguments)], - usage: Some(Usage::new(1, 1)), - origin: None, - }, - usage: Some(Usage::new(1, 1)), - finish_reason: Some("tool_calls".to_string()), - raw: None, - resolved_model: None, - continue_turn: None, - served_from_cache: false, - correlation: None, - resolved_route: None, - } + tinyagents_harness::testkit::tool_call_response(ToolCall::new(id, name, arguments)) + .with_usage(Usage::new(1, 1)) } fn text(body: &str) -> ModelResponse { diff --git a/crates/tinyagents-session/src/lib.rs b/crates/tinyagents-session/src/lib.rs index 5f8c7aee7..c4249de1a 100644 --- a/crates/tinyagents-session/src/lib.rs +++ b/crates/tinyagents-session/src/lib.rs @@ -85,6 +85,7 @@ mod context; pub mod entry_tree; mod migrations; pub mod ops; +mod paging; pub mod retention; pub mod run_ledger; mod store; diff --git a/crates/tinyagents-session/src/ops.rs b/crates/tinyagents-session/src/ops.rs index 7977ead76..0aff87502 100644 --- a/crates/tinyagents-session/src/ops.rs +++ b/crates/tinyagents-session/src/ops.rs @@ -353,59 +353,22 @@ pub fn list_sessions( ); with_connection(workspace_dir, |conn| { - let mut where_clauses: Vec = Vec::new(); - let mut param_values: Vec> = Vec::new(); - - if let Some(s) = status { - param_values.push(Box::new(s.to_string())); - where_clauses.push(format!("status = ?{}", param_values.len())); - } - if let Some(p) = parent_id { - param_values.push(Box::new(p.to_string())); - where_clauses.push(format!("parent_session_id = ?{}", param_values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - let lim = limit.unwrap_or(50).min(500) as i64; let off = offset.unwrap_or(0) as i64; - - let count_sql = format!("SELECT COUNT(*) FROM sessions {where_sql}"); - let total: u64 = { - let mut stmt = conn.prepare(&count_sql)?; - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - param_values.iter().map(|b| b.as_ref()).collect(); - stmt.query_row(params_ref.as_slice(), |r| r.get::<_, i64>(0))? as u64 - }; - - param_values.push(Box::new(lim)); - let lim_idx = param_values.len(); - param_values.push(Box::new(off)); - let off_idx = param_values.len(); - - let query_sql = format!( - "SELECT id, agent_definition_id, agent_definition_name, session_key, - parent_session_id, thread_id, source_channel, status, model, - turn_count, input_tokens, output_tokens, cached_input_tokens, - cost_usd, transcript_path, started_at, ended_at - FROM sessions {where_sql} - ORDER BY started_at DESC - LIMIT ?{lim_idx} OFFSET ?{off_idx}", - ); - - let mut stmt = conn.prepare(&query_sql)?; - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - param_values.iter().map(|b| b.as_ref()).collect(); - let rows = stmt.query_map(params_ref.as_slice(), map_session_row)?; - - let mut sessions = Vec::new(); - for row in rows { - sessions.push(row?); - } + let (sessions, total) = crate::paging::PagedQuery::default() + .eq("status", status.map(str::to_owned)) + .eq("parent_session_id", parent_id.map(str::to_owned)) + .page(lim, off)? + .fetch( + conn, + "sessions", + "id, agent_definition_id, agent_definition_name, session_key, + parent_session_id, thread_id, source_channel, status, model, + turn_count, input_tokens, output_tokens, cached_input_tokens, + cost_usd, transcript_path, started_at, ended_at", + "started_at DESC", + map_session_row, + )?; Ok(SessionSearchResult { sessions, total }) }) @@ -443,88 +406,36 @@ pub(super) fn search_sessions_inner( let lim = params.limit.unwrap_or(50).min(500) as i64; let off = params.offset.unwrap_or(0) as i64; - let mut where_clauses: Vec = Vec::new(); - let mut param_values: Vec> = Vec::new(); - + let mut query = crate::paging::PagedQuery::default(); if let Some(q) = params.query.as_ref().filter(|q| !q.trim().is_empty()) { - param_values.push(Box::new(fts_match_query(q))); - where_clauses.push(format!( - "s.id IN (SELECT session_id FROM sessions_fts WHERE sessions_fts MATCH ?{})", - param_values.len() - )); - } - - if let Some(ref agent) = params.agent_id { - param_values.push(Box::new(agent.clone())); - where_clauses.push(format!("s.agent_definition_id = ?{}", param_values.len())); + query.push(fts_match_query(q), |n| { + format!("s.id IN (SELECT session_id FROM sessions_fts WHERE sessions_fts MATCH ?{n})") + }); } - - if let Some(ref tool) = params.tool_name { - param_values.push(Box::new(tool.clone())); - where_clauses.push(format!( - "s.id IN (SELECT DISTINCT session_id FROM session_tool_calls WHERE tool_name = ?{})", - param_values.len() - )); - } - - if let Some(ref channel) = params.source_channel { - param_values.push(Box::new(channel.clone())); - where_clauses.push(format!("s.source_channel = ?{}", param_values.len())); - } - - if let Some(ref parent) = params.parent_session_id { - param_values.push(Box::new(parent.clone())); - where_clauses.push(format!("s.parent_session_id = ?{}", param_values.len())); - } - - if let Some(ref status) = params.status { - param_values.push(Box::new(status.clone())); - where_clauses.push(format!("s.status = ?{}", param_values.len())); - } - - if let Some(ref tid) = params.thread_id { - param_values.push(Box::new(tid.clone())); - where_clauses.push(format!("s.thread_id = ?{}", param_values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - - let count_sql = format!("SELECT COUNT(*) FROM sessions s {where_sql}"); - let total: u64 = { - let mut stmt = conn.prepare(&count_sql)?; - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - param_values.iter().map(|b| b.as_ref()).collect(); - stmt.query_row(params_ref.as_slice(), |r| r.get::<_, i64>(0))? as u64 - }; - - param_values.push(Box::new(lim)); - let lim_idx = param_values.len(); - param_values.push(Box::new(off)); - let off_idx = param_values.len(); - - let query = format!( - "SELECT s.id, s.agent_definition_id, s.agent_definition_name, s.session_key, - s.parent_session_id, s.thread_id, s.source_channel, s.status, s.model, - s.turn_count, s.input_tokens, s.output_tokens, s.cached_input_tokens, - s.cost_usd, s.transcript_path, s.started_at, s.ended_at - FROM sessions s {where_sql} - ORDER BY s.started_at DESC - LIMIT ?{lim_idx} OFFSET ?{off_idx}", - ); - - let mut stmt = conn.prepare(&query)?; - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - param_values.iter().map(|b| b.as_ref()).collect(); - let rows = stmt.query_map(params_ref.as_slice(), map_session_row)?; - - let mut sessions = Vec::new(); - for row in rows { - sessions.push(row?); + query.eq("s.agent_definition_id", params.agent_id.clone()); + if let Some(tool) = params.tool_name.clone() { + query.push(tool, |n| { + format!( + "s.id IN (SELECT DISTINCT session_id FROM session_tool_calls WHERE tool_name = ?{n})" + ) + }); } + let (sessions, total) = query + .eq("s.source_channel", params.source_channel.clone()) + .eq("s.parent_session_id", params.parent_session_id.clone()) + .eq("s.status", params.status.clone()) + .eq("s.thread_id", params.thread_id.clone()) + .page(lim, off)? + .fetch( + conn, + "sessions s", + "s.id, s.agent_definition_id, s.agent_definition_name, s.session_key, + s.parent_session_id, s.thread_id, s.source_channel, s.status, s.model, + s.turn_count, s.input_tokens, s.output_tokens, s.cached_input_tokens, + s.cost_usd, s.transcript_path, s.started_at, s.ended_at", + "s.started_at DESC", + map_session_row, + )?; Ok(SessionSearchResult { sessions, total }) } diff --git a/crates/tinyagents-session/src/paging/mod.rs b/crates/tinyagents-session/src/paging/mod.rs new file mode 100644 index 000000000..ff0647b4f --- /dev/null +++ b/crates/tinyagents-session/src/paging/mod.rs @@ -0,0 +1,118 @@ +//! One page of a filtered `SELECT`, plus the total number of matching rows. +//! +//! Every list endpoint over the session store does the same four things: +//! collect optional `WHERE` clauses with their bound values, count the matching +//! rows, bind `LIMIT`/`OFFSET` after the filter values, and map the page. +//! [`PagedQuery`] is that sequence once, so the placeholder numbering — the +//! part that silently breaks when one copy is edited and the others are not — +//! lives in one place. + +use rusqlite::Connection; +use rusqlite::types::ToSql; + +use tinyagents_harness::error::{Result, TinyAgentsError}; + +/// Filters for a paged query, each clause numbered after the values before it. +#[derive(Default)] +pub(crate) struct PagedQuery { + clauses: Vec, + values: Vec>, + limit: i64, + offset: i64, +} + +impl PagedQuery { + /// Binds `value` and adds the clause `clause` builds from its placeholder + /// number, for filters that are more than `column = value`. + pub(crate) fn push( + &mut self, + value: impl ToSql + 'static, + clause: impl FnOnce(usize) -> String, + ) -> &mut Self { + self.values.push(Box::new(value)); + self.clauses.push(clause(self.values.len())); + self + } + + /// Filters on `column = value` when `value` is `Some`. + pub(crate) fn eq(&mut self, column: &str, value: Option) -> &mut Self { + if let Some(value) = value { + self.push(value, |n| format!("{column} = ?{n}")); + } + self + } + + /// Filters on `column = value` when `value` is present and not blank. + pub(crate) fn eq_nonblank(&mut self, column: &str, value: Option<&str>) -> &mut Self { + self.eq( + column, + value + .filter(|value| !value.trim().is_empty()) + .map(str::to_owned), + ) + } + + /// Sets the page: at most `limit` rows, skipping the first `offset`. + pub(crate) fn page(&mut self, limit: i64, offset: i64) -> Result<&mut Self> { + if limit < 0 { + return Err(TinyAgentsError::Storage(format!( + "pagination limit must be non-negative: {limit}" + ))); + } + if offset < 0 { + return Err(TinyAgentsError::Storage(format!( + "pagination offset must be non-negative: {offset}" + ))); + } + self.limit = limit; + self.offset = offset; + Ok(self) + } + + /// Counts every row of `from` matching the filters, then maps `columns` + /// for the page of them in `order_by` order. Returns the page and the + /// total count. + pub(crate) fn fetch( + &mut self, + conn: &Connection, + from: &str, + columns: &str, + order_by: &str, + map: impl FnMut(&rusqlite::Row<'_>) -> rusqlite::Result, + ) -> Result<(Vec, u64)> { + let where_sql = if self.clauses.is_empty() { + String::new() + } else { + format!("WHERE {}", self.clauses.join(" AND ")) + }; + let total = conn.query_row( + &format!("SELECT COUNT(*) FROM {from} {where_sql}"), + self.params().as_slice(), + |row| row.get::<_, i64>(0), + )? as u64; + + let limit_idx = self.values.len() + 1; + let offset_idx = self.values.len() + 2; + let mut page_params = self.params(); + page_params.push(&self.limit); + page_params.push(&self.offset); + let mut stmt = conn.prepare(&format!( + "SELECT {columns} + FROM {from} {where_sql} + ORDER BY {order_by} + LIMIT ?{limit_idx} OFFSET ?{offset_idx}" + ))?; + let rows = stmt + .query_map(page_params.as_slice(), map)? + .collect::>>()?; + Ok((rows, total)) + } + + fn params(&self) -> Vec<&dyn ToSql> { + self.values.iter().map(|value| value.as_ref()).collect() + } +} + +#[cfg(test)] +#[path = "mod_tests.rs"] +mod tests; diff --git a/crates/tinyagents-session/src/paging/mod_tests.rs b/crates/tinyagents-session/src/paging/mod_tests.rs new file mode 100644 index 000000000..91f6c8583 --- /dev/null +++ b/crates/tinyagents-session/src/paging/mod_tests.rs @@ -0,0 +1,90 @@ +//! Unit tests for [`PagedQuery`]. + +use super::*; +use rusqlite::Connection; + +fn table() -> Connection { + let conn = Connection::open_in_memory().unwrap(); + conn.execute_batch( + "CREATE TABLE items (id INTEGER, kind TEXT, rank INTEGER); + INSERT INTO items VALUES (1, 'a', 3), (2, 'b', 1), (3, 'a', 2), (4, 'a', 5);", + ) + .unwrap(); + conn +} + +fn ids(query: &mut PagedQuery, limit: i64, offset: i64) -> (Vec, u64) { + query + .page(limit, offset) + .unwrap() + .fetch(&table(), "items", "id", "rank DESC", |row| row.get(0)) + .unwrap() +} + +#[test] +fn counts_every_match_and_returns_one_page_in_order() { + let mut query = PagedQuery::default(); + query.eq("kind", Some("a")); + + assert_eq!(ids(&mut query, 2, 0), (vec![4, 1], 3)); +} + +#[test] +fn the_offset_is_bound_after_the_filters() { + let mut query = PagedQuery::default(); + query.eq("kind", Some("a")); + + assert_eq!(ids(&mut query, 2, 2), (vec![3], 3)); +} + +#[test] +fn absent_and_blank_filters_add_no_clause() { + let mut query = PagedQuery::default(); + query + .eq("kind", None::) + .eq_nonblank("kind", Some(" ")) + .eq_nonblank("kind", None); + + assert_eq!(ids(&mut query, 10, 0), (vec![4, 1, 3, 2], 4)); +} + +#[test] +fn custom_clauses_use_their_placeholder_number() { + let mut query = PagedQuery::default(); + query + .eq_nonblank("kind", Some("a")) + .push(2_i64, |n| format!("rank > ?{n}")); + + assert_eq!(ids(&mut query, 10, 0), (vec![4, 1], 2)); +} + +#[test] +fn repeated_fetch_does_not_accumulate_pagination_parameters() { + let conn = table(); + let mut query = PagedQuery::default(); + query.eq("kind", Some("a")); + + query.page(2, 0).unwrap(); + assert_eq!( + query + .fetch(&conn, "items", "id", "rank DESC", |row| row.get(0)) + .unwrap(), + (vec![4, 1], 3) + ); + + query.page(1, 2).unwrap(); + assert_eq!( + query + .fetch(&conn, "items", "id", "rank DESC", |row| row.get(0)) + .unwrap(), + (vec![3], 3) + ); +} + +#[test] +fn rejects_negative_pagination_values_before_storing_them() { + let mut query = PagedQuery::default(); + + assert!(query.page(-1, 0).is_err()); + assert!(query.page(1, -1).is_err()); +} diff --git a/crates/tinyagents-session/src/run_ledger/ops.rs b/crates/tinyagents-session/src/run_ledger/ops.rs index a0c05b758..284d1b55c 100644 --- a/crates/tinyagents-session/src/run_ledger/ops.rs +++ b/crates/tinyagents-session/src/run_ledger/ops.rs @@ -645,70 +645,25 @@ pub fn list_agent_runs( ) -> Result { crate::store::with_connection(workspace_dir, |conn| { init_run_ledger_schema(conn)?; - let mut where_clauses = Vec::new(); - let mut values: Vec> = Vec::new(); - - if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { - values.push(Box::new(status.to_string())); - where_clauses.push(format!("status = ?{}", values.len())); - } - if let Some(kind) = request.kind.as_deref().filter(|s| !s.trim().is_empty()) { - values.push(Box::new(kind.to_string())); - where_clauses.push(format!("kind = ?{}", values.len())); - } - if let Some(parent) = request - .parent_run_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(parent.to_string())); - where_clauses.push(format!("parent_run_id = ?{}", values.len())); - } - if let Some(thread) = request - .parent_thread_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(thread.to_string())); - where_clauses.push(format!("parent_thread_id = ?{}", values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - let count_sql = format!("SELECT COUNT(*) FROM agent_runs {where_sql}"); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { - row.get::<_, i64>(0) - })? as usize; - let limit = request.limit.unwrap_or(50).min(500) as i64; let offset = request.offset.unwrap_or(0) as i64; - values.push(Box::new(limit)); - let limit_idx = values.len(); - values.push(Box::new(offset)); - let offset_idx = values.len(); - - let query_sql = format!( - "SELECT id, kind, parent_run_id, parent_thread_id, agent_id, status, - prompt_ref, worker_thread_id, - checkpoint_path, checkpoint_json, summary, error, metadata_json, - started_at, updated_at, completed_at - FROM agent_runs {where_sql} - ORDER BY updated_at DESC - LIMIT ?{limit_idx} OFFSET ?{offset_idx}" - ); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let mut stmt = conn.prepare(&query_sql)?; - let rows = stmt.query_map(params_ref.as_slice(), |row| map_agent_run_row(conn, row))?; - let mut runs = Vec::new(); - for row in rows { - runs.push(row?); - } + let (runs, count) = crate::paging::PagedQuery::default() + .eq_nonblank("status", request.status.as_deref()) + .eq_nonblank("kind", request.kind.as_deref()) + .eq_nonblank("parent_run_id", request.parent_run_id.as_deref()) + .eq_nonblank("parent_thread_id", request.parent_thread_id.as_deref()) + .page(limit, offset)? + .fetch( + conn, + "agent_runs", + "id, kind, parent_run_id, parent_thread_id, agent_id, status, + prompt_ref, worker_thread_id, + checkpoint_path, checkpoint_json, summary, error, metadata_json, + started_at, updated_at, completed_at", + "updated_at DESC", + |row| map_agent_run_row(conn, row), + )?; + let count = count as usize; Ok(AgentRunListResponse { runs, count }) }) } @@ -744,32 +699,57 @@ pub fn list_recent_run_events( }) } +/// Fetches one row by id inside its own connection, logging the lookup's +/// entry and whether it found anything — the shared body of every public +/// `get_*` in this module. +fn get_logged( + workspace_dir: &Path, + op: &str, + id: &str, + inner: impl FnOnce(&Connection, &str) -> Result>, +) -> Result> { + tracing::debug!("{LOG_PREFIX} {op}.entry id={id}"); + crate::store::with_connection(workspace_dir, |conn| { + init_run_ledger_schema(conn)?; + let found = inner(conn, id)?; + tracing::debug!("{LOG_PREFIX} {op}.exit id={id} found={}", found.is_some()); + Ok(found) + }) +} + +/// Runs a `SELECT … WHERE id = ?1` and maps the row, or `None` when no row +/// has that id — the shared body of every connection-scoped `*_inner` lookup. +fn get_row_by_id( + conn: &Connection, + sql: &str, + id: &str, + map: impl FnOnce(&rusqlite::Row<'_>) -> rusqlite::Result, +) -> Result> { + Ok(conn.prepare(sql)?.query_row(params![id], map).optional()?) +} + /// Connection-scoped workflow-run lookup, so an upsert can read its own write /// back inside the same transaction. fn get_workflow_run_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( + get_row_by_id( + conn, "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, child_run_ids_json, status, summary, started_at, updated_at, completed_at, revision, lease_owner, lease_expires_at FROM workflow_runs WHERE id = ?1", - )?; - Ok(stmt - .query_row(params![id], map_workflow_run_row) - .optional()?) + id, + map_workflow_run_row, + ) } /// Fetches a single [`WorkflowRun`] by id, or `None` if no row matches. pub fn get_workflow_run(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_workflow_run.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let run = get_workflow_run_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_workflow_run.exit id={id} found={}", - run.is_some() - ); - Ok(run) - }) + get_logged( + workspace_dir, + "get_workflow_run", + id, + get_workflow_run_inner, + ) } /// List durable workflow runs, most-recently-updated first, with optional @@ -789,68 +769,26 @@ pub fn list_workflow_runs( ); crate::store::with_connection(workspace_dir, |conn| { init_run_ledger_schema(conn)?; - let mut where_clauses = Vec::new(); - let mut values: Vec> = Vec::new(); - - if let Some(definition) = request - .definition_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(definition.to_string())); - where_clauses.push(format!("definition_id = ?{}", values.len())); - } - if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { - values.push(Box::new(status.to_string())); - where_clauses.push(format!("status = ?{}", values.len())); - } - if let Some(thread) = request - .parent_thread_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(thread.to_string())); - where_clauses.push(format!("parent_thread_id = ?{}", values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - let count_sql = format!("SELECT COUNT(*) FROM workflow_runs {where_sql}"); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { - row.get::<_, i64>(0) - })? as usize; - let limit = request.limit.unwrap_or(50).min(500) as i64; // `offset` is `u64`; convert checked so a value > i64::MAX surfaces a // clear error instead of wrapping negative and corrupting pagination. let offset = i64::try_from(request.offset.unwrap_or(0)) .storage_context("workflow run list offset exceeds i64::MAX")?; - values.push(Box::new(limit)); - let limit_idx = values.len(); - values.push(Box::new(offset)); - let offset_idx = values.len(); - - let query_sql = format!( - "SELECT id, definition_id, parent_thread_id, input_json, phase_states_json, - child_run_ids_json, status, summary, started_at, updated_at, completed_at, - revision, lease_owner, lease_expires_at - FROM workflow_runs {where_sql} - ORDER BY updated_at DESC - LIMIT ?{limit_idx} OFFSET ?{offset_idx}" - ); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let mut stmt = conn.prepare(&query_sql)?; - let rows = stmt.query_map(params_ref.as_slice(), map_workflow_run_row)?; - let mut runs = Vec::new(); - for row in rows { - runs.push(row?); - } + let (runs, count) = crate::paging::PagedQuery::default() + .eq_nonblank("definition_id", request.definition_id.as_deref()) + .eq_nonblank("status", request.status.as_deref()) + .eq_nonblank("parent_thread_id", request.parent_thread_id.as_deref()) + .page(limit, offset)? + .fetch( + conn, + "workflow_runs", + "id, definition_id, parent_thread_id, input_json, phase_states_json, + child_run_ids_json, status, summary, started_at, updated_at, completed_at, + revision, lease_owner, lease_expires_at", + "updated_at DESC", + map_workflow_run_row, + )?; + let count = count as usize; tracing::debug!( "{LOG_PREFIX} list_workflow_runs.exit count={count} returned={}", runs.len() @@ -911,16 +849,7 @@ pub fn upsert_agent_team(workspace_dir: &Path, upsert: AgentTeamUpsert) -> Resul /// Fetch a single team by id. pub fn get_agent_team(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let team = get_agent_team_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team.exit id={id} found={}", - team.is_some() - ); - Ok(team) - }) + get_logged(workspace_dir, "get_agent_team", id, get_agent_team_inner) } /// List teams, most-recently-updated first, with optional thread/status filters. @@ -937,59 +866,24 @@ pub fn list_agent_teams( ); crate::store::with_connection(workspace_dir, |conn| { init_run_ledger_schema(conn)?; - let mut where_clauses = Vec::new(); - let mut values: Vec> = Vec::new(); - - if let Some(thread) = request - .parent_thread_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(thread.to_string())); - where_clauses.push(format!("parent_thread_id = ?{}", values.len())); - } - if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { - values.push(Box::new(status.to_string())); - where_clauses.push(format!("status = ?{}", values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - let count_sql = format!("SELECT COUNT(*) FROM agent_teams {where_sql}"); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { - row.get::<_, i64>(0) - })? as usize; - let limit = request.limit.unwrap_or(50).min(500) as i64; // `offset` is `u64`; convert checked so a value > i64::MAX surfaces a // clear error instead of wrapping negative and corrupting pagination. let offset = i64::try_from(request.offset.unwrap_or(0)) .storage_context("agent team list offset exceeds i64::MAX")?; - values.push(Box::new(limit)); - let limit_idx = values.len(); - values.push(Box::new(offset)); - let offset_idx = values.len(); - - let query_sql = format!( - "SELECT id, parent_thread_id, lead_agent_id, status, summary, - created_at, updated_at, closed_at - FROM agent_teams {where_sql} - ORDER BY updated_at DESC - LIMIT ?{limit_idx} OFFSET ?{offset_idx}" - ); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let mut stmt = conn.prepare(&query_sql)?; - let rows = stmt.query_map(params_ref.as_slice(), map_agent_team_row)?; - let mut teams = Vec::new(); - for row in rows { - teams.push(row?); - } + let (teams, count) = crate::paging::PagedQuery::default() + .eq_nonblank("parent_thread_id", request.parent_thread_id.as_deref()) + .eq_nonblank("status", request.status.as_deref()) + .page(limit, offset)? + .fetch( + conn, + "agent_teams", + "id, parent_thread_id, lead_agent_id, status, summary, + created_at, updated_at, closed_at", + "updated_at DESC", + map_agent_team_row, + )?; + let count = count as usize; tracing::debug!( "{LOG_PREFIX} list_agent_teams.exit count={count} returned={}", teams.len() @@ -1057,16 +951,12 @@ pub fn upsert_agent_team_member( /// Fetch a single member by id. pub fn get_agent_team_member(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team_member.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let member = get_agent_team_member_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team_member.exit id={id} found={}", - member.is_some() - ); - Ok(member) - }) + get_logged( + workspace_dir, + "get_agent_team_member", + id, + get_agent_team_member_inner, + ) } /// List all members of a team, by creation order. @@ -1182,16 +1072,12 @@ pub fn upsert_agent_team_task( /// Fetch a single task by id. pub fn get_agent_team_task(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team_task.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let task = get_agent_team_task_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team_task.exit id={id} found={}", - task.is_some() - ); - Ok(task) - }) + get_logged( + workspace_dir, + "get_agent_team_task", + id, + get_agent_team_task_inner, + ) } /// List all tasks of a team, by `order_index` then creation order. @@ -1673,42 +1559,42 @@ pub fn release_agent_team_task(workspace_dir: &Path, team_id: &str, task_id: &st /// Connection-scoped team lookup, so an upsert can read its own write back /// inside the same transaction. fn get_agent_team_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( + get_row_by_id( + conn, "SELECT id, parent_thread_id, lead_agent_id, status, summary, created_at, updated_at, closed_at FROM agent_teams WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_row) - .optional() - .map_err(Into::into) + id, + map_agent_team_row, + ) } /// Connection-scoped member lookup, so an upsert can read its own write back /// inside the same transaction. fn get_agent_team_member_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( + get_row_by_id( + conn, "SELECT id, team_id, name, agent_id, member_status, current_task_id, worker_thread_id, run_id, created_at, updated_at FROM agent_team_members WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_member_row) - .optional() - .map_err(Into::into) + id, + map_agent_team_member_row, + ) } /// Connection-scoped task lookup, so a claim/completion transaction can read /// its own write back inside the same transaction. fn get_agent_team_task_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( + get_row_by_id( + conn, "SELECT id, team_id, title, objective, status, owner_member_id, claimed_by_member_id, claim_token, depends_on_json, gate_status, gate_reason, evidence_json, source_run_id, order_index, created_at, updated_at FROM agent_team_tasks WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_task_row) - .optional() - .map_err(Into::into) + id, + map_agent_team_task_row, + ) } fn map_agent_team_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { @@ -1763,16 +1649,16 @@ fn map_agent_team_task_row(row: &rusqlite::Row<'_>) -> rusqlite::Result Result> { - let mut stmt = conn.prepare( + get_row_by_id( + conn, "SELECT id, kind, parent_run_id, parent_thread_id, agent_id, status, prompt_ref, worker_thread_id, checkpoint_path, checkpoint_json, summary, error, metadata_json, started_at, updated_at, completed_at FROM agent_runs WHERE id = ?1", - )?; - stmt.query_row(params![id], |row| map_agent_run_row(conn, row)) - .optional() - .map_err(Into::into) + id, + |row| map_agent_run_row(conn, row), + ) } /// Connection-scoped telemetry lookup that errors when the row is absent — diff --git a/crates/tinyagents-session/src/run_ledger/ops/rows.rs b/crates/tinyagents-session/src/run_ledger/ops/rows.rs deleted file mode 100644 index 57eac3e30..000000000 --- a/crates/tinyagents-session/src/run_ledger/ops/rows.rs +++ /dev/null @@ -1,255 +0,0 @@ -//! SQLite row decoding and transaction-scoped lookup helpers. -//! -//! Keeping this detail outside the public ledger operations makes the -//! operation module focused on state transitions and query semantics. - -use chrono::{DateTime, Utc}; -use rusqlite::{Connection, OptionalExtension, params}; -use serde_json::{Value, json}; - -use tinyagents_harness::error::Result; - -use super::super::types::{ - AgentRun, AgentRunKind, AgentRunStatus, AgentTeam, AgentTeamMember, AgentTeamMemberStatus, - AgentTeamStatus, AgentTeamTask, AgentTeamTaskStatus, RunEvent, RunTelemetry, WorkflowRun, - WorkflowRunStatus, -}; -use crate::context::StorageContext; - -/// Connection-scoped team lookup, so an upsert can read its own write back -/// inside the same transaction. -pub(super) fn get_agent_team_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( - "SELECT id, parent_thread_id, lead_agent_id, status, summary, - created_at, updated_at, closed_at - FROM agent_teams WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_row) - .optional() - .map_err(Into::into) -} - -/// Connection-scoped member lookup, so an upsert can read its own write back -/// inside the same transaction. -pub(super) fn get_agent_team_member_inner( - conn: &Connection, - id: &str, -) -> Result> { - let mut stmt = conn.prepare( - "SELECT id, team_id, name, agent_id, member_status, - current_task_id, worker_thread_id, run_id, created_at, updated_at - FROM agent_team_members WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_member_row) - .optional() - .map_err(Into::into) -} - -/// Connection-scoped task lookup, so a claim/completion transaction can read -/// its own write back inside the same transaction. -pub(super) fn get_agent_team_task_inner( - conn: &Connection, - id: &str, -) -> Result> { - let mut stmt = conn.prepare( - "SELECT id, team_id, title, objective, status, owner_member_id, - claimed_by_member_id, claim_token, depends_on_json, gate_status, - gate_reason, evidence_json, source_run_id, order_index, - created_at, updated_at - FROM agent_team_tasks WHERE id = ?1", - )?; - stmt.query_row(params![id], map_agent_team_task_row) - .optional() - .map_err(Into::into) -} - -pub(super) fn map_agent_team_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(AgentTeam { - id: row.get(0)?, - parent_thread_id: row.get(1)?, - lead_agent_id: row.get(2)?, - status: AgentTeamStatus::parse(&row.get::<_, String>(3)?), - summary: row.get(4)?, - created_at: parse_rfc3339(&row.get::<_, String>(5)?)?, - updated_at: parse_rfc3339(&row.get::<_, String>(6)?)?, - closed_at: parse_rfc3339_opt(row.get(7)?)?, - }) -} - -pub(super) fn map_agent_team_member_row( - row: &rusqlite::Row<'_>, -) -> rusqlite::Result { - Ok(AgentTeamMember { - id: row.get(0)?, - team_id: row.get(1)?, - name: row.get(2)?, - agent_id: row.get(3)?, - member_status: AgentTeamMemberStatus::parse(&row.get::<_, String>(4)?), - current_task_id: row.get(5)?, - worker_thread_id: row.get(6)?, - run_id: row.get(7)?, - created_at: parse_rfc3339(&row.get::<_, String>(8)?)?, - updated_at: parse_rfc3339(&row.get::<_, String>(9)?)?, - }) -} - -pub(super) fn map_agent_team_task_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(AgentTeamTask { - id: row.get(0)?, - team_id: row.get(1)?, - title: row.get(2)?, - objective: row.get(3)?, - status: AgentTeamTaskStatus::parse(&row.get::<_, String>(4)?), - owner_member_id: row.get(5)?, - claimed_by_member_id: row.get(6)?, - claim_token: row.get(7)?, - depends_on: serde_json::from_str(&row.get::<_, String>(8)?).unwrap_or_default(), - gate_status: row.get(9)?, - gate_reason: row.get(10)?, - evidence: serde_json::from_str(&row.get::<_, String>(11)?).unwrap_or_default(), - source_run_id: row.get(12)?, - order_index: row.get(13)?, - created_at: parse_rfc3339(&row.get::<_, String>(14)?)?, - updated_at: parse_rfc3339(&row.get::<_, String>(15)?)?, - }) -} - -/// Connection-scoped agent-run lookup, so an upsert can read its own write -/// back inside the same transaction. -pub(super) fn get_agent_run_inner(conn: &Connection, id: &str) -> Result> { - let mut stmt = conn.prepare( - "SELECT id, kind, parent_run_id, parent_thread_id, agent_id, status, - prompt_ref, worker_thread_id, - checkpoint_path, checkpoint_json, summary, error, metadata_json, - started_at, updated_at, completed_at - FROM agent_runs WHERE id = ?1", - )?; - stmt.query_row(params![id], |row| map_agent_run_row(conn, row)) - .optional() - .map_err(Into::into) -} - -/// Connection-scoped telemetry lookup that errors when the row is absent — -/// used right after [`upsert_run_telemetry`] writes it, where a miss means -/// the write silently failed. -pub(super) fn get_run_telemetry_inner(conn: &Connection, run_id: &str) -> Result { - let mut stmt = conn.prepare( - "SELECT run_id, input_tokens, output_tokens, cached_input_tokens, cost_usd, - elapsed_ms, tool_count, model, provider, error, updated_at - FROM run_telemetry WHERE run_id = ?1", - )?; - stmt.query_row(params![run_id], map_run_telemetry_row) - .storage_context("run telemetry missing after upsert") -} - -/// Connection-scoped telemetry lookup used when joining telemetry onto an -/// [`AgentRun`], where no telemetry row yet existing is a normal `None` -/// rather than an error. -pub(super) fn get_optional_run_telemetry( - conn: &Connection, - run_id: &str, -) -> rusqlite::Result> { - let mut stmt = conn.prepare( - "SELECT run_id, input_tokens, output_tokens, cached_input_tokens, cost_usd, - elapsed_ms, tool_count, model, provider, error, updated_at - FROM run_telemetry WHERE run_id = ?1", - )?; - stmt.query_row(params![run_id], map_run_telemetry_row) - .optional() -} - -pub(super) fn map_agent_run_row( - conn: &Connection, - row: &rusqlite::Row<'_>, -) -> rusqlite::Result { - let id: String = row.get(0)?; - let checkpoint_json: Option = row.get(9)?; - let metadata_json: String = row.get(12)?; - Ok(AgentRun { - id: id.clone(), - kind: AgentRunKind::parse(&row.get::<_, String>(1)?), - parent_run_id: row.get(2)?, - parent_thread_id: row.get(3)?, - agent_id: row.get(4)?, - status: AgentRunStatus::parse(&row.get::<_, String>(5)?), - prompt_ref: row.get(6)?, - worker_thread_id: row.get(7)?, - checkpoint_path: row.get(8)?, - checkpoint: parse_json_opt(checkpoint_json), - summary: row.get(10)?, - error: row.get(11)?, - metadata: parse_json(metadata_json), - telemetry: get_optional_run_telemetry(conn, &id)?, - started_at: parse_rfc3339(&row.get::<_, String>(13)?)?, - updated_at: parse_rfc3339(&row.get::<_, String>(14)?)?, - completed_at: parse_rfc3339_opt(row.get(15)?)?, - }) -} - -pub(super) fn map_workflow_run_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(WorkflowRun { - id: row.get(0)?, - definition_id: row.get(1)?, - parent_thread_id: row.get(2)?, - input: parse_json(row.get(3)?), - phase_states: parse_json(row.get(4)?), - child_run_ids: serde_json::from_str(&row.get::<_, String>(5)?).unwrap_or_default(), - status: WorkflowRunStatus::parse(&row.get::<_, String>(6)?), - summary: row.get(7)?, - started_at: parse_rfc3339(&row.get::<_, String>(8)?)?, - updated_at: parse_rfc3339(&row.get::<_, String>(9)?)?, - completed_at: parse_rfc3339_opt(row.get(10)?)?, - revision: row.get::<_, i64>(11)? as u64, - lease_owner: row.get(12)?, - lease_expires_at: parse_rfc3339_opt(row.get(13)?)?, - }) -} - -pub(super) fn map_run_event_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(RunEvent { - run_id: row.get(0)?, - sequence: row.get::<_, i64>(1)? as u64, - event_type: row.get(2)?, - payload: parse_json(row.get(3)?), - timestamp: parse_rfc3339(&row.get::<_, String>(4)?)?, - }) -} - -pub(super) fn map_run_telemetry_row(row: &rusqlite::Row<'_>) -> rusqlite::Result { - Ok(RunTelemetry { - run_id: row.get(0)?, - input_tokens: row.get::<_, i64>(1)? as u64, - output_tokens: row.get::<_, i64>(2)? as u64, - cached_input_tokens: row.get::<_, i64>(3)? as u64, - cost_usd: row.get(4)?, - elapsed_ms: row.get::<_, Option>(5)?.map(|v| v as u64), - tool_count: row.get::<_, i64>(6)? as u64, - model: row.get(7)?, - provider: row.get(8)?, - error: row.get(9)?, - updated_at: Some(parse_rfc3339(&row.get::<_, String>(10)?)?), - }) -} - -pub(super) fn parse_json(raw: String) -> Value { - serde_json::from_str(&raw).unwrap_or_else(|_| json!({})) -} - -pub(super) fn parse_json_opt(raw: Option) -> Option { - raw.and_then(|value| serde_json::from_str(&value).ok()) -} - -pub(super) fn parse_rfc3339(raw: &str) -> rusqlite::Result> { - DateTime::parse_from_rfc3339(raw) - .map(|dt| dt.with_timezone(&Utc)) - .map_err(|err| { - rusqlite::Error::FromSqlConversionFailure(0, rusqlite::types::Type::Text, Box::new(err)) - }) -} - -pub(super) fn parse_rfc3339_opt(raw: Option) -> rusqlite::Result>> { - match raw { - Some(value) => parse_rfc3339(&value).map(Some), - None => Ok(None), - } -} diff --git a/crates/tinyagents-session/src/run_ledger/ops/team.rs b/crates/tinyagents-session/src/run_ledger/ops/team.rs deleted file mode 100644 index 5a68fa6fd..000000000 --- a/crates/tinyagents-session/src/run_ledger/ops/team.rs +++ /dev/null @@ -1,837 +0,0 @@ -//! Durable agent-team coordination operations. -//! -//! Teams, members, and tasks share the run-ledger database but have distinct -//! lifecycle and claim semantics from agent and workflow runs. - -use std::path::Path; - -use chrono::Utc; -use rusqlite::{Connection, OptionalExtension, params}; - -use tinyagents_harness::error::Result; - -use super::super::store::init_run_ledger_schema; -use super::super::types::{ - AgentTeam, AgentTeamListRequest, AgentTeamListResponse, AgentTeamMember, AgentTeamMemberUpsert, - AgentTeamTask, AgentTeamTaskStatus, AgentTeamTaskUpsert, AgentTeamUpsert, ClaimOutcome, - CompletionOutcome, -}; -use super::LOG_PREFIX; -use super::rows::{ - get_agent_team_inner, get_agent_team_member_inner, get_agent_team_task_inner, - map_agent_team_member_row, map_agent_team_row, map_agent_team_task_row, -}; -use crate::context::StorageContext; - -// --------------------------------------------------------------------------- -// Agent-team coordination (issue #3374) -// --------------------------------------------------------------------------- - -/// Insert or update a team row. -pub fn upsert_agent_team(workspace_dir: &Path, upsert: AgentTeamUpsert) -> Result { - let now = Utc::now(); - let created_at = upsert.created_at.unwrap_or(now); - tracing::debug!( - "{LOG_PREFIX} upsert_agent_team.entry id={} lead={} status={}", - upsert.id, - upsert.lead_agent_id, - upsert.status.as_str() - ); - // One transaction for the write *and* the read-back. Committing the insert - // on an autocommit connection, closing it, then re-opening to `get_*` hands - // the caller whatever a concurrent writer left behind rather than what this - // call wrote — an upsert that reports someone else's row. - let team = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - conn.execute( - "INSERT INTO agent_teams ( - id, parent_thread_id, lead_agent_id, status, summary, - created_at, updated_at, closed_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8) - ON CONFLICT(id) DO UPDATE SET - parent_thread_id = COALESCE(excluded.parent_thread_id, agent_teams.parent_thread_id), - lead_agent_id = excluded.lead_agent_id, - status = excluded.status, - summary = COALESCE(excluded.summary, agent_teams.summary), - updated_at = excluded.updated_at, - closed_at = COALESCE(excluded.closed_at, agent_teams.closed_at)", - params![ - upsert.id, - upsert.parent_thread_id, - upsert.lead_agent_id, - upsert.status.as_str(), - upsert.summary, - created_at.to_rfc3339(), - now.to_rfc3339(), - upsert.closed_at.map(|dt| dt.to_rfc3339()), - ], - ) - .storage_context("upsert agent team")?; - get_agent_team_inner(conn, &upsert.id)?.storage_context("agent team missing after upsert") - })?; - tracing::debug!("{LOG_PREFIX} upsert_agent_team.exit id={}", team.id); - Ok(team) -} - -/// Fetch a single team by id. -pub fn get_agent_team(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let team = get_agent_team_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team.exit id={id} found={}", - team.is_some() - ); - Ok(team) - }) -} - -/// List teams, most-recently-updated first, with optional thread/status filters. -pub fn list_agent_teams( - workspace_dir: &Path, - request: &AgentTeamListRequest, -) -> Result { - tracing::debug!( - "{LOG_PREFIX} list_agent_teams.entry parent_thread={:?} status={:?} limit={:?} offset={:?}", - request.parent_thread_id, - request.status, - request.limit, - request.offset - ); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let mut where_clauses = Vec::new(); - let mut values: Vec> = Vec::new(); - - if let Some(thread) = request - .parent_thread_id - .as_deref() - .filter(|s| !s.trim().is_empty()) - { - values.push(Box::new(thread.to_string())); - where_clauses.push(format!("parent_thread_id = ?{}", values.len())); - } - if let Some(status) = request.status.as_deref().filter(|s| !s.trim().is_empty()) { - values.push(Box::new(status.to_string())); - where_clauses.push(format!("status = ?{}", values.len())); - } - - let where_sql = if where_clauses.is_empty() { - String::new() - } else { - format!("WHERE {}", where_clauses.join(" AND ")) - }; - let count_sql = format!("SELECT COUNT(*) FROM agent_teams {where_sql}"); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let count = conn.query_row(&count_sql, params_ref.as_slice(), |row| { - row.get::<_, i64>(0) - })? as usize; - - let limit = request.limit.unwrap_or(50).min(500) as i64; - // `offset` is `u64`; convert checked so a value > i64::MAX surfaces a - // clear error instead of wrapping negative and corrupting pagination. - let offset = i64::try_from(request.offset.unwrap_or(0)) - .storage_context("agent team list offset exceeds i64::MAX")?; - values.push(Box::new(limit)); - let limit_idx = values.len(); - values.push(Box::new(offset)); - let offset_idx = values.len(); - - let query_sql = format!( - "SELECT id, parent_thread_id, lead_agent_id, status, summary, - created_at, updated_at, closed_at - FROM agent_teams {where_sql} - ORDER BY updated_at DESC - LIMIT ?{limit_idx} OFFSET ?{offset_idx}" - ); - let params_ref: Vec<&dyn rusqlite::types::ToSql> = - values.iter().map(|v| v.as_ref()).collect(); - let mut stmt = conn.prepare(&query_sql)?; - let rows = stmt.query_map(params_ref.as_slice(), map_agent_team_row)?; - let mut teams = Vec::new(); - for row in rows { - teams.push(row?); - } - tracing::debug!( - "{LOG_PREFIX} list_agent_teams.exit count={count} returned={}", - teams.len() - ); - Ok(AgentTeamListResponse { teams, count }) - }) -} - -/// Insert or update a team member. `UNIQUE(team_id, name)` enforces unique names. -pub fn upsert_agent_team_member( - workspace_dir: &Path, - upsert: AgentTeamMemberUpsert, -) -> Result { - let now = Utc::now(); - let created_at = upsert.created_at.unwrap_or(now); - tracing::debug!( - "{LOG_PREFIX} upsert_agent_team_member.entry id={} team={} name={} status={}", - upsert.id, - upsert.team_id, - upsert.name, - upsert.member_status.as_str() - ); - // One transaction for the write *and* the read-back. Committing the insert - // on an autocommit connection, closing it, then re-opening to `get_*` hands - // the caller whatever a concurrent writer left behind rather than what this - // call wrote — an upsert that reports someone else's row. - let member = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - conn.execute( - "INSERT INTO agent_team_members ( - id, team_id, name, agent_id, member_status, - current_task_id, worker_thread_id, run_id, created_at, updated_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10) - ON CONFLICT(id) DO UPDATE SET - name = excluded.name, - agent_id = COALESCE(excluded.agent_id, agent_team_members.agent_id), - member_status = excluded.member_status, - current_task_id = COALESCE(excluded.current_task_id, agent_team_members.current_task_id), - worker_thread_id = COALESCE(excluded.worker_thread_id, agent_team_members.worker_thread_id), - run_id = COALESCE(excluded.run_id, agent_team_members.run_id), - updated_at = excluded.updated_at", - params![ - upsert.id, - upsert.team_id, - upsert.name, - upsert.agent_id, - upsert.member_status.as_str(), - upsert.current_task_id, - upsert.worker_thread_id, - upsert.run_id, - created_at.to_rfc3339(), - now.to_rfc3339(), - ], - ) - .storage_context("upsert agent team member")?; - get_agent_team_member_inner(conn, &upsert.id)? - .storage_context("agent team member missing after upsert") - })?; - tracing::debug!( - "{LOG_PREFIX} upsert_agent_team_member.exit id={}", - member.id - ); - Ok(member) -} - -/// Fetch a single member by id. -pub fn get_agent_team_member(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team_member.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let member = get_agent_team_member_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team_member.exit id={id} found={}", - member.is_some() - ); - Ok(member) - }) -} - -/// List all members of a team, by creation order. -pub fn list_agent_team_members( - workspace_dir: &Path, - team_id: &str, -) -> Result> { - tracing::debug!("{LOG_PREFIX} list_agent_team_members.entry team={team_id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let mut stmt = conn.prepare( - "SELECT id, team_id, name, agent_id, member_status, - current_task_id, worker_thread_id, run_id, created_at, updated_at - FROM agent_team_members WHERE team_id = ?1 - ORDER BY created_at ASC", - )?; - let rows = stmt.query_map(params![team_id], map_agent_team_member_row)?; - let mut members = Vec::new(); - for row in rows { - members.push(row?); - } - tracing::debug!( - "{LOG_PREFIX} list_agent_team_members.exit team={team_id} count={}", - members.len() - ); - Ok(members) - }) -} - -/// Insert or update a team task. -pub fn upsert_agent_team_task( - workspace_dir: &Path, - upsert: AgentTeamTaskUpsert, -) -> Result { - let now = Utc::now(); - let created_at = upsert.created_at.unwrap_or(now); - let depends_on_json = - serde_json::to_string(&upsert.depends_on).storage_context("serialize task depends_on")?; - let evidence_json = - serde_json::to_string(&upsert.evidence).storage_context("serialize task evidence")?; - let gate_status = upsert.gate_status.unwrap_or_else(|| "pending".to_string()); - tracing::debug!( - "{LOG_PREFIX} upsert_agent_team_task.entry id={} team={} status={} deps={}", - upsert.id, - upsert.team_id, - upsert.status.as_str(), - upsert.depends_on.len() - ); - // One transaction for the write *and* the read-back. Committing the insert - // on an autocommit connection, closing it, then re-opening to `get_*` hands - // the caller whatever a concurrent writer left behind rather than what this - // call wrote — an upsert that reports someone else's row. - let task = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - conn.execute( - "INSERT INTO agent_team_tasks ( - id, team_id, title, objective, status, owner_member_id, - claimed_by_member_id, claim_token, depends_on_json, gate_status, - gate_reason, evidence_json, source_run_id, order_index, - created_at, updated_at - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, NULL, NULL, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) - ON CONFLICT(id) DO UPDATE SET - title = excluded.title, - objective = COALESCE(excluded.objective, agent_team_tasks.objective), - status = excluded.status, - -- A claim only means anything while the task is in_progress. - -- Editing a claimed task through this upsert (to change - -- dependencies, or to reset status during recovery) used to - -- leave claimed_by_member_id/claim_token set on a todo row, - -- which strands it: a fresh claim returns AlreadyClaimed, - -- completion returns NotClaimed, and release/shutdown skip it - -- because they only match in_progress. Drop the claim whenever - -- the new status is not in_progress; preserve it otherwise so - -- an unrelated edit does not steal a live claim. - claimed_by_member_id = CASE - WHEN excluded.status = 'in_progress' - THEN agent_team_tasks.claimed_by_member_id ELSE NULL END, - claim_token = CASE - WHEN excluded.status = 'in_progress' - THEN agent_team_tasks.claim_token ELSE NULL END, - owner_member_id = COALESCE(excluded.owner_member_id, agent_team_tasks.owner_member_id), - depends_on_json = excluded.depends_on_json, - gate_status = excluded.gate_status, - gate_reason = COALESCE(excluded.gate_reason, agent_team_tasks.gate_reason), - evidence_json = excluded.evidence_json, - source_run_id = COALESCE(excluded.source_run_id, agent_team_tasks.source_run_id), - order_index = excluded.order_index, - updated_at = excluded.updated_at", - params![ - upsert.id, - upsert.team_id, - upsert.title, - upsert.objective, - upsert.status.as_str(), - upsert.owner_member_id, - depends_on_json, - gate_status, - upsert.gate_reason, - evidence_json, - upsert.source_run_id, - upsert.order_index, - created_at.to_rfc3339(), - now.to_rfc3339(), - ], - ) - .storage_context("upsert agent team task")?; - get_agent_team_task_inner(conn, &upsert.id)? - .storage_context("agent team task missing after upsert") - })?; - tracing::debug!("{LOG_PREFIX} upsert_agent_team_task.exit id={}", task.id); - Ok(task) -} - -/// Fetch a single task by id. -pub fn get_agent_team_task(workspace_dir: &Path, id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} get_agent_team_task.entry id={id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let task = get_agent_team_task_inner(conn, id)?; - tracing::debug!( - "{LOG_PREFIX} get_agent_team_task.exit id={id} found={}", - task.is_some() - ); - Ok(task) - }) -} - -/// List all tasks of a team, by `order_index` then creation order. -pub fn list_agent_team_tasks(workspace_dir: &Path, team_id: &str) -> Result> { - tracing::debug!("{LOG_PREFIX} list_agent_team_tasks.entry team={team_id}"); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let mut stmt = conn.prepare( - "SELECT id, team_id, title, objective, status, owner_member_id, - claimed_by_member_id, claim_token, depends_on_json, gate_status, - gate_reason, evidence_json, source_run_id, order_index, - created_at, updated_at - FROM agent_team_tasks WHERE team_id = ?1 - ORDER BY order_index ASC, created_at ASC", - )?; - let rows = stmt.query_map(params![team_id], map_agent_team_task_row)?; - let mut tasks = Vec::new(); - for row in rows { - tasks.push(row?); - } - tracing::debug!( - "{LOG_PREFIX} list_agent_team_tasks.exit team={team_id} count={}", - tasks.len() - ); - Ok(tasks) - }) -} - -/// Atomically claim a task for a member. -/// -/// All steps run inside a single `with_connection` transaction so that the -/// dependency check and the compare-and-swap observe a consistent snapshot: -/// 1. Resolve the task by `(id, team_id)`; absent → [`ClaimOutcome::UnknownTask`]. -/// 2. For every dependency id, look up its status; collect those not `done` -/// into `unmet`. Non-empty → [`ClaimOutcome::Blocked`]. -/// 3. WHERE-guarded `UPDATE ... WHERE claimed_by_member_id IS NULL`: SQLite -/// serializes writers, so exactly one concurrent claimer flips the row from -/// unclaimed to claimed. `rows_affected == 0` → already taken -/// ([`ClaimOutcome::AlreadyClaimed`]); otherwise re-fetch and return -/// [`ClaimOutcome::Claimed`]. -pub fn claim_agent_team_task( - workspace_dir: &Path, - team_id: &str, - task_id: &str, - member_id: &str, - claim_token: &str, -) -> Result { - tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.entry team={team_id} task={task_id} member={member_id}" - ); - let outcome = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - - // 1. Resolve the task within this team. - let task = match get_agent_team_task_inner(conn, task_id)? { - Some(task) if task.team_id == team_id => task, - _ => { - tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.unknown team={team_id} task={task_id}" - ); - return Ok(ClaimOutcome::UnknownTask); - } - }; - - // 2. Dependency gate: every dep must be `done`. - let mut unmet = Vec::new(); - for dep_id in &task.depends_on { - let dep_status: Option = conn - .query_row( - "SELECT status FROM agent_team_tasks WHERE id = ?1 AND team_id = ?2", - params![dep_id, team_id], - |row| row.get(0), - ) - .optional()?; - let is_done = dep_status.as_deref() == Some(AgentTeamTaskStatus::Done.as_str()); - if !is_done { - unmet.push(dep_id.clone()); - } - } - if !unmet.is_empty() { - tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.blocked team={team_id} task={task_id} unmet={}", - unmet.len() - ); - return Ok(ClaimOutcome::Blocked { unmet }); - } - - // 3. Compare-and-swap on the unclaimed guard **and the status**. - // - // `claimed_by_member_id IS NULL` alone is not a guard: `upsert_agent_team_task` - // deliberately NULLs that column whenever the new status is not - // `in_progress`, so every `done` task also satisfies it. A stale worker - // re-claiming a finished task therefore flipped it straight back to - // `in_progress` — and stranded everything downstream, because the - // completion gate re-checks that each dependency is still `done` and now - // reports it unfinished. A terminal task is not claimable, whatever its - // claim column says. - let now = Utc::now(); - let rows_affected = conn - .execute( - "UPDATE agent_team_tasks - SET claimed_by_member_id = ?1, claim_token = ?2, status = 'in_progress', updated_at = ?3 - WHERE id = ?4 AND team_id = ?5 AND claimed_by_member_id IS NULL - AND status IN ('todo', 'ready', 'blocked')", - params![member_id, claim_token, now.to_rfc3339(), task_id, team_id], - ) - .storage_context("compare-and-swap claim agent team task")?; - if rows_affected == 0 { - tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.already_claimed team={team_id} task={task_id} \ - status={}", - task.status.as_str() - ); - return Ok(ClaimOutcome::AlreadyClaimed); - } - - let claimed = get_agent_team_task_inner(conn, task_id)? - .storage_context("claimed task missing after compare-and-swap")?; - Ok(ClaimOutcome::Claimed(Box::new(claimed))) - })?; - tracing::debug!( - "{LOG_PREFIX} claim_agent_team_task.exit team={team_id} task={task_id} outcome={}", - match &outcome { - ClaimOutcome::Claimed(_) => "claimed", - ClaimOutcome::AlreadyClaimed => "already_claimed", - ClaimOutcome::Blocked { .. } => "blocked", - ClaimOutcome::UnknownTask => "unknown", - } - ); - Ok(outcome) -} - -/// Quality-gate a task's completion and, on pass, transition it to `done`. -/// -/// Runs inside a single transaction so the gate evaluation and the status flip -/// observe one consistent snapshot: -/// 1. Resolve the task by `(id, team_id)`; absent → [`CompletionOutcome::UnknownTask`]. -/// 2. The completer must be the current claimant and the task must be -/// `in_progress`; otherwise [`CompletionOutcome::NotClaimed`]. -/// 3. Evaluate the quality gate (every dependency `done`, claimant matches any -/// pre-assigned owner, evidence present when `require_evidence`). Any unmet -/// invariant records `gate_status = "failed"` + the joined reasons and leaves -/// the task `in_progress` → [`CompletionOutcome::GateFailed`]. -/// 4. On pass, merge `evidence`, set `status = "done"`, `gate_status = "passed"`, -/// clear `gate_reason`, re-fetch → [`CompletionOutcome::Completed`]. -pub fn complete_agent_team_task( - workspace_dir: &Path, - team_id: &str, - task_id: &str, - member_id: &str, - evidence: &[String], - require_evidence: bool, -) -> Result { - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.entry team={team_id} task={task_id} member={member_id}" - ); - let outcome = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - - // 1. Resolve the task within this team. - let task = match get_agent_team_task_inner(conn, task_id)? { - Some(task) if task.team_id == team_id => task, - _ => { - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.unknown team={team_id} task={task_id}" - ); - return Ok(CompletionOutcome::UnknownTask); - } - }; - - // 2. Only the current claimant may complete, and only while in progress. - let is_claimant = task.claimed_by_member_id.as_deref() == Some(member_id); - let in_progress = task.status == AgentTeamTaskStatus::InProgress; - if !is_claimant || !in_progress { - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.not_claimed team={team_id} task={task_id} claimant={is_claimant} in_progress={in_progress}" - ); - return Ok(CompletionOutcome::NotClaimed); - } - - // Merge prior evidence with the newly-supplied links (de-duplicated, - // order-preserving) so a retry that adds evidence accumulates it. - let mut merged_evidence = task.evidence.clone(); - for link in evidence { - if !merged_evidence.iter().any(|e| e == link) { - merged_evidence.push(link.clone()); - } - } - - // 3. Quality gate. - let reasons = - evaluate_completion_gate(conn, team_id, &task, &merged_evidence, require_evidence)?; - let now = Utc::now(); - if !reasons.is_empty() { - let joined = reasons.join("; "); - // Persist the merged evidence even though the gate failed. Evidence - // accumulates across attempts (see `merged_evidence` above), so - // dropping it here punished a caller for an unrelated gate failure: - // after a dependency was fixed, a retry that did not resend the same - // links would fail `require_evidence` on evidence it had already - // submitted. Only the gate verdict is a failure; the submission is - // still real. - let evidence_json = serde_json::to_string(&merged_evidence) - .storage_context("serialize completion evidence")?; - conn.execute( - "UPDATE agent_team_tasks - SET gate_status = 'failed', gate_reason = ?1, - evidence_json = ?2, updated_at = ?3 - WHERE id = ?4 AND team_id = ?5", - params![joined, evidence_json, now.to_rfc3339(), task_id, team_id], - ) - .storage_context("record failed completion gate")?; - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.gate_failed team={team_id} task={task_id} reasons={}", - reasons.len() - ); - return Ok(CompletionOutcome::GateFailed { reasons }); - } - - // 4. Gate passed — flip to done. The WHERE clause is the real CAS: the - // `claimed_by_member_id` guard stops a concurrent shutdown/unclaim from - // completing a task it no longer holds, and the `status = 'in_progress'` - // guard stops a concurrent double-complete by the same member (the - // snapshot check above is a read, not part of the swap — only one of two - // racing UPDATEs flips `in_progress -> done`). - let evidence_json = serde_json::to_string(&merged_evidence) - .storage_context("serialize completion evidence")?; - let rows_affected = conn - .execute( - "UPDATE agent_team_tasks - SET status = 'done', gate_status = 'passed', gate_reason = NULL, - evidence_json = ?1, updated_at = ?2 - WHERE id = ?3 AND team_id = ?4 AND claimed_by_member_id = ?5 - AND status = 'in_progress'", - params![evidence_json, now.to_rfc3339(), task_id, team_id, member_id], - ) - .storage_context("complete agent team task")?; - if rows_affected == 0 { - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.lost_claim team={team_id} task={task_id}" - ); - return Ok(CompletionOutcome::NotClaimed); - } - - let done = get_agent_team_task_inner(conn, task_id)? - .storage_context("completed task missing after update")?; - Ok(CompletionOutcome::Completed(Box::new(done))) - })?; - tracing::debug!( - "{LOG_PREFIX} complete_agent_team_task.exit team={team_id} task={task_id} outcome={}", - match &outcome { - CompletionOutcome::Completed(_) => "completed", - CompletionOutcome::GateFailed { .. } => "gate_failed", - CompletionOutcome::NotClaimed => "not_claimed", - CompletionOutcome::UnknownTask => "unknown", - } - ); - Ok(outcome) -} - -/// Evaluate the quality-gate invariants for a completing task. Returns one -/// human-readable reason per unmet invariant (empty = gate passes). -fn evaluate_completion_gate( - conn: &Connection, - team_id: &str, - task: &AgentTeamTask, - merged_evidence: &[String], - require_evidence: bool, -) -> Result> { - let mut reasons = Vec::new(); - - // Every dependency must still be `done` (defends against a dependency that - // regressed after this task was claimed). - for dep_id in &task.depends_on { - let dep_status: Option = conn - .query_row( - "SELECT status FROM agent_team_tasks WHERE id = ?1 AND team_id = ?2", - params![dep_id, team_id], - |row| row.get(0), - ) - .optional()?; - if dep_status.as_deref() != Some(AgentTeamTaskStatus::Done.as_str()) { - reasons.push(format!("dependency {dep_id} is not done")); - } - } - - // No overlapping ownership: a pre-assigned owner must be the one completing. - if let Some(owner) = task - .owner_member_id - .as_ref() - .filter(|owner| Some(owner.as_str()) != task.claimed_by_member_id.as_deref()) - { - reasons.push(format!( - "task is owned by {owner} but claimed by {}", - task.claimed_by_member_id.as_deref().unwrap_or("nobody") - )); - } - - // Evidence gate. - if require_evidence && merged_evidence.is_empty() { - reasons.push("completion requires at least one evidence link".to_string()); - } - - Ok(reasons) -} - -/// Stop a team member and release any task it is actively working on. -/// -/// In one transaction: unclaim the member's `in_progress` tasks back to `todo` -/// (clearing claimant + token so another teammate can pick them up), then mark -/// the member `stopped` and clear its `current_task_id`. Returns the updated -/// member plus the ids of the tasks that were released, or `None` if the member -/// is not part of the team. -pub fn shutdown_agent_team_member( - workspace_dir: &Path, - team_id: &str, - member_id: &str, -) -> Result)>> { - tracing::debug!( - "{LOG_PREFIX} shutdown_agent_team_member.entry team={team_id} member={member_id}" - ); - let result = crate::store::with_transaction(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - - // Existence + team-membership check only; the row is intentionally not - // reused — the caller-facing member is re-read after the UPDATEs below so - // it reflects the stopped state. - match get_agent_team_member_inner(conn, member_id)? { - Some(found) if found.team_id == team_id => {} - _ => { - tracing::debug!( - "{LOG_PREFIX} shutdown_agent_team_member.unknown team={team_id} member={member_id}" - ); - return Ok(None); - } - } - - // Collect the ids first so the caller can report exactly what was freed. - let released: Vec = { - let mut stmt = conn.prepare( - "SELECT id FROM agent_team_tasks - WHERE team_id = ?1 AND claimed_by_member_id = ?2 AND status = 'in_progress'", - )?; - let ids = stmt.query_map(params![team_id, member_id], |row| row.get::<_, String>(0))?; - let mut out = Vec::new(); - for id in ids { - out.push(id?); - } - out - }; - - let now = Utc::now(); - conn.execute( - "UPDATE agent_team_tasks - SET claimed_by_member_id = NULL, claim_token = NULL, status = 'todo', updated_at = ?1 - WHERE team_id = ?2 AND claimed_by_member_id = ?3 AND status = 'in_progress'", - params![now.to_rfc3339(), team_id, member_id], - ) - .storage_context("release tasks on member shutdown")?; - conn.execute( - "UPDATE agent_team_members - SET member_status = 'stopped', current_task_id = NULL, updated_at = ?1 - WHERE id = ?2 AND team_id = ?3", - params![now.to_rfc3339(), member_id, team_id], - ) - .storage_context("stop agent team member")?; - - let member = get_agent_team_member_inner(conn, member_id)? - .storage_context("member missing after shutdown")?; - Ok(Some((member, released))) - })?; - tracing::debug!( - "{LOG_PREFIX} shutdown_agent_team_member.exit team={team_id} member={member_id} released={}", - result.as_ref().map(|(_, r)| r.len()).unwrap_or(0) - ); - Ok(result) -} - -/// Mark a member as actively running a task: status → `active`, with the -/// current task id and the worker/run identifiers of the spawned agent. Used by -/// the live runtime right after it claims a task and dispatches a worker. -/// Returns the updated member, or `None` if the member is not in the team. -pub fn mark_agent_team_member_running( - workspace_dir: &Path, - team_id: &str, - member_id: &str, - task_id: &str, - worker_thread_id: &str, - run_id: &str, -) -> Result> { - tracing::debug!( - "{LOG_PREFIX} mark_agent_team_member_running.entry team={team_id} member={member_id} task={task_id} run={run_id}" - ); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let now = Utc::now(); - let changed = conn - .execute( - "UPDATE agent_team_members - SET member_status = 'active', current_task_id = ?1, - worker_thread_id = ?2, run_id = ?3, updated_at = ?4 - WHERE id = ?5 AND team_id = ?6", - params![ - task_id, - worker_thread_id, - run_id, - now.to_rfc3339(), - member_id, - team_id - ], - ) - .storage_context("mark agent team member running")?; - if changed == 0 { - return Ok(None); - } - get_agent_team_member_inner(conn, member_id) - }) -} - -/// Mark a member idle: status → `idle`, clearing `current_task_id`. The -/// `worker_thread_id` / `run_id` are intentionally retained as a pointer to the -/// member's last run for history. Returns the updated member, or `None` if the -/// member is not in the team. Used when a worker run finishes (completed, -/// gate-failed, or failed) so the member is free to pick up new work. -pub fn mark_agent_team_member_idle( - workspace_dir: &Path, - team_id: &str, - member_id: &str, -) -> Result> { - tracing::debug!( - "{LOG_PREFIX} mark_agent_team_member_idle.entry team={team_id} member={member_id}" - ); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let now = Utc::now(); - let changed = conn - .execute( - "UPDATE agent_team_members - SET member_status = 'idle', current_task_id = NULL, updated_at = ?1 - WHERE id = ?2 AND team_id = ?3", - params![now.to_rfc3339(), member_id, team_id], - ) - .storage_context("mark agent team member idle")?; - if changed == 0 { - return Ok(None); - } - get_agent_team_member_inner(conn, member_id) - }) -} - -/// Release a single `in_progress` task back to `todo`, clearing its claim and -/// resetting the quality gate. Returns `true` if a row was actually released -/// (the task existed, belonged to the team, and was `in_progress`). Used by the -/// live runtime when a worker run fails or is aborted, so the task is free for -/// another teammate — the per-task analogue of the bulk release in -/// `shutdown_agent_team_member`. -pub fn release_agent_team_task(workspace_dir: &Path, team_id: &str, task_id: &str) -> Result { - tracing::debug!( - "{LOG_PREFIX} release_agent_team_task.entry team={team_id} task={task_id}" - ); - crate::store::with_connection(workspace_dir, |conn| { - init_run_ledger_schema(conn)?; - let now = Utc::now(); - let changed = conn - .execute( - "UPDATE agent_team_tasks - SET status = 'todo', claimed_by_member_id = NULL, claim_token = NULL, - gate_status = 'pending', gate_reason = NULL, updated_at = ?1 - WHERE id = ?2 AND team_id = ?3 AND status = 'in_progress'", - params![now.to_rfc3339(), task_id, team_id], - ) - .storage_context("release agent team task")?; - tracing::debug!( - "{LOG_PREFIX} release_agent_team_task.exit team={team_id} task={task_id} released={}", - changed > 0 - ); - Ok(changed > 0) - }) -} diff --git a/vendor/tinyinference b/vendor/tinyinference index c71e6b50b..a32eec813 160000 --- a/vendor/tinyinference +++ b/vendor/tinyinference @@ -1 +1 @@ -Subproject commit c71e6b50b40a10e3bf5d97c11ebb8fa2caa7cd60 +Subproject commit a32eec8138d7be0978dd232ec9f93af33945dbe6