6. Production
So far everything ran in one process with in-memory state. Now you split it:
- a worker that owns the pipeline, stores sessions in Postgres, and consumes tasks from NATS;
- an API server that authenticates callers, authorizes commands, and relays stream frames.
Prepare the infrastructure
nats-server &
createdb docs_assistant
psql docs_assistant -f ../riverbeam/migrations/0001_beam_schema.sql
export BEAM_DATABASE_URL=postgres://$USER@127.0.0.1:5432/docs_assistant
export NATS_URL=nats://127.0.0.1:4222
Apply the migration before the first start; see Database schema.
Share code between two binaries
Move the modules into a library so both binaries can use them, and move document_names from
main.rs into tools.rs as a pub fn:
pub mod agents;
pub mod knowledge;
pub mod tools;
In main.rs, replace mod agents; mod knowledge; mod tools; with
use docs_assistant::{agents, knowledge, tools}; so the one-process version keeps building.
[dependencies]
beam_store = { path = "../riverbeam/crates/beam_store" }
riverbase_task = { path = "../riverbase/crates/riverbase_task", features = ["postgres", "nats-io"] }
async-nats = "0.49"
The worker
use std::sync::Arc;
use beam_agents::Tool;
use beam_core::{BeamConfig, RiverBeam};
use beam_store::PgBeamStore;
use beam_stream::{BeamStream, NatsTransport, StreamTransport};
use beam_worker::{BeamWorker, NatsWorkerServer};
use docs_assistant::{agents, knowledge, tools};
use riverbase_task::{establish_tracker_dbpool, run_tracker_migrations_url, PostgresTrackerStore};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let nats = std::env::var("NATS_URL").unwrap_or_else(|_| "nats://127.0.0.1:4222".into());
let db = std::env::var("BEAM_DATABASE_URL")?;
let names = tools::document_names()?;
let tools: Vec<Arc<dyn Tool>> = vec![Arc::new(tools::ListDocuments { names: names.clone() })];
// Frames go to NATS, where the API server picks them up.
let transport: Arc<dyn StreamTransport> = Arc::new(NatsTransport::connect(&nats).await?);
let mut beam = RiverBeam::new(BeamConfig::default());
beam.set_stream(Arc::new(BeamStream::new(transport)));
beam.set_store(Arc::new(PgBeamStore::connect(&db).await?));
beam.set_knowledge(Arc::new(knowledge::load(&names).await?));
beam.register_agent(Arc::new(agents::assistant(tools)));
let tree = agents::working_tree("/srv/docs-assistant/work", "/srv/docs-remote.git")?;
beam.register_agent(Arc::new(agents::author(tree)));
let worker = Arc::new(BeamWorker::new(Arc::new(beam)));
worker.on_startup().await?; // writes the agents to ref_beam.agent
// Every task becomes a riverbase_task job.
run_tracker_migrations_url(&db)?;
let tracker = Arc::new(PostgresTrackerStore::new(establish_tracker_dbpool(&db)?));
NatsWorkerServer::connect(&nats, worker.clone(), tracker).await?.serve().await?;
worker.on_shutdown().await?;
Ok(())
}
NatsWorkerServer uses a plain NATS subscription on beam.task.*, not a queue group. If two
workers share a prefix, both run every task: each invoke runs twice, and pushes happen twice. Run
a single worker. Memory and usage counters also live in that worker's process.
The API server
The server never trusts identity or control fields from the request body. It takes the caller from
your authentication, overwrites auth_context, strips the metadata keys that change behaviour, and
then authorizes the command.
use std::convert::Infallible;
use std::sync::Arc;
use std::time::Duration;
use axum::extract::State;
use axum::http::{HeaderMap, StatusCode};
use axum::response::sse::{Event, Sse};
use axum::routing::post;
use axum::{Json, Router};
use beam_domain::{
command_activity, ActivityAuthorizer, ActivityRequest, BeamPolicyAuthorizer, BeamWorkerClient,
InvokePayload, NatsWorkerClient, DEFAULT_TASK_PREFIX,
};
use beam_stream::{BeamStream, NatsTransport};
use beam_types::{InvokeAuthContext, InvokeRequest, StreamChunk, StreamChunkType};
use futures_util::future;
use futures_util::stream::{BoxStream, StreamExt};
use riverbase_core::RiverbaseError;
use serde_json::{json, Value};
use uuid::Uuid;
struct AppState {
workers: NatsWorkerClient,
frames: NatsTransport,
authorizer: BeamPolicyAuthorizer,
}
type ApiError = (StatusCode, Json<Value>);
type EventStream = Sse<BoxStream<'static, Result<Event, Infallible>>>;
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)))
}
/// The authenticated caller's id. Replace this with your real authentication: here a gateway in
/// front of the server sets `x-user-id`, and the header must never be accepted from the internet.
fn caller(headers: &HeaderMap) -> Result<String, ApiError> {
headers
.get("x-user-id")
.and_then(|value| value.to_str().ok())
.map(str::to_string)
.ok_or_else(|| (StatusCode::UNAUTHORIZED, Json(json!({ "error": "not authenticated" }))))
}
/// Turn a client payload into a worker request the client cannot steer, and authorize it.
async fn prepare(
state: &AppState,
headers: &HeaderMap,
cmd: &str,
payload: InvokePayload,
stream: bool,
) -> Result<InvokeRequest, ApiError> {
let user_id = caller(headers)?;
let mut metadata = payload.metadata;
if let Some(map) = metadata.as_object_mut() {
for key in ["confirm", "sandbox_root", "protected_refs", "memory_scopes"] {
map.remove(key);
}
}
let activity = command_activity(cmd, &json!({ "agent_name": payload.agent_name }));
state
.authorizer
.enforce(&ActivityRequest::new(user_id.clone(), activity))
.await
.map_err(api_error)?;
Ok(InvokeRequest {
agent_name: payload.agent_name,
content: payload.content,
session_id: payload.session_id,
model_override: payload.model_override,
metadata,
interaction_id: None,
request_id: Some(Uuid::new_v4()),
stream,
actions: None,
auth_context: Some(InvokeAuthContext {
user_id: Some(user_id),
profile_id: None,
org_id: None,
mcp_token: None,
}),
})
}
async fn invoke(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(payload): Json<InvokePayload>,
) -> Result<Json<Value>, ApiError> {
let request = prepare(&state, &headers, "invoke", payload, false).await?;
let task = serde_json::to_value(&request).unwrap_or(Value::Null);
state.workers.request("invoke_agent", task).await.map(Json).map_err(api_error)
}
async fn astream(
State(state): State<Arc<AppState>>,
headers: HeaderMap,
Json(payload): Json<InvokePayload>,
) -> Result<EventStream, ApiError> {
let request = prepare(&state, &headers, "stream", payload, true).await?;
let request_id = request.request_id.unwrap_or_default().to_string();
let channel = format!(
"ws.channel.{}",
BeamStream::stream_key(&request.session_id.to_string(), &request_id)
);
// Subscribe before the worker can publish: NATS does not replay earlier frames.
let subscriber = state.frames.subscribe(&channel).await.map_err(api_error)?;
let task = serde_json::to_value(&request).unwrap_or(Value::Null);
state.workers.send("invoke_agent", task).await.map_err(api_error)?;
let events = subscriber
.filter_map(|message| future::ready(serde_json::from_slice::<StreamChunk>(&message.payload).ok()))
.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();
Ok(Sse::new(events))
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
let nats = std::env::var("NATS_URL").unwrap_or_else(|_| "nats://127.0.0.1:4222".into());
// A model-backed run can take minutes; the async-nats default request timeout is 10 seconds.
let client = async_nats::ConnectOptions::new()
.request_timeout(Some(Duration::from_secs(300)))
.connect(nats.as_str())
.await?;
let state = Arc::new(AppState {
workers: NatsWorkerClient::new(client, DEFAULT_TASK_PREFIX),
frames: NatsTransport::connect(&nats).await?,
authorizer: BeamPolicyAuthorizer::from_csv(&std::fs::read_to_string("policy.csv")?),
});
let app = Router::new()
.route("/invoke", post(invoke))
.route("/astream", post(astream))
.with_state(state);
let addr = std::env::var("BEAM_BIND_ADDR").unwrap_or_else(|_| "0.0.0.0:8080".into());
axum::serve(tokio::net::TcpListener::bind(&addr).await?, app).await?;
Ok(())
}
astream sends the task with its own request_id instead of calling BeamDomain::stream, so it can
subscribe to the channel before the worker starts. Port the /resume route from
step 5 the same way: it needs a PgBeamStore in the server to read pending actions.
The policy
Grant roles to user ids, never to names a caller could choose:
p, beam-user, flrs.beam, beam.invoke, agent, assistant, allow
p, beam-user, flrs.beam, beam.stream, agent, assistant, allow
p, beam-editor, flrs.beam, *, agent, *, allow
g, 6f1c2a90-0d4e-4b8e-9a51-2c7d3e9f0b11, beam-user
g, 0b7e5a62-3c1d-4f2a-8e9b-5d4c3b2a1f00, beam-editor
The authorizer also accepts a role name as the subject itself, so a caller whose id were
beam-editor would match the role's rows. Use opaque ids, such as UUIDs, for users.
Run it
cargo run --bin worker
cargo run --bin server
curl -N localhost:8080/astream \
-H 'content-type: application/json' \
-H 'x-user-id: 6f1c2a90-0d4e-4b8e-9a51-2c7d3e9f0b11' \
-d '{
"agent_name": "assistant",
"session_id": "00000000-0000-0000-0000-000000000003",
"content": [{"type": "text", "text": "How do I install the product?"}]
}'
Production checklist
| Check | Why |
|---|---|
| One worker per task prefix | Every worker on a subject runs every task |
Identity from authentication, written into auth_context | Authorization, quotas and audit all use auth_context.user_id |
| Opaque user ids | Role names also match as subjects |
| Your own policy rows | The bundled BEAM_POLICY_CSV makes alice an administrator |
Strip confirm, sandbox_root, protected_refs, memory_scopes | They skip approval, move the sandbox, or load other users' memory |
Check that a session_id belongs to the caller | Anyone who knows a session id can continue that conversation |
| Validate resume responses against the stored action | The runner runs whatever the response payload names |
| A request timeout above your longest run | The async-nats default is 10 seconds |
| A client timeout on streams | Failed runs publish no done frame |
| Only public documents in knowledge | Retrieval is not filtered by user |
| A sandbox root you control, and remotes that need no credentials | Push sends no credentials |
| Do not rely on the shipped security rules for prompt injection | Their scores never reach the block threshold |
The details behind each row are in Worker and API, Security, usage and audit, and Sessions and human approval.