capsule-runtime 0.8.0

//! The host side of `murmur:stream/events`: what a WASM component's mid-turn signals write to the
//! task's stream.
//!
//! [`invoke_tool_component`] registers one host function per WIT function, each a thin wrapper
//! around the matching function here, so the frames they write are testable without a component.
//! Every function is synchronous, returns nothing and never fails: a signal that has nowhere to go
//! is dropped, and the guest is never trapped for one.
//!
//! `emit-chunk` and `emit-thinking-chunk` stream text from any component. The two tool-call
//! functions write only when the dispatch carries a [`ToolCallProgress`], which only the agent
//! loop's driver turn does: from a tool, from a hook's `run-inference` and outside an A2A task they
//! do nothing.
//!
//! [`invoke_tool_component`]: crate::runtime::invoke_tool_component

use std::{
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc, Mutex,
    },
    time::Instant,
};

use crate::{
    streaming::{emit_frame, emit_text, emit_thinking, MurmurFrame, StreamSink, StreamTask},
    tool_call_progress::{SettledCall, ToolCallProgress},
};

/// Where one dispatch's `murmur:stream/events` calls go.
#[derive(Clone)]
pub(crate) struct StreamEventsTarget {
    pub(crate) sse: StreamSink,
    pub(crate) task: Option<StreamTask>,
    /// Set by every `emit-chunk`, so the agent loop knows the turn's text was streamed.
    pub(crate) chunks_emitted: Arc<AtomicBool>,
    /// The calls the session's driver has started this attempt. `None` on every dispatch that is
    /// not the agent loop's driver turn.
    pub(crate) tool_calls: Option<Arc<Mutex<ToolCallProgress>>>,
}

impl StreamEventsTarget {
    /// The task a frame goes to, when this dispatch serves an A2A task whose stream has somewhere
    /// to go.
    fn task(&self) -> Option<&StreamTask> {
        self.sse.as_ref()?;
        self.task.as_ref()
    }
}

/// `emit-chunk`: one streamed text fragment.
pub(crate) fn emit_chunk(target: &StreamEventsTarget, chunk: &str) {
    target.chunks_emitted.store(true, Ordering::Relaxed);
    if let Some(task) = target.task() {
        emit_text(&target.sse, task, chunk, false);
    }
}

/// `emit-thinking-chunk`: one streamed thinking fragment.
pub(crate) fn emit_thinking_chunk(target: &StreamEventsTarget, chunk: &str) {
    if let Some(task) = target.task() {
        emit_thinking(&target.sse, task, chunk);
    }
}

/// `tool-call-started`: the driver's model has begun writing call `id`, named `name`. Writes a
/// `tool-call-started` frame when the tracker accepts the start.
pub(crate) fn tool_call_started(target: &StreamEventsTarget, id: &str, name: &str) {
    let (Some(task), Some(tool_calls)) = (target.task(), target.tool_calls.as_ref()) else {
        return;
    };
    if !lock(tool_calls).start(id, name) {
        return;
    }
    emit_frame(
        &target.sse,
        task,
        MurmurFrame::ToolCallStarted {
            tool_call_id: id.to_string(),
            tool_name: name.to_string(),
        },
    );
}

/// `tool-call-input-bytes`: the driver's model has written `bytes` of call `id`'s input as of
/// `now`. Writes a `tool-call-progress` frame when one is due.
pub(crate) fn tool_call_input_bytes(
    target: &StreamEventsTarget,
    id: &str,
    bytes: u64,
    now: Instant,
) {
    let (Some(task), Some(tool_calls)) = (target.task(), target.tool_calls.as_ref()) else {
        return;
    };
    let Some(size) = lock(tool_calls).report(id, bytes, now) else {
        return;
    };
    write_progress(&target.sse, task, id, size);
}

/// The driver's dispatch returned at `now`: settle every call it started and has not settled,
/// writing the size still held for each. Returns each settled call's id and the name its start
/// carried, in the order they started. Outside an A2A task no call was started, so it writes and
/// returns nothing.
pub(crate) fn settle_started_calls(
    sse: &StreamSink,
    task: Option<&StreamTask>,
    tool_calls: &Mutex<ToolCallProgress>,
    now: Instant,
) -> Vec<(String, String)> {
    let settled = lock(tool_calls).complete_all(now);
    settled
        .into_iter()
        .map(|(id, SettledCall { name, held })| {
            if let (Some(task), Some(size)) = (task, held) {
                write_progress(sse, task, &id, size);
            }
            (id, name)
        })
        .collect()
}

fn write_progress(sse: &StreamSink, task: &StreamTask, tool_call_id: &str, input_bytes: u64) {
    emit_frame(
        sse,
        task,
        MurmurFrame::ToolCallProgress {
            tool_call_id: tool_call_id.to_string(),
            input_bytes,
        },
    );
}

