rutster/crates/rutster-tap/src/tap_client.rs

//! # TapClient — the async WSS pump loop (spec §4.2)
//!
//! Lives inside the `TapEngine` task (in the binary crate). Owns one WS
//! connection + the pump loop that shovels PCM between mpsc channels and
//! WS frames. The media loop never sees it — `TapAudioPipe` is the only
//! sync surface `loop_driver` touches.
//!
//! # Why TapClient never decides to reconnect
//!
//! Reconnect is the `TapEngine`'s job (spec §4.3). On any WS close / error,
//! `run_tap_client` returns; the engine rebuilds the client + applies
//! backoff. This keeps "the connection" (the wire) and "the reconnect
//! policy" (the backoff) as separate concerns — the one knows the wire,
//! the other knows the timing.
//!
//! # The `close` oneshot is shared, not owned
//!
//! `run_tap_client` takes `close: &mut oneshot::Receiver<()>` (a shared
//! borrow of the receiver), NOT `close: oneshot::Receiver<()>` by value.
//! Rationale: the `TapEngine` reconnect loop (Task 7) holds ONE close
//! oneshot per session and passes `&mut close` to each reconnect attempt
//! of `run_tap_client`. If this fn consumed it by value, each reconnect
//! would need a fresh oneshot — breaking the model where the binary holds
//! one close signal per session that survives across reconnects.

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

use futures_util::{SinkExt, StreamExt};
use rutster_call_model::ChannelId;
use rutster_media::PcmFrame;
use thiserror::Error;
use tokio::sync::{mpsc, oneshot};
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::tungstenite::Message;
use tracing::{debug, info, trace, warn};

use crate::metrics::TapMetrics;
use crate::protocol::{
    DecodedPayload, TapProtoError, decode_envelope, encode_audio_in, encode_hello,
    encode_session_end,
};

