chimera_core/

session.rs

use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

use futures_core::Stream;

use crate::{AgentError, AgentEvent, Input, TurnOptions, TurnResult, Usage};

pub trait Session: Send {
    /// Run a turn and return the buffered result.
    ///
    /// Default implementation calls `turn_stream` then `collect_turn`.
    fn turn(
        &mut self,
        input: Input,
        options: TurnOptions,
    ) -> Pin<Box<dyn Future<Output = Result<TurnOutput, AgentError>> + Send + '_>> {
        Box::pin(async move { self.turn_stream(input, options).await?.collect_turn().await })
    }

    /// Run a turn and return a stream of events.
    fn turn_stream(
        &mut self,
        input: Input,
        options: TurnOptions,
    ) -> Pin<Box<dyn Future<Output = Result<EventStream, AgentError>> + Send + '_>>;

    /// Session ID (populated after the first turn completes).
    fn session_id(&self) -> Option<&str>;

    /// Interrupt the current turn. No-op if idle.
    fn interrupt(&mut self) -> Pin<Box<dyn Future<Output = Result<(), AgentError>> + Send + '_>>;
}

#[must_use = "streams do nothing unless polled"]
pub struct EventStream {
    inner: Pin<Box<dyn Stream<Item = Result<AgentEvent, AgentError>> + Send>>,
}

impl EventStream {
    pub fn new(
        stream: impl Stream<Item = Result<AgentEvent, AgentError>> + Send + 'static,
    ) -> Self {
        Self {
            inner: Box::pin(stream),
        }
    }

    pub fn from_receiver(rx: tokio::sync::mpsc::Receiver<Result<AgentEvent, AgentError>>) -> Self {
        Self::new(ReceiverStream { inner: rx })
    }

    /// Consume the stream and collect into a `TurnOutput`.
    pub async fn collect_turn(mut self) -> Result<TurnOutput, AgentError> {
        let mut events = Vec::new();
        let mut response = None;
        let mut response_deltas = String::new();
        let mut saw_delta = false;
        let mut usage = None;
        let mut result = None;

        loop {
            match std::future::poll_fn(|cx| self.inner.as_mut().poll_next(cx)).await {
                Some(Ok(event)) => {
                    match &event {
                        AgentEvent::Message { text, .. } => {
                            response = Some(text.clone());
                        }
                        AgentEvent::TextDelta { delta, .. } => {
                            response_deltas.push_str(delta);
                            saw_delta = true;
                        }
                        AgentEvent::TurnCompleted {
                            usage: u,
                            result: r,
                        } => {
                            usage.clone_from(u);
                            result.clone_from(r);
                        }
                        AgentEvent::TurnFailed { message } => {
                            return Err(AgentError::TurnFailed {
                                message: message.clone(),
                            });
                        }
                        _ => {}
                    }
                    events.push(event);
                }
                Some(Err(e)) => return Err(e),
                None => break,
            }
        }

        if response.is_none() && saw_delta {
            response = Some(response_deltas);
        }

        Ok(TurnOutput {
            events,
            response,
            usage,
            result,
        })
    }
}

impl Stream for EventStream {
    type Item = Result<AgentEvent, AgentError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.as_mut().poll_next(cx)
    }
}

/// Adapts `mpsc::Receiver` to `Stream` without pulling in `tokio-stream`.
struct ReceiverStream {
    inner: tokio::sync::mpsc::Receiver<Result<AgentEvent, AgentError>>,
}

impl Stream for ReceiverStream {
    type Item = Result<AgentEvent, AgentError>;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.poll_recv(cx)
    }
}

