klieo-workflow-api 3.2.0

//! Shared fixtures for the run-service integration tests. Ollama-independent:
//! a stub LLM, in-process bus + memory, and a bearer authenticator that
//! grants both run + read scopes.

#![allow(dead_code)] // shared across test files; not every file uses every item

use std::collections::HashSet;
use std::sync::Arc;
use std::time::Duration;

use axum::body::Body;
use axum::http::{Request, Response};
use http_body_util::BodyExt;
use klieo::App;
use klieo_auth_common::{Authenticator, BearerTokenAuthenticator, Identity};
use klieo_bus_memory::MemoryBus;
use klieo_core::agent::{AgentContext, AgentEvent};
use klieo_core::test_utils::{
    FakeToolInvoker, InMemoryEpisodic, InMemoryLongTerm, InMemoryShortTerm,
};
use klieo_core::MemoryHandles;
use klieo_flows::{Flow, FlowError};
use klieo_provenance::{ProvenanceRepository, SqliteProvenanceRepository};
use klieo_workflow::test_support::stub_llm;
use klieo_workflow::Registry;
use klieo_workflow_api::{
    RunRecordView, RunStatus, RunStore, WorkflowRunRouterBuilder, WORKFLOW_READ_ALL_SCOPE,
    WORKFLOW_READ_SCOPE, WORKFLOW_RUN_SCOPE,
};
use serde_json::Value;

/// Test model id every fixture workflow references.
pub const TEST_MODEL: &str = "default";

/// Bearer token the fixtures use; the authenticator accepts any value.
pub const TEST_TOKEN: &str = "test-token";

/// A second, distinct caller identity — same base scopes as [`TEST_TOKEN`]
/// under [`distinct_author_authenticator`], but a different principal, so a
/// run it submits (or reads) has a different author.
pub const OTHER_TOKEN: &str = "other-token";

/// A caller identity granted the cross-author [`WORKFLOW_READ_ALL_SCOPE`]
/// under [`distinct_author_authenticator`].
pub const READ_ALL_TOKEN: &str = "read-all-token";

/// An App whose LLM is a stub (unused by runs) with in-process bus + memory.
pub async fn test_app() -> Arc<App> {
    let app = App::builder()
        .llm(stub_llm("STUB REPLY"))
        .memory(MemoryHandles::new(
            Arc::new(InMemoryShortTerm::default()),
            Arc::new(InMemoryLongTerm::default()),
            Arc::new(InMemoryEpisodic::default()),
        ))
        .bus(MemoryBus::new())
        .tools_invoker(Arc::new(FakeToolInvoker::new()))
        .build()
        .await
        .expect("test app builds");
    Arc::new(app)
}

/// A registry with the test model bound to a stub LLM.
pub fn test_registry() -> Registry {
    Registry::new().with_model(TEST_MODEL, stub_llm("STUB REPLY"))
}

/// An in-memory durable provenance repository (sqlite `:memory:`).
pub fn test_provenance() -> Arc<dyn ProvenanceRepository> {
    Arc::new(SqliteProvenanceRepository::open_in_memory().expect("in-memory provenance opens"))
}

/// A bearer authenticator that accepts any token and grants exactly `scopes`.
fn authenticator_granting(scopes: &'static [&'static str]) -> Arc<dyn Authenticator> {
    Arc::new(BearerTokenAuthenticator::new(move |token: &str| {
        let scopes = scopes.iter().map(|s| s.to_string()).collect();
        Ok(Identity::with_scopes(token, scopes))
    }))
}

/// A bearer authenticator that accepts any token and grants both scopes.
pub fn test_authenticator() -> Arc<dyn Authenticator> {
    authenticator_granting(&[WORKFLOW_RUN_SCOPE, WORKFLOW_READ_SCOPE])
}

/// A bearer authenticator granting ONLY the read scope (no run scope).
pub fn read_only_authenticator() -> Arc<dyn Authenticator> {
    authenticator_granting(&[WORKFLOW_READ_SCOPE])
}

/// A bearer authenticator granting ONLY the run scope (no read scope).
pub fn run_only_authenticator() -> Arc<dyn Authenticator> {
    authenticator_granting(&[WORKFLOW_RUN_SCOPE])
}

