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

Worker and API

In production you split Riverbeam into two kinds of process:

  • Workers own a RiverBeam and run the pipeline. They consume tasks from NATS.
  • API processes accept HTTP requests, authorize them, and send tasks to the workers.

Both connect to the same NATS server, and workers share one Postgres database.

Worker tasks​

BeamWorker wraps a RiverBeam and dispatches tasks by name with handle_task(task, payload):

TaskPayloadResult
invoke_agentAn InvokeRequestAgentResult JSON
ingest_knowledge_document{"text", "document_id"?, "source_url"?}IngestionResult
delete_knowledge_document{"document_id"}{"deleted": n}
remember_memory, forget_memoryNot implemented yet: BEM-092

An unknown task name fails with BEM-091. Call on_startup() once before serving: it writes the registered agents into the store's agent table.

A worker process​

worker/src/main.rs
use std::sync::Arc;

use beam_core::{BeamConfig, RiverBeam};
use beam_store::PgBeamStore;
use beam_stream::{BeamStream, NatsTransport, StreamTransport};
use beam_worker::{BeamWorker, NatsWorkerServer};
use riverbase_task::{establish_tracker_dbpool, 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")?;

// Frames go out over NATS for the API's SSE bridge.
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.register_echo_agent("assistant");

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

// Every task is recorded as a riverbase_task job.
let tracker = Arc::new(PostgresTrackerStore::new(establish_tracker_dbpool(&db)?));

let server = NatsWorkerServer::connect(&nats, worker.clone(), tracker).await?;
server.serve().await?;

worker.on_shutdown().await?;
Ok(())
}

NatsWorkerServer subscribes to beam.task.* and runs up to 16 tasks at once. Change them with with_prefix(prefix) and with_concurrency(n). For a request with a reply subject, it replies with {"ok": <result>} or {"error": <RiverbaseError>}. Every message is recorded in the riverbase_task tracker, so jobs show up with the rest of your Riverbase jobs; see riverbase_task.

Run one worker per prefix

NatsWorkerServer uses a plain NATS subscription, not a queue group. Every worker subscribed to beam.task.* receives every task, so two workers run each invoke twice, with duplicate side effects such as git pushes. Run one worker per prefix. Memory (InMemoryBackend) and usage counters (InMemoryUsageStore) also live inside that worker's process.

Calling workers from the API​

BeamWorkerClient is the API side of the task boundary:

MethodBehaviour
request(task, payload)Waits for the result
send(task, payload)Returns at once; for streaming. Errors from the task are not reported back.

NatsWorkerClient::connect uses the async-nats default request timeout of 10 seconds. A model-backed run can take longer: request then fails with BEM-096 while the worker keeps running. Connect with a longer timeout and wrap the connection with NatsWorkerClient::new(client, prefix):

use std::time::Duration;

use beam_domain::{NatsWorkerClient, DEFAULT_TASK_PREFIX};

let client = async_nats::ConnectOptions::new()
.request_timeout(Some(Duration::from_secs(300)))
.connect(nats_url.as_str())
.await?;
let workers = NatsWorkerClient::new(client, DEFAULT_TASK_PREFIX);
ClientUse
NatsWorkerClient::connect(nats_url)Separate processes. Publishes to beam.task.{task}. connect_with_prefix must match the worker's prefix.
LocalWorkerClient::new(worker)Same process, no broker; send spawns a task

BeamDomain puts typed commands on top of a client:

CommandPayloadReturns
invoke(InvokePayload)agent, content, session, optional model override, metadata, auth contextAgentResult JSON
stream(InvokePayload)the sameStreamAck (request_id, session_id, agent_name)
resume(ResumePayload)adds interaction_id and actionsStreamAck

Each command builds an InvokeRequest with a fresh request_id and sends the invoke_agent task.

An HTTP API​

An API process does four things for each request:

  1. Authenticate the caller with your own authentication.
  2. Rewrite the payload: set auth_context.user_id to the authenticated id, and remove the metadata keys that change behaviour (confirm, sandbox_root, protected_refs, memory_scopes).
  3. Authorize the command with a BeamPolicyAuthorizer.
  4. Forward it: request for invoke, send for streaming.