#[derive(Debug, Error)]
pub enum TapClientError {
    #[error("WebSocket error: {0}")]
    Ws(#[from] tokio_tungstenite::tungstenite::Error),
    #[error("Protocol error: {0}")]
    Proto(#[from] TapProtoError),
    #[error("hello handshake timed out")]
    HelloTimeout,
    #[error("close signal received")]
    Closed,
}

/// Run the pump loop on an already-connected `WebSocketStream`.
///
/// Returns `Ok(())` on graceful close (brain sent `bye` or the close
/// oneshot fired). Returns `Err(_)` on WS / protocol errors — the
/// `TapEngine` caller decides to retry.
///
/// `close` is a shared borrow (`&mut oneshot::Receiver<()>`) so the
/// `TapEngine` reconnect loop (Task 7) can share one close signal across
/// reconnect attempts of the same session (see module docs).
///
/// # slice-3 §5.2 — the tool-call side-channel
///
/// Two extra mpsc halves (relative to slice-2's audio-only pump) carry the
/// brain's `function_call` proposals out to the binary's poll task and the
/// binary's `function_call_output` replies back onto the wire:
///
/// - `tx_function_call` — the TapClient emits a `FunctionCallEvent` here
///   whenever it observes a tap `function_call` frame on its inbound WS
///   stream. The binary's poll task drains this in the same cycle it drains
///   the existing `flush_tx` side-channel (slice-2 §5.3 step 4 — one extra
///   channel, same cycle) and dispatches via `ToolRegistry::dispatch`.
/// - `rx_function_call_output` — the binary's poll task writes
///   `FunctionCallOutputEvent`s here (after a `ToolRegistry::dispatch` call
///   completes); the TapClient drains this in the same `tokio::select!`
///   pump as the audio + close arms and sends each as a `function_call_output`
///   tap WS frame to the brain.
///
/// Both halves are mpsc ends (not oneshot) because the brain may propose
/// multiple tool calls per session (one-of-many, not one-of-one) and the
/// binary may queue multiple replies before the TapClient's pump cycle
/// drains them.
//
// clippy::too_many_arguments: the slice-3 §5.2 design added two more mpsc
// halves to slice-2's already-5-arg pump signature for a total of 8. Each
// arg is a distinct channel end with a distinct lifetime owner (the WS
// stream, the session id, two sender/receiver pairs for audio, two for the
// tool-call side-channel, the metrics Arc, the close oneshot). Wrapping
// them in a struct would obscure that all-but-one are channel ends shared
// with the binary's poll task — the flat signature mirrors slice-2's
// precedent and keeps the call site readable. Suppress per AGENTS.md's
// "documented inline rationale" exception to the -D warnings bar.
#[allow(clippy::too_many_arguments)]
pub async fn run_tap_client<T>(
    mut ws: WebSocketStream<T>,
    session_id: ChannelId,
    rx_pcm_in: &mut mpsc::Receiver<PcmFrame>,
    tx_audio_out: mpsc::Sender<PcmFrame>,
    tx_function_call: mpsc::Sender<FunctionCallEvent>,
    rx_function_call_output: &mut mpsc::Receiver<FunctionCallOutputEvent>,
    metrics: Arc<TapMetrics>,
    close: &mut oneshot::Receiver<()>,
) -> Result<(), TapClientError>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    let session_start = Instant::now();
    let mut seq_egress: u64 = 0;

    // === Handshake: send hello, await brain hello (bounded 2s). ===
    let hello_str = encode_hello(&session_id.to_string(), seq_egress, 0)?;
    seq_egress += 1;
    ws.send(Message::Text(hello_str)).await?;
    info!(%session_id, "sent hello to brain");

    let hello_brain = tokio::time::timeout(
        std::time::Duration::from_secs(2),
        wait_for_brain_hello(&mut ws),
    )
    .await;
    // Every arm either returns or assigns last_seq_ingress — using a
    // match expression (rather than `let mut x = None; match { x = ... }`)
    // avoids the dead-store warning that an unconditional `None` initializer
    // would produce under `-D unused-assignments`.
    let mut last_seq_ingress: Option<u64> = match hello_brain {
        Ok(Ok(brain_seq)) => {
            info!(%session_id, "brain hello acked");
            Some(brain_seq)
        }
        Ok(Err(e)) => {
            warn!(error = ?e, %session_id, "brain hello failed");
            return Err(e);
        }
        Err(_) => return Err(TapClientError::HelloTimeout),
    };

    // === Pump loop. ===
    loop {
        tokio::select! {
            // Close signal from the binary (Channel::Closing). `close` is a
            // shared `&mut oneshot::Receiver<()>`; `&mut *close` reborrows
            // the inner receiver so the engine's reconnect loop can hand
            // us the same oneshot across reconnect attempts of one session.
            _ = &mut *close => {
                // slice-2 spec §5.1 step 5 + §5.2: on core-side teardown the
                // TapClient sends `session_end` (NOT `bye`) and awaits the
                // brain's `bye` ack, bounded by 500 ms. A brain that doesn't
                // `bye` back in time just gets a WS close — acceptable.
                info!(%session_id, "close signal; sending session_end + awaiting brain bye");
                let end_str = encode_session_end(
                    "core_shutdown",
                    seq_egress,
                    elapsed_ms(session_start),
                )?;
                // No `seq_egress += 1` here — the close path returns
                // immediately after the WS send, so the bump would be a
                // dead write (the brain never sees a higher seq from us).
                if let Err(e) = ws.send(Message::Text(end_str)).await {
                    warn!(error = ?e, %session_id, "ws send session_end failed; closing");
                    let _ = ws.close(None).await;
                    return Err(TapClientError::Closed);
                }
                // Bounded wait for brain `bye` (spec §5.2: 500 ms).
                // `wait_for_brain_bye` reads frames until it sees `bye` (or
                // the WS stream ends / errors). On timeout we just close the
                // WS — the brain didn't ack in time, the call is hanging up
                // anyway.
                match tokio::time::timeout(
                    Duration::from_millis(500),
                    wait_for_brain_bye(&mut ws),
                )
                .await
                {
                    Ok(Ok(())) => {
                        info!(%session_id, "brain bye acked; closing ws");
                    }
                    Ok(Err(e)) => {
                        warn!(error = ?e, %session_id, "brain bye wait errored; closing ws");
                    }
                    Err(_) => {
                        warn!(%session_id, "brain bye timed out (500ms); closing ws");
                    }
                }
                let _ = ws.close(None).await;
                return Err(TapClientError::Closed);
            }
            // Inbound PCM from peer → audio_in WS frame.
            frame = rx_pcm_in.recv() => {
                let Some(frame) = frame else {
                    // Peer side gone; just keep the WS open until close.
                    continue;
                };
                let ts = elapsed_ms(session_start);
                match encode_audio_in(&frame, seq_egress, ts) {
                    Ok(s) => {
                        seq_egress += 1;
                        if let Err(e) = ws.send(Message::Text(s)).await {
                            warn!(error = ?e, %session_id, "ws send audio_in failed");
                            return Err(e.into());
                        }
                    }
                    Err(e) => {
                        metrics.malformed_frames.fetch_add(1, Ordering::Relaxed);
                        warn!(error = ?e, "encode audio_in failed; dropping");
                    }
                }
            }
            // slice-3 §5.2: drain a `function_call_output` event the binary's
            // poll task wrote (after `ToolRegistry::dispatch` returned) + send
            // it as a `function_call_output` tap WS frame to the brain. The
            // `seq_egress` bump mirrors the audio arm — every egress frame
            // shares the same per-direction counter (spec §3.1).
            //
            // Like `rx_pcm_in.recv()`, this is one of many arms in the
            // select! — the engine's `tx_function_call_output` sender lives
            // in the binary; `run_tap_client` returns when the brain WS ends
            // or close fires, regardless of pending function_call_output
            // events (they're dropped on close — same posture as pending
            // audio_out frames on teardown).
            out = rx_function_call_output.recv() => {
                if let Some(out) = out {
                    let ts = elapsed_ms(session_start);
                    let result_str = out.0.result.to_string();
                    match crate::protocol::encode_function_call_output(
                        &out.0.id,
                        &out.0.status,
                        &result_str,
                        seq_egress,
                        ts,
                    ) {
                        Ok(s) => {
                            seq_egress += 1;
                            if let Err(e) = ws.send(Message::Text(s)).await {
                                warn!(error = ?e, %session_id, "ws send function_call_output failed");
                                return Err(e.into());
                            }
                            info!(%session_id, call_id = %out.0.id, status = %out.0.status, "sent function_call_output to brain");
                        }
                        Err(e) => {
                            metrics.malformed_frames.fetch_add(1, Ordering::Relaxed);
                            warn!(error = ?e, "encode function_call_output failed; dropping");
                        }
                    }
                }
            }
            // Inbound WS frame from brain.
            msg = ws.next() => {
                let Some(msg) = msg else {
                    info!(%session_id, "brain WS stream ended");
                    return Ok(());
                };
                let msg = msg?;
                // `Message::into_text` returns `Result<String>`; non-text
                // frames (Binary/Ping/Pong/Close) yield `Err(_)` and are
                // silently dropped (v1 is text-JSON only — spec §3.4).
                if let Ok(text) = msg.into_text() {
                    handle_brain_frame(
                        &text, &mut last_seq_ingress, &tx_audio_out,
                        &tx_function_call, &metrics, session_start,
                    ).await;
                }
            }
        }
    }
}

async fn wait_for_brain_hello<T>(ws: &mut WebSocketStream<T>) -> Result<u64, TapClientError>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    while let Some(msg) = ws.next().await {
        let msg = msg?;
        if let Ok(text) = msg.into_text() {
            let decoded = decode_envelope(&text)?;
            if let DecodedPayload::Hello(_) = decoded.payload {
                return Ok(decoded.seq);
            }
            // Non-hello frames before ack — log + continue (drop + observe).
            warn!("pre-hello frame from brain; ignoring");
        }
    }
    Err(TapClientError::Ws(
        tokio_tungstenite::tungstenite::Error::ConnectionClosed,
    ))
}

