rutster/crates/rutster-brain-realtime/src/main.rs

//! # rutster-brain-realtime binary (slice-3 spec §4.2 + §5.3)
//!
//! Standalone brain process the core dials out to via `RUTSTER_TAP_URL` (spec
//! §5.1). One process bridges many tap WS connections (one per call) to many
//! OpenAI Realtime sessions.
//!
//! # Modes
//!
//! - **`--features=mock`** (offline dev loop, slice-3 spec §7.3): starts an
//!   in-process `MockRealtimeBrain` WS server, dials it as the "OpenAI side."
//!   No API key, no network calls to OpenAI. Used by the integration test +
//!   the offline dev loop (`cargo run -p rutster-brain-realtime --features=mock`).
//! - **Default** (real OpenAI): loads `OPENAI_API_KEY` per
//!   [`api_key::load_api_key`], dials `wss://api.openai.com/v1/realtime?model=...`.
//!
//! Both modes share the same tap-side WS server + the same `run_openai_pump`
//! bridging logic — only the OpenAI-side dialer differs.

use std::net::SocketAddr;

use futures_util::{SinkExt, StreamExt};
#[cfg(feature = "mock")]
use rutster_brain_realtime::MockRealtimeBrain;
#[cfg(not(feature = "mock"))]
use rutster_brain_realtime::api_key;
use rutster_brain_realtime::openai_client;
use rutster_brain_realtime::translator::build_openai_session_update;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tracing::{info, warn};

/// Default bind addr for the brain process's tap-side WS server (spec §5.3).
/// Distinct from slice-2's echo brain (`:8081/echo`); the two coexist.
const DEFAULT_TAP_BIND: &str = "127.0.0.1:8082";
const DEFAULT_VOICE: &str = "alloy";
const DEFAULT_MODEL: &str = "gpt-4o-realtime";

#[tokio::main]
async fn main() {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "rutster_brain_realtime=info".into()),
        )
        .init();

    let tap_bind: SocketAddr = std::env::var("RUTSTER_TAP_BIND")
        .unwrap_or_else(|_| DEFAULT_TAP_BIND.to_string())
        .parse()
        .expect("RUTSTER_TAP_BIND must be a valid SocketAddr");
    let voice =
        std::env::var("OPENAI_REALTIME_VOICE").unwrap_or_else(|_| DEFAULT_VOICE.to_string());
    let model =
        std::env::var("OPENAI_REALTIME_MODEL").unwrap_or_else(|_| DEFAULT_MODEL.to_string());
    let _ = &model; // used only in the non-mock branch below.

    #[cfg(feature = "mock")]
    let openai_url: String;
    #[cfg(feature = "mock")]
    let openai_headers: Vec<(String, String)>;
    #[cfg(feature = "mock")]
    let _mock_guard;
    #[cfg(not(feature = "mock"))]
    let (openai_url, openai_headers) = {
        let api_key = match api_key::load_api_key() {
            Ok(k) => k,
            Err(e) => {
                eprintln!("failed to load OPENAI_API_KEY: {e}; refusing to start");
                std::process::exit(1);
            }
        };
        (
            openai_client::openai_realtime_url(&model),
            openai_client::openai_headers(&api_key),
        )
    };
    #[cfg(feature = "mock")]
    {
        let mock = MockRealtimeBrain::start().await.expect("mock bind ok");
        info!(mock_addr = %mock.addr, "MockRealtimeBrain started (--features=mock)");
        openai_url = mock.url();
        openai_headers = Vec::new();
        _mock_guard = mock;
    }
    let listener = tokio::net::TcpListener::bind(tap_bind)
        .await
        .expect("bind ok");
    info!(%tap_bind, %openai_url, "rutster-brain-realtime listening (core dials this)");

    loop {
        match listener.accept().await {
            Ok((stream, peer)) => {
                let url = openai_url.clone();
                let hdrs = openai_headers.clone();
                let voice = voice.clone();
                tokio::spawn(async move {
                    if let Err(e) = handle_tap_connection(stream, peer, &url, &hdrs, &voice).await {
                        warn!(error = ?e, %peer, "tap connection ended with error");
                    }
                });
            }
            Err(e) => warn!(error = %e, "accept failed; continuing"),
        }
    }
}

