use salvor_llm::{ContentBlock, Message, MessageRequest, Tool};
use serde_json::Value;
use crate::agent::Agent;
use crate::budgets::{BudgetExtensions, BudgetObservations};
use crate::compact::FailureTracker;
use crate::ctx::{Resumption, RunCtx, ToolCallResult};
use crate::error::RuntimeError;
use crate::runtime::ParkReason;
use crate::wire::content_string;
use salvor_core::Effect;
use time::OffsetDateTime;
#[derive(Debug, Clone)]
pub enum LoopOutcome {
Completed(Value),
Parked(ParkReason),
}
pub(crate) async fn drive(
ctx: &mut RunCtx,
agent: &Agent,
initial_input: &Value,
) -> Result<LoopOutcome, RuntimeError> {
let input = begin(ctx, agent, initial_input).await?;
let outcome = drive_loop(ctx, agent, &input).await?;
if let LoopOutcome::Completed(output) = &outcome {
ctx.complete_run(output).await?;
}
Ok(outcome)
}
pub(crate) async fn begin(
ctx: &mut RunCtx,
agent: &Agent,
initial_input: &Value,
) -> Result<Value, RuntimeError> {
ctx.begin(agent.def_hash(), initial_input).await
}
pub async fn drive_loop(
ctx: &mut RunCtx,
agent: &Agent,
input: &Value,
) -> Result<LoopOutcome, RuntimeError> {
let mut conversation: Vec<Message> = vec![Message::user(content_string(input))];
let llm_tools: Vec<Tool> = agent
.tools()
.descriptors()
.into_iter()
.map(|descriptor| Tool {
name: descriptor.name,
description: Some(descriptor.description),
input_schema: descriptor.input_schema,
})
.collect();
let mut steps: u64 = 0;
let mut input_tokens: u64 = 0;
let mut output_tokens: u64 = 0;
let mut started_at: Option<OffsetDateTime> = None;
let mut extensions = BudgetExtensions::default();
let mut failures = FailureTracker::new();
loop {
let now = ctx.now().await?;
let start = *started_at.get_or_insert(now);
loop {
let observations = BudgetObservations {
steps,
input_tokens,
output_tokens,
elapsed_seconds: (now - start).as_seconds_f64(),
};
let Some((budget, observed)) =
agent
.budgets()
.first_crossing(&extensions, agent.pricing(), &observations)
else {
break;
};
ctx.budget_exceeded(budget, observed).await?;
match ctx.await_resume().await? {
Resumption::Parked => {
return Ok(LoopOutcome::Parked(ParkReason::BudgetExceeded {
budget,
observed,
}));
}
Resumption::Resumed(resume_input) => extensions.absorb(&resume_input),
}
}
let mut request = MessageRequest::new(agent.model(), agent.max_response_tokens())
.with_messages(conversation.clone());
if let Some(system) = agent.system_prompt() {
request = request.with_system(system);
}
if !llm_tools.is_empty() {
request = request.with_tools(llm_tools.clone());
}
let turn = ctx.model_call(agent.client(), &request).await?;
steps += 1;
input_tokens = input_tokens.saturating_add(u64::from(turn.usage.input_tokens));
output_tokens = output_tokens.saturating_add(u64::from(turn.usage.output_tokens));
let tool_uses: Vec<(String, String, Value)> = turn
.response
.tool_uses()
.into_iter()
.map(|(id, name, tool_input)| (id.to_owned(), name.to_owned(), tool_input.clone()))
.collect();
conversation.push(Message::assistant_blocks(turn.response.content.clone()));
if tool_uses.is_empty() {
let output = Value::String(turn.response.text());
return Ok(LoopOutcome::Completed(output));
}
let mut result_blocks: Vec<ContentBlock> = Vec::with_capacity(tool_uses.len());
for (tool_use_id, name, tool_input) in tool_uses {
let Some(tool) = agent.tools().get(&name) else {
result_blocks.push(ContentBlock::tool_error(
tool_use_id,
format!("unknown tool `{name}`"),
));
continue;
};
let idempotency_key = match tool.effect() {
Effect::Idempotent => Some(format!("{:016x}", ctx.random().await?)),
Effect::Read | Effect::Write => None,
};
match ctx
.tool_call(tool, &tool_input, idempotency_key.as_deref())
.await?
{
ToolCallResult::Output(output) => {
failures.record_success();
result_blocks.push(ContentBlock::tool_result(
tool_use_id,
content_string(&output),
));
}
ToolCallResult::Failed(failure) => {
let content = failures.content_for_failure(&name, &failure.message);
result_blocks.push(ContentBlock::tool_error(tool_use_id, content));
}
ToolCallResult::Suspended(suspension) => {
ctx.suspend(&suspension.reason, &suspension.input_schema)
.await?;
match ctx.await_resume().await? {
Resumption::Parked => {
return Ok(LoopOutcome::Parked(ParkReason::Suspended {
reason: suspension.reason,
input_schema: suspension.input_schema,
}));
}
Resumption::Resumed(resume_input) => {
failures.record_success();
result_blocks.push(ContentBlock::tool_result(
tool_use_id,
content_string(&resume_input),
));
}
}
}
}
}
conversation.push(Message::user_blocks(result_blocks));
}
}