af_agent_runtime/
lib.rs

1//! The single default cloud Agent loop. It emits every model-visible fact and
2//! lifecycle transition before exposing it to projections or UI consumers.
3
4mod builder;
5mod compactor;
6mod error;
7mod extension;
8mod model;
9mod output;
10mod recovery;
11mod replay;
12mod tools;
13mod types;
14
15pub use af_agent::CancellationToken;
16pub use af_agent_session::{cancel_events, failure_events, recovery_events};
17pub use compactor::ModelCompactor;
18pub use error::RuntimeError;
19use model::ModelTurn;
20use output::assistant_content;
21use replay::{content_text, tool_message, transcript_from_events};
22use tools::PlannedCall;
23pub use types::{
24    ApproximateTokenMeter, CompactionResult, Compactor, EventWriter, RuntimeLimits, RuntimeOutcome,
25    TokenMeter, TurnRequest,
26};
27
28use std::{collections::HashSet, sync::Arc};
29
30use af_agent::{
31    ChatModel, ContextContributor, Hook, HookDecision, MountedPlugins, PromptRegistry, ToolRegistry,
32};
33use af_agent_session::{
34    DeliveryMode, Event, RecordedToolCall, RunStatus, SessionEvent, SessionProjection,
35    ToolAuthorizationStatus,
36};
37use af_llm::{ChatMessage, FinishReason, Role};
38use serde_json::{json, Value};
39
40pub struct AgentRuntime {
41    model: Arc<dyn ChatModel>,
42    model_name: String,
43    prompts: PromptRegistry,
44    contexts: Vec<Arc<dyn ContextContributor>>,
45    tools: ToolRegistry,
46    hooks: Vec<Arc<dyn Hook>>,
47    limits: RuntimeLimits,
48    meter: Arc<dyn TokenMeter>,
49    compactor: Arc<dyn Compactor>,
50    plugins: Option<MountedPlugins>,
51    reasoning_effort: Option<af_llm::ReasoningEffort>,
52    output_policy: Option<Value>,
53}
54
55struct InputDrain<'a> {
56    claimed: &'a mut HashSet<String>,
57    transcript: &'a mut Vec<ChatMessage>,
58    context_query: &'a mut String,
59    seen_seq: &'a mut u64,
60}
61
62impl AgentRuntime {
63    pub async fn run(
64        &self,
65        request: TurnRequest,
66        writer: &dyn EventWriter,
67        cancellation: CancellationToken,
68    ) -> Result<RuntimeOutcome, RuntimeError> {
69        let projection = SessionProjection::replay(&request.history)
70            .map_err(|error| RuntimeError::Invariant(error.to_string()))?;
71        if projection
72            .active_run_id
73            .as_deref()
74            .is_some_and(|active| active != request.run_id)
75        {
76            return Err(RuntimeError::SessionBusy);
77        }
78        let user_text = content_text(&request.content);
79        if user_text.trim().is_empty() {
80            return Err(RuntimeError::InvalidInput(
81                "text content is required".into(),
82            ));
83        }
84        if projection.active_run_id.is_none() {
85            writer
86                .append(vec![
87                    Event::InputClaimed {
88                        input_id: request.input_id.clone(),
89                        run_id: request.run_id.clone(),
90                    },
91                    Event::RunStarted {
92                        run_id: request.run_id.clone(),
93                        input_id: request.input_id.clone(),
94                    },
95                ])
96                .await?;
97        }
98        let resuming = projection.open_turn.is_some();
99        if !resuming {
100            writer
101                .append(vec![
102                    Event::TurnStarted {
103                        run_id: request.run_id.clone(),
104                        turn: 1,
105                    },
106                    Event::UserMessage {
107                        run_id: request.run_id.clone(),
108                        content: request.content.clone(),
109                    },
110                ])
111                .await?;
112        }
113
114        let mut transcript = transcript_from_events(&request.history);
115        let mut context_query = user_text.clone();
116        if !resuming {
117            transcript.push(ChatMessage::user(user_text));
118        }
119        let (mut prompt_tokens, mut completion_tokens) = projection.usage_for(&request.run_id);
120        let mut tool_calls = request
121            .history
122            .iter()
123            .filter(|event| {
124                matches!(&event.event, Event::ToolCall { run_id, .. } if run_id == &request.run_id)
125            })
126            .count() as u32;
127        let mut seen_seq = request.history.last().map_or(0, |event| event.seq);
128        let mut claimed_inputs = projection
129            .claimed_inputs
130            .keys()
131            .cloned()
132            .collect::<HashSet<_>>();
133
134        let recovery = self
135            .recover_open_surface(
136                &request,
137                &projection,
138                writer,
139                cancellation.clone(),
140                &mut transcript,
141                (prompt_tokens, completion_tokens),
142            )
143            .await?;
144        if let Some(outcome) = recovery.terminal {
145            return Ok(outcome);
146        }
147        let first_step = recovery.first_step;
148
149        Self::drain_inputs(
150            &request,
151            &request.history,
152            writer,
153            InputDrain {
154                claimed: &mut claimed_inputs,
155                transcript: &mut transcript,
156                context_query: &mut context_query,
157                seen_seq: &mut seen_seq,
158            },
159        )
160        .await?;
161
162        for step in first_step..=self.limits.max_steps {
163            if cancellation.is_cancelled() {
164                return self
165                    .finish_open(
166                        writer,
167                        &request.run_id,
168                        None,
169                        RunStatus::Cancelled,
170                        None,
171                        (prompt_tokens, completion_tokens),
172                    )
173                    .await;
174            }
175            let incoming = writer.load_after(seen_seq).await?;
176            if let Some(last) = incoming.last() {
177                seen_seq = last.seq;
178            }
179            Self::drain_inputs(
180                &request,
181                &incoming,
182                writer,
183                InputDrain {
184                    claimed: &mut claimed_inputs,
185                    transcript: &mut transcript,
186                    context_query: &mut context_query,
187                    seen_seq: &mut seen_seq,
188                },
189            )
190            .await?;
191            writer
192                .append(vec![Event::StepStarted {
193                    run_id: request.run_id.clone(),
194                    step,
195                }])
196                .await?;
197            let context_messages = match self
198                .context_messages(&request, step, &context_query, writer, cancellation.clone())
199                .await
200            {
201                Ok(messages) => messages,
202                Err(RuntimeError::Cancelled) => {
203                    return self
204                        .finish_open(
205                            writer,
206                            &request.run_id,
207                            Some(step),
208                            RunStatus::Cancelled,
209                            None,
210                            (prompt_tokens, completion_tokens),
211                        )
212                        .await
213                }
214                Err(error) => return Err(error),
215            };
216            let compacted = match self
217                .compact_if_needed(
218                    writer,
219                    &request.run_id,
220                    step,
221                    &mut transcript,
222                    &context_messages,
223                    cancellation.clone(),
224                )
225                .await
226            {
227                Ok(compacted) => compacted,
228                Err(RuntimeError::Cancelled) => {
229                    return self
230                        .finish_open(
231                            writer,
232                            &request.run_id,
233                            Some(step),
234                            RunStatus::Cancelled,
235                            None,
236                            (prompt_tokens, completion_tokens),
237                        )
238                        .await
239                }
240                Err(RuntimeError::CompactionConflict) => {
241                    writer
242                        .append(vec![Event::StepFinished {
243                            run_id: request.run_id.clone(),
244                            step,
245                        }])
246                        .await?;
247                    continue;
248                }
249                Err(error) => return Err(error),
250            };
251            prompt_tokens += compacted.0;
252            completion_tokens += compacted.1;
253            let completion = match self
254                .complete_with_retry(
255                    writer,
256                    ModelTurn {
257                        run_id: &request.run_id,
258                        step,
259                        transcript: &transcript,
260                        context: &context_messages,
261                        after_seq: seen_seq,
262                    },
263                    cancellation.clone(),
264                )
265                .await
266            {
267                Ok(completion) => completion,
268                Err(RuntimeError::Cancelled) => {
269                    return self
270                        .finish_open(
271                            writer,
272                            &request.run_id,
273                            Some(step),
274                            RunStatus::Cancelled,
275                            None,
276                            (prompt_tokens, completion_tokens),
277                        )
278                        .await
279                }
280                Err(RuntimeError::Steered) => {
281                    writer
282                        .append(vec![Event::StepFinished {
283                            run_id: request.run_id.clone(),
284                            step,
285                        }])
286                        .await?;
287                    continue;
288                }
289                Err(error) => return Err(error),
290            };
291            prompt_tokens += completion.prompt_tokens;
292            completion_tokens += completion.completion_tokens;
293            let response = completion.response;
294            if prompt_tokens + completion_tokens > self.limits.max_tokens {
295                return self
296                    .finish_open(
297                        writer,
298                        &request.run_id,
299                        Some(step),
300                        RunStatus::MaxStepsReached,
301                        None,
302                        (prompt_tokens, completion_tokens),
303                    )
304                    .await;
305            }
306            let choice = response
307                .choices
308                .into_iter()
309                .next()
310                .ok_or(RuntimeError::EmptyModelResponse)?;
311            let finish_reason = choice
312                .finish_reason
313                .ok_or_else(|| RuntimeError::Model("missing finish_reason".into()))?;
314            let message = choice.message;
315            let output_blocks = choice.output_blocks;
316            match finish_reason {
317                FinishReason::Length => {
318                    return self
319                        .finish_open(
320                            writer,
321                            &request.run_id,
322                            Some(step),
323                            RunStatus::MaxStepsReached,
324                            None,
325                            (prompt_tokens, completion_tokens),
326                        )
327                        .await
328                }
329                FinishReason::ContentFilter | FinishReason::Unknown(_) => {
330                    return Err(RuntimeError::FinishReason(finish_reason))
331                }
332                FinishReason::Stop => {
333                    let (content, answer) =
334                        assistant_content(message.content.as_deref(), output_blocks)?;
335                    if let Some(policy) = &self.output_policy {
336                        af_agent::validate_json_schema_value(
337                            policy,
338                            &serde_json::to_value(&content)
339                                .map_err(|error| RuntimeError::Invariant(error.to_string()))?,
340                            "assistant output",
341                        )
342                        .map_err(RuntimeError::Model)?;
343                    }
344                    let terminal = writer
345                        .append(vec![
346                            Event::AssistantMessage {
347                                run_id: request.run_id.clone(),
348                                step,
349                                attempt: completion.attempt,
350                                content,
351                            },
352                            Event::StepFinished {
353                                run_id: request.run_id.clone(),
354                                step,
355                            },
356                            Event::TurnFinished {
357                                run_id: request.run_id.clone(),
358                                turn: 1,
359                            },
360                            Event::RunFinished {
361                                run_id: request.run_id.clone(),
362                                status: RunStatus::Completed,
363                                error_code: None,
364                            },
365                        ])
366                        .await;
367                    if let Err(error) = terminal {
368                        let pending = writer.load_after(seen_seq).await?;
369                        if pending.iter().any(|event| matches!(
370                            &event.event,
371                            Event::InputQueued { run_id, mode: af_agent_session::DeliveryMode::Steer | af_agent_session::DeliveryMode::Inject, .. }
372                                if run_id == &request.run_id
373                        )) {
374                            writer
375                                .append(vec![Event::StepFinished {
376                                    run_id: request.run_id.clone(),
377                                    step,
378                                }])
379                                .await?;
380                            continue;
381                        }
382                        return Err(error);
383                    }
384                    return Ok(RuntimeOutcome {
385                        status: RunStatus::Completed.as_str().into(),
386                        final_text: answer,
387                        prompt_tokens,
388                        completion_tokens,
389                        waiting_interaction_id: None,
390                    });
391                }
392                FinishReason::ToolCalls => {}
393            }
394            if !output_blocks.is_empty() {
395                return Err(RuntimeError::Model(
396                    "structured output cannot accompany tool calls".into(),
397                ));
398            }
399            let calls = message.tool_calls.clone().unwrap_or_default();
400            if calls.is_empty() {
401                return Err(RuntimeError::Model(
402                    "finish_reason tool_calls without tool_calls".into(),
403                ));
404            }
405            if tool_calls + calls.len() as u32 > self.limits.max_tool_calls {
406                return self
407                    .finish_open(
408                        writer,
409                        &request.run_id,
410                        Some(step),
411                        RunStatus::MaxStepsReached,
412                        None,
413                        (prompt_tokens, completion_tokens),
414                    )
415                    .await;
416            }
417            let planned = calls
418                .into_iter()
419                .map(|call| {
420                    let id = if call.id.trim().is_empty() {
421                        uuid::Uuid::new_v4().to_string()
422                    } else {
423                        call.id.clone()
424                    };
425                    let canonical_name = self
426                        .tools
427                        .canonical_name(&call.function.name)
428                        .map(str::to_string);
429                    let (arguments, preflight_error) =
430                        match serde_json::from_str(&call.function.arguments) {
431                            Ok(arguments) => {
432                                let error = self
433                                    .tools
434                                    .validate_arguments(&call.function.name, &arguments)
435                                    .err();
436                                (arguments, error)
437                            }
438                            Err(error) => (
439                                json!({"_raw":call.function.arguments}),
440                                Some(format!("invalid JSON arguments: {error}")),
441                            ),
442                        };
443                    PlannedCall {
444                        transcript_id: id.clone(),
445                        id,
446                        name: canonical_name.unwrap_or(call.function.name),
447                        arguments,
448                        step,
449                        source_event_seq: 0,
450                        preflight_error,
451                    }
452                })
453                .collect::<Vec<_>>();
454            let mut durable_calls = vec![Event::AssistantToolCalls {
455                run_id: request.run_id.clone(),
456                step,
457                content: message.content.clone(),
458                calls: planned
459                    .iter()
460                    .map(|call| RecordedToolCall {
461                        call_id: call.id.clone(),
462                        tool: call.name.clone(),
463                        arguments: call.arguments.clone(),
464                    })
465                    .collect(),
466            }];
467            durable_calls.extend(planned.iter().map(|call| Event::ToolCall {
468                run_id: request.run_id.clone(),
469                step,
470                call_id: call.id.clone(),
471                tool: call.name.clone(),
472                arguments: call.arguments.clone(),
473            }));
474            let appended_calls = writer.append(durable_calls).await?;
475            transcript.push(ChatMessage {
476                role: Role::Assistant,
477                content: message.content,
478                tool_calls: Some(
479                    planned
480                        .iter()
481                        .map(|call| af_llm::ToolCall {
482                            id: call.id.clone(),
483                            kind: "function".into(),
484                            function: af_llm::FunctionCall {
485                                name: af_agent::model_tool_name(&call.name),
486                                arguments: call.arguments.to_string(),
487                            },
488                        })
489                        .collect(),
490                ),
491                tool_call_id: None,
492                name: None,
493            });
494
495            let mut planned = planned;
496            let mut decisions = Vec::with_capacity(planned.len());
497            for (call, envelope) in planned.iter_mut().zip(appended_calls.iter().skip(1)) {
498                call.source_event_seq = envelope.seq;
499                let decision = if let Some(reason) = &call.preflight_error {
500                    HookDecision::Deny {
501                        reason: reason.clone(),
502                    }
503                } else {
504                    self.authorize_call(&request, call, None, cancellation.clone())
505                        .await
506                };
507                decisions.push(decision);
508            }
509            tool_calls += planned.len() as u32;
510            if let Some(waiting_index) = decisions
511                .iter()
512                .position(|decision| matches!(decision, HookDecision::WaitForInput { .. }))
513            {
514                let mut events = Vec::with_capacity(planned.len() * 2 + 2);
515                for (index, call) in planned.iter().enumerate() {
516                    if index == waiting_index {
517                        events.push(Event::ToolAuthorization {
518                            run_id: request.run_id.clone(),
519                            step,
520                            call_id: call.id.clone(),
521                            status: ToolAuthorizationStatus::Waiting,
522                            reason: None,
523                        });
524                    } else {
525                        let reason = match &decisions[index] {
526                            HookDecision::Deny { reason } => reason.clone(),
527                            _ => "blocked_by_pending_interaction".into(),
528                        };
529                        events.extend([
530                            Event::ToolAuthorization {
531                                run_id: request.run_id.clone(),
532                                step,
533                                call_id: call.id.clone(),
534                                status: ToolAuthorizationStatus::Denied,
535                                reason: Some(reason.clone()),
536                            },
537                            Event::ToolResult {
538                                run_id: request.run_id.clone(),
539                                step,
540                                call_id: call.id.clone(),
541                                result: json!({"error":reason}),
542                                is_error: true,
543                            },
544                        ]);
545                    }
546                }
547                let call = &planned[waiting_index];
548                let HookDecision::WaitForInput { kind, mut payload } =
549                    decisions.swap_remove(waiting_index)
550                else {
551                    unreachable!("waiting decision selected above")
552                };
553                if let Value::Object(object) = &mut payload {
554                    object.insert("call_id".into(), Value::String(call.id.clone()));
555                    object.insert(
556                        "source_event_seq".into(),
557                        Value::from(call.source_event_seq),
558                    );
559                }
560                let interaction_id = uuid::Uuid::new_v4().to_string();
561                events.extend([
562                    Event::InteractionRequested {
563                        run_id: request.run_id.clone(),
564                        interaction_id: interaction_id.clone(),
565                        kind: if kind == "question" {
566                            af_agent_session::InteractionKind::UserQuestion
567                        } else {
568                            af_agent_session::InteractionKind::Action
569                        },
570                        payload,
571                    },
572                    Event::RunWaiting {
573                        run_id: request.run_id.clone(),
574                        interaction_id: interaction_id.clone(),
575                    },
576                ]);
577                writer.append(events).await?;
578                return Ok(RuntimeOutcome {
579                    status: "waiting_for_input".into(),
580                    final_text: None,
581                    prompt_tokens,
582                    completion_tokens,
583                    waiting_interaction_id: Some(interaction_id),
584                });
585            }
586
587            let mut executable = Vec::new();
588            for (call, decision) in planned.iter().zip(decisions) {
589                match decision {
590                    HookDecision::Continue => {
591                        writer
592                            .append(vec![Event::ToolAuthorization {
593                                run_id: request.run_id.clone(),
594                                step,
595                                call_id: call.id.clone(),
596                                status: ToolAuthorizationStatus::Allowed,
597                                reason: None,
598                            }])
599                            .await?;
600                        executable.push(call.clone());
601                    }
602                    HookDecision::Deny { reason } => {
603                        let value = json!({"error":reason});
604                        writer
605                            .append(vec![
606                                Event::ToolAuthorization {
607                                    run_id: request.run_id.clone(),
608                                    step,
609                                    call_id: call.id.clone(),
610                                    status: ToolAuthorizationStatus::Denied,
611                                    reason: Some(reason),
612                                },
613                                Event::ToolResult {
614                                    run_id: request.run_id.clone(),
615                                    step,
616                                    call_id: call.id.clone(),
617                                    result: value.clone(),
618                                    is_error: true,
619                                },
620                            ])
621                            .await?;
622                        transcript.push(tool_message(call, value));
623                    }
624                    HookDecision::WaitForInput { .. } => {
625                        writer
626                            .append(vec![Event::ToolAuthorization {
627                                run_id: request.run_id.clone(),
628                                step,
629                                call_id: call.id.clone(),
630                                status: ToolAuthorizationStatus::Waiting,
631                                reason: None,
632                            }])
633                            .await?;
634                        unreachable!("waiting decisions are handled as one atomic batch");
635                    }
636                }
637            }
638            let executed = self
639                .execute_tools(&request, &executable, writer, cancellation.clone(), None)
640                .await?;
641            let mut tool_outcome_unknown = false;
642            for (call, result) in executable.into_iter().zip(executed) {
643                let outcome_unknown = matches!(
644                    &result,
645                    Err(error) if error.contains("tool_outcome_unknown")
646                );
647                tool_outcome_unknown |= outcome_unknown;
648                let value = match result {
649                    Ok(value) => value,
650                    Err(error) => json!({"error": error}),
651                };
652                if !outcome_unknown {
653                    if let Err(error) = self
654                        .run_after_hooks(&request, &call, &value, None, cancellation.clone())
655                        .await
656                    {
657                        writer
658                            .append(vec![Event::Extension {
659                                run_id: request.run_id.clone(),
660                                plugin_id: "agentfactory.runtime".into(),
661                                event_type: "after_tool_failed".into(),
662                                payload: json!({"call_id":call.id,"error":error}),
663                            }])
664                            .await?;
665                    }
666                }
667                transcript.push(tool_message(&call, value));
668            }
669            writer
670                .append(vec![Event::StepFinished {
671                    run_id: request.run_id.clone(),
672                    step,
673                }])
674                .await?;
675            if tool_outcome_unknown {
676                return self
677                    .finish_open(
678                        writer,
679                        &request.run_id,
680                        None,
681                        RunStatus::Failed,
682                        None,
683                        (prompt_tokens, completion_tokens),
684                    )
685                    .await;
686            }
687        }
688        self.finish_open(
689            writer,
690            &request.run_id,
691            None,
692            RunStatus::MaxStepsReached,
693            None,
694            (prompt_tokens, completion_tokens),
695        )
696        .await
697    }
698
699    async fn drain_inputs(
700        request: &TurnRequest,
701        events: &[SessionEvent],
702        writer: &dyn EventWriter,
703        state: InputDrain<'_>,
704    ) -> Result<(), RuntimeError> {
705        for envelope in events {
706            let Event::InputQueued {
707                input_id,
708                run_id,
709                mode: DeliveryMode::Steer | DeliveryMode::Inject,
710                content,
711                ..
712            } = &envelope.event
713            else {
714                continue;
715            };
716            if run_id != &request.run_id || !state.claimed.insert(input_id.clone()) {
717                continue;
718            }
719            let appended = writer
720                .append(vec![
721                    Event::InputClaimed {
722                        input_id: input_id.clone(),
723                        run_id: request.run_id.clone(),
724                    },
725                    Event::UserMessage {
726                        run_id: request.run_id.clone(),
727                        content: content.clone(),
728                    },
729                ])
730                .await?;
731            *state.seen_seq = appended.last().map_or(*state.seen_seq, |event| event.seq);
732            *state.context_query = content_text(content);
733            state
734                .transcript
735                .push(ChatMessage::user(state.context_query.clone()));
736        }
737        Ok(())
738    }
739
740    async fn finish_open(
741        &self,
742        writer: &dyn EventWriter,
743        run_id: &str,
744        step: Option<u32>,
745        status: RunStatus,
746        final_text: Option<String>,
747        usage: (u64, u64),
748    ) -> Result<RuntimeOutcome, RuntimeError> {
749        let mut events = Vec::with_capacity(3);
750        if let Some(step) = step {
751            events.push(Event::StepFinished {
752                run_id: run_id.into(),
753                step,
754            });
755        }
756        events.push(Event::TurnFinished {
757            run_id: run_id.into(),
758            turn: 1,
759        });
760        events.push(Event::RunFinished {
761            run_id: run_id.into(),
762            status,
763            error_code: None,
764        });
765        writer.append(events).await?;
766        Ok(RuntimeOutcome {
767            status: status.as_str().into(),
768            final_text,
769            prompt_tokens: usage.0,
770            completion_tokens: usage.1,
771            waiting_interaction_id: None,
772        })
773    }
774}
af_agent_runtime/lib.rs

af_agent_runtime/
lib.rs