/// Bounded-wait helper: read frames from the brain until we see `bye`
/// (the brain's ack to our `session_end`), or the WS stream ends / errors.
/// slice-2 spec §5.2: the bound is enforced by the caller (`tokio::time::
/// timeout`); this fn just consumes frames until it sees `bye`.
///
/// Non-bye frames during the bye-wait (e.g. an in-flight `audio_out` the
/// brain had already queued) are silently dropped — the session is
/// hanging up, no more playout is desired.
async fn wait_for_brain_bye<T>(ws: &mut WebSocketStream<T>) -> Result<(), TapClientError>
where
    T: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
{
    while let Some(msg) = ws.next().await {
        let msg = msg?;
        if let Ok(text) = msg.into_text() {
            if let Ok(decoded) = decode_envelope(&text) {
                if matches!(decoded.payload, DecodedPayload::Bye(_)) {
                    return Ok(());
                }
                // Non-bye frame during teardown — ignore + continue.
                continue;
            }
        }
    }
    // WS stream ended without an explicit `bye` — surface as the same Ws
    // error `wait_for_brain_hello` returns; the caller logs + closes.
    Err(TapClientError::Ws(
        tokio_tungstenite::tungstenite::Error::ConnectionClosed,
    ))
}

async fn handle_brain_frame(
    text: &str,
    last_seq_ingress: &mut Option<u64>,
    tx_audio_out: &mpsc::Sender<PcmFrame>,
    tx_function_call: &mpsc::Sender<FunctionCallEvent>,
    metrics: &Arc<TapMetrics>,
    session_start: Instant,
) {
    let decoded = match decode_envelope(text) {
        Ok(d) => d,
        Err(e) => {
            metrics.malformed_frames.fetch_add(1, Ordering::Relaxed);
            warn!(error = ?e, "malformed brain frame; dropping");
            return;
        }
    };
    // seq gap detection (spec §3.1: gaps = loss; log + count; don't drop).
    if let Some(prev) = *last_seq_ingress {
        if decoded.seq > prev + 1 {
            let gap = decoded.seq - prev - 1;
            metrics.seq_gaps.fetch_add(gap, Ordering::Relaxed);
            debug!(gap, "seq gap detected");
        }
    }
    *last_seq_ingress = Some(decoded.seq);

    match decoded.payload {
        DecodedPayload::AudioOut(audio) => {
            match crate::protocol::decode_pcm(&audio.pcm, audio.samples) {
                Ok(frame) => {
                    if tx_audio_out.try_send(frame).is_err() {
                        metrics.outbound_dropped.fetch_add(1, Ordering::Relaxed);
                        trace!("outbound PCM dropped (playout ring full)");
                    }
                }
                Err(e) => {
                    metrics.malformed_frames.fetch_add(1, Ordering::Relaxed);
                    warn!(error = ?e, "failed to decode brain audio_out PCM");
                }
            }
        }
        DecodedPayload::Bye(p) => {
            info!(reason = %p.reason, "brain sent bye; closing");
            // Caller's pump loop sees Ok(()) and lets TapEngine decide to retry.
        }
        DecodedPayload::Error(p) => {
            warn!(code = %p.code, message = %p.message, "brain error frame");
        }
        DecodedPayload::Hello(_) => {
            // Re-hello mid-session (after a reconnect from the brain's POV).
            // Fine; we already tracked last_seq_ingress above.
        }
        DecodedPayload::Unknown => {
            metrics.unknown_frames.fetch_add(1, Ordering::Relaxed);
            warn!("unknown frame type from brain; dropping");
        }
        // SessionEnd is core→brain only; ignore if we receive one.
        DecodedPayload::SessionEnd(_) | DecodedPayload::AudioIn(_) => {
            metrics.unknown_frames.fetch_add(1, Ordering::Relaxed);
            warn!("unexpected frame direction from brain; dropping");
        }
        // Slice-3 spec §5.2: `function_call` flows through the side-channel
        // (NON-BLOCKING try_send — the binary's poll task drains on its own
        // cycle). The same "drop + observe" posture as audio_out applies if
        // the channel is full: a backed-up binary means we drop the proposal
        // and the brain gets no reply (the brain process knows no
        // function_call_output arrived → its OpenAI pump keeps going; the
        // model tolerates missing replies per OpenAI's design).
        DecodedPayload::FunctionCall(p) => {
            if tx_function_call.try_send(FunctionCallEvent(p)).is_err() {
                metrics.outbound_dropped.fetch_add(1, Ordering::Relaxed);
                warn!(
                    "function_call dropped (binary poll task not draining; brain will see no reply)"
                );
            }
        }
        // Slice-3 (spec §3.2): `function_call_output` is core→brain; ignore
        // if a brain sends one back (a misbehaving brain — same posture as
        // `SessionEnd`/`AudioIn` from brain above).
        DecodedPayload::FunctionCallOutput(_) => {
            metrics.unknown_frames.fetch_add(1, Ordering::Relaxed);
            warn!("unexpected function_call_output from brain; dropping");
        }
        // Slice-3 advisory — same "logged + counted, not forwarded" posture
        // as `Unknown`. The FOB reflex loop in step 4 will act on these;
        // slice-3 only pre-paves the wire event.
        DecodedPayload::SpeechStarted | DecodedPayload::SpeechStopped => {
            metrics.unknown_frames.fetch_add(1, Ordering::Relaxed);
            debug!("advisory interruption event observed; not acted on in slice-3");
        }
        DecodedPayload::ToolsUpdate(_) => {
            metrics.unknown_frames.fetch_add(1, Ordering::Relaxed);
            debug!(
                "tools.update observed; slice-3 dispatch keys off function_call by name, not catalog"
            );
        }
    }
    let _ = session_start; // used for ts computation if added later
}