#[derive(Debug, Clone)]
pub struct TurnOutput {
    pub events: Vec<AgentEvent>,
    pub response: Option<String>,
    pub usage: Option<Usage>,
    pub result: Option<TurnResult>,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn collect_turn_extracts_response() {
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tx.send(Ok(AgentEvent::TurnStarted)).await.unwrap();
        tx.send(Ok(AgentEvent::Message {
            id: "msg-1".into(),
            text: "Hello world".into(),
        }))
        .await
        .unwrap();
        tx.send(Ok(AgentEvent::TurnCompleted {
            usage: Some(Usage {
                input_tokens: Some(10),
                output_tokens: Some(5),
                ..Default::default()
            }),
            result: None,
        }))
        .await
        .unwrap();
        drop(tx);

        let stream = EventStream::from_receiver(rx);
        let output = stream.collect_turn().await.unwrap();

        assert_eq!(output.response.as_deref(), Some("Hello world"));
        assert_eq!(output.events.len(), 3);
        assert_eq!(output.usage.as_ref().unwrap().input_tokens, Some(10));
    }

    #[tokio::test]
    async fn collect_turn_fails_on_turn_failed() {
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tx.send(Ok(AgentEvent::TurnStarted)).await.unwrap();
        tx.send(Ok(AgentEvent::TurnFailed {
            message: "something broke".into(),
        }))
        .await
        .unwrap();
        drop(tx);

        let stream = EventStream::from_receiver(rx);
        let err = stream.collect_turn().await.unwrap_err();

        match err {
            AgentError::TurnFailed { message } => assert_eq!(message, "something broke"),
            _ => panic!("expected TurnFailed, got {err:?}"),
        }
    }

    #[tokio::test]
    async fn collect_turn_uses_last_message() {
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tx.send(Ok(AgentEvent::Message {
            id: "1".into(),
            text: "first".into(),
        }))
        .await
        .unwrap();
        tx.send(Ok(AgentEvent::Message {
            id: "2".into(),
            text: "second".into(),
        }))
        .await
        .unwrap();
        tx.send(Ok(AgentEvent::TurnCompleted {
            usage: None,
            result: None,
        }))
        .await
        .unwrap();
        drop(tx);

        let stream = EventStream::from_receiver(rx);
        let output = stream.collect_turn().await.unwrap();

        assert_eq!(output.response.as_deref(), Some("second"));
    }

    #[tokio::test]
    async fn collect_turn_uses_deltas_when_no_message() {
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tx.send(Ok(AgentEvent::TextDelta {
            id: "block_0".into(),
            delta: "Hello".into(),
        }))
        .await
        .unwrap();
        tx.send(Ok(AgentEvent::TextDelta {
            id: "block_0".into(),
            delta: " world".into(),
        }))
        .await
        .unwrap();
        tx.send(Ok(AgentEvent::TurnCompleted {
            usage: None,
            result: None,
        }))
        .await
        .unwrap();
        drop(tx);

        let stream = EventStream::from_receiver(rx);
        let output = stream.collect_turn().await.unwrap();

        assert_eq!(output.response.as_deref(), Some("Hello world"));
    }

    #[tokio::test]
    async fn collect_turn_propagates_stream_error() {
        let (tx, rx) = tokio::sync::mpsc::channel(16);

        tx.send(Ok(AgentEvent::TurnStarted)).await.unwrap();
        tx.send(Err(AgentError::ProcessFailed {
            code: 1,
            stderr: "crash".into(),
        }))
        .await
        .unwrap();
        drop(tx);

        let stream = EventStream::from_receiver(rx);
        let err = stream.collect_turn().await.unwrap_err();

        match err {
            AgentError::ProcessFailed { code, .. } => assert_eq!(code, 1),
            _ => panic!("expected ProcessFailed, got {err:?}"),
        }
    }

    #[tokio::test]
    async fn empty_stream_produces_empty_output() {
        let (_tx, rx) = tokio::sync::mpsc::channel::<Result<AgentEvent, AgentError>>(1);
        drop(_tx);

        let stream = EventStream::from_receiver(rx);
        let output = stream.collect_turn().await.unwrap();

        assert!(output.events.is_empty());
        assert!(output.response.is_none());
        assert!(output.usage.is_none());
    }
}