#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
use af_context::{InputId, InteractionId, ProfileRevisionId, RunId, SessionId, ToolCallId};
use std::collections::{BTreeMap, BTreeSet};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
pub const SESSION_EVENT_FORMAT_VERSION: u32 = 1;
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SessionEvent {
pub session_id: SessionId,
pub seq: u64,
pub occurred_at: DateTime<Utc>,
pub event: Event,
}
impl SessionEvent {
pub fn pending(session_id: impl Into<SessionId>, event: Event) -> Self {
Self {
session_id: session_id.into(),
seq: 0,
occurred_at: Utc::now(),
event,
}
}
pub fn format_version(&self) -> u32 {
self.event.format_version()
}
pub fn event_type(&self) -> &str {
self.event.event_type()
}
pub fn ignorable(&self) -> bool {
self.event.ignorable()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum Event {
SessionCreated {
profile_revision_id: ProfileRevisionId,
},
SessionForked {
parent_session_id: SessionId,
parent_seq: u64,
},
SessionDeleted {
reason: String,
},
InputQueued {
input_id: InputId,
run_id: RunId,
mode: DeliveryMode,
content: Vec<ContentBlock>,
explicit_skill: Option<String>,
},
InputClaimed {
input_id: InputId,
run_id: RunId,
},
InputCancelled {
input_id: InputId,
run_id: RunId,
error_code: String,
},
RunStarted {
run_id: RunId,
input_id: InputId,
},
RunWaiting {
run_id: RunId,
interaction_id: InteractionId,
},
RunResumed {
run_id: RunId,
interaction_id: InteractionId,
},
RunFinished {
run_id: RunId,
status: RunStatus,
error_code: Option<String>,
},
TurnStarted {
run_id: RunId,
turn: u32,
},
TurnFinished {
run_id: RunId,
turn: u32,
},
StepStarted {
run_id: RunId,
step: u32,
},
StepFinished {
run_id: RunId,
step: u32,
},
UserMessage {
run_id: RunId,
content: Vec<ContentBlock>,
},
AssistantDelta {
run_id: RunId,
step: u32,
attempt: u32,
content: String,
},
AssistantMessage {
run_id: RunId,
step: u32,
attempt: u32,
content: Vec<ContentBlock>,
},
AssistantToolCalls {
run_id: RunId,
step: u32,
content: Option<String>,
calls: Vec<RecordedToolCall>,
},
ToolCall {
run_id: RunId,
step: u32,
call_id: ToolCallId,
tool: String,
arguments: Value,
},
ToolAuthorization {
run_id: RunId,
step: u32,
call_id: ToolCallId,
status: ToolAuthorizationStatus,
reason: Option<String>,
},
ToolExecutionStarted {
run_id: RunId,
step: u32,
call_id: ToolCallId,
},
ToolResult {
run_id: RunId,
step: u32,
call_id: ToolCallId,
result: Value,
is_error: bool,
},
UsageRecorded {
run_id: RunId,
operation_id: String,
prompt_tokens: u64,
completion_tokens: u64,
#[serde(default)]
cost_units: u64,
},
RetryScheduled {
run_id: RunId,
attempt: u32,
delay_ms: u64,
reason: String,
},
ModelRequestPrepared {
run_id: RunId,
step: u32,
attempt: u32,
#[serde(default)]
provider_attempt_id: String,
#[serde(default)]
operation_id: String,
#[serde(default)]
reserved_prompt_tokens: u64,
#[serde(default)]
reserved_completion_tokens: u64,
request: Value,
prompt_sections: Value,
},
ContextInjected {
run_id: RunId,
step: u32,
contribution_id: String,
source: String,
version: String,
authority: String,
form: String,
content: Vec<ContentBlock>,
},
ModelAttemptFailed {
run_id: RunId,
step: u32,
attempt: u32,
error: String,
retryable: bool,
},
CompactionStarted {
run_id: RunId,
compaction_id: String,
source_through_seq: u64,
},
ToolResultsPruned {
run_id: RunId,
call_ids: Vec<ToolCallId>,
},
SummaryReplaced {
run_id: RunId,
through_seq: u64,
summary: String,
compactor: String,
model: String,
},
CompactionFinished {
run_id: RunId,
compaction_id: String,
status: String,
error: Option<String>,
},
InteractionRequested {
run_id: RunId,
interaction_id: InteractionId,
kind: InteractionKind,
payload: Value,
},
InteractionResolved {
run_id: RunId,
interaction_id: InteractionId,
resolution: InteractionResolution,
payload: Value,
},
ChildSessionLinked {
run_id: RunId,
child_session_id: SessionId,
provider: String,
},
Extension {
run_id: RunId,
plugin_id: String,
event_type: String,
payload: Value,
},
#[serde(skip)]
Opaque {
format_version: u32,
event_type: String,
ignorable: bool,
payload: Value,
},
}
impl Event {
pub fn format_version(&self) -> u32 {
match self {
Self::Opaque { format_version, .. } => *format_version,
_ => SESSION_EVENT_FORMAT_VERSION,
}
}
pub fn event_type(&self) -> &str {
match self {
Self::SessionCreated { .. } => "session_created",
Self::SessionForked { .. } => "session_forked",
Self::SessionDeleted { .. } => "session_deleted",
Self::InputQueued { .. } => "input_queued",
Self::InputClaimed { .. } => "input_claimed",
Self::InputCancelled { .. } => "input_cancelled",
Self::RunStarted { .. } => "run_started",
Self::RunWaiting { .. } => "run_waiting",
Self::RunResumed { .. } => "run_resumed",
Self::RunFinished { .. } => "run_finished",
Self::TurnStarted { .. } => "turn_started",
Self::TurnFinished { .. } => "turn_finished",
Self::StepStarted { .. } => "step_started",
Self::StepFinished { .. } => "step_finished",
Self::UserMessage { .. } => "user_message",
Self::AssistantDelta { .. } => "assistant_delta",
Self::AssistantMessage { .. } => "assistant_message",
Self::AssistantToolCalls { .. } => "assistant_tool_calls",
Self::ToolCall { .. } => "tool_call",
Self::ToolAuthorization { .. } => "tool_authorization",
Self::ToolExecutionStarted { .. } => "tool_execution_started",
Self::ToolResult { .. } => "tool_result",
Self::UsageRecorded { .. } => "usage_recorded",
Self::RetryScheduled { .. } => "retry_scheduled",
Self::ModelRequestPrepared { .. } => "model_request_prepared",
Self::ContextInjected { .. } => "context_injected",
Self::ModelAttemptFailed { .. } => "model_attempt_failed",
Self::CompactionStarted { .. } => "compaction_started",
Self::ToolResultsPruned { .. } => "tool_results_pruned",
Self::SummaryReplaced { .. } => "summary_replaced",
Self::CompactionFinished { .. } => "compaction_finished",
Self::InteractionRequested { .. } => "interaction_requested",
Self::InteractionResolved { .. } => "interaction_resolved",
Self::ChildSessionLinked { .. } => "child_session_linked",
Self::Extension { .. } => "extension",
Self::Opaque { event_type, .. } => event_type,
}
}
pub fn ignorable(&self) -> bool {
match self {
Self::Opaque { ignorable, .. } => *ignorable,
_ => false,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum DeliveryMode {
Followup,
Steer,
Inject,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
Completed,
Failed,
Cancelled,
MaxStepsReached,
}
impl RunStatus {
pub const fn as_str(self) -> &'static str {
match self {
Self::Completed => "completed",
Self::Failed => "failed",
Self::Cancelled => "cancelled",
Self::MaxStepsReached => "max_steps_reached",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InteractionKind {
Action,
UserQuestion,
}
impl InteractionKind {
pub const fn as_str(self) -> &'static str {
match self {
Self::Action => "action",
Self::UserQuestion => "user_question",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum InteractionResolution {
Confirmed,
Rejected,
Answered,
}
impl InteractionResolution {
pub const fn as_str(self) -> &'static str {
match self {
Self::Confirmed => "confirmed",
Self::Rejected => "rejected",
Self::Answered => "answered",
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ToolAuthorizationStatus {
Allowed,
Waiting,
Denied,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RecordedToolCall {
pub call_id: ToolCallId,
pub tool: String,
pub arguments: Value,
}
impl ToolAuthorizationStatus {
pub const fn as_str(self) -> &'static str {
match self {
Self::Allowed => "allowed",
Self::Waiting => "waiting",
Self::Denied => "denied",
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
Text {
text: String,
},
Resource {
resource_id: String,
media_type: String,
},
Data {
slot: String,
value: Value,
},
Citation {
resource_id: String,
label: String,
uri: String,
excerpt: Option<String>,
},
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct SessionProjection {
pub session_id: Option<SessionId>,
pub profile_revision_id: Option<ProfileRevisionId>,
pub deleted: bool,
pub last_seq: u64,
pub active_run_id: Option<RunId>,
pub waiting_interaction_id: Option<InteractionId>,
pub messages: Vec<ProjectedMessage>,
pub injected_context: Vec<ProjectedContext>,
pub run_status: BTreeMap<RunId, RunState>,
pub open_tool_calls: BTreeMap<ToolCallId, OpenToolCall>,
pub started_tool_calls: BTreeSet<ToolCallId>,
pub queued_inputs: BTreeMap<InputId, (RunId, DeliveryMode)>,
pub claimed_inputs: BTreeMap<InputId, RunId>,
pub open_turn: Option<u32>,
pub open_steps: BTreeSet<u32>,
pub next_step: u32,
pub open_compaction: Option<(String, u64)>,
pub summary: Option<String>,
usage_operations: BTreeMap<(RunId, String), (u64, u64, u64)>,
pending_provider_attempts: BTreeMap<(RunId, String), (u64, u64)>,
seen_tool_calls: BTreeSet<ToolCallId>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunState {
Running,
WaitingForInput,
Terminal(RunStatus),
}
impl RunState {
pub const fn as_str(self) -> &'static str {
match self {
Self::Running => "running",
Self::WaitingForInput => "waiting_for_input",
Self::Terminal(status) => status.as_str(),
}
}
pub const fn is_terminal(self) -> bool {
matches!(self, Self::Terminal(_))
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProjectedContext {
pub run_id: RunId,
pub step: u32,
pub contribution_id: String,
pub source: String,
pub version: String,
pub authority: String,
pub form: String,
pub content: Vec<ContentBlock>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct ProjectedMessage {
pub role: &'static str,
pub run_id: RunId,
pub content: Vec<ContentBlock>,
}
#[derive(Debug, Clone, PartialEq)]
pub struct OpenToolCall {
pub run_id: RunId,
pub step: u32,
pub tool: String,
pub arguments: Value,
pub source_event_seq: u64,
}
impl SessionProjection {
pub fn replay(events: &[SessionEvent]) -> Result<Self, EventError> {
let mut projection = Self::default();
for event in events {
projection.apply(event)?;
}
Ok(projection)
}
pub fn apply(&mut self, envelope: &SessionEvent) -> Result<(), EventError> {
if envelope.seq != self.last_seq + 1 {
return Err(EventError::Sequence {
expected: self.last_seq + 1,
actual: envelope.seq,
});
}
let session_id = self
.session_id
.get_or_insert_with(|| envelope.session_id.clone());
if *session_id != envelope.session_id {
return Err(EventError::SessionMismatch);
}
if envelope.format_version() != SESSION_EVENT_FORMAT_VERSION {
if envelope.ignorable() {
self.last_seq = envelope.seq;
return Ok(());
}
return Err(EventError::UnsupportedFormat(envelope.format_version()));
}
if self.deleted {
return Err(EventError::SessionClosed);
}
match &envelope.event {
Event::SessionCreated {
profile_revision_id,
} => {
if envelope.seq != 1 || self.profile_revision_id.is_some() {
return Err(EventError::DuplicateSession);
}
self.profile_revision_id = Some(profile_revision_id.clone());
}
Event::SessionDeleted { .. } => {
if let Some(run_id) = self.active_run_id.take() {
self.run_status
.insert(run_id, RunState::Terminal(RunStatus::Cancelled));
}
for (_, (run_id, _)) in std::mem::take(&mut self.queued_inputs) {
self.run_status
.insert(run_id, RunState::Terminal(RunStatus::Cancelled));
}
self.waiting_interaction_id = None;
self.open_turn = None;
self.open_steps.clear();
self.open_tool_calls.clear();
self.started_tool_calls.clear();
self.open_compaction = None;
self.deleted = true;
}
Event::RunStarted { run_id, input_id } => {
if self.active_run_id.is_some() {
return Err(EventError::ConcurrentRun);
}
if self.claimed_inputs.get(input_id) != Some(run_id) {
return Err(EventError::UnclaimedInput(input_id.clone()));
}
self.active_run_id = Some(run_id.clone());
self.next_step = 1;
self.run_status.insert(run_id.clone(), RunState::Running);
}
Event::RunWaiting {
run_id,
interaction_id,
} => {
self.require_active(run_id)?;
self.waiting_interaction_id = Some(interaction_id.clone());
self.run_status
.insert(run_id.clone(), RunState::WaitingForInput);
}
Event::RunResumed {
run_id,
interaction_id,
} => {
self.require_active(run_id)?;
if self.waiting_interaction_id.as_ref() != Some(interaction_id) {
return Err(EventError::InteractionMismatch);
}
self.waiting_interaction_id = None;
self.run_status.insert(run_id.clone(), RunState::Running);
}
Event::RunFinished { run_id, status, .. } => {
self.require_active(run_id)?;
if !self.open_tool_calls.is_empty()
|| !self.open_steps.is_empty()
|| self.open_turn.is_some()
|| self.open_compaction.is_some()
|| self
.queued_inputs
.values()
.any(|(target_run_id, _)| target_run_id == run_id)
{
return Err(EventError::OpenLifecycle);
}
self.run_status
.insert(run_id.clone(), RunState::Terminal(*status));
self.active_run_id = None;
self.waiting_interaction_id = None;
}
Event::UserMessage { run_id, content } => self.messages.push(ProjectedMessage {
role: "user",
run_id: run_id.clone(),
content: content.clone(),
}),
Event::AssistantMessage {
run_id, content, ..
} => self.messages.push(ProjectedMessage {
role: "assistant",
run_id: run_id.clone(),
content: content.clone(),
}),
Event::ContextInjected {
run_id,
step,
contribution_id,
source,
version,
authority,
form,
content,
} => {
self.require_active(run_id)?;
if !self.open_steps.contains(step) {
return Err(EventError::LifecycleMismatch);
}
self.injected_context.push(ProjectedContext {
run_id: run_id.clone(),
step: *step,
contribution_id: contribution_id.clone(),
source: source.clone(),
version: version.clone(),
authority: authority.clone(),
form: form.clone(),
content: content.clone(),
});
}
Event::InputQueued {
input_id,
run_id,
mode,
..
} => {
if *mode != DeliveryMode::Followup {
self.require_active(run_id)?;
}
if self
.queued_inputs
.insert(input_id.clone(), (run_id.clone(), *mode))
.is_some()
{
return Err(EventError::DuplicateInput(input_id.clone()));
}
}
Event::InputClaimed { input_id, run_id } => {
if self
.queued_inputs
.remove(input_id)
.map(|value| value.0)
.as_ref()
!= Some(run_id)
|| self
.claimed_inputs
.insert(input_id.clone(), run_id.clone())
.is_some()
{
return Err(EventError::UnqueuedInput(input_id.clone()));
}
}
Event::InputCancelled {
input_id,
run_id,
error_code,
} => {
if self
.queued_inputs
.remove(input_id)
.map(|value| value.0)
.as_ref()
!= Some(run_id)
{
return Err(EventError::UnqueuedInput(input_id.clone()));
}
self.run_status.insert(
run_id.clone(),
RunState::Terminal(if error_code == "cancelled" {
RunStatus::Cancelled
} else {
RunStatus::Failed
}),
);
}
Event::TurnStarted { run_id, turn } => {
self.require_active(run_id)?;
if self.open_turn.replace(*turn).is_some() {
return Err(EventError::ConcurrentTurn);
}
}
Event::TurnFinished { run_id, turn } => {
self.require_active(run_id)?;
if self.open_turn != Some(*turn)
|| !self.open_steps.is_empty()
|| !self.open_tool_calls.is_empty()
{
return Err(EventError::LifecycleMismatch);
}
self.open_turn = None;
}
Event::StepStarted { run_id, step } => {
self.require_active(run_id)?;
if self.open_turn.is_none()
|| !self.open_steps.is_empty()
|| *step != self.next_step
|| !self.open_steps.insert(*step)
{
return Err(EventError::ConcurrentStep);
}
}
Event::StepFinished { run_id, step } => {
self.require_active(run_id)?;
if self
.open_tool_calls
.values()
.any(|call| call.run_id == *run_id && call.step == *step)
|| !self.open_steps.remove(step)
{
return Err(EventError::LifecycleMismatch);
}
self.next_step = step.saturating_add(1);
}
Event::ToolCall {
run_id,
step,
call_id,
tool,
arguments,
} => {
self.require_active(run_id)?;
if !self.open_steps.contains(step) {
return Err(EventError::LifecycleMismatch);
}
if !self.seen_tool_calls.insert(call_id.clone())
|| self
.open_tool_calls
.insert(
call_id.clone(),
OpenToolCall {
run_id: run_id.clone(),
step: *step,
tool: tool.clone(),
arguments: arguments.clone(),
source_event_seq: envelope.seq,
},
)
.is_some()
{
return Err(EventError::DuplicateToolCall(call_id.clone()));
}
}
Event::ToolResult {
run_id,
step,
call_id,
..
} => {
self.require_active(run_id)?;
if !self.open_steps.contains(step) {
return Err(EventError::LifecycleMismatch);
}
let Some(call) = self.open_tool_calls.get(call_id) else {
return Err(EventError::OrphanToolResult(call_id.clone()));
};
if call.run_id != *run_id || call.step != *step {
return Err(EventError::ToolResultMismatch(call_id.clone()));
}
self.open_tool_calls.remove(call_id);
self.started_tool_calls.remove(call_id);
}
Event::ToolExecutionStarted {
run_id,
step,
call_id,
} => {
self.require_active(run_id)?;
let Some(call) = self.open_tool_calls.get(call_id) else {
return Err(EventError::OrphanToolResult(call_id.clone()));
};
if call.run_id != *run_id
|| call.step != *step
|| !self.started_tool_calls.insert(call_id.clone())
{
return Err(EventError::ToolResultMismatch(call_id.clone()));
}
}
Event::ModelRequestPrepared {
run_id,
operation_id,
reserved_prompt_tokens,
reserved_completion_tokens,
..
} if !operation_id.is_empty() => {
self.require_active(run_id)?;
let key = (run_id.clone(), operation_id.clone());
if !self.usage_operations.contains_key(&key) {
let reservation = (*reserved_prompt_tokens, *reserved_completion_tokens);
match self.pending_provider_attempts.get(&key) {
Some(existing) if *existing != reservation => {
return Err(EventError::UsageConflict(operation_id.clone()));
}
Some(_) => {}
None => {
self.pending_provider_attempts.insert(key, reservation);
}
}
}
}
Event::UsageRecorded {
run_id,
operation_id,
prompt_tokens,
completion_tokens,
cost_units,
} => {
self.require_active(run_id)?;
let key = (run_id.clone(), operation_id.clone());
let usage = (*prompt_tokens, *completion_tokens, *cost_units);
match self.usage_operations.get(&key) {
Some(existing) if *existing != usage => {
return Err(EventError::UsageConflict(operation_id.clone()));
}
Some(_) => {}
None => {
self.pending_provider_attempts.remove(&key);
self.usage_operations.insert(key, usage);
}
}
}
Event::CompactionStarted {
run_id,
compaction_id,
source_through_seq,
} => {
self.require_active(run_id)?;
if self.open_compaction.is_some() {
return Err(EventError::ConcurrentCompaction);
}
self.open_compaction = Some((compaction_id.clone(), *source_through_seq));
}
Event::CompactionFinished {
run_id,
compaction_id,
..
} => {
self.require_active(run_id)?;
if self.open_compaction.as_ref().map(|value| value.0.as_str())
!= Some(compaction_id.as_str())
{
return Err(EventError::CompactionMismatch);
}
self.open_compaction = None;
}
Event::SummaryReplaced { summary, .. } => self.summary = Some(summary.clone()),
Event::Opaque {
event_type,
ignorable: false,
..
} => return Err(EventError::UnknownRequired(event_type.clone())),
_ => {}
}
self.last_seq = envelope.seq;
Ok(())
}
fn require_active(&self, run_id: &RunId) -> Result<(), EventError> {
if self.active_run_id.as_ref() == Some(run_id) {
Ok(())
} else {
Err(EventError::RunMismatch)
}
}
pub fn usage_for(&self, run_id: &str) -> (u64, u64) {
self.usage_operations
.iter()
.filter(|((recorded_run_id, _), _)| recorded_run_id.as_str() == run_id)
.fold((0, 0), |total, (_, usage)| {
(total.0 + usage.0, total.1 + usage.1)
})
}
pub fn billable_units_for(&self, run_id: &str) -> u64 {
let recorded = self
.usage_operations
.iter()
.filter(|((recorded_run_id, _), _)| recorded_run_id.as_str() == run_id)
.map(|(_, usage)| usage.0 + usage.1 + usage.2)
.sum::<u64>();
recorded
+ self
.pending_provider_attempts
.iter()
.filter(|((recorded_run_id, _), _)| recorded_run_id.as_str() == run_id)
.map(|(_, usage)| usage.0 + usage.1)
.sum::<u64>()
}
}
#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum EventError {
#[error("unsupported session event format version {0}")]
UnsupportedFormat(u32),
#[error("unknown required session event type {0}")]
UnknownRequired(String),
#[error("event sequence mismatch: expected {expected}, got {actual}")]
Sequence {
expected: u64,
actual: u64,
},
#[error("event belongs to another session")]
SessionMismatch,
#[error("session creation must be the first and only creation event")]
DuplicateSession,
#[error("session already has an active run")]
ConcurrentRun,
#[error("session is closed")]
SessionClosed,
#[error("event does not match the active run")]
RunMismatch,
#[error("interaction does not match the waiting run")]
InteractionMismatch,
#[error("run cannot finish with an open turn, step or tool call")]
OpenLifecycle,
#[error("input was queued twice: {0}")]
DuplicateInput(InputId),
#[error("input was claimed before it was queued: {0}")]
UnqueuedInput(InputId),
#[error("run started from an unclaimed input: {0}")]
UnclaimedInput(InputId),
#[error("session already has an active turn")]
ConcurrentTurn,
#[error("turn already has this active step")]
ConcurrentStep,
#[error("turn or step lifecycle does not pair")]
LifecycleMismatch,
#[error("duplicate tool call {0}")]
DuplicateToolCall(ToolCallId),
#[error("tool result has no matching call {0}")]
OrphanToolResult(ToolCallId),
#[error("tool result does not match the call run and step: {0}")]
ToolResultMismatch(ToolCallId),
#[error("usage operation was recorded with different totals: {0}")]
UsageConflict(String),
#[error("session already has an active compaction")]
ConcurrentCompaction,
#[error("compaction lifecycle does not pair")]
CompactionMismatch,
#[error("event store conflict: {0}")]
Conflict(String),
#[error("event store unavailable: {0}")]
Unavailable(String),
}
#[async_trait]
pub trait SessionEventStore: Send + Sync {
async fn append(
&self,
tenant_id: &str,
session_id: &str,
expected_seq: u64,
events: Vec<Event>,
) -> Result<Vec<SessionEvent>, EventError>;
async fn load(
&self,
tenant_id: &str,
session_id: &str,
after_seq: u64,
) -> Result<Vec<SessionEvent>, EventError>;
}
pub fn text(value: impl Into<String>) -> Vec<ContentBlock> {
vec![ContentBlock::Text { text: value.into() }]
}
pub fn recovery_events(projection: &SessionProjection) -> Vec<Event> {
failure_events(projection, "worker_restarted")
}
pub fn failure_events(projection: &SessionProjection, error_code: &str) -> Vec<Event> {
termination_events(projection, RunStatus::Failed, error_code)
}
pub fn cancel_events(projection: &SessionProjection) -> Vec<Event> {
termination_events(projection, RunStatus::Cancelled, "cancelled")
}
pub fn session_deletion_events(projection: &SessionProjection, reason: &str) -> Vec<Event> {
let active = projection.active_run_id.as_deref();
let mut events = cancel_events(projection);
events.extend(
projection
.queued_inputs
.iter()
.filter(|(_, (run_id, _))| Some(run_id.as_str()) != active)
.map(|(input_id, (run_id, _))| Event::InputCancelled {
input_id: input_id.clone(),
run_id: run_id.clone(),
error_code: "cancelled".into(),
}),
);
events.push(Event::SessionDeleted {
reason: reason.into(),
});
events
}
fn termination_events(
projection: &SessionProjection,
status: RunStatus,
error_code: &str,
) -> Vec<Event> {
let Some(run_id) = &projection.active_run_id else {
return Vec::new();
};
let mut events = projection
.queued_inputs
.iter()
.filter(|(_, (target_run_id, _))| target_run_id == run_id)
.map(|(input_id, _)| Event::InputCancelled {
input_id: input_id.clone(),
run_id: run_id.clone(),
error_code: error_code.into(),
})
.collect::<Vec<_>>();
events.extend(
projection
.open_tool_calls
.iter()
.map(|(call_id, call)| Event::ToolResult {
run_id: run_id.clone(),
step: call.step,
call_id: call_id.clone(),
result: termination_result(error_code),
is_error: true,
}),
);
if let Some((compaction_id, _)) = &projection.open_compaction {
events.push(Event::CompactionFinished {
run_id: run_id.clone(),
compaction_id: compaction_id.clone(),
status: "failed".into(),
error: Some(error_code.into()),
});
}
events.extend(
projection
.open_steps
.iter()
.map(|step| Event::StepFinished {
run_id: run_id.clone(),
step: *step,
}),
);
if let Some(turn) = projection.open_turn {
events.push(Event::TurnFinished {
run_id: run_id.clone(),
turn,
});
}
events.push(Event::RunFinished {
run_id: run_id.clone(),
status,
error_code: Some(error_code.into()),
});
events
}
fn termination_result(error_code: &str) -> Value {
if error_code == "tool_outcome_unknown" || error_code == "worker_restarted" {
serde_json::json!({
"error": error_code,
"guidance": "The tool outcome is unknown. Verify external state before retrying any operation with side effects; ask the user when verification is unavailable."
})
} else {
serde_json::json!({"error":error_code})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn event(seq: u64, event: Event) -> SessionEvent {
SessionEvent {
session_id: "s".parse().unwrap(),
seq,
occurred_at: Utc::now(),
event,
}
}
#[test]
fn replay_enforces_single_run_and_tool_pairs() {
let events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::TurnStarted {
run_id: "r1".parse().unwrap(),
turn: 1,
},
),
event(
6,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 1,
},
),
event(
7,
Event::ToolCall {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
tool: "echo".into(),
arguments: serde_json::json!({"x":1}),
},
),
event(
8,
Event::ToolResult {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
result: serde_json::json!({"x":1}),
is_error: false,
},
),
event(
9,
Event::StepFinished {
run_id: "r1".parse().unwrap(),
step: 1,
},
),
event(
10,
Event::TurnFinished {
run_id: "r1".parse().unwrap(),
turn: 1,
},
),
event(
11,
Event::RunFinished {
run_id: "r1".parse().unwrap(),
status: RunStatus::Completed,
error_code: None,
},
),
];
let projection = SessionProjection::replay(&events).unwrap();
assert_eq!(projection.last_seq, 11);
assert!(projection.active_run_id.is_none());
}
#[test]
fn replay_rejects_orphan_tool_result() {
let events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::TurnStarted {
run_id: "r1".parse().unwrap(),
turn: 1,
},
),
event(
6,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 1,
},
),
event(
7,
Event::ToolResult {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "missing".parse().unwrap(),
result: Value::Null,
is_error: true,
},
),
];
assert_eq!(
SessionProjection::replay(&events).unwrap_err(),
EventError::OrphanToolResult("missing".parse().unwrap())
);
}
#[test]
fn queued_input_failure_is_not_projected_as_cancellation() {
let events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputCancelled {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
error_code: "profile_not_found".into(),
},
),
];
let projection = SessionProjection::replay(&events).unwrap();
assert_eq!(
projection.run_status.get("r1").map(|state| state.as_str()),
Some("failed")
);
}
#[test]
fn session_deletion_closes_active_and_queued_runs_before_tombstone() {
let mut events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("start"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::TurnStarted {
run_id: "r1".parse().unwrap(),
turn: 1,
},
),
event(
6,
Event::InputQueued {
input_id: "i2".parse().unwrap(),
run_id: "r2".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("later"),
explicit_skill: None,
},
),
];
let projection = SessionProjection::replay(&events).unwrap();
for event_value in session_deletion_events(&projection, "api_deleted") {
let seq = events.len() as u64 + 1;
events.push(event(seq, event_value));
}
let deleted = SessionProjection::replay(&events).unwrap();
assert!(deleted.deleted);
assert_eq!(
deleted.run_status.get("r1").map(|state| state.as_str()),
Some("cancelled")
);
assert_eq!(
deleted.run_status.get("r2").map(|state| state.as_str()),
Some("cancelled")
);
assert!(events.iter().any(|event| matches!(
&event.event,
Event::RunFinished { run_id, status: RunStatus::Cancelled, .. } if run_id == "r1"
)));
assert!(matches!(
events.last().map(|event| &event.event),
Some(Event::SessionDeleted { .. })
));
}
#[test]
fn usage_operations_are_idempotent_and_conflicts_fail_replay() {
let mut events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::UsageRecorded {
run_id: "r1".parse().unwrap(),
operation_id: "model:1:attempt:1".into(),
prompt_tokens: 7,
completion_tokens: 3,
cost_units: 5,
},
),
event(
6,
Event::UsageRecorded {
run_id: "r1".parse().unwrap(),
operation_id: "model:1:attempt:1".into(),
prompt_tokens: 7,
completion_tokens: 3,
cost_units: 5,
},
),
];
assert_eq!(
SessionProjection::replay(&events).unwrap().usage_for("r1"),
(7, 3)
);
assert_eq!(
SessionProjection::replay(&events)
.unwrap()
.billable_units_for("r1"),
15
);
events.push(event(
7,
Event::UsageRecorded {
run_id: "r1".parse().unwrap(),
operation_id: "model:1:attempt:1".into(),
prompt_tokens: 8,
completion_tokens: 3,
cost_units: 0,
},
));
assert_eq!(
SessionProjection::replay(&events).unwrap_err(),
EventError::UsageConflict("model:1:attempt:1".into())
);
}
#[test]
fn unresolved_provider_attempt_is_conservatively_billable_and_reconcilable() {
let mut events = vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::ModelRequestPrepared {
run_id: "r1".parse().unwrap(),
step: 1,
attempt: 1,
provider_attempt_id: "r1:model:1:attempt:1".into(),
operation_id: "model:1:attempt:1".into(),
reserved_prompt_tokens: 7,
reserved_completion_tokens: 11,
request: Value::Null,
prompt_sections: Value::Null,
},
),
];
assert_eq!(
SessionProjection::replay(&events)
.unwrap()
.billable_units_for("r1"),
18
);
events.push(event(
6,
Event::UsageRecorded {
run_id: "r1".parse().unwrap(),
operation_id: "model:1:attempt:1".into(),
prompt_tokens: 6,
completion_tokens: 2,
cost_units: 0,
},
));
assert_eq!(
SessionProjection::replay(&events)
.unwrap()
.billable_units_for("r1"),
8
);
}
fn open_step_events() -> Vec<SessionEvent> {
vec![
event(
1,
Event::SessionCreated {
profile_revision_id: "p1".parse().unwrap(),
},
),
event(
2,
Event::InputQueued {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
mode: DeliveryMode::Followup,
content: text("hi"),
explicit_skill: None,
},
),
event(
3,
Event::InputClaimed {
input_id: "i1".parse().unwrap(),
run_id: "r1".parse().unwrap(),
},
),
event(
4,
Event::RunStarted {
run_id: "r1".parse().unwrap(),
input_id: "i1".parse().unwrap(),
},
),
event(
5,
Event::TurnStarted {
run_id: "r1".parse().unwrap(),
turn: 1,
},
),
event(
6,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 1,
},
),
]
}
#[test]
fn replay_rejects_overlapping_or_out_of_order_steps() {
let mut overlapping = open_step_events();
overlapping.push(event(
7,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 2,
},
));
assert_eq!(
SessionProjection::replay(&overlapping).unwrap_err(),
EventError::ConcurrentStep
);
let mut skipped = open_step_events();
skipped[5] = event(
6,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 2,
},
);
assert_eq!(
SessionProjection::replay(&skipped).unwrap_err(),
EventError::ConcurrentStep
);
}
#[test]
fn replay_rejects_cross_step_results_and_dangling_calls() {
let mut events = open_step_events();
events.push(event(
7,
Event::ToolCall {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
tool: "echo".into(),
arguments: Value::Null,
},
));
events.push(event(
8,
Event::ToolResult {
run_id: "r1".parse().unwrap(),
step: 2,
call_id: "c1".parse().unwrap(),
result: Value::Null,
is_error: false,
},
));
assert_eq!(
SessionProjection::replay(&events).unwrap_err(),
EventError::LifecycleMismatch
);
let mut dangling = open_step_events();
dangling.push(event(
7,
Event::ToolCall {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
tool: "echo".into(),
arguments: Value::Null,
},
));
dangling.push(event(
8,
Event::StepFinished {
run_id: "r1".parse().unwrap(),
step: 1,
},
));
assert_eq!(
SessionProjection::replay(&dangling).unwrap_err(),
EventError::LifecycleMismatch
);
}
#[test]
fn replay_rejects_reused_tool_call_ids() {
let mut events = open_step_events();
events.extend([
event(
7,
Event::ToolCall {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
tool: "echo".into(),
arguments: Value::Null,
},
),
event(
8,
Event::ToolResult {
run_id: "r1".parse().unwrap(),
step: 1,
call_id: "c1".parse().unwrap(),
result: Value::Null,
is_error: false,
},
),
event(
9,
Event::StepFinished {
run_id: "r1".parse().unwrap(),
step: 1,
},
),
event(
10,
Event::StepStarted {
run_id: "r1".parse().unwrap(),
step: 2,
},
),
event(
11,
Event::ToolCall {
run_id: "r1".parse().unwrap(),
step: 2,
call_id: "c1".parse().unwrap(),
tool: "echo".into(),
arguments: Value::Null,
},
),
]);
assert_eq!(
SessionProjection::replay(&events).unwrap_err(),
EventError::DuplicateToolCall("c1".parse().unwrap())
);
}
#[test]
fn replay_keeps_injected_context_without_changing_run_state() {
let mut events = open_step_events();
events.push(event(
7,
Event::ContextInjected {
run_id: "r1".parse().unwrap(),
step: 1,
contribution_id: "memory:1".into(),
source: "memory".into(),
version: "v1".into(),
authority: "tenant".into(),
form: "message".into(),
content: text("tenant context"),
},
));
let projection = SessionProjection::replay(&events).unwrap();
assert_eq!(projection.active_run_id.as_deref(), Some("r1"));
assert_eq!(projection.open_steps, BTreeSet::from([1]));
assert_eq!(
projection.injected_context,
vec![ProjectedContext {
run_id: "r1".parse().unwrap(),
step: 1,
contribution_id: "memory:1".into(),
source: "memory".into(),
version: "v1".into(),
authority: "tenant".into(),
form: "message".into(),
content: text("tenant context"),
}]
);
}
#[test]
fn replay_skips_unknown_ignorable_formats_and_rejects_required_events() {
let optional = event(
1,
Event::Opaque {
format_version: SESSION_EVENT_FORMAT_VERSION,
event_type: "future_optional".into(),
ignorable: true,
payload: serde_json::json!({"answer": 42}),
},
);
let projection = SessionProjection::replay(&[optional]).unwrap();
assert_eq!(projection.last_seq, 1);
let required = event(
1,
Event::Opaque {
format_version: SESSION_EVENT_FORMAT_VERSION,
event_type: "future_required".into(),
ignorable: false,
payload: Value::Null,
},
);
assert_eq!(
SessionProjection::replay(&[required]).unwrap_err(),
EventError::UnknownRequired("future_required".into())
);
let unsupported = event(
1,
Event::Opaque {
format_version: SESSION_EVENT_FORMAT_VERSION + 1,
event_type: "future_optional".into(),
ignorable: true,
payload: Value::Null,
},
);
assert_eq!(
SessionProjection::replay(&[unsupported]).unwrap().last_seq,
1
);
let required_new_format = event(
1,
Event::Opaque {
format_version: SESSION_EVENT_FORMAT_VERSION + 1,
event_type: "future_required".into(),
ignorable: false,
payload: Value::Null,
},
);
assert_eq!(
SessionProjection::replay(&[required_new_format]).unwrap_err(),
EventError::UnsupportedFormat(SESSION_EVENT_FORMAT_VERSION + 1)
);
}
}