fn elapsed_ms(start: Instant) -> u64 {
    start.elapsed().as_millis() as u64
}

/// A `function_call` event the TapClient **observed** on its inbound WS
/// stream and forwarded to the binary's poll task via the
/// `tx_function_call` side-channel (spec §5.2). The binary's poll task
/// drains this (alongside the existing `flush_rx` side-channel — slice-2
/// §5.3 step 4) and dispatches each event through `ToolRegistry::dispatch`.
///
/// # Why a thin newtype over `FunctionCallPayload` (and not a bare alias)?
///
/// A type alias (`pub type FunctionCallEvent = FunctionCallPayload;`) would
/// let the binary pass a `FunctionCallPayload` where a `FunctionCallEvent`
/// is expected without surfacing the intent. The newtype draws a small but
/// real boundary: the wire-payload type (`FunctionCallPayload`) lives in
/// `protocol.rs` for (de)serialization; the side-channel event type
/// (`FunctionCallEvent`) lives here for dispatch. They share a shape but
/// carry different semantic weight — honoring the newtype-over-primitives
/// convention from AGENTS.md even at the message level.
#[derive(Debug, Clone)]
pub struct FunctionCallEvent(pub crate::protocol::FunctionCallPayload);

/// A `function_call_output` event the binary's poll task **emits** back to
/// the TapClient via the `tx_function_call_output` side-channel (spec §5.2
/// — the binary's poll task dispatches through `ToolRegistry::dispatch`,
/// serializes the `ToolResult`, and writes the output here). The TapClient
/// drains this in the same `tokio::select!` pump cycle as the audio + close
/// arms and sends each as a `function_call_output` tap WS frame.
///
/// Same newtype-over-payload rationale as `FunctionCallEvent`.
#[derive(Debug, Clone)]
pub struct FunctionCallOutputEvent(pub crate::protocol::FunctionCallOutputPayload);