use beam_domain::{command_activity, ActivityAuthorizer, ActivityRequest};
use serde_json::json;

// user_id comes from your authentication, never from the request body.
let activity = command_activity("invoke", &json!({ "agent_name": payload.agent_name }));
authorizer.enforce(&ActivityRequest::new(user_id.clone(), activity)).await?;

For streaming, subscribe to the frame channel before the worker starts publishing; NATS does not replay earlier frames. BeamDomain::stream picks the request_id itself, so build the InvokeRequest with your own request_id, subscribe, then send it.

The complete server, with /invoke, a streaming /astream that answers with SSE, and the error mapping, is in Tutorial step 6. Errors keep their BEM-* code and HTTP status; see Error codes.

Authorization​

Commands map to activities in the namespace flrs.beam on the resource agent:

CommandActivity
invokebeam.invoke
streambeam.stream
resumebeam.resume

The subject is auth_context.user_id, or anonymous when it is missing. The object is the agent name. A denial fails with BEM-067 (403).

BeamPolicyAuthorizer::from_csv reads Casbin-style rows. p rows allow a subject (or role) an activity; g rows assign roles:

api/policy.csv
p, beam-user, flrs.beam, beam.invoke, agent, *, allow
p, beam-user, flrs.beam, beam.stream, agent, *, allow
p, beam-user, flrs.beam, beam.resume, agent, *, allow
p, beam-admin, flrs.beam, *, *, *, allow
g, 6f1c2a90-0d4e-4b8e-9a51-2c7d3e9f0b11, beam-user
Do not ship the bundled policy

BEAM_POLICY_CSV and build_beam_authorizer() include the demonstration row g, alice, beam-admin, which makes any request with "user_id": "alice" an administrator. Load your own rows, and set auth_context.user_id from your authenticated session, never from the request body as sent by the client.

The authorizer also matches the subject itself against p rows, so a request whose user_id is a role name, such as beam-admin, gets that role's permissions. Use opaque user ids (UUIDs) and never let callers choose their own. A p row whose subject is * allows everyone. Roles resolve one level deep: a g row cannot point at another role.

BeamCommandEngine::with_authorizer applies the same check per command inside the domain, as a second layer.

Mounting in a Riverbase app​

BeamDomainBundle is a Riverbase Domain with the namespace flrs.beam. Its command engine forwards invoke, stream and resume to a worker client, and its query engine serves a catalog:

use std::sync::Arc;

use beam_domain::{BeamDomainBundle, BeamPolicyAuthorizer, Catalog, NatsWorkerClient};
use tokio::sync::RwLock;

let client = Arc::new(NatsWorkerClient::connect(&nats).await?);
let catalog = Arc::new(RwLock::new(Catalog::new()));
let authorizer = Arc::new(BeamPolicyAuthorizer::from_csv(&policy_csv));

let beam_domain = BeamDomainBundle::with_authorizer(client, catalog, authorizer);
app.mount_domain(Arc::new(beam_domain)).await?;
Query resourceRowsItem key
agentAgentInfo: name, description, supported_modalities, versionname
modelModelInfo: name, provider, modalitiesname
agent_sessionSessionInfo: session_id, namesession_id

The catalog is filled by your code with set_agents, set_models and upsert_session; it does not read the store. Routes follow the Riverbase HTTP conventions, for example GET /api/flrs.beam/agent.list.

Through the generic command route, stream and resume forward the payload unchanged and return {"status": "accepted"}. No request_id is added and stream is not forced, so nothing is published unless the client sends both, and the client never learns a request_id it did not choose. Prefer your own routes for streaming.

Current limitations​

  • Workers do not share tasks: see Run one worker per prefix.
  • send does not report task errors, and a failed streaming run publishes no error frame. Clients of streaming routes need a timeout.
  • remember_memory and forget_memory are not implemented (BEM-092).