Skip to main content

orbit_runtime/
server.rs

1//! HTTP API (RFC-0001 §11.2) and the HTTP/JSON binding of the worker
2//! protocol (§8). See `protocol/README.md` for the wire format.
3
4use crate::engine::{Engine, EngineError, RegisterWorker, TaskAssignment};
5use crate::identity::{Identities, WorkerIdentity};
6use crate::model::TaskError;
7use axum::extract::{DefaultBodyLimit, Path, Query, Request, State};
8use axum::http::{HeaderMap, HeaderValue, StatusCode};
9use axum::middleware::{self, Next};
10use axum::response::{IntoResponse, Response};
11use axum::routing::{get, post};
12use axum::{Extension, Json, Router};
13use serde::Deserialize;
14use serde_json::{json, Value};
15use std::net::SocketAddr;
16use std::sync::{Arc, Mutex};
17use std::time::Duration;
18use tokio::sync::Notify;
19
20pub const API_VERSION: &str = "orbit-api/1";
21
22pub struct AppState {
23    pub engine: Mutex<Engine>,
24    pub changed: Arc<Notify>,
25    pub auth_token: Option<String>,
26    /// Optional credential restricted to the worker protocol.
27    pub worker_token: Option<String>,
28    /// Per-worker identities (`--workers-file`).
29    pub identities: Identities,
30}
31
32pub type Shared = Arc<AppState>;
33
34pub struct ApiError(pub(crate) EngineError);
35
36impl From<EngineError> for ApiError {
37    fn from(e: EngineError) -> Self {
38        ApiError(e)
39    }
40}
41
42impl IntoResponse for ApiError {
43    fn into_response(self) -> Response {
44        let status =
45            StatusCode::from_u16(self.0.status).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
46        (
47            status,
48            Json(json!({ "error": { "code": self.0.code, "message": self.0.message } })),
49        )
50            .into_response()
51    }
52}
53
54fn bad_request(msg: impl Into<String>) -> ApiError {
55    ApiError(EngineError::new(400, "INVALID_REQUEST", msg))
56}
57
58type ApiResult<T> = Result<T, ApiError>;
59
60pub(crate) fn engine(state: &Shared) -> std::sync::MutexGuard<'_, Engine> {
61    state.engine.lock().unwrap_or_else(|p| p.into_inner())
62}
63
64pub fn router(state: Shared, body_limit: usize) -> Router {
65    Router::new()
66        .route("/v1/health", get(health))
67        // Answered by the cluster gate before authentication; the route
68        // exists so the gate wraps it even when other routers are merged in
69        // (the web UI), which would otherwise replace the gated fallback.
70        .route("/v1/ready", get(ready))
71        .route("/v1/metrics", get(metrics))
72        .route("/v1/workflows", post(register_workflow).get(list_workflows))
73        .route("/v1/workflows/{name}", get(get_workflow_latest))
74        .route("/v1/workflows/{name}/{version}", get(get_workflow))
75        .route("/v1/executions", post(start_execution).get(list_executions))
76        .route("/v1/executions/{id}", get(get_execution))
77        .route("/v1/executions/{id}/history", get(get_history))
78        .route("/v1/executions/{id}/cancel", post(cancel_execution))
79        .route("/v1/executions/{id}/signals/{name}", post(send_signal))
80        .route("/v1/executions/{id}/queries/{name}", get(run_query))
81        .route("/v1/executions/{id}/wait", get(wait_execution))
82        .route("/v1/schedules", post(create_schedule).get(list_schedules))
83        .route(
84            "/v1/schedules/{id}",
85            get(get_schedule)
86                .patch(update_schedule)
87                .delete(delete_schedule),
88        )
89        .route("/v1/schedules/{id}/pause", post(pause_schedule))
90        .route("/v1/schedules/{id}/resume", post(resume_schedule))
91        .route("/v1/schedules/{id}/trigger", post(trigger_schedule))
92        .route("/v1/workers", get(list_workers))
93        .route("/v1/workers/register", post(register_worker))
94        .route("/v1/workers/heartbeat", post(heartbeat))
95        .route("/v1/workers/poll", post(poll))
96        .route("/v1/workers/unregister", post(unregister))
97        .route("/v1/tasks/complete", post(complete_task))
98        .route("/v1/tasks/fail", post(fail_task))
99        .layer(middleware::from_fn_with_state(state.clone(), auth))
100        .layer(DefaultBodyLimit::max(body_limit))
101        .with_state(state)
102}
103
104/// Endpoints a worker credential may call.
105fn is_worker_path(path: &str) -> bool {
106    path.starts_with("/v1/workers/") || path.starts_with("/v1/tasks/")
107}
108
109/// Who is calling, as established by [`auth`].
110#[derive(Clone)]
111pub enum Caller {
112    /// No authentication is configured (loopback development).
113    Open,
114    Admin,
115    /// The shared, restricted worker token (`--worker-token`).
116    SharedWorker,
117    /// A worker from the workers file, by token or client certificate.
118    Worker(Arc<WorkerIdentity>),
119    /// An application from the workers file (`[[client]]`): the execution
120    /// API for the workflows and actions it was granted.
121    Client(Arc<crate::identity::ClientIdentity>),
122}
123
124impl Caller {
125    fn identity(&self) -> Option<&WorkerIdentity> {
126        match self {
127            Caller::Worker(w) => Some(w),
128            _ => None,
129        }
130    }
131}
132
133fn forbidden(msg: impl Into<String>) -> ApiError {
134    ApiError(EngineError::new(403, "FORBIDDEN", msg))
135}
136
137/// A client identity may only act on the workflows and actions it was granted.
138fn allow(caller: &Caller, action: &str, workflow: &str) -> ApiResult<()> {
139    if let Caller::Client(c) = caller {
140        if !c.allows(action, workflow) {
141            return Err(forbidden(format!(
142                "client `{}` may not {action} executions of `{workflow}`",
143                c.id
144            )));
145        }
146    }
147    Ok(())
148}
149
150/// A client may read a workflow definition only if it may read its executions.
151fn allow_definition(caller: &Caller, workflow: &str) -> ApiResult<()> {
152    if let Caller::Client(c) = caller {
153        if !c.can_read(workflow) {
154            return Err(forbidden(format!(
155                "client `{}` may not read the definition of `{workflow}`",
156                c.id
157            )));
158        }
159    }
160    Ok(())
161}
162
163/// `allow` for an existing execution, looked up by id.
164fn allow_execution(state: &Shared, caller: &Caller, action: &str, id: &str) -> ApiResult<()> {
165    if matches!(caller, Caller::Client(_)) {
166        let workflow = engine(state).execution(id)?.workflow;
167        allow(caller, action, &workflow)?;
168    }
169    Ok(())
170}
171
172/// Authentication.
173///
174/// When any credential is configured, every endpoint except `/v1/health`
175/// requires one. The admin token can call anything. Worker credentials (the
176/// shared `--worker-token`, a per-worker token from the workers file, or a
177/// client certificate under mutual TLS) can only call the worker protocol
178/// endpoints, so a compromised worker host cannot start, cancel or read
179/// executions. Per-worker identities are further limited by the handlers to
180/// their own id, tasks, sessions and task tokens.
181/// Establish the caller from a bearer token and/or a verified client
182/// certificate. Shared by the HTTP API and the gRPC binding. `Ok(None)` means
183/// no credential was presented while authentication is required.
184pub(crate) fn identify(
185    state: &AppState,
186    bearer: Option<&str>,
187    cert_der: Option<&[u8]>,
188) -> Result<Option<Caller>, ApiError> {
189    let secured = state.auth_token.is_some() || !state.identities.is_empty();
190    if !secured {
191        return Ok(Some(Caller::Open));
192    }
193    if let Some(t) = bearer {
194        if state
195            .auth_token
196            .as_ref()
197            .is_some_and(|a| constant_time_eq(t.as_bytes(), a.as_bytes()))
198        {
199            return Ok(Some(Caller::Admin));
200        }
201        if let Some(w) = state.identities.by_token(t) {
202            return Ok(Some(Caller::Worker(Arc::new(w.clone()))));
203        }
204        if let Some(c) = state.identities.client_by_token(t) {
205            return Ok(Some(Caller::Client(Arc::new(c.clone()))));
206        }
207        if state
208            .worker_token
209            .as_ref()
210            .is_some_and(|w| constant_time_eq(t.as_bytes(), w.as_bytes()))
211        {
212            return Ok(Some(Caller::SharedWorker));
213        }
214    }
215    if let Some(der) = cert_der {
216        if let Some(w) = state.identities.by_certificate(der) {
217            return Ok(Some(Caller::Worker(Arc::new(w.clone()))));
218        }
219        return match state.identities.client_by_certificate(der) {
220            Some(c) => Ok(Some(Caller::Client(Arc::new(c.clone())))),
221            None => Err(forbidden(
222                "the client certificate is valid but matches no worker or client in the workers file",
223            )),
224        };
225    }
226    Ok(None)
227}
228
229async fn auth(State(state): State<Shared>, mut req: Request, next: Next) -> Response {
230    let path = req.uri().path().to_string();
231    let bearer = req
232        .headers()
233        .get("authorization")
234        .and_then(|v| v.to_str().ok())
235        .and_then(|v| v.strip_prefix("Bearer "))
236        .map(str::to_string);
237    let cert = req
238        .extensions()
239        .get::<Arc<axum_server::tls::TlsInfo>>()
240        .and_then(|info| info.peer_certificates.first().cloned());
241    let caller = match identify(&state, bearer.as_deref(), cert.as_deref()) {
242        Ok(Some(c)) => c,
243        Ok(None) if path == "/v1/health" => Caller::Open,
244        Ok(None) => {
245            return ApiError(EngineError::new(
246                401,
247                "UNAUTHENTICATED",
248                "missing or invalid credentials",
249            ))
250            .into_response()
251        }
252        Err(e) => return e.into_response(),
253    };
254    if matches!(caller, Caller::SharedWorker | Caller::Worker(_))
255        && !is_worker_path(&path)
256        && path != "/v1/health"
257    {
258        return forbidden("worker credentials may only call /v1/workers/* and /v1/tasks/*")
259            .into_response();
260    }
261    if let Caller::Client(c) = &caller {
262        let admin_only = (path == "/v1/workflows" && req.method() == axum::http::Method::POST)
263            || path.starts_with("/v1/workers")
264            || path == "/v1/metrics";
265        if is_worker_path(&path) || admin_only {
266            return forbidden(format!(
267                "client `{}` may only use the execution API (start, signal, query, read, cancel) and schedules",
268                c.id
269            ))
270            .into_response();
271        }
272    }
273    req.extensions_mut().insert(caller);
274    let mut resp = next.run(req).await;
275    resp.headers_mut()
276        .insert("x-orbit-api-version", HeaderValue::from_static(API_VERSION));
277    resp
278}
279
280/// A per-worker identity may only use sessions it opened.
281pub(crate) fn authorize_session(state: &Shared, caller: &Caller, session: &str) -> ApiResult<()> {
282    if let Some(id) = caller.identity() {
283        if let Some(owner) = engine(state).session_worker_lookup(session) {
284            if owner != id.id {
285                return Err(forbidden(format!(
286                    "session belongs to worker `{owner}`, not `{}`",
287                    id.id
288                )));
289            }
290        }
291    }
292    Ok(())
293}
294
295/// A per-worker identity may only report results for tokens issued to it.
296pub(crate) fn authorize_token(state: &Shared, caller: &Caller, token: &str) -> ApiResult<()> {
297    if let Some(id) = caller.identity() {
298        if let Some(owner) = engine(state).token_worker_lookup(token) {
299            if owner != id.id {
300                return Err(forbidden(format!(
301                    "task token was issued to worker `{owner}`, not `{}`",
302                    id.id
303                )));
304            }
305        }
306    }
307    Ok(())
308}
309
310fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
311    a.len() == b.len() && a.iter().zip(b).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
312}
313
314async fn ready() -> Json<Value> {
315    Json(json!({ "role": "single", "ready": true }))
316}
317
318async fn health(State(state): State<Shared>) -> Json<Value> {
319    let e = engine(&state);
320    Json(json!({
321        "status": "ok",
322        "api_version": API_VERSION,
323        "runtime_version": env!("CARGO_PKG_VERSION"),
324        "ir_versions": orbit_ir::SUPPORTED_IR_VERSIONS,
325        "active_executions": e.active_count(),
326        "worker_sessions": e.session_count(),
327    }))
328}
329
330async fn metrics(State(state): State<Shared>) -> Response {
331    let e = engine(&state);
332    let body = e.metrics.render(e.active_count(), e.session_count());
333    let mut headers = HeaderMap::new();
334    headers.insert(
335        "content-type",
336        HeaderValue::from_static("text/plain; version=0.0.4"),
337    );
338    (headers, body).into_response()
339}
340
341async fn register_workflow(State(state): State<Shared>, body: String) -> ApiResult<Response> {
342    let ir = orbit_ir::WorkflowIr::from_json(&body)
343        .map_err(|e| ApiError(EngineError::new(400, "INVALID_IR", e.0)))?;
344    let (name, version, hash) = (
345        ir.name.clone(),
346        ir.version.clone(),
347        ir.definition_hash.clone(),
348    );
349    let created = engine(&state).register_definition(ir)?;
350    let status = if created {
351        StatusCode::CREATED
352    } else {
353        StatusCode::OK
354    };
355    Ok((status, Json(json!({ "name": name, "version": version, "definition_hash": hash, "created": created }))).into_response())
356}
357
358fn summary(d: &orbit_ir::WorkflowIr) -> Value {
359    json!({
360        "name": d.name,
361        "version": d.version,
362        "description": d.description,
363        "definition_hash": d.definition_hash,
364        "steps": d.steps.iter().map(|s| s.id.clone()).collect::<Vec<_>>(),
365        "tasks": d.contracts.keys().collect::<Vec<_>>(),
366    })
367}
368
369async fn list_workflows(
370    State(state): State<Shared>,
371    Extension(caller): Extension<Caller>,
372) -> Json<Value> {
373    let defs = engine(&state).definitions();
374    let visible = |name: &str| match &caller {
375        Caller::Client(c) => c.can_read(name),
376        _ => true,
377    };
378    Json(
379        json!({ "workflows": defs.iter().filter(|d| visible(&d.name)).map(|d| summary(d)).collect::<Vec<_>>() }),
380    )
381}
382
383async fn get_workflow(
384    State(state): State<Shared>,
385    Extension(caller): Extension<Caller>,
386    Path((name, version)): Path<(String, String)>,
387) -> ApiResult<Json<Value>> {
388    allow_definition(&caller, &name)?;
389    let d = engine(&state)
390        .definition(&name, Some(&version))
391        .ok_or_else(|| {
392            ApiError(EngineError::new(
393                404,
394                "WORKFLOW_NOT_FOUND",
395                "workflow not found",
396            ))
397        })?;
398    Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
399}
400
401async fn get_workflow_latest(
402    State(state): State<Shared>,
403    Extension(caller): Extension<Caller>,
404    Path(name): Path<String>,
405) -> ApiResult<Json<Value>> {
406    allow_definition(&caller, &name)?;
407    let d = engine(&state).definition(&name, None).ok_or_else(|| {
408        ApiError(EngineError::new(
409            404,
410            "WORKFLOW_NOT_FOUND",
411            "workflow not found",
412        ))
413    })?;
414    Ok(Json(serde_json::to_value(&*d).unwrap_or_default()))
415}
416
417#[derive(Deserialize)]
418struct StartRequest {
419    workflow: String,
420    #[serde(default)]
421    version: Option<String>,
422    #[serde(default)]
423    input: Value,
424    #[serde(default)]
425    execution_id: Option<String>,
426}
427
428async fn start_execution(
429    State(state): State<Shared>,
430    Extension(caller): Extension<Caller>,
431    Json(req): Json<StartRequest>,
432) -> ApiResult<Response> {
433    allow(&caller, "start", &req.workflow)?;
434    let r = engine(&state).start(
435        &req.workflow,
436        req.version.as_deref(),
437        req.input,
438        req.execution_id,
439    )?;
440    let status = if r.created {
441        StatusCode::CREATED
442    } else {
443        StatusCode::OK
444    };
445    Ok((status, Json(serde_json::to_value(r).unwrap_or_default())).into_response())
446}
447
448#[derive(Deserialize)]
449struct ListQuery {
450    status: Option<String>,
451    workflow: Option<String>,
452    limit: Option<u32>,
453}
454
455async fn list_executions(
456    State(state): State<Shared>,
457    Extension(caller): Extension<Caller>,
458    Query(q): Query<ListQuery>,
459) -> ApiResult<Json<Value>> {
460    let limit = q.limit.unwrap_or(50).clamp(1, 1000);
461    // A client sees only the workflows it may read: look further back so the
462    // page still holds `limit` executions it is allowed to see.
463    let fetch = if matches!(caller, Caller::Client(_)) {
464        1000
465    } else {
466        limit
467    };
468    let mut list = engine(&state).list_executions(
469        q.status.as_deref().map(str::to_uppercase).as_deref(),
470        q.workflow.as_deref(),
471        fetch,
472    )?;
473    if let Caller::Client(c) = &caller {
474        list.retain(|e| c.can_read(&e.workflow));
475        list.truncate(limit as usize);
476    }
477    Ok(Json(json!({ "executions": list })))
478}
479
480async fn get_execution(
481    State(state): State<Shared>,
482    Extension(caller): Extension<Caller>,
483    Path(id): Path<String>,
484) -> ApiResult<Json<Value>> {
485    allow_execution(&state, &caller, "read", &id)?;
486    let s = engine(&state).execution(&id)?;
487    Ok(Json(serde_json::to_value(s).unwrap_or_default()))
488}
489
490#[derive(Deserialize)]
491struct HistoryQuery {
492    run: Option<u32>,
493}
494
495/// `GET /v1/executions/{id}/history[?run=N]`: the latest run's history, or
496/// run `N`'s (continue-as-new).
497async fn get_history(
498    State(state): State<Shared>,
499    Extension(caller): Extension<Caller>,
500    Path(id): Path<String>,
501    Query(q): Query<HistoryQuery>,
502) -> ApiResult<Json<Value>> {
503    allow_execution(&state, &caller, "read", &id)?;
504    let e = engine(&state);
505    let h = match q.run {
506        Some(run) => e.run_history(&id, run)?,
507        None => e.history(&id)?,
508    };
509    let run = crate::model::ExecutionState::replay(&id, &h[..h.len().min(1)]).run;
510    let mut body = json!({ "execution_id": id, "events": h });
511    if run > 1 || q.run.is_some() {
512        body["run"] = json!(run);
513    }
514    Ok(Json(body))
515}
516
517// --------------------------------------------------------------- schedules
518
519/// A client identity needs the `schedule` action for the schedule's workflow.
520fn allow_schedule(caller: &Caller, workflow: &str) -> ApiResult<()> {
521    if let Caller::Client(c) = caller {
522        if !c.allows("schedule", workflow) {
523            return Err(forbidden(format!(
524                "client `{}` may not manage schedules of `{workflow}`",
525                c.id
526            )));
527        }
528    }
529    Ok(())
530}
531
532fn allow_schedule_id(state: &Shared, caller: &Caller, id: &str) -> ApiResult<()> {
533    if matches!(caller, Caller::Client(_)) {
534        let workflow = engine(state).schedule_workflow(id)?;
535        allow_schedule(caller, &workflow)?;
536    }
537    Ok(())
538}
539
540async fn create_schedule(
541    State(state): State<Shared>,
542    Extension(caller): Extension<Caller>,
543    Json(req): Json<crate::schedule::CreateSchedule>,
544) -> ApiResult<Response> {
545    allow_schedule(&caller, req.spec["workflow"].as_str().unwrap_or(""))?;
546    let v = engine(&state).create_schedule(req)?;
547    Ok((StatusCode::CREATED, Json(v)).into_response())
548}
549
550async fn list_schedules(
551    State(state): State<Shared>,
552    Extension(caller): Extension<Caller>,
553) -> ApiResult<Json<Value>> {
554    let mut list = engine(&state).list_schedules()?;
555    if let Caller::Client(c) = &caller {
556        list.retain(|s| c.allows("schedule", s["workflow"].as_str().unwrap_or("")));
557    }
558    Ok(Json(json!({ "schedules": list })))
559}
560
561async fn get_schedule(
562    State(state): State<Shared>,
563    Extension(caller): Extension<Caller>,
564    Path(id): Path<String>,
565) -> ApiResult<Json<Value>> {
566    allow_schedule_id(&state, &caller, &id)?;
567    Ok(Json(engine(&state).get_schedule(&id)?))
568}
569
570async fn update_schedule(
571    State(state): State<Shared>,
572    Extension(caller): Extension<Caller>,
573    Path(id): Path<String>,
574    Json(patch): Json<Value>,
575) -> ApiResult<Json<Value>> {
576    allow_schedule_id(&state, &caller, &id)?;
577    if let Some(w) = patch.get("workflow").and_then(Value::as_str) {
578        allow_schedule(&caller, w)?;
579    }
580    Ok(Json(engine(&state).update_schedule(&id, patch)?))
581}
582
583async fn delete_schedule(
584    State(state): State<Shared>,
585    Extension(caller): Extension<Caller>,
586    Path(id): Path<String>,
587) -> ApiResult<Json<Value>> {
588    allow_schedule_id(&state, &caller, &id)?;
589    engine(&state).delete_schedule(&id)?;
590    Ok(Json(json!({ "schedule_id": id, "deleted": true })))
591}
592
593async fn pause_schedule(
594    State(state): State<Shared>,
595    Extension(caller): Extension<Caller>,
596    Path(id): Path<String>,
597) -> ApiResult<Json<Value>> {
598    allow_schedule_id(&state, &caller, &id)?;
599    Ok(Json(engine(&state).set_schedule_paused(&id, true)?))
600}
601
602async fn resume_schedule(
603    State(state): State<Shared>,
604    Extension(caller): Extension<Caller>,
605    Path(id): Path<String>,
606) -> ApiResult<Json<Value>> {
607    allow_schedule_id(&state, &caller, &id)?;
608    Ok(Json(engine(&state).set_schedule_paused(&id, false)?))
609}
610
611/// `POST /v1/schedules/{id}/trigger`: start an execution now (the overlap
612/// policy applies; the timing is not changed).
613async fn trigger_schedule(
614    State(state): State<Shared>,
615    Extension(caller): Extension<Caller>,
616    Path(id): Path<String>,
617) -> ApiResult<Json<Value>> {
618    allow_schedule_id(&state, &caller, &id)?;
619    let t = engine(&state).trigger_schedule(&id)?;
620    Ok(Json(serde_json::to_value(t).unwrap_or_default()))
621}
622
623#[derive(Deserialize, Default)]
624struct CancelRequest {
625    #[serde(default)]
626    reason: Option<String>,
627}
628
629async fn cancel_execution(
630    State(state): State<Shared>,
631    Extension(caller): Extension<Caller>,
632    Path(id): Path<String>,
633    body: Option<Json<CancelRequest>>,
634) -> ApiResult<Json<Value>> {
635    allow_execution(&state, &caller, "cancel", &id)?;
636    let reason = body
637        .and_then(|b| b.0.reason)
638        .unwrap_or_else(|| "requested via API".into());
639    let s = engine(&state).cancel(&id, &reason)?;
640    Ok(Json(serde_json::to_value(s).unwrap_or_default()))
641}
642
643/// `POST /v1/executions/{id}/signals/{name}` with the payload as the body
644/// (an empty body sends `{}`).
645async fn send_signal(
646    State(state): State<Shared>,
647    Extension(caller): Extension<Caller>,
648    Path((id, name)): Path<(String, String)>,
649    body: Option<Json<Value>>,
650) -> ApiResult<Json<Value>> {
651    allow_execution(&state, &caller, "signal", &id)?;
652    let payload = body.map(|b| b.0).unwrap_or_else(|| json!({}));
653    let seq = engine(&state).signal(&id, &name, payload)?;
654    Ok(Json(json!({
655        "execution_id": id,
656        "signal": name,
657        "accepted": true,
658        "seq": seq,
659    })))
660}
661
662/// `GET /v1/executions/{id}/queries/{name}`: evaluate a declared query.
663async fn run_query(
664    State(state): State<Shared>,
665    Extension(caller): Extension<Caller>,
666    Path((id, name)): Path<(String, String)>,
667) -> ApiResult<Json<Value>> {
668    allow_execution(&state, &caller, "query", &id)?;
669    let value = engine(&state).query(&id, &name)?;
670    Ok(Json(
671        json!({ "execution_id": id, "query": name, "value": value }),
672    ))
673}
674
675#[derive(Deserialize)]
676struct WaitQuery {
677    timeout_ms: Option<u64>,
678}
679
680/// Long-poll until the execution reaches a terminal status (or the timeout).
681async fn wait_execution(
682    State(state): State<Shared>,
683    Extension(caller): Extension<Caller>,
684    Path(id): Path<String>,
685    Query(q): Query<WaitQuery>,
686) -> ApiResult<Json<Value>> {
687    allow_execution(&state, &caller, "read", &id)?;
688    let deadline = tokio::time::Instant::now()
689        + Duration::from_millis(q.timeout_ms.unwrap_or(30_000).min(300_000));
690    loop {
691        let notified = state.changed.notified();
692        tokio::pin!(notified);
693        notified.as_mut().enable();
694        let s = engine(&state).execution(&id)?;
695        if s.status.is_terminal() || tokio::time::Instant::now() >= deadline {
696            return Ok(Json(serde_json::to_value(s).unwrap_or_default()));
697        }
698        let _ = tokio::time::timeout_at(deadline, notified).await;
699    }
700}
701
702async fn list_workers(State(state): State<Shared>) -> ApiResult<Json<Value>> {
703    Ok(Json(json!({ "workers": engine(&state).workers()? })))
704}
705
706async fn register_worker(
707    State(state): State<Shared>,
708    Extension(caller): Extension<Caller>,
709    Json(req): Json<RegisterWorker>,
710) -> ApiResult<Json<Value>> {
711    if let Some(id) = caller.identity() {
712        if req.worker_id != id.id {
713            return Err(forbidden(format!(
714                "this credential belongs to worker `{}`, not `{}`",
715                id.id, req.worker_id
716            )));
717        }
718        if let Some(rt) = &id.runtime {
719            if !rt.eq_ignore_ascii_case(&req.runtime) {
720                return Err(forbidden(format!(
721                    "worker `{}` must announce runtime `{rt}`",
722                    id.id
723                )));
724            }
725        }
726        let denied: Vec<&str> = req
727            .contracts
728            .iter()
729            .map(|c| c.task_name.as_str())
730            .filter(|t| !id.may_serve(t))
731            .collect();
732        if !denied.is_empty() {
733            return Err(forbidden(format!(
734                "worker `{}` may not serve: {} (allowed: {})",
735                id.id,
736                denied.join(", "),
737                id.tasks.join(", ")
738            )));
739        }
740    }
741    Ok(Json(
742        serde_json::to_value(engine(&state).register_worker(req)?).unwrap_or_default(),
743    ))
744}
745
746#[derive(Deserialize)]
747struct SessionRequest {
748    worker_session_id: String,
749    #[serde(default)]
750    wait_ms: Option<u64>,
751}
752
753async fn heartbeat(
754    State(state): State<Shared>,
755    Extension(caller): Extension<Caller>,
756    Json(req): Json<SessionRequest>,
757) -> ApiResult<Json<Value>> {
758    authorize_session(&state, &caller, &req.worker_session_id)?;
759    Ok(Json(
760        serde_json::to_value(engine(&state).heartbeat(&req.worker_session_id)?).unwrap_or_default(),
761    ))
762}
763
764async fn unregister(
765    State(state): State<Shared>,
766    Extension(caller): Extension<Caller>,
767    Json(req): Json<SessionRequest>,
768) -> ApiResult<Json<Value>> {
769    authorize_session(&state, &caller, &req.worker_session_id)?;
770    engine(&state).unregister_worker(&req.worker_session_id)?;
771    Ok(Json(json!({ "accepted": true })))
772}
773
774/// Long-poll for a task. Returns 204 when nothing arrived within `wait_ms`.
775async fn poll(
776    State(state): State<Shared>,
777    Extension(caller): Extension<Caller>,
778    Json(req): Json<SessionRequest>,
779) -> ApiResult<Response> {
780    authorize_session(&state, &caller, &req.worker_session_id)?;
781    let wait = Duration::from_millis(req.wait_ms.unwrap_or(20_000).min(60_000));
782    let deadline = tokio::time::Instant::now() + wait;
783    loop {
784        let notified = state.changed.notified();
785        tokio::pin!(notified);
786        notified.as_mut().enable();
787        let task: Option<TaskAssignment> = crate::shared::poll(&state, &req.worker_session_id)?;
788        if let Some(t) = task {
789            return Ok(Json(serde_json::to_value(t).unwrap_or_default()).into_response());
790        }
791        if tokio::time::Instant::now() >= deadline {
792            return Ok(StatusCode::NO_CONTENT.into_response());
793        }
794        let _ = tokio::time::timeout_at(deadline, notified).await;
795    }
796}
797
798#[derive(Deserialize)]
799struct CompleteRequest {
800    token: String,
801    #[serde(default)]
802    output: Value,
803}
804
805async fn complete_task(
806    State(state): State<Shared>,
807    Extension(caller): Extension<Caller>,
808    Json(req): Json<CompleteRequest>,
809) -> ApiResult<Json<Value>> {
810    if req.token.is_empty() {
811        return Err(bad_request("token is required"));
812    }
813    authorize_token(&state, &caller, &req.token)?;
814    Ok(Json(
815        serde_json::to_value(engine(&state).complete(&req.token, req.output)?).unwrap_or_default(),
816    ))
817}
818
819#[derive(Deserialize)]
820struct FailRequest {
821    token: String,
822    error_code: String,
823    #[serde(default)]
824    error_message: String,
825    #[serde(default = "default_true")]
826    retryable: bool,
827    #[serde(default)]
828    details: Option<Value>,
829}
830
831fn default_true() -> bool {
832    true
833}
834
835async fn fail_task(
836    State(state): State<Shared>,
837    Extension(caller): Extension<Caller>,
838    Json(req): Json<FailRequest>,
839) -> ApiResult<Json<Value>> {
840    if req.error_code.is_empty() {
841        return Err(bad_request("error_code is required"));
842    }
843    authorize_token(&state, &caller, &req.token)?;
844    let error = TaskError {
845        code: req.error_code,
846        message: req.error_message,
847        retryable: req.retryable,
848        details: req.details,
849    };
850    Ok(Json(
851        serde_json::to_value(engine(&state).fail(&req.token, error)?).unwrap_or_default(),
852    ))
853}
854
855/// Background loop: timers, retries, deadlines, liveness and built-in tasks.
856pub async fn run_ticker(state: Shared, interval: Duration) {
857    let mut ticker = tokio::time::interval(interval);
858    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
859    loop {
860        ticker.tick().await;
861        let jobs = match engine(&state).tick() {
862            Ok(j) => j,
863            Err(e) => {
864                tracing::error!(error = %e, "tick failed");
865                continue;
866            }
867        };
868        for job in jobs {
869            let st = state.clone();
870            tokio::spawn(async move {
871                let input = job.input.clone();
872                let timeout = job.timeout_ms;
873                let result = tokio::task::spawn_blocking(move || {
874                    crate::http_adapter::execute(&input, timeout)
875                })
876                .await
877                .unwrap_or_else(|e| Err(TaskError::new("INTERNAL", e.to_string(), true)));
878                let mut e = engine(&st);
879                let r = match result {
880                    Ok(out) => e.complete(&job.token, out),
881                    Err(err) => e.fail(&job.token, err),
882                };
883                if let Err(err) = r {
884                    tracing::error!(error = %err, "recording built-in task result failed");
885                }
886            });
887        }
888    }
889}
890
891/// PEM files for serving HTTPS directly.
892pub struct TlsConfig {
893    pub cert: std::path::PathBuf,
894    pub key: std::path::PathBuf,
895    /// Enables mutual TLS: client certificates signed by this CA are verified
896    /// and identify workers. Clients without a certificate can still connect
897    /// and authenticate with a token.
898    pub client_ca: Option<std::path::PathBuf>,
899}
900
901fn io_err(msg: String) -> std::io::Error {
902    std::io::Error::new(std::io::ErrorKind::InvalidInput, msg)
903}
904
905fn read_pem_certs(
906    path: &std::path::Path,
907) -> std::io::Result<Vec<rustls::pki_types::CertificateDer<'static>>> {
908    let pem =
909        std::fs::read(path).map_err(|e| io_err(format!("cannot read {}: {e}", path.display())))?;
910    let certs = rustls_pemfile::certs(&mut pem.as_slice())
911        .collect::<Result<Vec<_>, _>>()
912        .map_err(|e| io_err(format!("{}: {e}", path.display())))?;
913    if certs.is_empty() {
914        return Err(io_err(format!(
915            "{} contains no certificates",
916            path.display()
917        )));
918    }
919    Ok(certs)
920}
921
922/// Build the rustls server configuration (with client verification for mTLS).
923fn rustls_config(tls: &TlsConfig) -> std::io::Result<rustls::ServerConfig> {
924    let chain = read_pem_certs(&tls.cert)?;
925    let key_pem = std::fs::read(&tls.key)
926        .map_err(|e| io_err(format!("cannot read {}: {e}", tls.key.display())))?;
927    let key = rustls_pemfile::private_key(&mut key_pem.as_slice())
928        .map_err(|e| io_err(format!("{}: {e}", tls.key.display())))?
929        .ok_or_else(|| io_err(format!("{} contains no private key", tls.key.display())))?;
930    let provider = Arc::new(rustls::crypto::aws_lc_rs::default_provider());
931    let builder = rustls::ServerConfig::builder_with_provider(provider.clone())
932        .with_safe_default_protocol_versions()
933        .map_err(|e| io_err(e.to_string()))?;
934    let builder = match &tls.client_ca {
935        Some(ca) => {
936            let mut roots = rustls::RootCertStore::empty();
937            for c in read_pem_certs(ca)? {
938                roots
939                    .add(c)
940                    .map_err(|e| io_err(format!("{}: {e}", ca.display())))?;
941            }
942            let verifier = rustls::server::WebPkiClientVerifier::builder_with_provider(
943                Arc::new(roots),
944                provider,
945            )
946            .allow_unauthenticated()
947            .build()
948            .map_err(|e| io_err(format!("client CA: {e}")))?;
949            builder.with_client_cert_verifier(verifier)
950        }
951        None => builder.with_no_client_auth(),
952    };
953    let mut config = builder
954        .with_single_cert(chain, key)
955        .map_err(|e| io_err(format!("TLS certificate/key: {e}")))?;
956    config.alpn_protocols = vec![b"h2".to_vec(), b"http/1.1".to_vec()];
957    Ok(config)
958}
959
960pub struct ServerConfig {
961    pub addr: SocketAddr,
962    pub auth_token: Option<String>,
963    pub worker_token: Option<String>,
964    pub identities: Identities,
965    pub tls: Option<TlsConfig>,
966    /// Also serve the gRPC binding of the worker protocol on this address.
967    pub grpc_addr: Option<SocketAddr>,
968    /// Pre-bound listeners (socket activation, tests): used instead of
969    /// binding `addr` / `grpc_addr`.
970    pub listener: Option<std::net::TcpListener>,
971    pub grpc_listener: Option<std::net::TcpListener>,
972    pub body_limit: usize,
973    pub tick_interval: Duration,
974    /// Serve the web UI at `/ui/` ([`crate::ui`]); `orbit server --no-ui` turns it off.
975    pub ui: bool,
976}
977
978pub async fn serve(
979    engine: Engine,
980    config: ServerConfig,
981    shutdown: impl std::future::Future<Output = ()> + Send + 'static,
982) -> std::io::Result<()> {
983    serve_node(engine, config, None, shutdown).await
984}
985
986/// [`serve`] as a member of a cluster (`orbit server --cluster`): `engine`
987/// is the node's standby engine ([`crate::cluster::Node::standby_engine`]);
988/// the node's election thread swaps in a leader engine once elected.
989pub async fn serve_node(
990    engine: Engine,
991    config: ServerConfig,
992    cluster: Option<Arc<crate::cluster::Node>>,
993    shutdown: impl std::future::Future<Output = ()> + Send + 'static,
994) -> std::io::Result<()> {
995    // Several rustls crypto backends are compiled in (via different
996    // dependencies); choose one before any TLS configuration is built.
997    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
998    let changed = engine.changed.clone();
999    // Shared mode (`--cluster-mode shared`): the engine carries its context.
1000    let shared = engine.shared_context();
1001    let state = Arc::new(AppState {
1002        engine: Mutex::new(engine),
1003        changed,
1004        auth_token: config.auth_token,
1005        worker_token: config.worker_token,
1006        identities: config.identities,
1007    });
1008    let app = router(state.clone(), config.body_limit);
1009    let app = match &shared {
1010        Some(ctx) => crate::shared::layer(app, ctx.clone(), state.clone()),
1011        None => crate::cluster::layer(app, cluster.clone(), state.clone()),
1012    };
1013    // Outside the authentication and cluster layers: the static assets need
1014    // no credential and every node serves them.
1015    let app = if config.ui {
1016        app.merge(crate::ui::router())
1017    } else {
1018        app
1019    };
1020    if let Some(node) = &cluster {
1021        crate::cluster::start(node.clone(), state.clone())?;
1022    }
1023    if let Some(ctx) = &shared {
1024        crate::shared::start(ctx.clone(), state.clone())?;
1025    }
1026    let ticker = tokio::spawn(run_ticker(state.clone(), config.tick_interval));
1027    // One shutdown signal fanned out to every listener.
1028    let (stop_tx, stop_rx) = tokio::sync::watch::channel(false);
1029    tokio::spawn(async move {
1030        shutdown.await;
1031        let _ = stop_tx.send(true);
1032    });
1033    let stopped = |mut rx: tokio::sync::watch::Receiver<bool>| async move {
1034        let _ = rx.wait_for(|v| *v).await;
1035    };
1036    let bind = |addr: SocketAddr| {
1037        std::net::TcpListener::bind(addr)
1038            .map_err(|e| std::io::Error::new(e.kind(), format!("cannot listen on {addr}: {e}")))
1039    };
1040    let grpc_listener = match (config.grpc_listener, config.grpc_addr) {
1041        (Some(l), _) => Some(l),
1042        (None, Some(addr)) => Some(bind(addr)?),
1043        (None, None) => None,
1044    };
1045    let http_listener = match config.listener {
1046        Some(l) => l,
1047        None => bind(config.addr)?,
1048    };
1049    http_listener.set_nonblocking(true)?;
1050    let http_addr = http_listener.local_addr()?;
1051    let grpc = match grpc_listener {
1052        Some(std_listener) => {
1053            std_listener.set_nonblocking(true)?;
1054            let addr = std_listener.local_addr()?;
1055            let listener = tokio::net::TcpListener::from_std(std_listener)?;
1056            let incoming =
1057                tonic::transport::server::TcpIncoming::from_listener(listener, true, None)
1058                    .map_err(|e| io_err(format!("gRPC listener: {e}")))?;
1059            let mut builder = tonic::transport::Server::builder();
1060            if let Some(tls) = &config.tls {
1061                let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1062                let read = |p: &std::path::Path| {
1063                    std::fs::read(p)
1064                        .map_err(|e| io_err(format!("cannot read {}: {e}", p.display())))
1065                };
1066                let mut t = tonic::transport::ServerTlsConfig::new().identity(
1067                    tonic::transport::Identity::from_pem(read(&tls.cert)?, read(&tls.key)?),
1068                );
1069                if let Some(ca) = &tls.client_ca {
1070                    t = t
1071                        .client_ca_root(tonic::transport::Certificate::from_pem(read(ca)?))
1072                        .client_auth_optional(true);
1073                }
1074                builder = builder
1075                    .tls_config(t)
1076                    .map_err(|e| io_err(format!("gRPC TLS: {e}")))?;
1077            }
1078            let service = crate::grpc::GrpcWorkers::new(state.clone());
1079            let rx = stop_rx.clone();
1080            tracing::info!(addr = %addr, tls = config.tls.is_some(), "gRPC worker protocol listening");
1081            let router = match &cluster {
1082                Some(node) => {
1083                    builder.add_service(tonic::service::interceptor::InterceptedService::new(
1084                        service,
1085                        crate::cluster::grpc_gate(node.clone()),
1086                    ))
1087                }
1088                None => builder.add_service(service),
1089            };
1090            Some(tokio::spawn(async move {
1091                if let Err(e) = router
1092                    .serve_with_incoming_shutdown(incoming, stopped(rx))
1093                    .await
1094                {
1095                    tracing::error!(error = %e, "gRPC server failed");
1096                }
1097            }))
1098        }
1099        None => None,
1100    };
1101    let shutdown = stopped(stop_rx);
1102    let result = match config.tls {
1103        None => {
1104            let listener = tokio::net::TcpListener::from_std(http_listener)?;
1105            tracing::info!(addr = %http_addr, tls = false, ui = config.ui, "ORBIT runtime listening");
1106            axum::serve(listener, app)
1107                .with_graceful_shutdown(shutdown)
1108                .await
1109        }
1110        Some(tls) => {
1111            // Several rustls crypto backends are compiled in (via different
1112            // dependencies); pick one explicitly so the choice is not ambiguous.
1113            let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
1114            let rustls =
1115                axum_server::tls_rustls::RustlsConfig::from_config(Arc::new(rustls_config(&tls)?));
1116            let handle = axum_server::Handle::new();
1117            let h = handle.clone();
1118            tokio::spawn(async move {
1119                shutdown.await;
1120                h.graceful_shutdown(Some(Duration::from_secs(10)));
1121            });
1122            tracing::info!(addr = %http_addr, tls = true, mtls = tls.client_ca.is_some(), ui = config.ui, "ORBIT runtime listening");
1123            axum_server::from_tcp_rustls(http_listener, rustls)?
1124                .handle(handle)
1125                .serve(app.into_make_service())
1126                .await
1127        }
1128    };
1129    ticker.abort();
1130    if let Some(ctx) = &shared {
1131        ctx.stop();
1132    }
1133    if let Some(g) = grpc {
1134        let _ = g.await;
1135    }
1136    result
1137}