#[cfg(test)]
mod tests {
    // TapClient is heavily async; its real behavior is exercised in the
    // integration test (Task 8) against the in-process EchoServer. Unit
    // tests here cover the pure helpers.

    use super::*;
    use crate::protocol::encode_function_call;

    #[test]
    fn elapsed_ms_is_monotonic_nonneg() {
        let start = Instant::now();
        let ms = elapsed_ms(start);
        // First call ~0; just assert it's a valid u64.
        assert_eq!(ms, ms); // tautology but clippy-clean
    }

    /// slice-3 spec §5.2: when the TapClient observes a tap `function_call`
    /// frame on its inbound WS stream it emits a `FunctionCallEvent` on
    /// the `tx_function_call` side-channel. The binary's poll task drains
    /// that and dispatches via `ToolRegistry::dispatch`. This test pins the
    /// contract end-to-end through the pure helper (`handle_brain_frame`):
    /// hand it a wire-encoded function_call frame + a fresh mpsc pair and
    /// assert the receiver observes the forwarded event.
    #[tokio::test]
    async fn handle_brain_frame_forwards_function_call_to_side_channel() {
        let (tx_fc, mut rx_fc) = mpsc::channel::<FunctionCallEvent>(8);
        let (tx_audio_out, _rx_audio_out) = mpsc::channel::<PcmFrame>(8);
        let metrics = Arc::new(TapMetrics::new());

        // Build a wire function_call frame: id="call-1", name="hangup", args={}.
        let wire = encode_function_call("call-1", "hangup", "{}", 1, 100).unwrap();
        let mut last_seq: Option<u64> = None;

        handle_brain_frame(
            &wire,
            &mut last_seq,
            &tx_audio_out,
            &tx_fc,
            &metrics,
            Instant::now(),
        )
        .await;

        // The side-channel must have observed exactly one FunctionCallEvent
        // carrying the wire's id/name/args.
        let event = tokio::time::timeout(Duration::from_millis(200), rx_fc.recv())
            .await
            .expect("tx_function_call drained within 200ms")
            .expect("channel not closed");
        assert_eq!(event.0.id, "call-1");
        assert_eq!(event.0.name, "hangup");
        assert_eq!(event.0.args, serde_json::json!({}));
        // seq tracking still updates for the side-channeled event.
        assert_eq!(last_seq, Some(1));
    }

