4. An HTTP API with streaming

You put the assistant behind HTTP. The worker still runs in the same process: an in-process client replaces NATS, and a broadcast transport carries the stream frames.

Add the crates

docs-assistant/Cargo.toml

[dependencies]
beam_domain = { path = "../riverbeam/crates/beam_domain" }
beam_stream = { path = "../riverbeam/crates/beam_stream" }
beam_worker = { path = "../riverbeam/crates/beam_worker" }

axum = "0.8"
futures-util = "0.3"
tokio = { version = "1", features = ["rt-multi-thread", "macros", "net", "sync"] }
tokio-stream = { version = "0.1", features = ["sync"] }

The API

Two routes:

• POST /invoke waits for the agent and returns the AgentResult as JSON.
• POST /astream starts a streaming run and answers with the frames as server-sent events (SSE).

/astream subscribes to the transport before it starts the run. The in-process worker starts publishing frames within milliseconds, and a broadcast channel does not replay frames sent before a subscription.

docs-assistant/src/api.rs

use std::convert::Infallible;
use std::sync::Arc;

use axum::http::StatusCode;
use axum::response::sse::{Event, Sse};
use axum::extract::State;
use axum::routing::post;
use axum::{Json, Router};
use beam_domain::{BeamDomain, InvokePayload, LocalWorkerClient, StreamAck};
use beam_stream::{BeamStream, BroadcastTransport, PublishedFrame};
use beam_types::StreamChunkType;
use futures_util::future;
use futures_util::stream::{BoxStream, StreamExt};
use riverbase_core::RiverbaseError;
use serde_json::Value;
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;

pub struct AppState {
    pub domain: BeamDomain<LocalWorkerClient>,
    pub frames: Arc<BroadcastTransport>,
}

type ApiError = (StatusCode, Json<Value>);
type EventStream = Sse<BoxStream<'static, Result<Event, Infallible>>>;

/// Keep the BEM-* code and HTTP status of Riverbeam errors.
pub fn api_error(err: RiverbaseError) -> ApiError {
    let status = StatusCode::from_u16(err.http_status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    (status, Json(serde_json::to_value(&err).unwrap_or(Value::Null)))
}

pub fn router(state: Arc<AppState>) -> Router {
    Router::new()
        .route("/invoke", post(invoke))
        .route("/astream", post(astream))
        .with_state(state)
}

async fn invoke(
    State(state): State<Arc<AppState>>,
    Json(payload): Json<InvokePayload>,
) -> Result<Json<Value>, ApiError> {
    state.domain.invoke(payload).await.map(Json).map_err(api_error)
}

async fn astream(
    State(state): State<Arc<AppState>>,
    Json(payload): Json<InvokePayload>,
) -> Result<EventStream, ApiError> {
    let receiver = state.frames.subscribe(); // before the run starts
    let ack = state.domain.stream(payload).await.map_err(api_error)?;
    Ok(frames_for(receiver, &ack))
}

/// The frames of one request as SSE, ending after the `done` frame.
pub fn frames_for(receiver: broadcast::Receiver<PublishedFrame>, ack: &StreamAck) -> EventStream {
    let channel = format!(
        "ws.channel.{}",
        BeamStream::stream_key(&ack.session_id.to_string(), &ack.request_id.to_string())
    );
    let events = BroadcastStream::new(receiver)
        .filter_map(move |received| {
            let frame = received
                .ok()
                .filter(|published| published.channel == channel)
                .map(|published| published.frame);
            future::ready(frame)
        })
        .scan(false, |finished, frame| {
            if *finished {
                return future::ready(None);
            }
            *finished = matches!(frame.chunk_type, StreamChunkType::Done | StreamChunkType::Error);
            future::ready(Some(frame))
        })
        .map(|frame| {
            let event = Event::default()
                .json_data(&frame)
                .unwrap_or_else(|_| Event::default().data("serialize-error"));
            Ok::<_, Infallible>(event)
        })
        .boxed();
    Sse::new(events)
}

Wire it up

The pipeline now runs inside a BeamWorker, and the API reaches it through BeamDomain:

docs-assistant/src/main.rs

mod agents;
mod api;
mod knowledge;
mod tools;

use std::sync::Arc;

use beam_agents::Tool;
use beam_core::{BeamConfig, RiverBeam};
use beam_domain::{BeamDomain, LocalWorkerClient};
use beam_stream::{BeamStream, BroadcastTransport};
use beam_worker::BeamWorker;

// document_names() as in step 2

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let names = document_names()?;
    let tools: Vec<Arc<dyn Tool>> = vec![Arc::new(tools::ListDocuments { names: names.clone() })];

    let frames = Arc::new(BroadcastTransport::new(1024));

    let mut beam = RiverBeam::new(BeamConfig::default());
    beam.set_stream(Arc::new(BeamStream::new(frames.clone())));
    beam.set_knowledge(Arc::new(knowledge::load(&names).await?));
    beam.register_agent(Arc::new(agents::assistant(tools)));

    let worker = Arc::new(BeamWorker::new(Arc::new(beam)));
    worker.on_startup().await?;

    let state = Arc::new(api::AppState {
        domain: BeamDomain::new(LocalWorkerClient::new(worker)),
        frames,
    });

    let listener = tokio::net::TcpListener::bind("127.0.0.1:8080").await?;
    axum::serve(listener, api::router(state)).await?;
    Ok(())
}

Try it

cargo run

curl -s localhost:8080/invoke -H 'content-type: application/json' -d '{
  "agent_name": "assistant",
  "session_id": "00000000-0000-0000-0000-000000000001",
  "content": [{"type": "text", "text": "Which documents are there?"}]
}'

curl -N localhost:8080/astream -H 'content-type: application/json' -d '{
  "agent_name": "assistant",
  "session_id": "00000000-0000-0000-0000-000000000001",
  "content": [{"type": "text", "text": "And which one covers installation?"}]
}'

data: {"type":"placeholder", …}
data: {"type":"thinking", …}
data: {"type":"chunk","data":{"role":"ai", …}, …}
data: {"type":"done","result":{"content":[…],"agent_name":"assistant", …}, …}

Both calls use the same session_id, so the second question is answered with the first in the history.

If a streaming run fails

The pipeline does not publish an error frame when a run fails. LocalWorkerClient discards the error, and over NATS it is only recorded in the job tracker. A failed stream simply never sends done, so give clients a timeout.

The API has no authentication yet. You add it in step 6. Next: let the assistant edit documents.