agent-runtime-http-server 0.1.0

//! Axum transport for the Runtime Open API.

use std::{convert::Infallible, sync::Arc, time::Duration};

use async_trait::async_trait;
use axum::{
    Json, Router,
    extract::{DefaultBodyLimit, Path, Query, State},
    http::{HeaderMap, StatusCode, header},
    response::{
        IntoResponse, Response, Sse,
        sse::{Event, KeepAlive},
    },
    routing::{get, post},
};
use futures_util::Stream;
use runtime_api_contract::{
    ApiErrorBody, CreateExecutionRequest, SubmitInputRequest, openapi_document,
};
use runtime_link_api::{LinkError, RuntimeLink};
use runtime_types::{CallerScope, CredentialHandle, ExecutionId, RequestAuthority};
use serde::Deserialize;
use serde_json::json;
use sha2::{Digest, Sha256};

#[derive(Debug)]
pub struct AuthError {
    pub message: String,
}

#[async_trait]
pub trait IngressAuthenticator: Send + Sync + 'static {
    async fn authenticate(&self, headers: &HeaderMap) -> Result<RequestAuthority, AuthError>;
}

/// Safe single-tenant authenticator for local/container deployments.
/// Production multi-tenant hosts should inject a JWT/introspection adapter.
pub struct StaticBearerAuthenticator {
    expected_token: String,
    caller: CallerScope,
}

impl StaticBearerAuthenticator {
    pub fn new(expected_token: impl Into<String>, caller: CallerScope) -> Result<Self, AuthError> {
        let expected_token = expected_token.into();
        if expected_token.is_empty() {
            return Err(AuthError {
                message: "static bearer token must not be empty".into(),
            });
        }
        caller.validate().map_err(|error| AuthError {
            message: error.to_string(),
        })?;
        Ok(Self {
            expected_token,
            caller,
        })
    }
}

#[async_trait]
impl IngressAuthenticator for StaticBearerAuthenticator {
    async fn authenticate(&self, headers: &HeaderMap) -> Result<RequestAuthority, AuthError> {
        let value = headers
            .get(header::AUTHORIZATION)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.strip_prefix("Bearer "))
            .ok_or_else(|| AuthError {
                message: "missing bearer credential".into(),
            })?;
        if value != self.expected_token {
            return Err(AuthError {
                message: "invalid bearer credential".into(),
            });
        }
        Ok(RequestAuthority {
            caller: self.caller.clone(),
            credential: CredentialHandle::new(value).map_err(|error| AuthError {
                message: error.to_string(),
            })?,
        })
    }
}

/// Authenticator for a Runtime listener reachable only through a trusted
/// identity-aware gateway. The gateway owns credential verification and
/// injects the caller scope headers; Runtime forwards the bearer through an
/// opaque credential handle to its scoped Agent Infra client.
pub struct TrustedGatewayAuthenticator {
    shared_secret_digest: [u8; 32],
}

impl TrustedGatewayAuthenticator {
    pub fn new(shared_secret: &str) -> Result<Self, AuthError> {
        if shared_secret.len() < 32 || shared_secret.len() > 4096 {
            return Err(AuthError {
                message: "trusted gateway secret must contain 32..=4096 bytes".into(),
            });
        }
        Ok(Self {
            shared_secret_digest: Sha256::digest(shared_secret.as_bytes()).into(),
        })
    }
}

#[async_trait]
impl IngressAuthenticator for TrustedGatewayAuthenticator {
    async fn authenticate(&self, headers: &HeaderMap) -> Result<RequestAuthority, AuthError> {
        fn required(headers: &HeaderMap, name: &'static str) -> Result<String, AuthError> {
            headers
                .get(name)
                .and_then(|value| value.to_str().ok())
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned)
                .ok_or_else(|| AuthError {
                    message: format!("trusted gateway header {name} is missing"),
                })
        }
        let supplied_secret = required(headers, "x-runtime-gateway-secret")?;
        let supplied_digest: [u8; 32] = Sha256::digest(supplied_secret.as_bytes()).into();
        let matches = self
            .shared_secret_digest
            .iter()
            .zip(supplied_digest)
            .fold(0_u8, |difference, (expected, actual)| {
                difference | (*expected ^ actual)
            })
            == 0;
        if !matches {
            return Err(AuthError {
                message: "trusted gateway authentication failed".into(),
            });
        }
        let token = headers
            .get(header::AUTHORIZATION)
            .and_then(|value| value.to_str().ok())
            .and_then(|value| value.strip_prefix("Bearer "))
            .ok_or_else(|| AuthError {
                message: "trusted gateway bearer credential is missing".into(),
            })?;
        let caller = CallerScope {
            subject: required(headers, "x-runtime-subject")?,
            tenant_id: required(headers, "x-runtime-tenant-id")?,
            project_id: required(headers, "x-runtime-project-id")?,
            capabilities: headers
                .get("x-runtime-capabilities")
                .and_then(|value| value.to_str().ok())
                .unwrap_or_default()
                .split(',')
                .map(str::trim)
                .filter(|value| !value.is_empty())
                .map(ToOwned::to_owned)
                .collect(),
        };
        caller.validate().map_err(|error| AuthError {
            message: error.to_string(),
        })?;
        Ok(RequestAuthority {
            caller,
            credential: CredentialHandle::new(token).map_err(|error| AuthError {
                message: error.to_string(),
            })?,
        })
    }
}