    /// slice-3 spec §5.2 — the *advisory* interrupt events (`speech_started`
    /// /`speech_stopped`) and `tools.update` are observed (logged + counted)
    /// but do NOT flow through the function_call side-channel (only
    /// `function_call` does — that's the only event with a binary-side
    /// disposal). This pins that boundary: an advisory event must NOT
    /// produce a `FunctionCallEvent` even with the channel plumbed.
    #[tokio::test]
    async fn advisory_events_are_logged_not_forwarded_to_function_call_channel() {
        let (tx_fc, mut rx_fc) = mpsc::channel::<FunctionCallEvent>(8);
        let (tx_audio_out, _rx_audio_out) = mpsc::channel::<PcmFrame>(8);
        let metrics = Arc::new(TapMetrics::new());

        let wire = crate::protocol::encode_speech_started(2, 200).unwrap();
        let mut last_seq: Option<u64> = None;
        handle_brain_frame(
            &wire,
            &mut last_seq,
            &tx_audio_out,
            &tx_fc,
            &metrics,
            Instant::now(),
        )
        .await;

        // No FunctionCallEvent forwarded — the channel stays empty. Pick a
        // tight bounded receive so the test fails fast if a future refactor
        // starts forwarding advisory events here.
        assert!(
            tokio::time::timeout(Duration::from_millis(50), rx_fc.recv())
                .await
                .is_err(),
            "no FunctionCallEvent expected for advisory events"
        );
        // The advisory event IS still observed via metrics (seq gap tracking
        // + the unknown-slot counter remains 0 — speech_started is now a
        // known payload variant).
        assert_eq!(last_seq, Some(2));
    }
}