fix: harden interceptor lifecycle contracts

crates/agen/tests/annotated_history_test.rs

@@ -1,8 +1,15 @@
 mod common;
 
+use agen::interceptor::{
+    AssistantTurnEndContext, Interceptor, InterceptorCallId, InterceptorInvocation,
+    InterceptorPhase, InterceptorResult, PendingHistoryAppendsContext, PreLlmRequestContext,
+    PreRequestAction, PromptAction, PromptSubmitContext, RunExitContext, TurnEndAction,
+};
 use agen::llm_client::event::{Event, ResponseStatus, StatusEvent};
 use agen::{Engine, EngineError, History, HistoryEntry, Item, Role};
+use async_trait::async_trait;
 use common::MockLlmClient;
+use std::sync::{Arc, Mutex};
 
 fn completed_text_events(text: &str) -> Vec<Event> {
     vec![
@@ -47,6 +54,125 @@ async fn run_preserves_item_annotations_without_projecting_them() {
     assert_eq!(history.items_cloned().len(), 2);
 }
 
+#[derive(Clone)]
+struct AnnotationObservingInterceptor {
+    observed: Arc<Mutex<Vec<(InterceptorInvocation, Vec<String>)>>>,
+}
+
+impl AnnotationObservingInterceptor {
+    fn record(&self, invocation: &InterceptorInvocation, history: &[HistoryEntry<String>]) {
+        self.observed.lock().unwrap().push((
+            invocation.clone(),
+            history
+                .iter()
+                .map(|entry| entry.annotation.clone())
+                .collect(),
+        ));
+    }
+}
+
+#[async_trait]
+impl Interceptor<String> for AnnotationObservingInterceptor {
+    async fn on_prompt_submit(
+        &self,
+        context: PromptSubmitContext<'_, String>,
+    ) -> InterceptorResult<PromptAction> {
+        self.record(&context.invocation, context.history);
+        Ok(PromptAction::Continue)
+    }
+
+    async fn pending_history_appends(
+        &self,
+        context: PendingHistoryAppendsContext<'_, String>,
+    ) -> InterceptorResult<Vec<Item>> {
+        self.record(&context.invocation, context.history);
+        Ok(Vec::new())
+    }
+
+    async fn pre_llm_request(
+        &self,
+        context: PreLlmRequestContext<'_, String>,
+    ) -> InterceptorResult<PreRequestAction> {
+        self.record(&context.invocation, context.history);
+        Ok(PreRequestAction::Continue)
+    }
+
+    async fn on_assistant_turn_end(
+        &self,
+        context: AssistantTurnEndContext<'_, String>,
+    ) -> InterceptorResult<TurnEndAction> {
+        assert_eq!(context.assistant_entries.len(), 1);
+        assert_eq!(context.assistant_entries[0].annotation, "2:assistant");
+        self.record(&context.invocation, context.history);
+        Ok(TurnEndAction::Finish)
+    }
+
+    async fn on_run_exit(&self, context: RunExitContext<'_, String>) -> InterceptorResult<()> {
+        self.record(&context.invocation, context.history);
+        Ok(())
+    }
+}
+
+#[tokio::test]
+async fn interceptor_contexts_preserve_annotations_and_typed_lifecycle_identity() {
+    let client = MockLlmClient::new(completed_text_events("assistant reply"));
+    let mut engine = Engine::<_, agen::state::Mutable, String>::new_annotated(client);
+    let observed = Arc::new(Mutex::new(Vec::new()));
+    engine.set_interceptor(AnnotationObservingInterceptor {
+        observed: observed.clone(),
+    });
+    let mut history = History::<String>::new();
+    let mut next = 0usize;
+    let mut annotate = |item: &Item| {
+        next += 1;
+        let kind = if item.is_assistant_message() {
+            "assistant"
+        } else {
+            "user"
+        };
+        Ok(format!("{next}:{kind}"))
+    };
+
+    let output = engine
+        .run_with_annotation(&mut history, "hello", &mut annotate)
+        .await;
+    assert!(matches!(output.result, agen::EngineRunExit::Finished));
+
+    let observed = observed.lock().unwrap();
+    let phases: Vec<_> = observed
+        .iter()
+        .map(|(invocation, _)| invocation.phase)
+        .collect();
+    assert_eq!(
+        phases,
+        [
+            InterceptorPhase::PromptSubmit,
+            InterceptorPhase::PendingHistoryAppends,
+            InterceptorPhase::PreLlmRequest,
+            InterceptorPhase::AssistantTurnEnd,
+            InterceptorPhase::RunExit,
+        ]
+    );
+    assert!(
+        observed
+            .iter()
+            .all(|(invocation, _)| invocation.run_id == observed[0].0.run_id)
+    );
+    assert_eq!(
+        observed
+            .iter()
+            .map(|(invocation, _)| invocation.counters.invocation.get())
+            .collect::<Vec<_>>(),
+        [0, 1, 2, 3, 4]
+    );
+    assert_eq!(observed[2].0.call_id, Some(InterceptorCallId::Llm(0)));
+    assert_eq!(observed[3].0.call_id, Some(InterceptorCallId::Llm(0)));
+    assert_eq!(observed[1].1, ["1:user"]);
+    assert_eq!(observed[2].1, ["1:user"]);
+    assert_eq!(observed[3].1, ["1:user", "2:assistant"]);
+    assert_eq!(observed[4].1, ["1:user", "2:assistant"]);
+}
+
 #[test]
 fn append_failure_does_not_make_item_live() {
     let client = MockLlmClient::new(vec![]);