use std::time::Duration;
use af_agent::{ChatModel, ModelFailure, RequestErrorAction, StepContext};
use af_agent_session::Event;
use af_context::RunId;
use af_llm::{ChatMessage, CompletionRequest, StreamOptions};
use crate::{AgentRuntime, CancellationToken, EventWriter, RuntimeError};
pub(super) struct ModelCompletion {
pub(super) response: af_llm::CompletionResponse,
pub(super) attempt: u32,
pub(super) prompt_tokens: u64,
pub(super) completion_tokens: u64,
}
pub(super) struct ModelTurn<'a> {
pub(super) run_id: &'a RunId,
pub(super) step: u32,
pub(super) step_context: &'a StepContext,
pub(super) transcript: &'a [ChatMessage],
pub(super) context: &'a [ChatMessage],
pub(super) after_seq: u64,
}
impl AgentRuntime {
pub(super) async fn complete_with_retry(
&self,
writer: &dyn EventWriter,
turn: ModelTurn<'_>,
cancellation: CancellationToken,
) -> Result<ModelCompletion, RuntimeError> {
let ModelTurn {
run_id,
step,
step_context,
transcript,
context,
after_seq,
} = turn;
let policy = self.request_policy(step_context).await;
let mut request = CompletionRequest::new(
&policy.model,
std::iter::once(ChatMessage::system(self.prompts.render()))
.chain(transcript.iter().cloned())
.chain(context.iter().cloned())
.collect(),
)
.temperature(policy.temperature)
.stream(true);
if let Some(max_tokens) = policy.max_tokens {
request = request.max_tokens(max_tokens);
}
if let Some(effort) = policy.reasoning_effort {
request = request.reasoning_effort(effort);
}
request.stream_options = Some(StreamOptions {
include_usage: true,
});
if !self.tools.is_empty() {
request = request.tools(self.tools.specs());
}
let mut last = None;
let estimated_prompt_tokens = self.meter.count(&policy.model, &request.messages);
let mut total_prompt_tokens = 0;
let mut total_completion_tokens = 0;
for attempt in 1..=self.limits.provider_attempts.max(1) {
let operation_id = format!("model:{step}:attempt:{attempt}");
let provider_attempt_id = format!("{run_id}:{operation_id}");
request.provider_attempt_id = Some(provider_attempt_id.clone());
writer
.append(vec![Event::ModelRequestPrepared {
run_id: run_id.clone(),
step,
attempt,
provider_attempt_id,
operation_id: operation_id.clone(),
reserved_prompt_tokens: estimated_prompt_tokens,
reserved_completion_tokens: u64::from(request.max_tokens),
request: serde_json::to_value(&request)
.map_err(|error| RuntimeError::Invariant(error.to_string()))?,
prompt_sections: serde_json::to_value(self.prompts.sections())
.map_err(|error| RuntimeError::Invariant(error.to_string()))?,
}])
.await?;
let (delta_tx, mut delta_rx) = tokio::sync::mpsc::unbounded_channel();
let model_call = self.model.complete_streaming(&request, delta_tx);
tokio::pin!(model_call);
let deadline = tokio::time::sleep(self.limits.provider_deadline);
tokio::pin!(deadline);
let mut input_poll = tokio::time::interval(Duration::from_millis(100));
let mut streamed = String::new();
let result = loop {
tokio::select! {
response = &mut model_call => {
while let Ok((content, _has_tool_calls)) = delta_rx.try_recv() {
let delta = content.strip_prefix(&streamed).unwrap_or(&content).to_string();
streamed = content;
if !delta.is_empty() {
writer.append(vec![Event::AssistantDelta {
run_id: run_id.clone(), step, attempt, content: delta,
}]).await?;
}
}
break Some(response)
},
update = delta_rx.recv() => {
if let Some((content, _has_tool_calls)) = update {
let delta = content.strip_prefix(&streamed).unwrap_or(&content).to_string();
streamed = content;
if !delta.is_empty() {
writer.append(vec![Event::AssistantDelta {
run_id: run_id.clone(), step, attempt, content: delta,
}]).await?;
}
}
}
_ = &mut deadline => break None,
_ = cancellation.cancelled() => {
let completion_tokens = self.estimated_completion_tokens(&streamed);
writer.append(vec![
Event::ModelAttemptFailed {
run_id: run_id.clone(), step, attempt,
error: "cancelled".into(), retryable: false,
},
Event::UsageRecorded {
run_id: run_id.clone(),
operation_id: operation_id.clone(),
prompt_tokens: estimated_prompt_tokens,
completion_tokens,
cost_units: 0,
},
]).await?;
return Err(RuntimeError::Cancelled);
}
_ = input_poll.tick() => {
if writer.load_after(after_seq).await?.iter().any(|event| matches!(
&event.event,
Event::InputQueued { run_id: target, mode: af_agent_session::DeliveryMode::Steer, .. }
if target == run_id
)) {
let completion_tokens = self.estimated_completion_tokens(&streamed);
writer.append(vec![
Event::ModelAttemptFailed {
run_id: run_id.clone(), step, attempt,
error: "steered".into(), retryable: false,
},
Event::UsageRecorded {
run_id: run_id.clone(),
operation_id: operation_id.clone(),
prompt_tokens: estimated_prompt_tokens,
completion_tokens,
cost_units: 0,
},
]).await?;
return Err(RuntimeError::Steered);
}
}
}
};
let retryable = match result {
Some(Ok(response)) => {
let (prompt_tokens, completion_tokens) = response.usage.map_or_else(
|| {
(
estimated_prompt_tokens,
response
.first_content()
.map_or(1, |content| self.estimated_completion_tokens(content)),
)
},
|usage| {
(
u64::from(usage.prompt_tokens),
u64::from(usage.completion_tokens),
)
},
);
writer
.append(vec![Event::UsageRecorded {
run_id: run_id.clone(),
operation_id: operation_id.clone(),
prompt_tokens,
completion_tokens,
cost_units: 0,
}])
.await?;
total_prompt_tokens += prompt_tokens;
total_completion_tokens += completion_tokens;
return Ok(ModelCompletion {
response,
attempt,
prompt_tokens: total_prompt_tokens,
completion_tokens: total_completion_tokens,
});
}
Some(Err(error)) => {
let retryable = error.is_retryable();
last = Some(error.to_string());
retryable
}
None => {
last = Some("deadline exceeded".into());
true
}
};
writer
.append(vec![
Event::ModelAttemptFailed {
run_id: run_id.clone(),
step,
attempt,
error: last.clone().unwrap_or_default(),
retryable,
},
Event::UsageRecorded {
run_id: run_id.clone(),
operation_id,
prompt_tokens: estimated_prompt_tokens,
completion_tokens: self.estimated_completion_tokens(&streamed),
cost_units: 0,
},
])
.await?;
total_prompt_tokens += estimated_prompt_tokens;
total_completion_tokens += self.estimated_completion_tokens(&streamed);
let failure = ModelFailure {
attempt,
error: last.clone().unwrap_or_default(),
retryable,
attempts_remaining: attempt < self.limits.provider_attempts,
};
match self.request_error(step_context, &failure).await {
RequestErrorAction::Terminal => break,
RequestErrorAction::Retry => {}
RequestErrorAction::Default if !retryable => break,
RequestErrorAction::Default => {}
}
if attempt < self.limits.provider_attempts {
let delay_ms = 100_u64.saturating_mul(2_u64.pow(attempt - 1));
writer
.append(vec![Event::RetryScheduled {
run_id: run_id.clone(),
attempt,
delay_ms,
reason: last.clone().unwrap_or_default(),
}])
.await?;
tokio::time::sleep(Duration::from_millis(delay_ms)).await;
}
}
Err(RuntimeError::Model(
last.unwrap_or_else(|| "model failed".into()),
))
}
fn estimated_completion_tokens(&self, content: &str) -> u64 {
if content.is_empty() {
return 0;
}
self.meter
.count(&self.model_name, &[ChatMessage::assistant(content)])
.max(1)
}
}