#[async_trait]
pub trait ReadinessProbe: Send + Sync + 'static {
    async fn ready(&self) -> Result<(), String>;
}

pub struct AlwaysReady;
#[async_trait]
impl ReadinessProbe for AlwaysReady {
    async fn ready(&self) -> Result<(), String> {
        Ok(())
    }
}

#[derive(Clone)]
struct AppState {
    link: Arc<dyn RuntimeLink>,
    authenticator: Arc<dyn IngressAuthenticator>,
    readiness: Arc<dyn ReadinessProbe>,
}

pub fn router(
    link: Arc<dyn RuntimeLink>,
    authenticator: Arc<dyn IngressAuthenticator>,
    readiness: Arc<dyn ReadinessProbe>,
) -> Router {
    Router::new()
        .route("/livez", get(live))
        .route("/readyz", get(ready))
        .route("/openapi.json", get(openapi))
        .route("/v1/executions", post(create_execution))
        .route("/v1/executions/{execution_id}", get(get_execution))
        .route("/v1/executions/{execution_id}/events", get(get_events))
        .route("/v1/executions/{execution_id}/stream", get(stream_events))
        .route("/v1/executions/{execution_id}/inputs", post(submit_input))
        .route(
            "/v1/executions/{execution_id}/cancel",
            post(cancel_execution),
        )
        .layer(DefaultBodyLimit::max(1024 * 1024))
        .with_state(AppState {
            link,
            authenticator,
            readiness,
        })
}

async fn live() -> Json<serde_json::Value> {
    Json(json!({"status":"ok","service":"runtime-host"}))
}

async fn openapi() -> Json<serde_json::Value> {
    Json(openapi_document())
}

async fn ready(State(state): State<AppState>) -> Response {
    match state.readiness.ready().await {
        Ok(()) => Json(json!({"status":"ready"})).into_response(),
        Err(message) => (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(json!({"status":"not_ready","message":message})),
        )
            .into_response(),
    }
}

async fn create_execution(
    State(state): State<AppState>,
    headers: HeaderMap,
    Json(request): Json<CreateExecutionRequest>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let key = match headers
        .get("idempotency-key")
        .and_then(|value| value.to_str().ok())
    {
        Some(value) => value,
        None => {
            return api_error(
                StatusCode::BAD_REQUEST,
                "IDEMPOTENCY_KEY_REQUIRED",
                "Idempotency-Key header is required",
            );
        }
    };
    match state.link.create_execution(&authority, key, request).await {
        Ok(value) => (StatusCode::ACCEPTED, Json(value)).into_response(),
        Err(error) => link_error(error),
    }
}