/// A bearer authenticator whose granted scopes depend on the token: every
/// token (including [`TEST_TOKEN`] and [`OTHER_TOKEN`]) gets run + read
/// scope, so submits and reads under different tokens produce different
/// run authors; [`READ_ALL_TOKEN`] additionally gets the cross-author
/// [`WORKFLOW_READ_ALL_SCOPE`].
pub fn distinct_author_authenticator() -> Arc<dyn Authenticator> {
    Arc::new(BearerTokenAuthenticator::new(|token: &str| {
        let mut scopes: HashSet<String> = [WORKFLOW_RUN_SCOPE, WORKFLOW_READ_SCOPE]
            .into_iter()
            .map(String::from)
            .collect();
        if token == READ_ALL_TOKEN {
            scopes.insert(WORKFLOW_READ_ALL_SCOPE.to_string());
        }
        Ok(Identity::with_scopes(token, scopes))
    }))
}

/// Assemble a run-service router over the given registry, with a fresh
/// in-process run store and the supplied timeout + concurrency cap.
pub async fn build_router(
    registry: Registry,
    run_timeout: Duration,
    max_concurrent: usize,
) -> axum::Router {
    build_router_with_auth(test_authenticator(), registry, run_timeout, max_concurrent).await
}

/// Like [`build_router`] but with a caller-supplied authenticator, for
/// exercising the per-route scope-deny branches.
pub async fn build_router_with_auth(
    authenticator: Arc<dyn Authenticator>,
    registry: Registry,
    run_timeout: Duration,
    max_concurrent: usize,
) -> axum::Router {
    WorkflowRunRouterBuilder::new()
        .with_authenticator(authenticator)
        .with_app(test_app().await)
        .with_registry(registry)
        .with_run_store(RunStore::new(MemoryBus::new().kv))
        .with_provenance(test_provenance())
        .with_run_timeout(run_timeout)
        .with_max_concurrent_runs(max_concurrent)
        .build()
        .expect("router builds")
}

/// Like [`build_router`] but with caller-owned run store + provenance, so a
/// test can inspect the audit chain after a run.
pub async fn build_router_with_provenance(
    registry: Registry,
    run_store: RunStore,
    provenance: Arc<dyn ProvenanceRepository>,
    run_timeout: Duration,
    max_concurrent: usize,
) -> axum::Router {
    WorkflowRunRouterBuilder::new()
        .with_authenticator(test_authenticator())
        .with_app(test_app().await)
        .with_registry(registry)
        .with_run_store(run_store)
        .with_provenance(provenance)
        .with_run_timeout(run_timeout)
        .with_max_concurrent_runs(max_concurrent)
        .build()
        .expect("router builds")
}

/// All chain entries recorded under `scope` (the run id).
pub async fn chain_entries(
    provenance: &Arc<dyn ProvenanceRepository>,
    scope: &str,
) -> Vec<klieo_provenance::ChainEntry> {
    provenance
        .list_range(scope, 0, u64::MAX)
        .await
        .expect("chain range query succeeds")
}

/// A subflow that panics — exercises the executor's panic → `Failed` path.
pub fn panicking_subflow() -> Arc<dyn Flow> {
    struct Boom;
    #[async_trait::async_trait]
    impl Flow for Boom {
        fn name(&self) -> &str {
            "boom"
        }
        async fn run(&self, _ctx: AgentContext, _input: Value) -> Result<Value, FlowError> {
            panic!("intentional test panic");
        }
    }
    Arc::new(Boom)
}

