pi-oxidized 0.1.1

package/.agent-tasks/pi-rust-rewrite/tests/provider_drain_isolation.rs

@@ -0,0 +1,760 @@
+//! Verification check 14: ProviderDrain isolation.
+//!
+//! Fixture: `fixtures/provider-drain/isolation-expectations.json`.
+//!
+//! Block one extension event consumer during a long mock provider stream,
+//! fill its capacity-64 queue, and assert:
+//! - exactly one `extension_error` isolation signal (`ExtensionEvent::Lagged`)
+//! - only that extension disconnects
+//! - the provider drain + native UI partial watch keep advancing
+//! - the agent turn reaches exactly one final settle
+//! - session persistence records the final assistant and tool results
+
+use std::collections::BTreeMap;
+use std::path::PathBuf;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::time::Duration;
+
+use futures::future::BoxFuture;
+use futures::stream::{self, BoxStream, StreamExt};
+use pi::core::agent_session::prompt::PromptOptions;
+use pi::core::agent_session::{AgentSession, AgentSessionConfig, AgentSessionEvent};
+use pi::core::sessions::SessionEntry;
+use pi_agent::{
+    AgentEvent, AgentTool, AgentToolResult, EXTENSION_EVENT_CAPACITY, ExtensionEvent, ToolError,
+    ToolUpdates,
+};
+use pi_ai::{
+    AssistantContent, AssistantMessage, AssistantMessageEvent, Context, DoneReason, Message, Model,
+    ModelCost, ModelInput, Provider, ProviderError, StopReason, StreamOptions, TextContent,
+    ToolCall, ToolResultContent,
+};
+use serde_json::{Map, Value, json};
+use tokio::sync::{Notify, mpsc, oneshot};
+use tokio::time::timeout;
+
+type TestResult<T = ()> = Result<T, String>;
+
+#[derive(Debug)]
+struct IsolationExpectations {
+    extension_event_capacity: usize,
+    first_turn_partial_count: usize,
+    tool_call_id: String,
+    tool_name: String,
+    final_assistant_text: String,
+    tool_result_text: String,
+    extension_error_count: usize,
+    healthy_extension_lag_count: usize,
+    agent_settled_count: usize,
+    min_stream_starts: usize,
+    min_partial_watch_versions: usize,
+    tool_execution_count: usize,
+}
+
+fn fixture_path() -> PathBuf {
+    PathBuf::from(env!("CARGO_MANIFEST_DIR"))
+        .join("fixtures/provider-drain/isolation-expectations.json")
+}
+
+fn load_expectations() -> TestResult<IsolationExpectations> {
+    let path = fixture_path();
+    let raw = std::fs::read_to_string(&path)
+        .map_err(|err| format!("read fixture {}: {err}", path.display()))?;
+    let value: Value =
+        serde_json::from_str(&raw).map_err(|err| format!("parse fixture: {err}"))?;
+
+    let capacity = value
+        .get("extensionEventCapacity")
+        .and_then(Value::as_u64)
+        .ok_or_else(|| "fixture missing extensionEventCapacity".to_owned())?
+        as usize;
+    let stream = value
+        .get("stream")
+        .ok_or_else(|| "fixture missing stream".to_owned())?;
+    let assertions = value
+        .get("assertions")
+        .ok_or_else(|| "fixture missing assertions".to_owned())?;
+
+    Ok(IsolationExpectations {
+        extension_event_capacity: capacity,
+        first_turn_partial_count: stream
+            .get("firstTurnPartialCount")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing firstTurnPartialCount".to_owned())?
+            as usize,
+        tool_call_id: stream
+            .get("toolCallId")
+            .and_then(Value::as_str)
+            .ok_or_else(|| "fixture missing toolCallId".to_owned())?
+            .to_owned(),
+        tool_name: stream
+            .get("toolName")
+            .and_then(Value::as_str)
+            .ok_or_else(|| "fixture missing toolName".to_owned())?
+            .to_owned(),
+        final_assistant_text: stream
+            .get("finalAssistantText")
+            .and_then(Value::as_str)
+            .ok_or_else(|| "fixture missing finalAssistantText".to_owned())?
+            .to_owned(),
+        tool_result_text: stream
+            .get("toolResultText")
+            .and_then(Value::as_str)
+            .ok_or_else(|| "fixture missing toolResultText".to_owned())?
+            .to_owned(),
+        extension_error_count: assertions
+            .get("extensionErrorCount")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing extensionErrorCount".to_owned())?
+            as usize,
+        healthy_extension_lag_count: assertions
+            .get("healthyExtensionLagCount")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing healthyExtensionLagCount".to_owned())?
+            as usize,
+        agent_settled_count: assertions
+            .get("agentSettledCount")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing agentSettledCount".to_owned())?
+            as usize,
+        min_stream_starts: assertions
+            .get("minStreamStarts")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing minStreamStarts".to_owned())?
+            as usize,
+        min_partial_watch_versions: assertions
+            .get("minPartialWatchVersions")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing minPartialWatchVersions".to_owned())?
+            as usize,
+        tool_execution_count: assertions
+            .get("toolExecutionCount")
+            .and_then(Value::as_u64)
+            .ok_or_else(|| "fixture missing toolExecutionCount".to_owned())?
+            as usize,
+    })
+}
+
+fn test_model() -> Model {
+    Model {
+        id: "m".to_owned(),
+        name: "m".to_owned(),
+        api: "test-api".to_owned(),
+        provider: "test-provider".to_owned(),
+        base_url: String::new(),
+        reasoning: false,
+        thinking_level_map: None,
+        input: vec![ModelInput::Text],
+        cost: ModelCost::default(),
+        context_window: 8_192,
+        max_tokens: 1_024,
+        headers: None,
+        compat: None,
+        extra: BTreeMap::new(),
+    }
+}
+
+fn assistant_with_text(text: &str) -> AssistantMessage {
+    let mut message = AssistantMessage::new("test-api", "test-provider", "m", 1);
+    message
+        .content
+        .push(AssistantContent::Text(TextContent::new(text)));
+    message
+}
+
+fn start_event() -> AssistantMessageEvent {
+    AssistantMessageEvent::Start {
+        partial: AssistantMessage::new("test-api", "test-provider", "m", 1),
+    }
+}
+
+fn text_delta_event(text: &str) -> AssistantMessageEvent {
+    AssistantMessageEvent::TextDelta {
+        content_index: 0,
+        delta: text.to_owned(),
+        partial: assistant_with_text(text),
+    }
+}
+
+fn done_text_event(text: &str) -> AssistantMessageEvent {
+    let mut message = assistant_with_text(text);
+    message.stop_reason = StopReason::Stop;
+    AssistantMessageEvent::Done {
+        reason: DoneReason::Stop,
+        message,
+    }
+}
+
+fn done_tool_event(call: ToolCall) -> AssistantMessageEvent {
+    let mut message = AssistantMessage::new("test-api", "test-provider", "m", 1);
+    message.content = vec![AssistantContent::ToolCall(call)];
+    message.stop_reason = StopReason::ToolUse;
+    AssistantMessageEvent::Done {
+        reason: DoneReason::ToolUse,
+        message,
+    }
+}
+
+fn long_isolation_scripts(
+    expectations: &IsolationExpectations,
+) -> Vec<Vec<Result<AssistantMessageEvent, ProviderError>>> {
+    // Enough deltas to force well past EXTENSION_EVENT_CAPACITY agent events
+    // once lifecycle + message_update fan-out is included.
+    let mut first = Vec::with_capacity(expectations.first_turn_partial_count.saturating_add(2));
+    first.push(Ok(start_event()));
+    for i in 0..expectations.first_turn_partial_count {
+        let text = format!("partial-{i}");
+        first.push(Ok(text_delta_event(&text)));
+    }
+    first.push(Ok(done_tool_event(ToolCall::new(
+        expectations.tool_call_id.clone(),
+        expectations.tool_name.clone(),
+        Map::new(),
+    ))));
+
+    let second = vec![
+        Ok(start_event()),
+        Ok(done_text_event(&expectations.final_assistant_text)),
+    ];
+    vec![first, second]
+}
+
+#[derive(Clone)]
+struct ScriptedProvider {
+    scripts: Arc<std::sync::Mutex<Vec<Vec<Result<AssistantMessageEvent, ProviderError>>>>>,
+    stream_starts: Arc<AtomicUsize>,
+    items_emitted: Arc<AtomicUsize>,
+}
+
+impl ScriptedProvider {
+    fn new(scripts: Vec<Vec<Result<AssistantMessageEvent, ProviderError>>>) -> Self {
+        Self {
+            scripts: Arc::new(std::sync::Mutex::new(scripts)),
+            stream_starts: Arc::new(AtomicUsize::new(0)),
+            items_emitted: Arc::new(AtomicUsize::new(0)),
+        }
+    }
+}
+
+impl Provider for ScriptedProvider {
+    fn stream(
+        &self,
+        _model: &Model,
+        _context: Context,
+        _options: StreamOptions,
+    ) -> BoxStream<'static, Result<AssistantMessageEvent, ProviderError>> {
+        self.stream_starts.fetch_add(1, Ordering::SeqCst);
+        let script = self
+            .scripts
+            .lock()
+            .ok()
+            .and_then(|mut guard| {
+                if guard.is_empty() {
+                    None
+                } else {
+                    Some(guard.remove(0))
+                }
+            })
+            .unwrap_or_default();
+        let items = Arc::new(script);
+        let items_emitted = Arc::clone(&self.items_emitted);
+        stream::unfold(0usize, move |index| {
+            let items = Arc::clone(&items);
+            let items_emitted = Arc::clone(&items_emitted);
+            async move {
+                if index >= items.len() {
+                    return None;
+                }
+                items_emitted.fetch_add(1, Ordering::SeqCst);
+                Some((items[index].clone(), index + 1))
+            }
+        })
+        .boxed()
+    }
+}
+
+struct EchoTool {
+    name: String,
+    result_text: String,
+    executed: Arc<AtomicUsize>,
+}
+
+impl EchoTool {
+    fn new(name: impl Into<String>, result_text: impl Into<String>) -> Self {
+        Self {
+            name: name.into(),
+            result_text: result_text.into(),
+            executed: Arc::new(AtomicUsize::new(0)),
+        }
+    }
+}
+
+impl AgentTool for EchoTool {
+    fn name(&self) -> &str {
+        &self.name
+    }
+
+    fn label(&self) -> &str {
+        &self.name
+    }
+
+    fn description(&self) -> &str {
+        "echo tool for isolation contract"
+    }
+
+    fn parameters(&self) -> &Value {
+        static SCHEMA: std::sync::LazyLock<Value> =
+            std::sync::LazyLock::new(|| json!({"type":"object","properties":{}}));
+        &SCHEMA
+    }
+
+    fn validate_arguments(
+        &self,
+        args: &Map<String, Value>,
+    ) -> Result<Map<String, Value>, ToolError> {
+        Ok(args.clone())
+    }
+
+    fn execute(
+        &self,
+        _tool_call_id: &str,
+        _args: Map<String, Value>,
+        _cancel: tokio_util::sync::CancellationToken,
+        _updates: ToolUpdates,
+    ) -> BoxFuture<'static, Result<AgentToolResult, ToolError>> {
+        let executed = Arc::clone(&self.executed);
+        let result_text = self.result_text.clone();
+        Box::pin(async move {
+            executed.fetch_add(1, Ordering::SeqCst);
+            Ok(AgentToolResult {
+                content: vec![ToolResultContent::Text(TextContent::new(result_text))],
+                details: json!({}),
+                added_tool_names: None,
+                terminate: None,
+            })
+        })
+    }
+}
+
+fn is_agent_end(event: &AgentEvent) -> bool {
+    matches!(event, AgentEvent::AgentEnd { .. })
+}
+
+fn message_role(entry: &SessionEntry) -> Option<&str> {
+    match entry {
+        SessionEntry::Message(message) => Some(message.message.role()),
+        _ => None,
+    }
+}
+
+fn assistant_entry_text(entry: &SessionEntry) -> Option<String> {
+    let SessionEntry::Message(message) = entry else {
+        return None;
+    };
+    let Some(Message::Assistant(assistant)) = message.message.as_llm() else {
+        return None;
+    };
+    let mut text = String::new();
+    for block in &assistant.content {
+        if let AssistantContent::Text(content) = block {
+            text.push_str(&content.text);
+        }
+    }
+    if text.is_empty() {
+        None
+    } else {
+        Some(text)
+    }
+}
+
+fn tool_result_entry_text(entry: &SessionEntry) -> Option<String> {
+    let SessionEntry::Message(message) = entry else {
+        return None;
+    };
+    let Some(Message::ToolResult(result)) = message.message.as_llm() else {
+        return None;
+    };
+    let mut text = String::new();
+    for block in &result.content {
+        if let ToolResultContent::Text(content) = block {
+            text.push_str(&content.text);
+        }
+    }
+    if text.is_empty() {
+        None
+    } else {
+        Some(text)
+    }
+}
+
+/// Drain a healthy extension until the agent run notifies completion.
+async fn drain_healthy_extension(
+    mut extension: pi_agent::ExtensionSubscription,
+    done: Arc<Notify>,
+    events_tx: mpsc::UnboundedSender<AgentEvent>,
+    lag_tx: oneshot::Sender<usize>,
+) {
+    let mut lagged = 0usize;
+    loop {
+        tokio::select! {
+            () = done.notified() => {
+                loop {
+                    match extension.try_recv() {
+                        Ok(ExtensionEvent::Event(event)) => {
+                            let _ = events_tx.send(*event);
+                        }
+                        Ok(ExtensionEvent::Lagged) => lagged = lagged.saturating_add(1),
+                        Err(mpsc::error::TryRecvError::Empty) => break,
+                        Err(mpsc::error::TryRecvError::Disconnected) => break,
+                    }
+                }
+                let _ = lag_tx.send(lagged);
+                return;
+            }
+            item = extension.recv() => {
+                match item {
+                    Some(ExtensionEvent::Event(event)) => {
+                        let _ = events_tx.send(*event);
+                    }
+                    Some(ExtensionEvent::Lagged) => {
+                        lagged = lagged.saturating_add(1);
+                    }
+                    None => {
+                        let _ = lag_tx.send(lagged);
+                        return;
+                    }
+                }
+            }
+        }
+    }
+}
+
+/// Collect hung extension buffered events + the single Lagged terminal.
+async fn collect_hung_lag(
+    hung: &mut pi_agent::ExtensionSubscription,
+    capacity: usize,
+) -> TestResult<(usize, usize)> {
+    let mut hung_events = 0usize;
+    let mut hung_lagged = 0usize;
+
+    // Non-blocking drain first.
+    for _ in 0..capacity.saturating_add(8) {
+        match hung.try_recv() {
+            Ok(ExtensionEvent::Event(_)) => hung_events = hung_events.saturating_add(1),
+            Ok(ExtensionEvent::Lagged) => {
+                hung_lagged = hung_lagged.saturating_add(1);
+                return Ok((hung_events, hung_lagged));
+            }
+            Err(mpsc::error::TryRecvError::Empty) => break,
+            Err(mpsc::error::TryRecvError::Disconnected) => break,
+        }
+    }
+
+    // Timeout-guarded final recv for the lag terminal only.
+    if hung_lagged == 0 {
+        match timeout(Duration::from_millis(200), hung.recv()).await {
+            Ok(Some(ExtensionEvent::Event(_))) => {
+                hung_events = hung_events.saturating_add(1);
+                match timeout(Duration::from_millis(200), hung.recv()).await {
+                    Ok(Some(ExtensionEvent::Lagged)) => hung_lagged = 1,
+                    Ok(Some(ExtensionEvent::Event(_))) => {
+                        return Err(
+                            "hung extension produced events after overflow without Lagged".into(),
+                        );
+                    }
+                    Ok(None) | Err(_) => {}
+                }
+            }
+            Ok(Some(ExtensionEvent::Lagged)) => hung_lagged = 1,
+            Ok(None) | Err(_) => {}
+        }
+    }
+
+    Ok((hung_events, hung_lagged))
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+async fn provider_drain_isolates_lagging_extension_consumer() -> TestResult {
+    let expectations = load_expectations()?;
+    if expectations.extension_event_capacity != EXTENSION_EVENT_CAPACITY {
+        return Err(format!(
+            "fixture capacity {} != EXTENSION_EVENT_CAPACITY {EXTENSION_EVENT_CAPACITY}",
+            expectations.extension_event_capacity
+        ));
+    }
+
+    let provider = Arc::new(ScriptedProvider::new(long_isolation_scripts(&expectations)));
+    let tool = Arc::new(EchoTool::new(
+        expectations.tool_name.clone(),
+        expectations.tool_result_text.clone(),
+    ));
+
+    let mut config = AgentSessionConfig::test_config(
+        Arc::clone(&provider) as Arc<dyn Provider>,
+        test_model(),
+    )
+    .map_err(|err| err.to_string())?;
+    config.system_prompt = "isolation".to_owned();
+    config.tools = vec![Arc::clone(&tool) as Arc<dyn AgentTool>];
+    let session = AgentSession::new(config).map_err(|err| err.to_string())?;
+
+    // Session-level settle signal from the product event pump.
+    let settled = Arc::new(AtomicUsize::new(0));
+    let settled_c = Arc::clone(&settled);
+    let _unsub = session.subscribe(move |event| {
+        if matches!(event, AgentSessionEvent::AgentSettled) {
+            settled_c.fetch_add(1, Ordering::SeqCst);
+        }
+    });
+
+    // Hung extension: never reads during the long stream.
+    let mut hung = session.agent.subscribe_extension();
+    // Healthy extension: continuously drained.
+    let healthy = session.agent.subscribe_extension();
+    // Native UI partial watch.
+    let mut partial_rx = session.agent.partial();
+    partial_rx.borrow_and_update();
+
+    let run_done = Arc::new(Notify::new());
+    let healthy_done = Arc::clone(&run_done);
+
+    let (healthy_events_tx, mut healthy_events_rx) = mpsc::unbounded_channel();
+    let (healthy_lag_tx, healthy_lag_rx) = oneshot::channel();
+    let healthy_task = tokio::spawn(drain_healthy_extension(
+        healthy,
+        healthy_done,
+        healthy_events_tx,
+        healthy_lag_tx,
+    ));
+
+    // Track native partial watch advancement without sleeps in the happy path.
+    let partial_versions = Arc::new(AtomicUsize::new(0));
+    let partial_versions_bg = Arc::clone(&partial_versions);
+    let partial_done = Arc::clone(&run_done);
+    let partial_task = tokio::spawn(async move {
+        loop {
+            tokio::select! {
+                () = partial_done.notified() => break,
+                changed = partial_rx.changed() => {
+                    if changed.is_err() {
+                        break;
+                    }
+                    if partial_rx.borrow_and_update().is_some() {
+                        partial_versions_bg.fetch_add(1, Ordering::SeqCst);
+                    }
+                }
+            }
+        }
+        if partial_rx.has_changed().unwrap_or(false) && partial_rx.borrow_and_update().is_some() {
+            partial_versions_bg.fetch_add(1, Ordering::SeqCst);
+        }
+    });
+
+    // Run the product session prompt under a hard timeout guard only.
+    let mut options = PromptOptions::new();
+    options.expand_prompt_templates = false;
+    timeout(Duration::from_secs(5), session.prompt("stream-me", options))
+        .await
+        .map_err(|_| "session.prompt timed out under hung extension".to_owned())?
+        .map_err(|err| err.to_string())?;
+    timeout(Duration::from_secs(2), session.wait_for_idle())
+        .await
+        .map_err(|_| "session.wait_for_idle timed out".to_owned())?;
+
+    // Wake background collectors.
+    run_done.notify_waiters();
+    run_done.notify_waiters();
+
+    timeout(Duration::from_secs(2), partial_task)
+        .await
+        .map_err(|_| "partial watcher timed out".to_owned())?
+        .map_err(|err| err.to_string())?;
+    timeout(Duration::from_secs(2), healthy_task)
+        .await
+        .map_err(|_| "healthy extension drain timed out".to_owned())?
+        .map_err(|err| err.to_string())?;
+    let healthy_lagged = timeout(Duration::from_secs(1), healthy_lag_rx)
+        .await
+        .map_err(|_| "healthy lag channel timed out".to_owned())?
+        .map_err(|err| err.to_string())?;
+
+    // Drain remaining healthy events channel.
+    let mut healthy_events = Vec::new();
+    while let Ok(event) = healthy_events_rx.try_recv() {
+        healthy_events.push(event);
+    }
+
+    // Hung extension must have overflowed (capacity 64) and be disconnected.
+    if !hung.is_lagged() {
+        return Err("hung extension did not disconnect after queue overflow".to_owned());
+    }
+
+    let (hung_events, hung_lagged) =
+        collect_hung_lag(&mut hung, expectations.extension_event_capacity).await?;
+
+    if hung_events > expectations.extension_event_capacity {
+        return Err(format!(
+            "hung extension retained more than capacity events: {hung_events}"
+        ));
+    }
+    if hung_lagged != expectations.extension_error_count {
+        return Err(format!(
+            "expected exactly {} extension_error/Lagged, got {hung_lagged} (events={hung_events})",
+            expectations.extension_error_count
+        ));
+    }
+    // Second lag poll stays empty/disconnected.
+    match hung.try_recv() {
+        Ok(ExtensionEvent::Lagged) => {
+            return Err("duplicate extension_error/Lagged after first signal".to_owned());
+        }
+        Ok(ExtensionEvent::Event(_)) => {
+            return Err("event after extension_error/Lagged".to_owned());
+        }
+        Err(_) => {}
+    }
+
+    // Only the hung consumer disconnects.
+    if healthy_lagged != expectations.healthy_extension_lag_count {
+        return Err(format!(
+            "healthy extension lag count mismatch: expected {}, got {healthy_lagged}",
+            expectations.healthy_extension_lag_count
+        ));
+    }
+    if healthy_events.is_empty() {
+        return Err("healthy extension received no events".to_owned());
+    }
+    if !healthy_events.iter().any(is_agent_end) {
+        return Err("healthy extension missed agent_end".to_owned());
+    }
+
+    // Provider drain advanced through both streams (tool turn + final turn).
+    let stream_starts = provider.stream_starts.load(Ordering::SeqCst);
+    if stream_starts < expectations.min_stream_starts {
+        return Err(format!(
+            "provider drain did not start enough streams: {stream_starts}"
+        ));
+    }
+    let items_emitted = provider.items_emitted.load(Ordering::SeqCst);
+    if items_emitted <= expectations.extension_event_capacity {
+        return Err(format!(
+            "provider stream too short for isolation pressure: {items_emitted}"
+        ));
+    }
+
+    // Native UI watch advanced during the long stream.
+    let partials = partial_versions.load(Ordering::SeqCst);
+    if partials < expectations.min_partial_watch_versions {
+        return Err(format!(
+            "native partial watch did not advance enough: {partials}"
+        ));
+    }
+
+    // Exactly one session-level settle.
+    let settled_count = settled.load(Ordering::SeqCst);
+    if settled_count != expectations.agent_settled_count {
+        return Err(format!(
+            "expected exactly {} agent_settled, got {settled_count}",
+            expectations.agent_settled_count
+        ));
+    }
+
+    // Session persistence records the final assistant and tool results.
+    let entries: Vec<SessionEntry> = {
+        let manager = session.session_manager();
+        let sm = manager.lock().await;
+        sm.get_entries().into_iter().cloned().collect()
+    };
+    let roles: Vec<&str> = entries.iter().filter_map(message_role).collect();
+    let assistant_count = roles.iter().filter(|role| **role == "assistant").count();
+    let tool_result_count = roles.iter().filter(|role| **role == "toolResult").count();
+    let user_count = roles.iter().filter(|role| **role == "user").count();
+
+    if user_count < 1 {
+        return Err(format!(
+            "session persistence missing user message; roles={roles:?}"
+        ));
+    }
+    // Tool-turn assistant + final-answer assistant.
+    if assistant_count < 2 {
+        return Err(format!(
+            "session persistence missing assistant finals; roles={roles:?}"
+        ));
+    }
+    if tool_result_count != 1 {
+        return Err(format!(
+            "session persistence missing tool result; roles={roles:?}"
+        ));
+    }
+
+    let final_texts: Vec<String> = entries.iter().filter_map(assistant_entry_text).collect();
+    if !final_texts
+        .iter()
+        .any(|text| text == &expectations.final_assistant_text)
+    {
+        return Err(format!(
+            "session persistence missing final assistant text {:?}; got {final_texts:?}",
+            expectations.final_assistant_text
+        ));
+    }
+
+    let tool_texts: Vec<String> = entries.iter().filter_map(tool_result_entry_text).collect();
+    if !tool_texts
+        .iter()
+        .any(|text| text == &expectations.tool_result_text)
+    {
+        return Err(format!(
+            "session persistence missing tool result text {:?}; got {tool_texts:?}",
+            expectations.tool_result_text
+        ));
+    }
+
+    // Agent transcript also holds the final assistant and tool result.
+    let transcript = session.agent.transcript();
+    let transcript_roles: Vec<&str> = transcript.iter().map(pi_agent::AgentMessage::role).collect();
+    if !transcript_roles.iter().any(|role| *role == "toolResult") {
+        return Err(format!(
+            "agent transcript missing tool result; roles={transcript_roles:?}"
+        ));
+    }
+    let last_assistant = session
+        .agent
+        .last_assistant()
+        .ok_or_else(|| "agent transcript missing final assistant".to_owned())?;
+    let last_text = last_assistant
+        .content
+        .iter()
+        .find_map(|block| match block {
+            AssistantContent::Text(content) => Some(content.text.as_str()),
+            _ => None,
+        })
+        .unwrap_or("");
+    if last_text != expectations.final_assistant_text {
+        return Err(format!(
+            "agent last assistant text mismatch: {last_text:?}"
+        ));
+    }
+
+    let executed = tool.executed.load(Ordering::SeqCst);
+    if executed != expectations.tool_execution_count {
+        return Err(format!(
+            "expected tool to run {} time(s), got {executed}",
+            expectations.tool_execution_count
+        ));
+    }
+
+    Ok(())
+}
+
+#[test]
+fn extension_event_capacity_matches_fixture_budget() -> TestResult {
+    let expectations = load_expectations()?;
+    assert_eq!(EXTENSION_EVENT_CAPACITY, 64);
+    assert_eq!(
+        expectations.extension_event_capacity,
+        EXTENSION_EVENT_CAPACITY
+    );
+    assert_eq!(expectations.extension_error_count, 1);
+    Ok(())
+}