async fn get_execution(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let id = match parse_id(id) {
        Ok(value) => value,
        Err(error) => return invalid_execution_id(error),
    };
    match state.link.execution(&authority, &id).await {
        Ok(value) => Json(value).into_response(),
        Err(error) => link_error(error),
    }
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
struct EventQuery {
    after: Option<u64>,
    limit: Option<usize>,
}

async fn get_events(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
    Query(query): Query<EventQuery>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let id = match parse_id(id) {
        Ok(value) => value,
        Err(error) => return invalid_execution_id(error),
    };
    match state
        .link
        .events(&authority, &id, query.after, query.limit.unwrap_or(100))
        .await
    {
        Ok(value) => Json(value).into_response(),
        Err(error) => link_error(error),
    }
}

async fn stream_events(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
    Query(query): Query<EventQuery>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let id = match parse_id(id) {
        Ok(value) => value,
        Err(error) => return invalid_execution_id(error),
    };
    if let Err(error) = state.link.execution(&authority, &id).await {
        return link_error(error);
    }
    let link = state.link.clone();
    let stream = event_stream(link, authority, id, query.after);
    Sse::new(stream)
        .keep_alive(
            KeepAlive::new()
                .interval(Duration::from_secs(15))
                .text("keepalive"),
        )
        .into_response()
}

fn event_stream(
    link: Arc<dyn RuntimeLink>,
    authority: RequestAuthority,
    id: ExecutionId,
    initial_after: Option<u64>,
) -> impl Stream<Item = Result<Event, Infallible>> {
    async_stream::stream! {
        let mut after = initial_after;
        loop {
            match link.events(&authority, &id, after, 100).await {
                Ok(page) => {
                    let page_was_empty = page.items.is_empty();
                    for item in page.items {
                        after = Some(item.sequence);
                        let event = Event::default().id(item.sequence.to_string()).event("runtime").json_data(&item).unwrap_or_else(|_| Event::default().event("error").data("serialization failure"));
                        yield Ok(event);
                    }
                    if page_was_empty {
                        match link.execution(&authority, &id).await {
                            Ok(view) if view.state.is_terminal() => break,
                            Err(_) => break,
                            _ => {}
                        }
                    }
                }
                Err(error) => { yield Ok(Event::default().event("error").data(error.to_string())); break; }
            }
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }
}

async fn submit_input(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
    Json(request): Json<SubmitInputRequest>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let id = match parse_id(id) {
        Ok(value) => value,
        Err(error) => return invalid_execution_id(error),
    };
    match state.link.submit_input(&authority, &id, request).await {
        Ok(value) => Json(value).into_response(),
        Err(error) => link_error(error),
    }
}

async fn cancel_execution(
    State(state): State<AppState>,
    headers: HeaderMap,
    Path(id): Path<String>,
) -> Response {
    let authority = match authenticate(&state, &headers).await {
        Ok(value) => value,
        Err(response) => return response,
    };
    let id = match parse_id(id) {
        Ok(value) => value,
        Err(error) => return invalid_execution_id(error),
    };
    match state.link.cancel(&authority, &id).await {
        Ok(value) => Json(value).into_response(),
        Err(error) => link_error(error),
    }
}

async fn authenticate(state: &AppState, headers: &HeaderMap) -> Result<RequestAuthority, Response> {
    state
        .authenticator
        .authenticate(headers)
        .await
        .map_err(|_| {
            api_error(
                StatusCode::UNAUTHORIZED,
                "UNAUTHENTICATED",
                "request authentication failed",
            )
        })
}

fn parse_id(value: String) -> Result<ExecutionId, runtime_types::IdentifierError> {
    ExecutionId::new(value)
}

fn invalid_execution_id(error: runtime_types::IdentifierError) -> Response {
    api_error(
        StatusCode::BAD_REQUEST,
        "EXECUTION_ID_INVALID",
        &error.to_string(),
    )
}

fn link_error(error: LinkError) -> Response {
    match error {
        LinkError::Invalid(message) => {
            api_error(StatusCode::BAD_REQUEST, "INVALID_REQUEST", &message)
        }
        LinkError::NotFound => api_error(
            StatusCode::NOT_FOUND,
            "EXECUTION_NOT_FOUND",
            "execution not found",
        ),
        LinkError::Forbidden => api_error(
            StatusCode::FORBIDDEN,
            "FORBIDDEN",
            "caller cannot access this execution",
        ),
        LinkError::Conflict(message) => {
            api_error(StatusCode::CONFLICT, "EXECUTION_CONFLICT", &message)
        }
        LinkError::Overloaded => api_error(
            StatusCode::TOO_MANY_REQUESTS,
            "RUNTIME_OVERLOADED",
            "runtime is at capacity",
        ),
        LinkError::Unavailable(message) => api_error(
            StatusCode::SERVICE_UNAVAILABLE,
            "DEPENDENCY_UNAVAILABLE",
            &message,
        ),
        LinkError::Internal(message) => {
            tracing::error!(%message, "runtime internal error");
            api_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                "INTERNAL",
                "runtime internal failure",
            )
        }
    }
}

fn api_error(status: StatusCode, code: &str, message: &str) -> Response {
    (
        status,
        Json(ApiErrorBody {
            code: code.into(),
            message: message.into(),
            request_id: None,
            details: None,
        }),
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use axum::{body::Body, http::Request};
    use http_body_util::BodyExt;
    use runtime_api_contract::{CreateExecutionResponse, EventPage, ExecutionState, ExecutionView};
    use runtime_types::RuntimeInstanceId;
    use tower::ServiceExt;

    use super::*;

    struct FakeLink;

    fn view() -> ExecutionView {
        ExecutionView {
            id: ExecutionId::new("execution-1").unwrap(),
            runtime_instance_id: RuntimeInstanceId::new("runtime-1").unwrap(),
            conversation_id: None,
            state: ExecutionState::Queued,
            outcome: None,
            failure: None,
            created_at_ms: 1,
            updated_at_ms: 1,
        }
    }

    #[async_trait]
    impl RuntimeLink for FakeLink {
        async fn create_execution(
            &self,
            _: &RequestAuthority,
            _: &str,
            _: CreateExecutionRequest,
        ) -> Result<CreateExecutionResponse, LinkError> {
            Ok(CreateExecutionResponse {
                execution: view(),
                replayed: false,
            })
        }
        async fn execution(
            &self,
            _: &RequestAuthority,
            _: &ExecutionId,
        ) -> Result<ExecutionView, LinkError> {
            Ok(view())
        }
        async fn events(
            &self,
            _: &RequestAuthority,
            _: &ExecutionId,
            _: Option<u64>,
            _: usize,
        ) -> Result<EventPage, LinkError> {
            Ok(EventPage {
                items: Vec::new(),
                next_after: None,
            })
        }
        async fn submit_input(
            &self,
            _: &RequestAuthority,
            _: &ExecutionId,
            _: SubmitInputRequest,
        ) -> Result<ExecutionView, LinkError> {
            Ok(view())
        }
        async fn cancel(
            &self,
            _: &RequestAuthority,
            _: &ExecutionId,
        ) -> Result<ExecutionView, LinkError> {
            Ok(view())
        }
    }

    fn app() -> Router {
        let auth = StaticBearerAuthenticator::new(
            "token",
            CallerScope {
                subject: "user".into(),
                tenant_id: "tenant".into(),
                project_id: "project".into(),
                capabilities: Vec::new(),
            },
        )
        .unwrap();
        router(Arc::new(FakeLink), Arc::new(auth), Arc::new(AlwaysReady))
    }

    fn create_request(authenticated: bool) -> Request<Body> {
        let mut request = Request::post("/v1/executions")
            .header("content-type", "application/json")
            .header("idempotency-key", "key-1");
        if authenticated {
            request = request.header("authorization", "Bearer token");
        }
        request.body(Body::from(r#"{"runtimeInstanceId":"runtime-1","input":{"type":"user_message","text":"hello"},"options":{"deadlineSeconds":60,"maxModelTurns":4,"maxToolCalls":8}}"#)).unwrap()
    }

    #[tokio::test]
    async fn create_requires_authentication_and_returns_contract_body() {
        let unauthorized = app().oneshot(create_request(false)).await.unwrap();
        assert_eq!(unauthorized.status(), StatusCode::UNAUTHORIZED);

        let response = app().oneshot(create_request(true)).await.unwrap();
        assert_eq!(response.status(), StatusCode::ACCEPTED);
        let body = response.into_body().collect().await.unwrap().to_bytes();
        let body: serde_json::Value = serde_json::from_slice(&body).unwrap();
        assert_eq!(body["execution"]["id"], "execution-1");
    }

    #[tokio::test]
    async fn trusted_gateway_scope_requires_the_shared_ingress_secret() {
        let authenticator =
            TrustedGatewayAuthenticator::new("a-runtime-gateway-secret-with-at-least-32-bytes")
                .unwrap();
        let mut headers = HeaderMap::new();
        headers.insert("authorization", "Bearer delegated-token".parse().unwrap());
        headers.insert("x-runtime-subject", "user".parse().unwrap());
        headers.insert("x-runtime-tenant-id", "tenant".parse().unwrap());
        headers.insert("x-runtime-project-id", "project".parse().unwrap());
        assert!(authenticator.authenticate(&headers).await.is_err());
        headers.insert(
            "x-runtime-gateway-secret",
            "a-runtime-gateway-secret-with-at-least-32-bytes"
                .parse()
                .unwrap(),
        );
        let authority = authenticator.authenticate(&headers).await.unwrap();
        assert_eq!(authority.caller.subject, "user");
    }
}