/// A subflow that sleeps `delay` (so a live subscriber can attach), then emits
/// a tagged `ToolCallStarted` + `Completed` through `ctx.progress`, then echoes.
pub fn progress_emitting_subflow(tag: &'static str, delay: Duration) -> Arc<dyn Flow> {
    struct Emit(&'static str, Duration);
    #[async_trait::async_trait]
    impl Flow for Emit {
        fn name(&self) -> &str {
            self.0
        }
        async fn run(&self, ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            tokio::time::sleep(self.1).await;
            if let Some(progress) = &ctx.progress {
                let _ = progress.send(AgentEvent::ToolCallStarted {
                    name: self.0.to_string(),
                });
                let _ = progress.send(AgentEvent::Completed);
            }
            Ok(input)
        }
    }
    Arc::new(Emit(tag, delay))
}

/// A subflow that sleeps `delay` then echoes — exercises the timeout path.
pub fn slow_subflow(delay: Duration) -> Arc<dyn Flow> {
    struct Slow(Duration);
    #[async_trait::async_trait]
    impl Flow for Slow {
        fn name(&self) -> &str {
            "slow"
        }
        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            tokio::time::sleep(self.0).await;
            Ok(input)
        }
    }
    Arc::new(Slow(delay))
}

/// A single-subflow workflow definition JSON referencing `subflow_ref`.
pub fn single_subflow_def(subflow_ref: &str) -> Value {
    serde_json::json!({
        "id": "test-workflow",
        "entry": "n",
        "nodes": [{ "id": "n", "kind": "subflow", "ref": subflow_ref }],
        "edges": []
    })
}

/// A `POST /runs` request carrying `definition` + `input`.
pub fn submit_request(definition: Value, input: Value) -> Request<Body> {
    let body = serde_json::json!({ "definition": definition, "input": input });
    Request::builder()
        .method("POST")
        .uri("/runs")
        .header("authorization", format!("Bearer {TEST_TOKEN}"))
        .header("content-type", "application/json")
        .body(Body::from(serde_json::to_vec(&body).unwrap()))
        .unwrap()
}

/// An authenticated `GET` request to `uri`, bearing `token`.
pub fn authed_get_as(uri: &str, token: &str) -> Request<Body> {
    Request::builder()
        .method("GET")
        .uri(uri)
        .header("authorization", format!("Bearer {token}"))
        .body(Body::empty())
        .unwrap()
}

/// An authenticated `GET` request to `uri`, as [`TEST_TOKEN`].
pub fn authed_get(uri: &str) -> Request<Body> {
    authed_get_as(uri, TEST_TOKEN)
}

/// A `GET /runs/{id}` request, as [`TEST_TOKEN`].
pub fn get_request(run_id: &str) -> Request<Body> {
    authed_get(&format!("/runs/{run_id}"))
}

/// A `GET /runs/{id}` request, as `token`.
pub fn get_request_as(run_id: &str, token: &str) -> Request<Body> {
    authed_get_as(&format!("/runs/{run_id}"), token)
}

/// Collect a response body into JSON.
pub async fn body_json(response: Response<Body>) -> Value {
    let bytes = response.into_body().collect().await.unwrap().to_bytes();
    serde_json::from_slice(&bytes).unwrap_or(Value::Null)
}

/// Collect a response body into a UTF-8 string (for SSE streams). Drives a
/// streaming body to completion.
pub async fn body_text(response: Response<Body>) -> String {
    let bytes = response.into_body().collect().await.unwrap().to_bytes();
    String::from_utf8(bytes.to_vec()).expect("body is valid UTF-8")
}

/// Read the `run_id` from a `202` submit response.
pub async fn run_id_of(response: Response<Body>) -> String {
    body_json(response).await["run_id"]
        .as_str()
        .expect("202 body carries run_id")
        .to_string()
}

/// True once a status is terminal (not `Pending`/`Running`).
fn is_terminal(status: RunStatus) -> bool {
    !matches!(status, RunStatus::Pending | RunStatus::Running)
}

/// Poll `GET /runs/{id}` until the run reaches a terminal status, bounded.
pub async fn poll_until_terminal(router: axum::Router, run_id: &str) -> RunRecordView {
    use tower::ServiceExt;
    const MAX_POLLS: u32 = 100;
    const POLL_INTERVAL: Duration = Duration::from_millis(20);
    for _ in 0..MAX_POLLS {
        let response = router.clone().oneshot(get_request(run_id)).await.unwrap();
        let view: RunRecordView = serde_json::from_value(body_json(response).await).unwrap();
        if is_terminal(view.status) {
            return view;
        }
        tokio::time::sleep(POLL_INTERVAL).await;
    }
    panic!("run {run_id} did not reach a terminal status in time");
}