feat(grid-agent): add generic LLM tool loop (#119)
Some checks failed
CI / rust-skia (Rust only) (push) Successful in 2m50s
CI / required (push) Failing after 2m46s

This commit is contained in:
2026-08-17 20:45:56 +00:00
parent 1e1e95a58a
commit a46bc42a8f
11 changed files with 2365 additions and 6 deletions

View File

@@ -6,7 +6,9 @@
pub mod backend;
pub mod config;
pub mod llm;
pub mod service;
pub mod tool_loop;
pub mod types;
#[cfg(feature = "live-grid")]
@@ -17,7 +19,15 @@ pub use config::{
GridConnection, Limits, LlmConnection, MapEnvironment, OperatingMode, SecretString,
StdEnvironment, Timeouts,
};
pub use llm::{
Completion, CompletionMessage, ContentPart, ImageDetail, LlmClient, LlmError,
LlmTransportLimits, ToolDefinition, ToolSchema, Usage,
};
pub use service::{AgentService, ServiceError, ServiceHandle, ServiceState};
pub use tool_loop::{
HistorySummarizer, SessionGeneration, ToolExecution, ToolExecutor, ToolFuture, ToolLoop,
ToolLoopError, ToolLoopLimits, ToolLoopOutcome,
};
pub use types::{
BoundaryError, BoundedText, BoundedVec, ControlCommand, Conversation, ConversationMessage,
GridEvent, GridEventKind, LlmRequest, LlmResult, MessageRole, ObservableEvent, PolicyDecision,

View File

@@ -0,0 +1,764 @@
//! Provider-neutral exact-endpoint LLM transport and bounded wire envelope.
#![allow(clippy::missing_errors_doc)] // Public methods share the exhaustive typed LlmError boundary.
use crate::config::LlmConnection;
use crate::types::{
BoundedText, BoundedVec, MAX_BODY_BYTES, MAX_CONVERSATION_MESSAGES, MAX_IDENTIFIER_BYTES,
MAX_MESSAGE_BYTES, MAX_TOOL_CALLS, MessageRole, ProposedToolCall,
};
use libremetaverse_types::compat::CancellationToken;
use serde::Deserialize;
use serde_json::{Map, Value, json};
use std::collections::{BTreeMap, BTreeSet};
use std::error::Error;
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::Semaphore;
const MAX_TOOLS: usize = 64;
const MAX_SCHEMA_PROPERTIES: usize = 128;
const MAX_SCHEMA_DEPTH: usize = 16;
const MAX_SCHEMA_NODES: usize = 1_024;
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ContentPart {
Text(BoundedText<MAX_MESSAGE_BYTES>),
Image {
url: BoundedText<MAX_BODY_BYTES>,
detail: ImageDetail,
},
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ImageDetail {
Auto,
Low,
High,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct CompletionMessage {
pub role: MessageRole,
pub content: BoundedVec<ContentPart, 16>,
pub tool_call_id: Option<BoundedText<MAX_IDENTIFIER_BYTES>>,
pub proposed_calls: BoundedVec<ProposedToolCall, MAX_TOOL_CALLS>,
}
impl CompletionMessage {
pub fn text(role: MessageRole, text: impl Into<String>) -> Result<Self, LlmError> {
let mut content = BoundedVec::new();
content
.try_push(
"completion_message.content",
ContentPart::Text(BoundedText::new("completion_message.text", text)?),
)
.map_err(LlmError::Boundary)?;
Ok(Self {
role,
content,
tool_call_id: None,
proposed_calls: BoundedVec::new(),
})
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToolSchema {
Object {
properties: BTreeMap<String, ToolSchema>,
required: BTreeSet<String>,
additional_properties: bool,
},
String,
Integer,
Number,
Boolean,
Array {
items: Box<ToolSchema>,
max_items: usize,
},
}
impl ToolSchema {
pub fn validate_schema(&self) -> Result<(), LlmError> {
let mut remaining_nodes = MAX_SCHEMA_NODES;
self.validate_schema_inner(0, &mut remaining_nodes)
}
fn validate_schema_inner(
&self,
depth: usize,
remaining_nodes: &mut usize,
) -> Result<(), LlmError> {
if depth > MAX_SCHEMA_DEPTH || *remaining_nodes == 0 {
return Err(LlmError::InvalidToolSchema);
}
*remaining_nodes -= 1;
match self {
Self::Object {
properties,
required,
..
} => {
if properties.len() > MAX_SCHEMA_PROPERTIES
|| required.len() > MAX_SCHEMA_PROPERTIES
|| required.iter().any(|name| !properties.contains_key(name))
|| properties
.keys()
.any(|name| name.is_empty() || name.len() > MAX_IDENTIFIER_BYTES)
{
return Err(LlmError::InvalidToolSchema);
}
for schema in properties.values() {
schema.validate_schema_inner(depth + 1, remaining_nodes)?;
}
Ok(())
}
Self::Array { items, max_items } => {
if *max_items == 0 || *max_items > 256 {
return Err(LlmError::InvalidToolSchema);
}
items.validate_schema_inner(depth + 1, remaining_nodes)
}
_ => Ok(()),
}
}
pub fn validate_value(&self, value: &Value) -> Result<(), LlmError> {
match self {
Self::Object {
properties,
required,
additional_properties,
} => {
let object = value.as_object().ok_or(LlmError::InvalidToolArguments)?;
if required.iter().any(|name| !object.contains_key(name))
|| (!additional_properties
&& object.keys().any(|name| !properties.contains_key(name)))
{
return Err(LlmError::InvalidToolArguments);
}
for (name, value) in object {
if let Some(schema) = properties.get(name) {
schema.validate_value(value)?;
}
}
Ok(())
}
Self::String if value.is_string() => Ok(()),
Self::Integer if value.as_i64().is_some() || value.as_u64().is_some() => Ok(()),
Self::Number if value.is_number() => Ok(()),
Self::Boolean if value.is_boolean() => Ok(()),
Self::Array { items, max_items } => {
let values = value.as_array().ok_or(LlmError::InvalidToolArguments)?;
if values.len() > *max_items {
return Err(LlmError::InvalidToolArguments);
}
for value in values {
items.validate_value(value)?;
}
Ok(())
}
_ => Err(LlmError::InvalidToolArguments),
}
}
fn wire_value(&self) -> Value {
match self {
Self::Object {
properties,
required,
additional_properties,
} => {
let properties = properties
.iter()
.map(|(name, schema)| (name.clone(), schema.wire_value()))
.collect::<Map<_, _>>();
json!({
"type": "object",
"properties": properties,
"required": required,
"additionalProperties": additional_properties
})
}
Self::String => json!({"type":"string"}),
Self::Integer => json!({"type":"integer"}),
Self::Number => json!({"type":"number"}),
Self::Boolean => json!({"type":"boolean"}),
Self::Array { items, max_items } => {
json!({"type":"array", "items":items.wire_value(), "maxItems":max_items})
}
}
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolDefinition {
pub name: BoundedText<MAX_IDENTIFIER_BYTES>,
pub description: BoundedText<MAX_MESSAGE_BYTES>,
pub schema: ToolSchema,
pub mutating: bool,
}
impl ToolDefinition {
pub fn validate(&self) -> Result<(), LlmError> {
if !self
.name
.chars()
.all(|character| character.is_ascii_alphanumeric() || matches!(character, '_' | '-'))
{
return Err(LlmError::InvalidToolSchema);
}
self.schema.validate_schema()
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Usage {
pub prompt_tokens: Option<u64>,
pub completion_tokens: Option<u64>,
pub total_tokens: Option<u64>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Completion {
pub request_id: u64,
pub correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
pub message: CompletionMessage,
pub usage: Option<Usage>,
pub latency: Duration,
pub attempts: usize,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct LlmTransportLimits {
pub connect_timeout: Duration,
pub request_timeout: Duration,
pub read_idle_timeout: Duration,
pub pool_idle_timeout: Duration,
pub total_timeout: Duration,
pub max_prompt_bytes: usize,
pub max_response_bytes: usize,
pub max_concurrent_requests: usize,
pub max_retries: usize,
pub max_retry_delay: Duration,
}
impl Default for LlmTransportLimits {
fn default() -> Self {
Self {
connect_timeout: Duration::from_secs(10),
request_timeout: Duration::from_mins(1),
read_idle_timeout: Duration::from_secs(15),
pool_idle_timeout: Duration::from_secs(30),
total_timeout: Duration::from_secs(90),
max_prompt_bytes: 1024 * 1024,
max_response_bytes: 1024 * 1024,
max_concurrent_requests: 8,
max_retries: 2,
max_retry_delay: Duration::from_secs(5),
}
}
}
impl LlmTransportLimits {
pub fn validate(&self) -> Result<(), LlmError> {
if self.connect_timeout.is_zero()
|| self.request_timeout.is_zero()
|| self.read_idle_timeout.is_zero()
|| self.pool_idle_timeout.is_zero()
|| self.total_timeout.is_zero()
|| self.max_prompt_bytes == 0
|| self.max_prompt_bytes > MAX_BODY_BYTES
|| self.max_response_bytes == 0
|| self.max_response_bytes > MAX_BODY_BYTES
|| self.max_concurrent_requests == 0
|| self.max_concurrent_requests > 256
|| self.max_retries > 8
|| self.max_retry_delay > Duration::from_mins(1)
{
return Err(LlmError::UnsafeLimits);
}
Ok(())
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum LlmError {
Boundary(crate::types::BoundaryError),
UnsafeLimits,
InvalidToolSchema,
InvalidToolArguments,
PromptTooLarge,
ResponseTooLarge,
Cancelled,
Timeout,
Transport,
RedirectRefused,
HttpStatus(u16),
MalformedJson,
UnsupportedResponse,
DuplicateToolCallId,
}
impl fmt::Display for LlmError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Boundary(error) => write!(formatter, "LLM boundary rejected input: {error}"),
Self::UnsafeLimits => formatter.write_str("unsafe LLM transport limits"),
Self::InvalidToolSchema => formatter.write_str("invalid registered tool schema"),
Self::InvalidToolArguments => formatter.write_str("tool arguments do not match schema"),
Self::PromptTooLarge => formatter.write_str("LLM prompt exceeds its byte bound"),
Self::ResponseTooLarge => formatter.write_str("LLM response exceeds its byte bound"),
Self::Cancelled => formatter.write_str("LLM request cancelled"),
Self::Timeout => formatter.write_str("LLM request timed out"),
Self::Transport => formatter.write_str("LLM transport failed"),
Self::RedirectRefused => formatter.write_str("LLM endpoint redirect refused"),
Self::HttpStatus(status) => write!(formatter, "LLM endpoint returned HTTP {status}"),
Self::MalformedJson => formatter.write_str("LLM endpoint returned malformed JSON"),
Self::UnsupportedResponse => formatter.write_str("unsupported LLM response shape"),
Self::DuplicateToolCallId => {
formatter.write_str("LLM response repeated a tool-call ID")
}
}
}
}
impl Error for LlmError {}
impl From<crate::types::BoundaryError> for LlmError {
fn from(value: crate::types::BoundaryError) -> Self {
Self::Boundary(value)
}
}
pub struct LlmClient {
connection: LlmConnection,
client: reqwest::Client,
limits: LlmTransportLimits,
slots: Arc<Semaphore>,
next_request_id: AtomicU64,
}
impl fmt::Debug for LlmClient {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("LlmClient")
.field("endpoint", &self.connection.endpoint_url)
.field("limits", &self.limits)
.field("available_slots", &self.slots.available_permits())
.finish_non_exhaustive()
}
}
impl LlmClient {
pub fn new(connection: LlmConnection, limits: LlmTransportLimits) -> Result<Self, LlmError> {
limits.validate()?;
let client = reqwest::Client::builder()
.redirect(reqwest::redirect::Policy::none())
.connect_timeout(limits.connect_timeout)
.timeout(limits.request_timeout)
.read_timeout(limits.read_idle_timeout)
.pool_idle_timeout(limits.pool_idle_timeout)
.pool_max_idle_per_host(limits.max_concurrent_requests)
.build()
.map_err(|_| LlmError::Transport)?;
Ok(Self {
connection,
client,
slots: Arc::new(Semaphore::new(limits.max_concurrent_requests)),
limits,
next_request_id: AtomicU64::new(1),
})
}
pub async fn complete(
&self,
messages: &[CompletionMessage],
tools: &[ToolDefinition],
cancellation: &CancellationToken,
) -> Result<Completion, LlmError> {
if messages.is_empty()
|| messages.len() > MAX_CONVERSATION_MESSAGES
|| tools.len() > MAX_TOOLS
{
return Err(LlmError::PromptTooLarge);
}
for tool in tools {
tool.validate()?;
}
let body = request_body(messages, tools)?;
if body.len() > self.limits.max_prompt_bytes {
return Err(LlmError::PromptTooLarge);
}
let request_id = self.next_request_id.fetch_add(1, Ordering::Relaxed);
let correlation_id = BoundedText::new(
"llm.correlation_id",
format!("agent-request-{request_id:016x}"),
)?;
let started = Instant::now();
let operation = async {
let permit = tokio::select! {
() = cancellation.cancelled() => return Err(LlmError::Cancelled),
permit = Arc::clone(&self.slots).acquire_owned() => permit.map_err(|_| LlmError::Cancelled)?,
};
let result = self
.complete_with_retries(request_id, correlation_id, body, cancellation, started)
.await;
drop(permit);
result
};
tokio::time::timeout(self.limits.total_timeout, operation)
.await
.map_err(|_| LlmError::Timeout)?
}
async fn complete_with_retries(
&self,
request_id: u64,
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
body: Vec<u8>,
cancellation: &CancellationToken,
started: Instant,
) -> Result<Completion, LlmError> {
for attempt in 0..=self.limits.max_retries {
match self
.complete_once(
request_id,
correlation_id.clone(),
&body,
cancellation,
started,
attempt + 1,
)
.await
{
Ok(completion) => return Ok(completion),
Err(AttemptError { error, retry_after })
if attempt < self.limits.max_retries && transient(&error) =>
{
let delay = retry_after
.unwrap_or_else(|| deterministic_backoff(request_id, attempt))
.min(self.limits.max_retry_delay);
tokio::select! {
() = cancellation.cancelled() => return Err(LlmError::Cancelled),
() = tokio::time::sleep(delay) => {}
}
}
Err(AttemptError { error, .. }) => return Err(error),
}
}
Err(LlmError::Transport)
}
async fn complete_once(
&self,
request_id: u64,
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
body: &[u8],
cancellation: &CancellationToken,
started: Instant,
attempts: usize,
) -> Result<Completion, AttemptError> {
let response = tokio::select! {
() = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
response = self.client
.post(self.connection.endpoint_url.expose_url())
.bearer_auth(self.connection.api_key.expose_secret())
.header("content-type", "application/json")
.header("accept", "application/json")
.header("x-correlation-id", correlation_id.as_str())
.body(body.to_vec())
.send() => response.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
};
let status = response.status();
if status.is_redirection() {
return Err(AttemptError::new(LlmError::RedirectRefused));
}
if !status.is_success() {
let retry_after = response
.headers()
.get("retry-after")
.and_then(|value| value.to_str().ok())
.and_then(|value| value.parse::<u64>().ok())
.map(Duration::from_secs);
return Err(AttemptError {
error: LlmError::HttpStatus(status.as_u16()),
retry_after,
});
}
let bytes = read_bounded(response, self.limits.max_response_bytes, cancellation).await?;
parse_completion(
request_id,
correlation_id,
&bytes,
started.elapsed(),
attempts,
)
.map_err(AttemptError::new)
}
}
struct AttemptError {
error: LlmError,
retry_after: Option<Duration>,
}
impl AttemptError {
const fn new(error: LlmError) -> Self {
Self {
error,
retry_after: None,
}
}
}
async fn read_bounded(
mut response: reqwest::Response,
maximum: usize,
cancellation: &CancellationToken,
) -> Result<Vec<u8>, AttemptError> {
if response
.content_length()
.is_some_and(|length| usize::try_from(length).map_or(true, |length| length > maximum))
{
return Err(AttemptError::new(LlmError::ResponseTooLarge));
}
let mut body = Vec::with_capacity(
response
.content_length()
.and_then(|value| usize::try_from(value).ok())
.unwrap_or(0)
.min(maximum),
);
loop {
let chunk = tokio::select! {
() = cancellation.cancelled() => return Err(AttemptError::new(LlmError::Cancelled)),
chunk = response.chunk() => chunk.map_err(|error| AttemptError::new(classify_reqwest(&error)))?,
};
let Some(chunk) = chunk else { break };
if body.len().saturating_add(chunk.len()) > maximum {
return Err(AttemptError::new(LlmError::ResponseTooLarge));
}
body.extend_from_slice(&chunk);
}
Ok(body)
}
fn request_body(
messages: &[CompletionMessage],
tools: &[ToolDefinition],
) -> Result<Vec<u8>, LlmError> {
let messages = messages.iter().map(message_wire_value).collect::<Vec<_>>();
let tools = tools
.iter()
.map(|tool| {
json!({
"type":"function",
"function": {
"name": tool.name.as_str(),
"description": tool.description.as_str(),
"parameters": tool.schema.wire_value()
}
})
})
.collect::<Vec<_>>();
serde_json::to_vec(&json!({"messages":messages, "tools":tools}))
.map_err(|_| LlmError::MalformedJson)
}
fn message_wire_value(message: &CompletionMessage) -> Value {
let role = match message.role {
MessageRole::System => "system",
MessageRole::Avatar => "user",
MessageRole::Agent => "assistant",
MessageRole::Tool => "tool",
};
let content = message.content.as_slice().iter().map(|part| match part {
ContentPart::Text(text) => json!({"type":"text", "text":text.as_str()}),
ContentPart::Image { url, detail } => json!({
"type":"image_url",
"image_url":{"url":url.as_str(), "detail":match detail { ImageDetail::Auto=>"auto", ImageDetail::Low=>"low", ImageDetail::High=>"high"}}
}),
}).collect::<Vec<_>>();
let mut value = json!({"role":role,"content":content});
if let Some(call_id) = &message.tool_call_id {
value["tool_call_id"] = Value::String(call_id.as_str().to_owned());
}
if !message.proposed_calls.is_empty() {
value["tool_calls"] = Value::Array(message.proposed_calls.as_slice().iter().map(|call| json!({
"id":call.call_id.as_str(), "type":"function", "function":{"name":call.name.as_str(),"arguments":call.arguments_json.as_str()}
})).collect());
}
value
}
#[derive(Deserialize)]
struct WireResponse {
choices: Vec<WireChoice>,
usage: Option<WireUsage>,
}
#[derive(Deserialize)]
struct WireChoice {
message: WireMessage,
}
#[derive(Deserialize)]
struct WireMessage {
content: Option<String>,
#[serde(default)]
tool_calls: Vec<WireToolCall>,
}
#[derive(Deserialize)]
struct WireToolCall {
id: String,
#[serde(rename = "type")]
kind: Option<String>,
function: WireFunction,
}
#[derive(Deserialize)]
struct WireFunction {
name: String,
arguments: String,
}
#[derive(Deserialize)]
struct WireUsage {
#[serde(rename = "prompt_tokens")]
prompt: Option<u64>,
#[serde(rename = "completion_tokens")]
completion: Option<u64>,
#[serde(rename = "total_tokens")]
total: Option<u64>,
}
fn parse_completion(
request_id: u64,
correlation_id: BoundedText<MAX_IDENTIFIER_BYTES>,
body: &[u8],
latency: Duration,
attempts: usize,
) -> Result<Completion, LlmError> {
let wire: WireResponse = serde_json::from_slice(body).map_err(|_| LlmError::MalformedJson)?;
let choice = wire
.choices
.into_iter()
.next()
.ok_or(LlmError::UnsupportedResponse)?;
let mut proposed_calls = BoundedVec::new();
let mut ids = BTreeSet::new();
for call in choice.message.tool_calls {
if call.kind.as_deref().is_some_and(|kind| kind != "function") {
return Err(LlmError::UnsupportedResponse);
}
if !ids.insert(call.id.clone()) {
return Err(LlmError::DuplicateToolCallId);
}
proposed_calls
.try_push(
"completion.tool_calls",
ProposedToolCall::from_model_output(
call.id,
call.function.name,
call.function.arguments,
)?,
)
.map_err(LlmError::Boundary)?;
}
let mut content = BoundedVec::new();
if let Some(text) = choice.message.content {
content
.try_push(
"completion.content",
ContentPart::Text(BoundedText::new_allow_empty("completion.text", text)?),
)
.map_err(LlmError::Boundary)?;
}
if content.is_empty() && proposed_calls.is_empty() {
return Err(LlmError::UnsupportedResponse);
}
Ok(Completion {
request_id,
correlation_id,
message: CompletionMessage {
role: MessageRole::Agent,
content,
tool_call_id: None,
proposed_calls,
},
usage: wire.usage.map(|usage| Usage {
prompt_tokens: usage.prompt,
completion_tokens: usage.completion,
total_tokens: usage.total,
}),
latency,
attempts,
})
}
fn classify_reqwest(error: &reqwest::Error) -> LlmError {
if error.is_timeout() {
LlmError::Timeout
} else {
LlmError::Transport
}
}
fn transient(error: &LlmError) -> bool {
matches!(error, LlmError::Transport | LlmError::Timeout)
|| matches!(
error,
LlmError::HttpStatus(408 | 425 | 429 | 500 | 502 | 503 | 504)
)
}
fn deterministic_backoff(request_id: u64, attempt: usize) -> Duration {
let exponent = u32::try_from(attempt).unwrap_or(8).min(8);
let base = 100_u64.saturating_mul(1_u64 << exponent);
let jitter = request_id
.wrapping_mul(6_364_136_223_846_793_005)
.rotate_left(exponent)
% 97;
Duration::from_millis(base.saturating_add(jitter))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::AgentConfig;
#[tokio::test]
async fn total_timeout_includes_waiting_for_a_concurrency_slot() {
let limits = LlmTransportLimits {
total_timeout: Duration::from_millis(20),
max_concurrent_requests: 1,
..LlmTransportLimits::default()
};
let connection = AgentConfig::offline("http://127.0.0.1:9/exact", "test-key")
.expect("test config")
.llm;
let client = LlmClient::new(connection, limits).expect("test client");
let permit = Arc::clone(&client.slots)
.acquire_owned()
.await
.expect("semaphore open");
assert_eq!(
client
.complete(
&[CompletionMessage::text(MessageRole::Avatar, "hello").expect("message")],
&[],
&CancellationToken::default(),
)
.await
.expect_err("slot wait must time out"),
LlmError::Timeout
);
drop(permit);
}
}

View File

@@ -0,0 +1,534 @@
//! Bounded multi-turn tool orchestration above the provider-neutral transport.
use crate::llm::{CompletionMessage, ContentPart, LlmClient, LlmError, ToolDefinition, Usage};
use crate::types::{
BoundedText, BoundedVec, MAX_BODY_BYTES, MAX_CONVERSATION_MESSAGES, MAX_MESSAGE_BYTES,
MAX_TOOL_CALLS, MessageRole, ProposedToolCall,
};
use libremetaverse_types::compat::{CancellationToken, CancellationTokenSource};
use serde_json::Value;
use std::error::Error;
use std::fmt;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
const MAX_REGISTERED_TOOLS: usize = 64;
const MAX_LOOP_TURNS: usize = 32;
const MAX_SESSION_TOOL_CALLS: usize = 256;
pub type ToolFuture<'a> = Pin<Box<dyn Future<Output = ToolExecution> + Send + 'a>>;
/// Downstream execution boundary. Issue #120 can implement policy evaluation
/// here; this loop guarantees its input already passed name/schema checks.
pub trait ToolExecutor: Send + Sync {
fn execute<'a>(
&'a self,
definition: &'a ToolDefinition,
call: &'a ProposedToolCall,
arguments: &'a Value,
cancellation: &'a CancellationToken,
) -> ToolFuture<'a>;
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToolExecution {
Completed(BoundedText<MAX_BODY_BYTES>),
Rejected(BoundedText<MAX_MESSAGE_BYTES>),
Failed(BoundedText<MAX_MESSAGE_BYTES>),
/// A mutating operation may have reached the world, so it must never be
/// repeated or converted into another model request automatically.
AmbiguousMutation,
}
/// Optional deterministic history summarizer. Failure is explicitly safe: the
/// loop inserts a bounded truncation marker instead.
pub trait HistorySummarizer: Send + Sync {
/// Returns a deterministic bounded summary, or `None` to request the safe
/// truncation marker fallback.
fn summarize(&self, omitted: &[CompletionMessage]) -> Option<CompletionMessage>;
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolLoopLimits {
pub max_turns: usize,
pub max_tool_calls_per_turn: usize,
pub max_tool_calls_per_session: usize,
pub max_history_messages: usize,
pub max_history_bytes: usize,
pub wall_clock_timeout: Duration,
}
impl Default for ToolLoopLimits {
fn default() -> Self {
Self {
max_turns: 8,
max_tool_calls_per_turn: 8,
max_tool_calls_per_session: 32,
max_history_messages: 64,
max_history_bytes: 512 * 1024,
wall_clock_timeout: Duration::from_mins(2),
}
}
}
impl ToolLoopLimits {
/// # Errors
///
/// Rejects zero or above-hard-ceiling loop limits.
pub fn validate(&self) -> Result<(), ToolLoopError> {
if self.max_turns == 0
|| self.max_turns > MAX_LOOP_TURNS
|| self.max_tool_calls_per_turn == 0
|| self.max_tool_calls_per_turn > MAX_TOOL_CALLS
|| self.max_tool_calls_per_session == 0
|| self.max_tool_calls_per_session > MAX_SESSION_TOOL_CALLS
|| self.max_history_messages < 2
|| self.max_history_messages > MAX_CONVERSATION_MESSAGES
|| self.max_history_bytes < 1024
|| self.max_history_bytes > MAX_BODY_BYTES
|| self.wall_clock_timeout.is_zero()
|| self.wall_clock_timeout > Duration::from_mins(10)
{
return Err(ToolLoopError::UnsafeLimits);
}
Ok(())
}
}
/// Monotonic session epoch. Disconnect, expiry, or operator replacement calls
/// `supersede`; late inference and tool results then fail closed.
#[derive(Debug)]
struct SessionState {
generation: u64,
cancellation: CancellationTokenSource,
}
/// Session generation and its owned cancellation source are updated under one
/// lock so a caller can never obtain an uncancelled token for a stale epoch.
#[derive(Debug)]
pub struct SessionGeneration(Mutex<SessionState>);
impl Default for SessionGeneration {
fn default() -> Self {
Self(Mutex::new(SessionState {
generation: 0,
cancellation: CancellationTokenSource::new(),
}))
}
}
impl SessionGeneration {
#[must_use]
pub fn current(&self) -> u64 {
self.state().generation
}
pub fn supersede(&self) -> u64 {
let mut state = self.state();
state.cancellation.cancel();
state.generation = state.generation.saturating_add(1);
state.cancellation = CancellationTokenSource::new();
state.generation
}
#[must_use]
pub fn is_current(&self, generation: u64) -> bool {
self.current() == generation
}
fn cancellation_for(&self, generation: u64) -> Option<CancellationToken> {
let state = self.state();
(state.generation == generation).then(|| state.cancellation.token())
}
fn state(&self) -> std::sync::MutexGuard<'_, SessionState> {
self.0
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ToolLoopOutcome {
pub final_message: CompletionMessage,
pub turns: usize,
pub tool_calls: usize,
pub usage: BoundedVec<Usage, MAX_LOOP_TURNS>,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum ToolLoopError {
Transport(LlmError),
UnsafeLimits,
DuplicateToolName,
DuplicateToolCallId,
TooManyTools,
EmptyHistory,
HistoryLimit,
EndlessToolLoop,
ToolCallLimit,
AmbiguousMutation,
Cancelled,
WallClockTimeout,
Superseded,
Boundary(crate::types::BoundaryError),
}
impl fmt::Display for ToolLoopError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Transport(error) => write!(formatter, "LLM tool loop transport failed: {error}"),
Self::UnsafeLimits => formatter.write_str("unsafe tool-loop limits"),
Self::DuplicateToolName => formatter.write_str("duplicate registered tool name"),
Self::DuplicateToolCallId => {
formatter.write_str("model repeated a tool-call ID in the session")
}
Self::TooManyTools => formatter.write_str("too many registered tools"),
Self::EmptyHistory => formatter.write_str("tool loop requires initial history"),
Self::HistoryLimit => formatter.write_str("tool-loop history exceeds its hard bound"),
Self::EndlessToolLoop => formatter.write_str("model exceeded the bounded tool turns"),
Self::ToolCallLimit => {
formatter.write_str("model exceeded the session tool-call bound")
}
Self::AmbiguousMutation => formatter.write_str("mutating tool outcome is ambiguous"),
Self::Cancelled => formatter.write_str("tool loop cancelled"),
Self::WallClockTimeout => formatter.write_str("tool loop wall-clock bound elapsed"),
Self::Superseded => formatter.write_str("tool loop session was superseded"),
Self::Boundary(error) => write!(formatter, "tool loop boundary failed: {error}"),
}
}
}
impl Error for ToolLoopError {}
impl From<LlmError> for ToolLoopError {
fn from(value: LlmError) -> Self {
if value == LlmError::Cancelled {
Self::Cancelled
} else {
Self::Transport(value)
}
}
}
impl From<crate::types::BoundaryError> for ToolLoopError {
fn from(value: crate::types::BoundaryError) -> Self {
Self::Boundary(value)
}
}
pub struct ToolLoop {
client: Arc<LlmClient>,
tools: BoundedVec<ToolDefinition, MAX_REGISTERED_TOOLS>,
limits: ToolLoopLimits,
summarizer: Option<Arc<dyn HistorySummarizer>>,
}
impl fmt::Debug for ToolLoop {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ToolLoop")
.field("client", &self.client)
.field("tool_count", &self.tools.len())
.field("limits", &self.limits)
.field("has_summarizer", &self.summarizer.is_some())
.finish()
}
}
impl ToolLoop {
/// # Errors
///
/// Rejects invalid limits, schemas, counts, or duplicate tool names.
pub fn new(
client: Arc<LlmClient>,
tools: Vec<ToolDefinition>,
limits: ToolLoopLimits,
) -> Result<Self, ToolLoopError> {
limits.validate()?;
if tools.len() > MAX_REGISTERED_TOOLS {
return Err(ToolLoopError::TooManyTools);
}
let mut names = std::collections::BTreeSet::new();
for tool in &tools {
tool.validate()?;
if !names.insert(tool.name.as_str()) {
return Err(ToolLoopError::DuplicateToolName);
}
}
Ok(Self {
client,
tools: BoundedVec::try_from_vec("tool_loop.tools", tools)?,
limits,
summarizer: None,
})
}
#[must_use]
pub fn with_summarizer(mut self, summarizer: Arc<dyn HistorySummarizer>) -> Self {
self.summarizer = Some(summarizer);
self
}
/// # Errors
///
/// Returns typed transport, bound, cancellation, supersession, endless-loop,
/// or ambiguous-mutation failures.
pub async fn run(
&self,
mut history: Vec<CompletionMessage>,
generation_owner: &SessionGeneration,
generation: u64,
cancellation: &CancellationToken,
executor: &dyn ToolExecutor,
) -> Result<ToolLoopOutcome, ToolLoopError> {
if history.is_empty() {
return Err(ToolLoopError::EmptyHistory);
}
if history.len() > MAX_CONVERSATION_MESSAGES {
return Err(ToolLoopError::HistoryLimit);
}
let session_cancellation = generation_owner
.cancellation_for(generation)
.ok_or(ToolLoopError::Superseded)?;
let linked_cancellation =
CancellationTokenSource::new_linked(&[cancellation.clone(), session_cancellation]);
let linked_token = linked_cancellation.token();
let future = self.run_inner(
&mut history,
generation_owner,
generation,
&linked_token,
executor,
);
let result = tokio::time::timeout(self.limits.wall_clock_timeout, future)
.await
.map_err(|_| ToolLoopError::WallClockTimeout)?;
drop(linked_cancellation);
result
}
async fn run_inner(
&self,
history: &mut Vec<CompletionMessage>,
generation_owner: &SessionGeneration,
generation: u64,
cancellation: &CancellationToken,
executor: &dyn ToolExecutor,
) -> Result<ToolLoopOutcome, ToolLoopError> {
let mut total_calls = 0;
let mut usage = BoundedVec::new();
let mut seen_call_ids = std::collections::BTreeSet::new();
for turn in 1..=self.limits.max_turns {
ensure_live(generation_owner, generation, cancellation)?;
compact_history(history, &self.limits, self.summarizer.as_deref())?;
let completion_result = self
.client
.complete(history, self.tools.as_slice(), cancellation)
.await;
ensure_live(generation_owner, generation, cancellation)?;
let completion = completion_result?;
if let Some(item) = completion.usage.clone() {
usage.try_push("tool_loop.usage", item)?;
}
let calls = completion.message.proposed_calls.clone().into_inner();
if calls.len() > self.limits.max_tool_calls_per_turn {
return Err(ToolLoopError::ToolCallLimit);
}
if calls
.iter()
.any(|call| !seen_call_ids.insert(call.call_id.as_str().to_owned()))
{
return Err(ToolLoopError::DuplicateToolCallId);
}
history.push(completion.message.clone());
if calls.is_empty() {
return Ok(ToolLoopOutcome {
final_message: completion.message,
turns: turn,
tool_calls: total_calls,
usage,
});
}
total_calls = total_calls.saturating_add(calls.len());
if total_calls > self.limits.max_tool_calls_per_session {
return Err(ToolLoopError::ToolCallLimit);
}
for call in calls {
ensure_live(generation_owner, generation, cancellation)?;
let observation = self
.evaluate_call(&call, generation_owner, generation, cancellation, executor)
.await?;
history.push(observation);
}
}
Err(ToolLoopError::EndlessToolLoop)
}
async fn evaluate_call(
&self,
call: &ProposedToolCall,
generation_owner: &SessionGeneration,
generation: u64,
cancellation: &CancellationToken,
executor: &dyn ToolExecutor,
) -> Result<CompletionMessage, ToolLoopError> {
let Some(definition) = self
.tools
.as_slice()
.iter()
.find(|definition| definition.name.as_str() == call.name.as_str())
else {
return tool_observation(call, "unknown tool; no execution occurred");
};
let Ok(arguments) = serde_json::from_str::<Value>(call.arguments_json.as_str()) else {
return tool_observation(call, "arguments are malformed JSON; no execution occurred");
};
if definition.schema.validate_value(&arguments).is_err() {
return tool_observation(call, "arguments rejected by registered schema");
}
ensure_live(generation_owner, generation, cancellation)?;
let execution = tokio::select! {
() = cancellation.cancelled() => {
ensure_live(generation_owner, generation, cancellation)?;
return Err(ToolLoopError::Cancelled);
}
execution = executor.execute(definition, call, &arguments, cancellation) => execution,
};
ensure_live(generation_owner, generation, cancellation)?;
match execution {
ToolExecution::Completed(result) => tool_observation(call, result.as_str()),
ToolExecution::Rejected(reason) => {
tool_observation(call, &format!("tool rejected: {}", reason.as_str()))
}
ToolExecution::Failed(reason) => {
tool_observation(call, &format!("tool failed safely: {}", reason.as_str()))
}
ToolExecution::AmbiguousMutation => Err(ToolLoopError::AmbiguousMutation),
}
}
}
fn ensure_live(
owner: &SessionGeneration,
generation: u64,
cancellation: &CancellationToken,
) -> Result<(), ToolLoopError> {
if !owner.is_current(generation) {
Err(ToolLoopError::Superseded)
} else if cancellation.is_cancellation_requested() {
Err(ToolLoopError::Cancelled)
} else {
Ok(())
}
}
fn tool_observation(
call: &ProposedToolCall,
body: &str,
) -> Result<CompletionMessage, ToolLoopError> {
const TRUNCATION_MARKER: &str = "\n[tool observation truncated]";
let body = if body.len() <= MAX_MESSAGE_BYTES {
body.to_owned()
} else {
let mut end = MAX_MESSAGE_BYTES - TRUNCATION_MARKER.len();
while !body.is_char_boundary(end) {
end -= 1;
}
format!("{}{TRUNCATION_MARKER}", &body[..end])
};
let mut message = CompletionMessage::text(MessageRole::Tool, body)?;
message.tool_call_id = Some(call.call_id.clone());
Ok(message)
}
fn compact_history(
history: &mut Vec<CompletionMessage>,
limits: &ToolLoopLimits,
summarizer: Option<&dyn HistorySummarizer>,
) -> Result<(), ToolLoopError> {
let mut omitted = Vec::new();
while history.len() > 1
&& (history.len() >= limits.max_history_messages
|| history_bytes(history) > limits.max_history_bytes)
{
omitted.push(history.remove(0));
}
if omitted.is_empty() {
return (history_bytes(history) <= limits.max_history_bytes)
.then_some(())
.ok_or(ToolLoopError::HistoryLimit);
}
if history_bytes(history) > limits.max_history_bytes {
return Err(ToolLoopError::HistoryLimit);
}
let fallback = CompletionMessage::text(
MessageRole::System,
format!(
"[history safely truncated: {} earlier messages omitted]",
omitted.len()
),
)?;
let summary = summarizer
.and_then(|summarizer| summarizer.summarize(&omitted))
.filter(valid_summary)
.filter(|summary| {
history_bytes(history).saturating_add(message_bytes(summary))
<= limits.max_history_bytes
})
.or_else(|| {
(history_bytes(history).saturating_add(message_bytes(&fallback))
<= limits.max_history_bytes)
.then_some(fallback)
});
if let Some(summary) = summary {
history.insert(0, summary);
}
while history.len() > limits.max_history_messages {
history.remove(1);
}
Ok(())
}
fn history_bytes(history: &[CompletionMessage]) -> usize {
history
.iter()
.map(message_bytes)
.fold(0, usize::saturating_add)
}
fn message_bytes(message: &CompletionMessage) -> usize {
let content = message
.content
.as_slice()
.iter()
.map(|part| match part {
ContentPart::Text(text) => text.len(),
ContentPart::Image { url, .. } => url.len(),
})
.fold(0, usize::saturating_add);
let call_id = message
.tool_call_id
.as_ref()
.map_or(0, |call_id| call_id.len());
message
.proposed_calls
.as_slice()
.iter()
.map(|call| {
call.call_id
.len()
.saturating_add(call.name.len())
.saturating_add(call.arguments_json.len())
})
.fold(content.saturating_add(call_id), usize::saturating_add)
}
fn valid_summary(message: &CompletionMessage) -> bool {
message.role == MessageRole::System
&& message.tool_call_id.is_none()
&& message.proposed_calls.is_empty()
}

View File

@@ -330,6 +330,21 @@ impl ProposedToolCall {
arguments_json,
})
}
/// Admits bounded model output before schema/JSON validation in the tool
/// loop. This is crate-private so application-created calls still use
/// [`Self::new`].
pub(crate) fn from_model_output(
call_id: impl Into<String>,
name: impl Into<String>,
arguments_json: impl Into<String>,
) -> Result<Self, BoundaryError> {
Ok(Self {
call_id: BoundedText::new("tool_call.call_id", call_id)?,
name: BoundedText::new("tool_call.name", name)?,
arguments_json: BoundedText::new("tool_call.arguments_json", arguments_json)?,
})
}
}
#[derive(Clone, Debug, Eq, PartialEq)]