use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use async_trait::async_trait;
use serde_json::json;
use tokio::time::Instant;
use crate::audio::AudioRecorder;
use crate::brain::AgentBrain;
use crate::codec::Resampler;
use crate::error::{FlowcatError, Result};
use crate::processor::frame::{AudioFrame, Frame, FrameClass, StartParams};
use crate::processor::{Envelope, FrameProcessor, Link, ProcessorSetup, StopReason};
use crate::realtime::{RealtimeKickoff, RealtimeLlm};
use crate::session::SessionSource;
use crate::transcript::Transcript;
use crate::transport::{MediaIn, MediaTransport};
use crate::types::{AudioChunk, BrainAction, Finalize, RealtimeEvent, ToolDecl, Usage};
use super::{Pipeline, PipelineTask, PipelineTaskParams, SourceHandle, SourcePump};
const GEMINI_OUTPUT_RATE: u32 = 24_000;
#[derive(Debug, Clone)]
pub(crate) struct ModelToolCall {
pub(crate) id: String,
pub(crate) name: String,
pub(crate) args: serde_json::Value,
}
impl crate::processor::frame::CustomFrame for ModelToolCall {
fn frame_class(&self) -> FrameClass {
FrameClass::System
}
fn name(&self) -> &'static str {
"ModelToolCall"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[derive(Debug, Clone)]
pub(crate) struct ToolResult {
pub(crate) id: String,
pub(crate) result: serde_json::Value,
}
impl crate::processor::frame::CustomFrame for ToolResult {
fn frame_class(&self) -> FrameClass {
FrameClass::System
}
fn uninterruptible(&self) -> bool {
true
}
fn name(&self) -> &'static str {
"ToolResult"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[derive(Debug, Clone)]
pub(crate) struct Reprompt {
pub(crate) prompt: String,
pub(crate) tools: Vec<ToolDecl>,
}
impl crate::processor::frame::CustomFrame for Reprompt {
fn frame_class(&self) -> FrameClass {
FrameClass::System
}
fn uninterruptible(&self) -> bool {
true
}
fn name(&self) -> &'static str {
"Reprompt"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
pub(crate) struct LiveState {
pub(crate) recorder: AudioRecorder,
pub(crate) transcript: Transcript,
pub(crate) disposition: Option<String>,
pub(crate) collected_vars: serde_json::Value,
usage: Usage,
pub(crate) tool_calls: u64,
error: Option<FlowcatError>,
pub(crate) transport_dead: bool,
}
impl LiveState {
pub(crate) fn new(carrier_rate: u32) -> Self {
Self {
recorder: AudioRecorder::new(carrier_rate),
transcript: Transcript::new(),
disposition: None,
collected_vars: serde_json::Value::Null,
usage: Usage::default(),
tool_calls: 0,
error: None,
transport_dead: false,
}
}
pub(crate) fn record_error(&mut self, e: FlowcatError) {
let is_transport = matches!(e, FlowcatError::Transport(_));
if is_transport && self.transport_dead {
return;
}
tracing::warn!(error = %e, "live-phase error (s2s pipeline)");
if is_transport {
self.transport_dead = true;
}
if self.error.is_none() {
self.error = Some(e);
}
}
fn accumulate_usage(&mut self, u: &Usage) {
self.usage.input_tokens = add_opt(self.usage.input_tokens, u.input_tokens);
self.usage.output_tokens = add_opt(self.usage.output_tokens, u.output_tokens);
self.usage.total_tokens = add_opt(self.usage.total_tokens, u.total_tokens);
if u.extra.is_some() {
self.usage.extra = u.extra.clone();
}
}
fn usage_json(&self) -> serde_json::Value {
let mut v = serde_json::to_value(&self.usage).unwrap_or(serde_json::Value::Null);
if let serde_json::Value::Object(ref mut map) = v {
map.entry("duration_seconds".to_string())
.or_insert_with(|| serde_json::json!(self.recorder.duration_seconds()));
map.entry("user_turns".to_string())
.or_insert_with(|| serde_json::json!(self.transcript.user_turns()));
map.entry("bot_turns".to_string())
.or_insert_with(|| serde_json::json!(self.transcript.bot_turns()));
map.entry("tool_calls".to_string())
.or_insert_with(|| serde_json::json!(self.tool_calls));
}
v
}
}
fn add_opt(a: Option<u64>, b: Option<u64>) -> Option<u64> {
match (a, b) {
(Some(a), Some(b)) => Some(a + b),
(Some(a), None) => Some(a),
(None, b) => b,
}
}
pub(crate) fn merge_disposition(
vars: serde_json::Value,
disposition: Option<String>,
) -> serde_json::Value {
match (vars, disposition) {
(serde_json::Value::Object(mut map), Some(d)) => {
map.entry("call_disposition".to_string())
.or_insert(serde_json::Value::String(d));
serde_json::Value::Object(map)
}
(other, _) => other,
}
}
pub(crate) type SharedState = Arc<Mutex<LiveState>>;
pub struct TransportInput {
name: &'static str,
}
impl TransportInput {
pub(crate) fn new() -> Self {
Self {
name: "TransportInput",
}
}
}
#[async_trait]
impl FrameProcessor for TransportInput {
fn name(&self) -> &str {
self.name
}
}
pub(crate) struct SharedTransport<T: MediaTransport>(Arc<tokio::sync::Mutex<T>>);
impl<T: MediaTransport> Clone for SharedTransport<T> {
fn clone(&self) -> Self {
SharedTransport(self.0.clone())
}
}
impl<T: MediaTransport> SharedTransport<T> {
pub(crate) fn new(t: T) -> Self {
SharedTransport(Arc::new(tokio::sync::Mutex::new(t)))
}
}
#[async_trait]
impl<T: MediaTransport> MediaTransport for SharedTransport<T> {
async fn recv(&mut self) -> Option<MediaIn> {
self.0.lock().await.recv().await
}
async fn send_audio(&mut self, chunk: AudioChunk) -> std::result::Result<(), FlowcatError> {
self.0.lock().await.send_audio(chunk).await
}
async fn send_clear(&mut self) -> std::result::Result<(), FlowcatError> {
self.0.lock().await.send_clear().await
}
fn carrier_rate(&self) -> u32 {
0
}
}
pub(crate) fn spawn_transport_pump<T: MediaTransport + 'static>(
mut transport: SharedTransport<T>,
head: tokio::sync::mpsc::UnboundedSender<Frame>,
) -> SourcePump {
SourcePump::spawn(head, move |h: SourceHandle| async move {
loop {
match transport.recv().await {
Some(MediaIn::StreamStart { call_id }) => {
tracing::debug!(call_id, "media stream started (s2s pipeline)");
let _ = h.emit(Frame::ClientConnected);
}
Some(MediaIn::Audio(chunk)) => {
let frame = Frame::InputAudio(Arc::new(AudioFrame::from(&chunk)));
if h.emit(frame).is_err() {
break; }
}
Some(MediaIn::Stop) | None => {
h.end();
break;
}
}
}
})
}
pub struct RealtimeServiceProcessor<R: RealtimeLlm + RealtimeKickoff> {
realtime: Arc<tokio::sync::Mutex<R>>,
in_resampler: Option<Resampler>,
carrier_rate: u32,
input_rate: u32,
transport_started: bool,
kicked_off: bool,
reader_spawned: bool,
state: SharedState,
end_tx: tokio::sync::mpsc::UnboundedSender<Frame>,
reader_handle: Option<tokio::task::AbortHandle>,
}
impl<R: RealtimeLlm + RealtimeKickoff + 'static> RealtimeServiceProcessor<R> {
fn new(
realtime: R,
carrier_rate: u32,
input_rate: u32,
state: SharedState,
end_tx: tokio::sync::mpsc::UnboundedSender<Frame>,
) -> Self {
Self {
realtime: Arc::new(tokio::sync::Mutex::new(realtime)),
in_resampler: None,
carrier_rate,
input_rate,
transport_started: false,
kicked_off: false,
reader_spawned: false,
state,
end_tx,
reader_handle: None,
}
}
fn spawn_reader(&mut self, link: Link) {
if self.reader_spawned {
return;
}
self.reader_spawned = true;
let realtime = self.realtime.clone();
let end_tx = self.end_tx.clone();
let handle = tokio::spawn(async move {
let notify = { realtime.lock().await.event_notify() };
let mut user_partial = String::new();
let mut bot_partial = String::new();
macro_rules! flush_bot {
() => {
if !bot_partial.is_empty() {
link.push_down(Frame::Transcription {
text: std::mem::take(&mut bot_partial),
user_id: Arc::from("bot"),
language: None,
final_: true,
})
.await;
link.push_down(Frame::BotStoppedSpeaking).await;
}
};
}
loop {
let event = if let Some(notify) = ¬ify {
loop {
let polled = {
let mut rt = realtime.lock().await;
rt.poll_event().await
};
match polled {
crate::realtime::PollEvent::Ready(ev) => break ev,
crate::realtime::PollEvent::Pending => notify.notified().await,
}
}
} else {
let mut rt = realtime.lock().await;
rt.next_event().await
};
let Some(event) = event else {
flush_bot!();
let _ = end_tx.send(Frame::End { reason: None });
break;
};
match event {
RealtimeEvent::AudioOut(chunk) => {
let af = Arc::new(AudioFrame::from(&chunk));
link.push_down(Frame::OutputAudio(af)).await;
}
RealtimeEvent::UserInterimText(delta) => {
user_partial.push_str(&delta);
link.push_down(Frame::InterimTranscription {
text: user_partial.clone(),
user_id: Arc::from("user"),
language: None,
})
.await;
}
RealtimeEvent::UserText(text) => {
user_partial.clear();
link.push_down(Frame::Transcription {
text,
user_id: Arc::from("user"),
language: None,
final_: true,
})
.await;
}
RealtimeEvent::BotText(text) => {
bot_partial.push_str(&text);
}
RealtimeEvent::ToolCall { id, name, args } => {
link.push_down(Frame::Custom(Arc::new(ModelToolCall { id, name, args })))
.await;
}
RealtimeEvent::Interrupted => {
flush_bot!();
link.broadcast(Frame::Interruption).await;
}
RealtimeEvent::Usage(u) => {
flush_bot!();
link.push_down(Frame::Custom(Arc::new(UsageReport(u))))
.await;
}
RealtimeEvent::Closed => {
flush_bot!();
let _ = end_tx.send(Frame::End { reason: None });
break;
}
}
}
});
self.reader_handle = Some(handle.abort_handle());
}
}
#[derive(Debug, Clone)]
pub(crate) struct UsageReport(pub(crate) Usage);
impl crate::processor::frame::CustomFrame for UsageReport {
fn frame_class(&self) -> FrameClass {
FrameClass::Data
}
fn name(&self) -> &'static str {
"UsageReport"
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
}
#[async_trait]
impl<R: RealtimeLlm + RealtimeKickoff + 'static> FrameProcessor for RealtimeServiceProcessor<R> {
fn name(&self) -> &str {
"RealtimeService"
}
async fn start(&mut self, _setup: &ProcessorSetup, _params: &StartParams) -> Result<()> {
self.in_resampler = Some(Resampler::new(self.carrier_rate, self.input_rate)?);
Ok(())
}
async fn stop(&mut self, _reason: StopReason) -> Result<()> {
if let Some(handle) = self.reader_handle.take() {
handle.abort();
}
Ok(())
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::ClientConnected => {
self.transport_started = true;
self.spawn_reader(link.clone());
if self.transport_started && !self.kicked_off {
let mut rt = self.realtime.lock().await;
if let Err(e) = rt.kickoff().await {
drop(rt);
self.state.lock().unwrap().record_error(e);
let _ = self.end_tx.send(Frame::End { reason: None });
return Ok(());
}
self.kicked_off = true;
}
link.push(env.meta, env.frame, env.direction).await;
}
Frame::InputAudio(audio) => {
self.spawn_reader(link.clone());
let chunk = AudioChunk::from(audio.as_ref());
let resampler = self
.in_resampler
.as_mut()
.expect("in_resampler set in start()");
match resampler.process(&chunk) {
Ok(up) => {
if !up.is_empty() {
let mut rt = self.realtime.lock().await;
if let Err(e) = rt.send_audio(up).await {
drop(rt);
self.state.lock().unwrap().record_error(e);
let _ = self.end_tx.send(Frame::End { reason: None });
return Ok(());
}
}
}
Err(e) => {
self.state.lock().unwrap().record_error(e);
let _ = self.end_tx.send(Frame::End { reason: None });
return Ok(());
}
}
link.push(env.meta, env.frame, env.direction).await;
}
Frame::Custom(c) if c.as_any().is::<ToolResult>() => {
let tr = c.as_any().downcast_ref::<ToolResult>().unwrap();
let mut rt = self.realtime.lock().await;
if let Err(e) = rt.send_tool_result(tr.id.clone(), tr.result.clone()).await {
drop(rt);
self.state.lock().unwrap().record_error(e);
}
}
Frame::Custom(c) if c.as_any().is::<Reprompt>() => {
let rp = c.as_any().downcast_ref::<Reprompt>().unwrap();
let mut rt = self.realtime.lock().await;
if let Err(e) = rt.update_system(rp.prompt.clone(), rp.tools.clone()).await {
drop(rt);
self.state.lock().unwrap().record_error(e);
}
}
_ => {
link.push(env.meta, env.frame, env.direction).await;
}
}
Ok(())
}
}
#[async_trait]
pub trait ToolRelay: Send + Sync {
async fn node_tools(&self, node_id: &str) -> Vec<ToolDecl>;
async fn relay(&self, node_id: &str, tool_name: &str, args: &serde_json::Value) -> String;
}
pub struct SessionToolRelay<S: SessionSource> {
session: Arc<S>,
run_id: i64,
token: String,
}
impl<S: SessionSource> SessionToolRelay<S> {
pub(crate) fn new(session: Arc<S>, run_id: i64, token: String) -> Self {
Self {
session,
run_id,
token,
}
}
}
#[async_trait]
impl<S: SessionSource> ToolRelay for SessionToolRelay<S> {
async fn node_tools(&self, node_id: &str) -> Vec<ToolDecl> {
self.session
.node_tools(self.run_id, &self.token, node_id)
.await
.unwrap_or_default()
}
async fn relay(&self, node_id: &str, tool_name: &str, args: &serde_json::Value) -> String {
match self
.session
.tool_call(self.run_id, &self.token, node_id, tool_name, args)
.await
{
Ok(c) => c,
Err(_) => "the tool is temporarily unavailable".to_string(),
}
}
}
fn mcp_result_struct(content: &str) -> serde_json::Value {
match serde_json::from_str::<serde_json::Value>(content) {
Ok(v @ serde_json::Value::Object(_)) => v,
Ok(other) => json!({ "result": other }),
Err(_) => json!({ "result": content }),
}
}
pub struct BrainProcessor<B: AgentBrain, Rl: ToolRelay> {
brain: B,
relay: Arc<Rl>,
mcp_names: std::collections::HashSet<String>,
state: SharedState,
end_tx: tokio::sync::mpsc::UnboundedSender<Frame>,
}
impl<B: AgentBrain, Rl: ToolRelay> BrainProcessor<B, Rl> {
pub(crate) fn new(
brain: B,
relay: Arc<Rl>,
mcp_names: std::collections::HashSet<String>,
state: SharedState,
end_tx: tokio::sync::mpsc::UnboundedSender<Frame>,
) -> Self {
Self {
brain,
relay,
mcp_names,
state,
end_tx,
}
}
async fn resolve_tools(&mut self) -> Vec<ToolDecl> {
let node_id = self.brain.current_node_id();
let mcp = self.relay.node_tools(&node_id).await;
self.mcp_names = mcp.iter().map(|t| t.name.clone()).collect();
let mut tools = self.brain.tools();
tools.extend(mcp);
tools
}
}
#[async_trait]
impl<B: AgentBrain + 'static, Rl: ToolRelay + 'static> FrameProcessor for BrainProcessor<B, Rl> {
fn name(&self) -> &str {
"Brain"
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
let Frame::Custom(c) = &env.frame else {
link.push(env.meta, env.frame, env.direction).await;
return Ok(());
};
let Some(tc) = c.as_any().downcast_ref::<ModelToolCall>() else {
link.push(env.meta, env.frame, env.direction).await;
return Ok(());
};
let id = tc.id.clone();
let name = tc.name.clone();
let args = tc.args.clone();
if self.mcp_names.contains(&name) {
self.state.lock().unwrap().tool_calls += 1;
let node_id = self.brain.current_node_id();
let content = self.relay.relay(&node_id, &name, &args).await;
link.push_up(Frame::Custom(Arc::new(ToolResult {
id,
result: mcp_result_struct(&content),
})))
.await;
return Ok(());
}
match self.brain.on_tool_call(&name, &args) {
BrainAction::Transition {
system_prompt, say, ..
} => {
let tools = self.resolve_tools().await;
link.push_up(Frame::Custom(Arc::new(Reprompt {
prompt: system_prompt,
tools,
})))
.await;
link.push_up(Frame::Custom(Arc::new(ToolResult {
id,
result: json!({ "status": "moved" }),
})))
.await;
let marker = match &say {
Some(s) if !s.is_empty() => format!("[transition: {name}] {s}"),
_ => format!("[transition: {name}]"),
};
self.state.lock().unwrap().transcript.push_bot(&marker);
}
BrainAction::Stay => {
link.push_up(Frame::Custom(Arc::new(ToolResult {
id,
result: json!({ "status": "stay" }),
})))
.await;
}
BrainAction::End { disposition } => {
link.push_up(Frame::Custom(Arc::new(ToolResult {
id,
result: json!({ "status": "ended" }),
})))
.await;
self.state.lock().unwrap().disposition = disposition;
let vars = self.brain.collected_vars();
self.state.lock().unwrap().collected_vars = vars;
let _ = self.end_tx.send(Frame::End { reason: None });
}
}
Ok(())
}
}
pub(super) const MAX_PLAYOUT_DRAIN: Duration = Duration::from_secs(15);
pub(super) fn advance_playout(
until: Option<Instant>,
now: Instant,
samples: usize,
rate: u32,
) -> Instant {
let dur = Duration::from_secs_f64(samples as f64 / rate.max(1) as f64);
let base = match until {
Some(t) if t > now => t,
_ => now,
};
base + dur
}
pub struct TransportOutput<T: MediaTransport> {
transport: T,
out_resampler: Option<Resampler>,
carrier_rate: u32,
bot_audio_until: Option<Instant>,
state: SharedState,
}
impl<T: MediaTransport> TransportOutput<T> {
fn new(transport: T, carrier_rate: u32, state: SharedState) -> Self {
Self {
transport,
out_resampler: None,
carrier_rate,
bot_audio_until: None,
state,
}
}
}
#[async_trait]
impl<T: MediaTransport + 'static> FrameProcessor for TransportOutput<T> {
fn name(&self) -> &str {
"TransportOutput"
}
async fn start(&mut self, _setup: &ProcessorSetup, _params: &StartParams) -> Result<()> {
self.out_resampler = Some(Resampler::new(GEMINI_OUTPUT_RATE, self.carrier_rate)?);
Ok(())
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::OutputAudio(audio) | Frame::TtsAudio { audio, .. } => {
let chunk = AudioChunk::from(audio.as_ref());
self.state.lock().unwrap().recorder.push_outbound(&chunk);
let resampler = self
.out_resampler
.as_mut()
.expect("out_resampler set in start()");
match resampler.process(&chunk) {
Ok(down) => {
if !down.is_empty() {
let samples = down.pcm.len();
if let Err(e) = self.transport.send_audio(down).await {
self.state.lock().unwrap().record_error(e);
} else {
self.bot_audio_until = Some(advance_playout(
self.bot_audio_until,
Instant::now(),
samples,
self.carrier_rate,
));
}
}
}
Err(e) => {
self.state.lock().unwrap().record_error(e);
}
}
}
Frame::Interruption => {
if let Err(e) = self.transport.send_clear().await {
self.state.lock().unwrap().record_error(e);
}
self.bot_audio_until = None;
link.push(env.meta, env.frame, env.direction).await;
}
Frame::End { .. } => {
if let Some(until) = self.bot_audio_until.take() {
let now = Instant::now();
if until > now {
tokio::time::sleep((until - now).min(MAX_PLAYOUT_DRAIN)).await;
}
}
link.push(env.meta, env.frame, env.direction).await;
}
_ => {
link.push(env.meta, env.frame, env.direction).await;
}
}
Ok(())
}
}
pub struct RecorderProcessor {
state: SharedState,
}
impl RecorderProcessor {
pub(crate) fn new(state: SharedState) -> Self {
Self { state }
}
}
#[async_trait]
impl FrameProcessor for RecorderProcessor {
fn name(&self) -> &str {
"Recorder"
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
match &env.frame {
Frame::InputAudio(audio) => {
let chunk = AudioChunk::from(audio.as_ref());
self.state.lock().unwrap().recorder.push_inbound(&chunk);
}
Frame::Custom(c) if c.as_any().is::<UsageReport>() => {
let u = c.as_any().downcast_ref::<UsageReport>().unwrap();
self.state.lock().unwrap().accumulate_usage(&u.0);
}
_ => {}
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
pub struct TranscriptProcessor {
state: SharedState,
}
impl TranscriptProcessor {
pub(crate) fn new(state: SharedState) -> Self {
Self { state }
}
}
#[async_trait]
impl FrameProcessor for TranscriptProcessor {
fn name(&self) -> &str {
"Transcript"
}
async fn process_frame(&mut self, env: Envelope, link: &Link) -> Result<()> {
if let Frame::Transcription { text, user_id, .. } = &env.frame {
let mut st = self.state.lock().unwrap();
if user_id.as_ref() == "bot" {
st.transcript.push_bot(text);
} else {
st.transcript.push_user(text);
}
}
link.push(env.meta, env.frame, env.direction).await;
Ok(())
}
}
pub struct FinalizeProcessor<S: SessionSource> {
session: Arc<S>,
run_id: i64,
token: String,
state: SharedState,
finalized: Arc<AtomicBool>,
}
impl<S: SessionSource> FinalizeProcessor<S> {
pub(crate) fn new(session: Arc<S>, run_id: i64, token: String, state: SharedState) -> Self {
Self {
session,
run_id,
token,
state,
finalized: Arc::new(AtomicBool::new(false)),
}
}
async fn finalize(&self) {
let (recording_bytes, transcript_bytes, usage_json, collected_vars) = {
let mut st = self.state.lock().unwrap();
let vars = merge_disposition(st.collected_vars.clone(), st.disposition.clone());
st.collected_vars = vars.clone();
let recording = st.recorder.render_wav();
let transcript = st.transcript.render();
(recording, transcript, st.usage_json(), vars)
};
let recording_url = match recording_bytes {
Ok(bytes) => self.upload_artifact("recording", bytes, "audio/wav").await,
Err(e) => {
tracing::warn!(error = %e, "render recording wav failed; skipping recording artifact");
None
}
};
let transcript_url = self
.upload_artifact("transcript", transcript_bytes, "application/json")
.await;
let fin = Finalize {
usage: usage_json,
collected_vars,
recording_url,
transcript_url,
};
if let Err(e) = self.session.complete(self.run_id, &self.token, fin).await {
tracing::warn!(error = %e, run_id = self.run_id, "session.complete failed");
}
}
async fn upload_artifact(
&self,
kind: &str,
bytes: Vec<u8>,
content_type: &str,
) -> Option<String> {
let target = match self
.session
.artifact_upload_url(self.run_id, &self.token, kind)
.await
{
Ok(t) => t,
Err(e) => {
tracing::warn!(error = %e, kind, "artifact_upload_url failed; skipping {kind}");
return None;
}
};
let ct = if target.content_type.is_empty() {
content_type.to_string()
} else {
target.content_type.clone()
};
match self.session.put_bytes(&target.url, bytes, &ct).await {
Ok(()) => Some(target.key),
Err(e) => {
tracing::warn!(error = %e, kind, "put_bytes failed; skipping {kind} url");
None
}
}
}
}
#[async_trait]
impl<S: SessionSource + 'static> FrameProcessor for FinalizeProcessor<S> {
fn name(&self) -> &str {
"Finalize"
}
async fn stop(&mut self, reason: StopReason) -> Result<()> {
if matches!(reason, StopReason::EndOfTask | StopReason::Cancelled)
&& !self.finalized.swap(true, Ordering::SeqCst)
{
self.finalize().await;
}
Ok(())
}
}
pub struct S2sTask {
pub task: PipelineTask,
pump: SourcePump,
}
impl S2sTask {
pub async fn run(self) -> Result<()> {
let res = self.task.run().await;
self.pump.abort();
res
}
}
#[allow(clippy::too_many_arguments)]
pub async fn build_s2s_task<T, R, B, S>(
transport: T,
realtime: R,
brain: B,
session: S,
run_id: i64,
token: String,
model: String,
) -> Result<S2sTask>
where
T: MediaTransport + 'static,
R: RealtimeLlm + RealtimeKickoff + 'static,
B: AgentBrain + 'static,
S: SessionSource + 'static,
{
build_s2s_task_with_observers(
transport,
realtime,
brain,
session,
run_id,
token,
model,
vec![],
)
.await
}
#[allow(clippy::too_many_arguments)]
pub async fn build_s2s_task_with_observers<T, R, B, S>(
transport: T,
mut realtime: R,
brain: B,
session: S,
run_id: i64,
token: String,
model: String,
observers: Vec<Arc<dyn crate::observer::FrameObserver>>,
) -> Result<S2sTask>
where
T: MediaTransport + 'static,
R: RealtimeLlm + RealtimeKickoff + 'static,
B: AgentBrain + 'static,
S: SessionSource + 'static,
{
let carrier_rate = transport.carrier_rate();
let session = Arc::new(session);
let relay = Arc::new(SessionToolRelay::new(
session.clone(),
run_id,
token.clone(),
));
let node_id = brain.current_node_id();
let mcp = relay.node_tools(&node_id).await;
let mcp_names: std::collections::HashSet<String> = mcp.iter().map(|t| t.name.clone()).collect();
let mut initial_tools = brain.tools();
initial_tools.extend(mcp);
let input_rate = realtime.input_sample_rate();
let setup = crate::types::RealtimeSetup {
model,
system_prompt: brain.system_prompt(),
tools: initial_tools,
input_sample_rate: input_rate,
output_sample_rate: GEMINI_OUTPUT_RATE,
};
if let Err(e) = realtime.connect(setup).await {
let state: SharedState = Arc::new(Mutex::new(LiveState::new(carrier_rate)));
state.lock().unwrap().record_error(e);
let fin = FinalizeProcessor::new(session, run_id, token, state);
fin.finalize().await;
return Err(FlowcatError::Realtime(
"realtime connect failed; finalized".into(),
));
}
let state: SharedState = Arc::new(Mutex::new(LiveState::new(carrier_rate)));
let shared = SharedTransport::new(transport);
let (end_tx, mut end_rx) = tokio::sync::mpsc::unbounded_channel::<Frame>();
let processors: Vec<Box<dyn FrameProcessor>> = vec![
Box::new(TransportInput::new()),
Box::new(RealtimeServiceProcessor::new(
realtime,
carrier_rate,
input_rate,
state.clone(),
end_tx.clone(),
)),
Box::new(BrainProcessor::new(
brain,
relay,
mcp_names,
state.clone(),
end_tx.clone(),
)),
Box::new(RecorderProcessor::new(state.clone())),
Box::new(TranscriptProcessor::new(state.clone())),
Box::new(TransportOutput::new(
shared.clone(),
carrier_rate,
state.clone(),
)),
Box::new(FinalizeProcessor::new(
session,
run_id,
token,
state.clone(),
)),
];
let pipeline = Pipeline::new(processors);
let params = PipelineTaskParams {
audio_in_sample_rate: input_rate,
audio_out_sample_rate: GEMINI_OUTPUT_RATE,
idle_timeout: None,
..Default::default()
};
let task = PipelineTask::new(pipeline, params, observers);
let head = task.queue_sender();
tokio::spawn(async move {
while let Some(f) = end_rx.recv().await {
if head.send(f).is_err() {
break;
}
}
});
let pump = spawn_transport_pump(shared, task.queue_sender());
Ok(S2sTask { task, pump })
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pipeline::s2s_test_mocks::{
Captured, McpMockRealtime, MockBrain, MockRealtime, MockSession, MockSocket, CARRIER_RATE,
};
use crate::serializer::PlivoSerializer;
use crate::transport::WsCarrierTransport;
use crate::types::WsOut;
use std::sync::atomic::AtomicBool;
use std::time::Duration;
const TEST_MODEL: &str = "models/test-realtime";
#[tokio::test]
async fn stop_aborts_the_reader_task() {
let reader = tokio::spawn(async {
loop {
tokio::task::yield_now().await;
}
});
let (end_tx, _end_rx) = tokio::sync::mpsc::unbounded_channel::<Frame>();
let state: SharedState = Arc::new(Mutex::new(LiveState::new(CARRIER_RATE)));
let mock = MockRealtime::new(
Arc::new(AtomicBool::new(false)),
Arc::new(Mutex::new(0)),
Arc::new(Mutex::new(Vec::new())),
);
let mut proc = RealtimeServiceProcessor::new(mock, CARRIER_RATE, 16_000, state, end_tx);
proc.reader_handle = Some(reader.abort_handle());
assert!(!reader.is_finished(), "reader runs before stop()");
proc.stop(StopReason::EndOfTask).await.unwrap();
let joined = reader.await;
assert!(
joined.is_err() && joined.unwrap_err().is_cancelled(),
"stop() must abort the reader task"
);
assert!(
proc.reader_handle.is_none(),
"the handle is taken on stop()"
);
}
#[derive(Debug, PartialEq)]
struct Captures {
play_audio: Vec<String>,
tool_results: Vec<(String, String)>,
recording_url: Option<String>,
transcript_url: Option<String>,
collected_vars: serde_json::Value,
total_tokens: Option<u64>,
relayed: Vec<(String, String, serde_json::Value)>,
raw_results: Vec<(String, serde_json::Value)>,
}
fn play_audio_frames(sent: &[WsOut]) -> Vec<String> {
sent.iter()
.filter_map(|o| match o {
WsOut::Text(t) if t.contains("playAudio") => Some(t.clone()),
_ => None,
})
.collect()
}
fn collect(
sent: &Arc<Mutex<Vec<WsOut>>>,
tool_results: &Arc<Mutex<Vec<(String, String)>>>,
raw_results: &Arc<Mutex<Vec<(String, serde_json::Value)>>>,
captured: &Arc<Mutex<Captured>>,
) -> Captures {
let c = captured.lock().unwrap();
let fin = c.finalize.as_ref();
Captures {
play_audio: play_audio_frames(&sent.lock().unwrap()),
tool_results: tool_results.lock().unwrap().clone(),
recording_url: fin.and_then(|f| f.recording_url.clone()),
transcript_url: fin.and_then(|f| f.transcript_url.clone()),
collected_vars: fin
.map(|f| f.collected_vars.clone())
.unwrap_or(serde_json::Value::Null),
total_tokens: fin.and_then(|f| f.usage.get("total_tokens").and_then(|v| v.as_u64())),
relayed: c.tool_calls.clone(),
raw_results: raw_results.lock().unwrap().clone(),
}
}
async fn run_pipeline_scenario1() -> Captures {
let sent = Arc::new(Mutex::new(Vec::<WsOut>::new()));
let connected = Arc::new(AtomicBool::new(false));
let audio_received = Arc::new(Mutex::new(0usize));
let tool_results = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let seen_tools = Arc::new(Mutex::new(Vec::<String>::new()));
let captured = Arc::new(Mutex::new(Captured::default()));
let raw_results = Arc::new(Mutex::new(Vec::<(String, serde_json::Value)>::new()));
let transport = WsCarrierTransport::new(
MockSocket::new(sent.clone()),
PlivoSerializer::new(CARRIER_RATE),
);
let s2s = build_s2s_task(
transport,
MockRealtime::new(
connected.clone(),
audio_received.clone(),
tool_results.clone(),
),
MockBrain::new(seen_tools.clone()),
MockSession::new(captured.clone()),
4242,
"tok-abc".into(),
TEST_MODEL.into(),
)
.await
.expect("build_s2s_task");
tokio::time::timeout(Duration::from_secs(5), s2s.run())
.await
.expect("S2S pipeline timed out")
.expect("S2S pipeline errored");
collect(&sent, &tool_results, &raw_results, &captured)
}
#[tokio::test]
async fn s2s_pipeline_matches_golden_audio_and_finalize() {
let pipeline = run_pipeline_scenario1().await;
assert_eq!(
pipeline.play_audio.len(),
1,
"expected exactly one playAudio frame for the scripted bot audio"
);
let frame: serde_json::Value =
serde_json::from_str(&pipeline.play_audio[0]).expect("playAudio frame is valid JSON");
assert_eq!(frame["event"], "playAudio");
assert!(
frame["media"]["payload"]
.as_str()
.is_some_and(|p| !p.is_empty()),
"playAudio frame must carry a non-empty base64 payload"
);
assert_eq!(
pipeline.tool_results,
vec![("fc-end-1".into(), "ended".into())]
);
assert_eq!(
pipeline.recording_url.as_deref(),
Some("runs/4242/recording")
);
assert_eq!(
pipeline.transcript_url.as_deref(),
Some("runs/4242/transcript")
);
assert_eq!(pipeline.collected_vars["call_disposition"], "completed");
assert_eq!(pipeline.collected_vars["name"], "Ada");
assert_eq!(pipeline.collected_vars["intent"], "support");
assert_eq!(pipeline.total_tokens, Some(15));
}
#[derive(Default)]
struct UserTranscriptRecorder {
seen: Mutex<Vec<(bool, String)>>,
}
#[async_trait::async_trait]
impl crate::observer::FrameObserver for UserTranscriptRecorder {
async fn on_push(&self, e: &crate::observer::FramePushEvent<'_>) {
if e.source != "RealtimeService" {
return;
}
match e.frame {
Frame::InterimTranscription { text, user_id, .. } if user_id.as_ref() == "user" => {
self.seen.lock().unwrap().push((false, text.clone()));
}
Frame::Transcription {
text,
user_id,
final_,
..
} if user_id.as_ref() == "user" => {
self.seen.lock().unwrap().push((*final_, text.clone()));
}
_ => {}
}
}
}
#[tokio::test]
async fn user_transcription_streams_one_interim_line_then_finalizes() {
let sent = Arc::new(Mutex::new(Vec::<WsOut>::new()));
let recorder = Arc::new(UserTranscriptRecorder::default());
let transport = WsCarrierTransport::new(
MockSocket::new(sent.clone()),
PlivoSerializer::new(CARRIER_RATE),
);
let s2s = build_s2s_task_with_observers(
transport,
MockRealtime::new(
Arc::new(AtomicBool::new(false)),
Arc::new(Mutex::new(0)),
Arc::new(Mutex::new(Vec::new())),
),
MockBrain::new(Arc::new(Mutex::new(Vec::new()))),
MockSession::new(Arc::new(Mutex::new(Captured::default()))),
4242,
"tok-abc".into(),
TEST_MODEL.into(),
vec![recorder.clone()],
)
.await
.expect("build_s2s_task_with_observers");
tokio::time::timeout(Duration::from_secs(5), s2s.run())
.await
.expect("S2S pipeline timed out")
.expect("S2S pipeline errored");
let seen = recorder.seen.lock().unwrap().clone();
assert_eq!(
seen,
vec![
(false, "I'd ".to_string()),
(false, "I'd like to end now".to_string()),
(true, "I'd like to end now".to_string()),
],
"expected two growing interim lines then one final, got {seen:?}"
);
}
#[derive(Default)]
struct BotTranscriptRecorder {
lines: Mutex<Vec<String>>,
stopped: std::sync::atomic::AtomicUsize,
}
#[async_trait::async_trait]
impl crate::observer::FrameObserver for BotTranscriptRecorder {
async fn on_push(&self, e: &crate::observer::FramePushEvent<'_>) {
if e.source != "RealtimeService" {
return;
}
match e.frame {
Frame::Transcription { text, user_id, .. } if user_id.as_ref() == "bot" => {
self.lines.lock().unwrap().push(text.clone());
}
Frame::BotStoppedSpeaking => {
self.stopped
.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
}
_ => {}
}
}
}
#[tokio::test]
async fn bot_transcript_coalesces_into_one_final_and_signals_stopped() {
let sent = Arc::new(Mutex::new(Vec::<WsOut>::new()));
let recorder = Arc::new(BotTranscriptRecorder::default());
let transport = WsCarrierTransport::new(
MockSocket::new(sent.clone()),
PlivoSerializer::new(CARRIER_RATE),
);
let s2s = build_s2s_task_with_observers(
transport,
MockRealtime::new(
Arc::new(AtomicBool::new(false)),
Arc::new(Mutex::new(0)),
Arc::new(Mutex::new(Vec::new())),
),
MockBrain::new(Arc::new(Mutex::new(Vec::new()))),
MockSession::new(Arc::new(Mutex::new(Captured::default()))),
4242,
"tok-abc".into(),
TEST_MODEL.into(),
vec![recorder.clone()],
)
.await
.expect("build_s2s_task_with_observers");
tokio::time::timeout(Duration::from_secs(5), s2s.run())
.await
.expect("S2S pipeline timed out")
.expect("S2S pipeline errored");
let lines = recorder.lines.lock().unwrap().clone();
assert_eq!(
lines,
vec!["Hello, how can I help?".to_string()],
"bot transcript should be one coalesced final line, got {lines:?}"
);
assert!(
recorder.stopped.load(std::sync::atomic::Ordering::SeqCst) >= 1,
"a BotStoppedSpeaking must be emitted so the UI finalizes the bubble"
);
}
const MCP_TOOL: &str = "list_doctors";
fn mcp_node_tools() -> Vec<ToolDecl> {
vec![ToolDecl {
name: MCP_TOOL.into(),
description: "List available doctors.".into(),
params: json!({
"type": "object",
"properties": { "specialty": { "type": "string" } }
}),
}]
}
async fn run_pipeline_scenario2() -> (Captures, Vec<Vec<String>>) {
let sent = Arc::new(Mutex::new(Vec::<WsOut>::new()));
let seen_tools = Arc::new(Mutex::new(Vec::<String>::new()));
let raw_results = Arc::new(Mutex::new(Vec::<(String, serde_json::Value)>::new()));
let advertised = Arc::new(Mutex::new(Vec::<Vec<String>>::new()));
let captured = Arc::new(Mutex::new(Captured::default()));
let tool_results = Arc::new(Mutex::new(Vec::<(String, String)>::new()));
let session =
MockSession::with_node_tools(captured.clone(), mcp_node_tools(), "Dr. Smith, Dr. Lee");
let transport = WsCarrierTransport::new(
MockSocket::new(sent.clone()),
PlivoSerializer::new(CARRIER_RATE),
);
let s2s = build_s2s_task(
transport,
McpMockRealtime::new(raw_results.clone(), advertised.clone(), MCP_TOOL),
MockBrain::new(seen_tools.clone()),
session,
4242,
"tok-abc".into(),
TEST_MODEL.into(),
)
.await
.expect("build_s2s_task");
tokio::time::timeout(Duration::from_secs(5), s2s.run())
.await
.expect("S2S pipeline timed out")
.expect("S2S pipeline errored");
let caps = collect(&sent, &tool_results, &raw_results, &captured);
let adv = advertised.lock().unwrap().clone();
(caps, adv)
}
#[tokio::test]
async fn s2s_pipeline_matches_golden_mcp_relay_not_transition() {
let (pipeline, pipe_adv) = run_pipeline_scenario2().await;
assert_eq!(pipeline.relayed.len(), 1, "exactly one MCP tool relayed");
let (node_id, tool_name, args) = &pipeline.relayed[0];
assert_eq!(node_id, "start");
assert_eq!(tool_name, MCP_TOOL);
assert_eq!(args["specialty"], "dentistry");
let mcp = pipeline
.raw_results
.iter()
.find(|(id, _)| id == "fc-mcp-1")
.expect("MCP tool result sent back to the model");
assert_eq!(mcp.1, json!({ "result": "Dr. Smith, Dr. Lee" }));
let end = pipeline
.raw_results
.iter()
.find(|(id, _)| id == "fc-end-1")
.expect("endCall result sent");
assert_eq!(end.1["status"], "ended");
assert_eq!(
pipe_adv.first(),
Some(&vec!["end_call".to_string(), MCP_TOOL.to_string()]),
"transitions first, then MCP tools"
);
}
#[test]
fn merge_disposition_inserts_into_object() {
let merged = merge_disposition(json!({ "a": 1 }), Some("busy".into()));
assert_eq!(merged["call_disposition"], "busy");
assert_eq!(merged["a"], 1);
}
#[test]
fn merge_disposition_does_not_clobber_existing() {
let merged = merge_disposition(
json!({ "call_disposition": "preset" }),
Some("override".into()),
);
assert_eq!(merged["call_disposition"], "preset", "existing value wins");
}
#[test]
fn merge_disposition_none_and_non_object_passthrough() {
assert_eq!(
merge_disposition(json!({ "a": 1 }), None),
json!({ "a": 1 })
);
assert_eq!(
merge_disposition(json!("scalar"), Some("x".into())),
json!("scalar")
);
}
#[test]
fn add_opt_sums_and_treats_none_as_zero() {
assert_eq!(add_opt(Some(2), Some(3)), Some(5));
assert_eq!(add_opt(Some(2), None), Some(2));
assert_eq!(add_opt(None, Some(3)), Some(3));
assert_eq!(add_opt(None, None), None);
}
#[test]
fn live_state_accumulates_usage_across_reports() {
let mut st = LiveState::new(CARRIER_RATE);
st.accumulate_usage(&Usage {
input_tokens: Some(10),
output_tokens: Some(2),
total_tokens: Some(12),
extra: None,
});
st.accumulate_usage(&Usage {
input_tokens: Some(5),
output_tokens: Some(1),
total_tokens: Some(6),
extra: Some(json!({ "cached": 3 })),
});
st.transcript.push_user("hi");
st.transcript.push_bot("hello");
st.tool_calls = 2;
let u = st.usage_json();
assert_eq!(u["input_tokens"], 15);
assert_eq!(u["output_tokens"], 3);
assert_eq!(u["total_tokens"], 18);
assert_eq!(u["extra"]["cached"], 3, "last non-null extra kept");
assert!(
u["duration_seconds"].is_number(),
"usage_json must surface duration_seconds"
);
assert_eq!(u["user_turns"], 1, "one user turn surfaced");
assert_eq!(u["bot_turns"], 1, "one bot turn surfaced");
assert_eq!(u["tool_calls"], 2, "tool calls tallied");
}
#[test]
fn mcp_result_object_passes_through_as_struct() {
let v = mcp_result_struct(r#"{"a2ui":{"kind":"card"},"prompt":"hi"}"#);
assert!(v.is_object(), "object result must stay an object");
assert_eq!(v["a2ui"]["kind"], "card");
assert_eq!(v["prompt"], "hi");
}
#[test]
fn mcp_result_plain_text_is_wrapped() {
let v = mcp_result_struct("Dr. Smith, Dr. Lee");
assert_eq!(v, json!({ "result": "Dr. Smith, Dr. Lee" }));
assert!(v.is_object());
}
#[test]
fn mcp_result_json_non_object_is_wrapped() {
assert_eq!(mcp_result_struct("[1,2,3]"), json!({ "result": [1, 2, 3] }));
assert_eq!(mcp_result_struct("42"), json!({ "result": 42 }));
assert_eq!(mcp_result_struct("true"), json!({ "result": true }));
assert!(mcp_result_struct("[1,2,3]").is_object());
}
#[test]
fn advance_playout_tracks_realtime_tail() {
let now = Instant::now();
let t = advance_playout(None, now, 8000, 8000);
assert_eq!(t, now + Duration::from_secs(1));
let t2 = advance_playout(Some(t), now, 8000, 8000);
assert_eq!(t2, now + Duration::from_secs(2));
let gap = advance_playout(Some(now - Duration::from_secs(5)), now, 8000, 8000);
assert_eq!(gap, now + Duration::from_secs(1));
let frame = advance_playout(None, now, 160, 8000);
assert_eq!(frame, now + Duration::from_millis(20));
}
}