/// The tracker, recovered from a poisoned lock: a panic elsewhere must not trap the guest here.
fn lock(tool_calls: &Mutex<ToolCallProgress>) -> std::sync::MutexGuard<'_, ToolCallProgress> {
    tool_calls
        .lock()
        .unwrap_or_else(std::sync::PoisonError::into_inner)
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use serde_json::Value;

    use super::*;
    use crate::streaming::{buffered, event_kind, frame_of, SseEventBuffer};

    /// A target writing into a real buffer, carrying a tracker when `tracked`, and that buffer.
    fn target(tracked: bool) -> (StreamEventsTarget, Arc<Mutex<SseEventBuffer>>) {
        let (tx, _rx) = tokio::sync::broadcast::channel(16);
        let frames = Arc::new(Mutex::new(SseEventBuffer::new(64)));
        let target = StreamEventsTarget {
            sse: Some((tx, Arc::clone(&frames))),
            task: Some(StreamTask::new("tsk_test", "ctx_test")),
            chunks_emitted: Arc::new(AtomicBool::new(false)),
            tool_calls: tracked.then(|| Arc::new(Mutex::new(ToolCallProgress::default()))),
        };
        (target, frames)
    }

    /// The kind and murmur frame of every event the buffer holds, in order.
    fn frames_in(frames: &Arc<Mutex<SseEventBuffer>>) -> Vec<(String, Value)> {
        buffered(frames)
            .iter()
            .map(|result| {
                crate::a2a::assert_bare_status(
                    &result["statusUpdate"]["status"],
                    "TASK_STATE_WORKING",
                );
                assert_eq!(result["statusUpdate"]["taskId"], "tsk_test");
                (event_kind(result), frame_of(result))
            })
            .collect()
    }

    fn kinds(frames: &Arc<Mutex<SseEventBuffer>>) -> Vec<String> {
        frames_in(frames)
            .into_iter()
            .map(|(kind, _)| kind)
            .collect()
    }

    #[test]
    fn stream_events_without_a_tracker_a_tool_call_writes_no_frame() {
        let (target, frames) = target(false);
        tool_call_started(&target, "c1", "echo-tool");
        tool_call_input_bytes(&target, "c1", 12, Instant::now());
        assert!(kinds(&frames).is_empty());
        assert!(!target.chunks_emitted.load(Ordering::Relaxed));
    }

    #[test]
    fn stream_events_with_a_tracker_a_start_writes_one_tool_call_started_frame() {
        let (target, frames) = target(true);
        tool_call_started(&target, "c1", "echo-tool");
        tool_call_started(&target, "c1", "echo-tool");
        tool_call_started(&target, "", "echo-tool");
        assert_eq!(
            frames_in(&frames),
            [(
                "tool-call-started".to_string(),
                serde_json::json!({
                    "frame": "tool-call-started",
                    "toolCallId": "c1",
                    "toolName": "echo-tool",
                })
            )]
        );
        // A tool-call signal is not streamed text.
        assert!(!target.chunks_emitted.load(Ordering::Relaxed));
    }

    #[test]
    fn stream_events_progress_is_coalesced_and_the_rest_is_flushed_on_return() {
        let (target, frames) = target(true);
        let now = Instant::now();
        tool_call_input_bytes(&target, "ghost", 5, now);
        tool_call_started(&target, "c1", "echo-tool");
        tool_call_input_bytes(&target, "c1", 15, now);
        tool_call_input_bytes(&target, "c1", 30, now + Duration::from_millis(1));
        tool_call_input_bytes(&target, "c1", 10, now + Duration::from_millis(2));
        let settled = settle_started_calls(
            &target.sse,
            target.task.as_ref(),
            target.tool_calls.as_ref().unwrap(),
            now + Duration::from_millis(3),
        );
        assert_eq!(settled, [("c1".to_string(), "echo-tool".to_string())]);
        let progress: Vec<Value> = frames_in(&frames)
            .into_iter()
            .filter(|(kind, _)| kind == "tool-call-progress")
            .map(|(_, data)| data["inputBytes"].clone())
            .collect();
        assert_eq!(progress, [15, 30]);
        // Settled: a late report writes nothing.
        tool_call_input_bytes(&target, "c1", 90, now + Duration::from_secs(1));
        assert_eq!(
            kinds(&frames),
            [
                "tool-call-started",
                "tool-call-progress",
                "tool-call-progress"
            ]
        );
    }

    #[test]
    fn stream_events_without_sse_nothing_is_written_or_tracked() {
        let tool_calls = Arc::new(Mutex::new(ToolCallProgress::default()));
        let target = StreamEventsTarget {
            sse: None,
            task: None,
            chunks_emitted: Arc::new(AtomicBool::new(false)),
            tool_calls: Some(Arc::clone(&tool_calls)),
        };
        emit_chunk(&target, "hello");
        emit_thinking_chunk(&target, "hmm");
        tool_call_started(&target, "c1", "echo-tool");
        tool_call_input_bytes(&target, "c1", 12, Instant::now());
        assert!(settle_started_calls(&None, None, &tool_calls, Instant::now()).is_empty());
        // The start never reached the tracker, so the same id may still start in a task.
        assert!(tool_calls.lock().unwrap().start("c1", "echo-tool"));
    }

    #[test]
    fn stream_events_emit_chunk_without_a_tracker_still_writes_text() {
        let (target, frames) = target(false);
        emit_chunk(&target, "hello");
        emit_thinking_chunk(&target, "hmm");
        assert_eq!(
            frames_in(&frames),
            [
                (
                    "text".to_string(),
                    serde_json::json!({"frame": "text", "text": "hello", "final": false})
                ),
                (
                    "thinking".to_string(),
                    serde_json::json!({"frame": "thinking", "text": "hmm", "final": false})
                ),
            ]
        );
        assert!(target.chunks_emitted.load(Ordering::Relaxed));
    }
}