/// Handle one tap WS connection from the core. Spawns two mpsc forwarders
/// (tap→pump, pump→tap) + dials the OpenAI side, then runs `run_openai_pump`
/// on the OpenAI side + awaits both forwarder tasks.
///
/// The tap-side handshake (hello/session_end) is handled here; the OpenAI
/// pump sees only audio_in + function_call_output on its inbound + emits
/// only audio_out/speech_*/function_call/error on its outbound.
async fn handle_tap_connection(
    stream: tokio::net::TcpStream,
    peer: SocketAddr,
    openai_url: &str,
    openai_headers: &[(String, String)],
    voice: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let mut tap_ws = tokio_tungstenite::accept_async(stream).await?;
    info!(%peer, "tap WS connection accepted");

    // === Tap handshake: first frame must be hello; send hello ack. ===
    let hello_in = tap_ws
        .next()
        .await
        .ok_or("tap connection closed before hello")??;
    let hello_text = hello_in.into_text().map_err(|_| "hello not text")?;
    let decoded = rutster_tap::decode_envelope(&hello_text)?;
    let session_id = match decoded.payload {
        rutster_tap::DecodedPayload::Hello(p) => p.session_id,
        _ => return Err("first tap frame not hello".into()),
    };
    info!(%peer, %session_id, "tap hello received");
    let ack = rutster_tap::encode_hello(&session_id, 0, 0)?;
    tap_ws.send(Message::Text(ack)).await?;

    // === Bridge channels between tap WS and OpenAI pump. ===
    // tap_via (Sender): the core's tap frames flow IN here — the forwarder
    //   reads text from tap_ws and forwards via this channel. The OpenAI
    //   pump's `tap_rx: mpsc::Receiver<String>` is the receiving end.
    // pump_tap_tx (Sender): the pump's outbound tap frames (audio_out,
    //   speech_started, function_call, error) flow OUT here. The forwarder
    //   reads from the paired Receiver and writes to tap_ws.
    let (tap_via, pump_tap_rx) = mpsc::channel::<String>(64);
    let (pump_tap_tx, mut tap_out_rx) = mpsc::channel::<String>(64);

    // Split the tap WS into sink + stream so the two forwarder tasks can
    // own each half independently. `tokio_tungstenite::WebSocketStream::split`
    // yields `(SplitSink<..., Message>, SplitStream<...>)` — the canonical
    // tokio pattern for fan-out into two parallel async tasks.
    let (mut tap_sink, mut tap_stream) = tap_ws.split();

    // === Forwarder: tap_ws_stream → pump_tap_rx ===
    let in_fwd = tokio::spawn(async move {
        while let Some(msg_res) = tap_stream.next().await {
            match msg_res {
                Ok(m) => {
                    if let Ok(text) = m.into_text() {
                        if tap_via.send(text).await.is_err() {
                            break; // pump dropped
                        }
                    }
                }
                Err(e) => {
                    warn!(error = ?e, "tap_ws recv error");
                    break;
                }
            }
        }
    });

    // === Forwarder: pump_tap_tx → tap_ws_sink ===
    let out_fwd = tokio::spawn(async move {
        while let Some(text) = tap_out_rx.recv().await {
            if let Err(e) = tap_sink.send(Message::Text(text)).await {
                warn!(error = ?e, "tap_ws send error");
                break;
            }
        }
    });

    // === Dial the OpenAI side (real or mock). ===
    let request = openai_url.into_client_request()?;
    let mut request = request;
    for (k, v) in openai_headers {
        request.headers_mut().insert(
            http::HeaderName::from_bytes(k.as_bytes())?,
            http::HeaderValue::from_str(v)?,
        );
    }
    let (openai_ws, _resp) = tokio_tungstenite::connect_async(request).await?;
    info!(%peer, %openai_url, "OpenAI side connected");

    // `run_openai_pump` sends its own session.update on startup per
    // `build_openai_session_update`; don't pre-send here.
    let _ = build_openai_session_update;
    info!(%peer, voice = %voice, "starting OpenAI pump");

    let pump_result =
        openai_client::run_openai_pump(openai_ws, pump_tap_rx, pump_tap_tx, voice.to_string())
            .await;
    match pump_result {
        Ok(()) => info!(%peer, "OpenAI pump exited cleanly"),
        Err(e) => warn!(%peer, error = ?e, "OpenAI pump exited with error"),
    }
    in_fwd.abort();
    out_fwd.abort();
    Ok(())
}