Chuyển tới nội dung chính

Streaming

A streaming request returns an acknowledgement at once. The agent's frames arrive on a channel named after the session and the request:

ws.channel.{session_id}:{request_id}

beam_stream publishes the frames through a pluggable transport. Clients usually read them as server-sent events (SSE) from the API process.

Turn streaming on​

Attach a BeamStream to the pipeline:

use std::sync::Arc;

use beam_stream::{BeamStream, NatsTransport, StreamTransport};

let transport: Arc<dyn StreamTransport> = Arc::new(NatsTransport::connect(&nats_url).await?);
beam.set_stream(Arc::new(BeamStream::new(transport)));

A request is streamed when it has stream: true and a request_id. The stream and resume commands set both. Without an attached stream, the request still runs and returns its result, but nothing is published.

Frames​

Each frame is a StreamChunk. Its type is one of:

TypeMeaning
placeholderThe agent has started
thinkingProgress text
toolA tool call; data names the tool and its arguments
chunkReply text in data.content
doneThe final frame; result holds the full AgentResult
errorReserved for agents that report a failure as a frame; error holds the message

Every frame also carries agent_name, session_id, interaction_id and run_id. Stop reading a channel after done or error.

Failed runs send no final frame

When a streaming run fails, for example a security block, a quota denial, or an agent error, the pipeline returns the error to the worker but publishes nothing to the channel. send does not report it back either. Give stream readers a timeout.

Transports​

TransportUseSubscribe with
NatsTransportWorker and API in separate processestransport.subscribe(&channel), a NATS subscriber
BroadcastTransportOne process, many listenerstransport.subscribe(), a tokio::sync::broadcast::Receiver; filter by channel
InMemoryTransportTeststransport.frames() returns every published frame

NatsTransport publishes JSON to the subject beam.stream.{channel} with : replaced by .. For example, the channel ws.channel.s1:r2 maps to beam.stream.ws.channel.s1.r2. Use NatsTransport::with_prefix(bus, prefix) to change the beam.stream prefix.

Implement StreamTransport to publish elsewhere:

#[async_trait]
pub trait StreamTransport: Send + Sync {
async fn publish(&self, channel: &str, frame: &StreamChunk) -> RiverbaseResult<String>;
}

Relaying frames as SSE​

In the API process, subscribe to the channel and forward each frame. An axum handler:

api/src/sse.rs
use std::convert::Infallible;
use std::sync::Arc;

use axum::extract::{Path, State};
use axum::response::sse::{Event, Sse};
use beam_stream::{BeamStream, NatsTransport};
use beam_types::StreamChunk;
use futures_util::stream::{self, BoxStream, StreamExt};

async fn astream_sse(
State(transport): State<Arc<NatsTransport>>,
Path((session_id, request_id)): Path<(String, String)>,
) -> Sse<BoxStream<'static, Result<Event, Infallible>>> {
let channel = format!("ws.channel.{}", BeamStream::stream_key(&session_id, &request_id));
let events = match transport.subscribe(&channel).await {
Ok(subscriber) => subscriber
.map(|msg| {
let event = serde_json::from_slice::<StreamChunk>(&msg.payload)
.ok()
.and_then(|frame| Event::default().json_data(&frame).ok())
.unwrap_or_else(|| Event::default().data("serialize-error"));
Ok::<_, Infallible>(event)
})
.boxed(),
Err(err) => {
let message = err.to_string();
stream::once(async move { Ok(Event::default().event("error").data(message)) }).boxed()
}
};
Sse::new(events)
}

Subscribe before the worker starts publishing: NATS does not replay frames published before the subscription. With the two-call flow below, a fast run can finish before the second call subscribes. The tutorial's single-call /astream subscribes first, then sends the task, and is the safer pattern.

# 1. start the stream and read request_id from the acknowledgement
curl -s localhost:8080/astream -H 'content-type: application/json' -d '{
"agent_name": "assistant",
"session_id": "00000000-0000-0000-0000-000000000001",
"content": [{"type": "text", "text": "stream please"}]
}'

# 2. read the frames
curl -N localhost:8080/astream/00000000-0000-0000-0000-000000000001/<request_id>

Here /astream returns the acknowledgement and /astream/{session_id}/{request_id} is the SSE handler above. Authenticate and authorize both routes as described in Worker and API.

Merging text deltas​

BeamStream::with_config(transport, StreamConfig { .. }) can merge consecutive chunk frames before publishing, to send fewer, larger frames:

FieldEffectDefault
publish_chunk_token_thresholdPublish once the buffered text reaches this many tokens (about four bytes per token). 0 publishes every frame immediately.0
publish_chunk_max_charsAlso publish once the buffered text reaches this many bytes (only when the threshold is above 0)0

Frames merge only when they share the message, agent, run and tool call. Any other frame flushes the buffer first. The pipeline calls flush at the end of each request, so no text is lost.

Current limitations​

  • Agents return all their frames when they finish, so frames reach subscribers in a burst after the agent is done, not as the model generates them.
  • Session and request ids that contain . or NATS wildcards produce unexpected subjects. Use UUIDs.
  • Security checks the reply only after the frames are published; see Security.