fabro/lib/components/fabro-petri/tests/model.rs

//! A model call from a Petri run authenticates through Fabro's vault, and
//! the skills step reads the Fabro home the runtime was given.
//!
//! The `hello` bundle's agent stage calls the OpenAI twin through a model
//! client built over a vault that holds the key; the twin requires a
//! bearer token and logs requests under it, so a request logged under the
//! vault's key proves the key came from the vault. The Host scope runs in
//! process.

mod support;

use std::collections::HashMap;
use std::sync::Arc;

use fabro_auth::VaultCredentialSource;
use fabro_llm::test_support::test_catalog_with_provider_base_url;
use fabro_petri::check::Launch;
use fabro_petri::engine::{self, RunStatus};
use fabro_petri::runtime::{self, RuntimeSpec};
use fabro_test::{TwinScenario, TwinScenarios, twin_openai};
use fabro_types::{RunId, SecretType};
use fabro_vault::Vault;
use lithos_llm::catalog::ProviderId;
use petri_store::MemoryRunStore;
use support::{Silent, all_records, hello_bundle, no_questions, run_request};
use tokio::fs;
use tokio::sync::RwLock as AsyncRwLock;

const OPENAI_MODEL: &str = "gpt-5.4";

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn a_model_call_authenticates_through_the_vault_and_skills_read_the_home() {
    let twin = twin_openai().await;
    let namespace = format!("{}::{}", module_path!(), line!());
    TwinScenarios::new(&namespace)
        .scenario(TwinScenario::responses(OPENAI_MODEL).text("A haiku, added."))
        .load(twin)
        .await;
    let root = tempfile::tempdir().expect("a temp dir");
    let home = root.path().join("fabro-home");
    std::fs::create_dir_all(home.join("skills")).expect("the skills dir creates");

    // The vault holds the key; nothing in the environment does.
    let mut vault = Vault::from_entries(HashMap::new());
    vault
        .set("OPENAI_API_KEY", &namespace, SecretType::Token, None)
        .expect("a detached vault takes an entry");
    let credentials = Arc::new(VaultCredentialSource::vault_only(Arc::new(
        AsyncRwLock::new(vault),
    )));
    let catalog = test_catalog_with_provider_base_url("openai", &twin.base_url);
    let client = runtime::model_client(catalog, credentials, None, &[ProviderId::new("openai")])
        .expect("the model client builds")
        .expect("openai is eligible");
    let runtime = RuntimeSpec {
        model_client: Some(client),
        fabro_home: Some(home.clone()),
        ..RuntimeSpec::default()
    };

    let workflow = fs::read_to_string(hello_bundle().join("workflow.fabro"))
        .await
        .expect("the hello workflow is checked in");
    let settings = fs::read_to_string(hello_bundle().join("workflow.toml"))
        .await
        .expect("the hello settings are checked in");
    let graphs = support::admit(
        &[("workflow.fabro", &workflow), ("workflow.toml", &settings)],
        Launch {
            default_model: Some(OPENAI_MODEL.to_string()),
            ..Launch::default()
        },
        &runtime,
    );
    let store = Arc::new(MemoryRunStore::new());
    let request = run_request(
        "hello",
        &root.path().join("run"),
        graphs,
        store.clone(),
        runtime,
        no_questions(Arc::new(Silent)),
    );

    let outcome = engine::run(request).await.expect("the run ends");

    assert_eq!(outcome.status, RunStatus::Success, "{outcome:?}");
    let logs = twin.request_logs(&namespace).await;
    let requests = logs["requests"]
        .as_array()
        .expect("twin request logs are an array");
    assert!(
        requests
            .iter()
            .any(|request| request["model"] == OPENAI_MODEL),
        "the stage should have called the twin with the vault's key, got {logs}"
    );
    let records = all_records(store.as_ref(), "hello").await;
    let resolved = records
        .iter()
        .find(|record| record.to_string().contains("\"attractor.skills\""))
        .unwrap_or_else(|| panic!("the skills step recorded what it searched: {records:?}"));
    let configured = home.join("skills").display().to_string();
    assert!(
        resolved.to_string().contains(&configured),
        "the configured home is searched: {resolved}"
    );
}

/// The engine binds the durable identity for both native step kinds and both
/// entry modes, without changing a shared client used outside a run.
#[tokio::test]
async fn model_session_headers_follow_run_identity_on_start_and_resume() {
    use std::sync::Mutex;

    use fabro_llm::test_support::{self, ScriptedAdapter};
    use lithos_llm::Client;
    use lithos_llm::middleware::{Call, Observer, ObserverMiddleware};
    use petri_store::{Access, OwnerId, RunKey, RunStore as _};

    struct Sessions(Arc<Mutex<Vec<Option<String>>>>);
    impl Observer for Sessions {
        fn on_start(&self, call: &Call) {
            self.0.lock().unwrap().push(
                call.context()
                    .default_headers()
                    .get("x-session-id")
                    .map(|value| value.to_str().unwrap().to_owned()),
            );
        }
    }
    let seen = Arc::new(Mutex::new(Vec::new()));
    let client = Client::builder()
        .catalog(test_support::test_catalog())
        .adapter(
            ProviderId::new("openai"),
            ScriptedAdapter::new(vec![test_support::text_response(
                "openai",
                OPENAI_MODEL,
                "done",
            )]),
        )
        .middleware(ObserverMiddleware::new(Sessions(Arc::clone(&seen))))
        .build()
        .unwrap()
        .client;
    let runtime = RuntimeSpec {
        model_client: Some(client.clone()),
        ..RuntimeSpec::default()
    };
    let workflow = r#"digraph Headers {
        graph [backend="api", default_max_retries=0]
        start [shape=Mdiamond]
        exit [shape=Msquare]
        agent [shape=box, prompt="Say done"]
        prompt [shape=tab, prompt="Say done"]
        start -> agent -> prompt -> exit
    }"#;
    for resume in [false, true] {
        let run_id = RunId::new().to_string();
        let root = tempfile::tempdir().unwrap();
        let store = Arc::new(MemoryRunStore::new());
        let graphs = support::admit(
            &[
                ("workflow.fabro", workflow),
                ("workflow.toml", support::SETTINGS),
            ],
            Launch {
                default_model: Some(format!("openai/{OPENAI_MODEL}")),
                ..Launch::default()
            },
            &runtime,
        );
        let mut request = run_request(
            &run_id,
            root.path(),
            graphs,
            store.clone(),
            runtime.clone(),
            no_questions(Arc::new(Silent)),
        );
        if resume {
            // Recover the same durable key after creation was interrupted.
            drop(
                store
                    .open(&RunKey::new(&run_id), Access::Create {
                        owner: OwnerId::new("interrupted"),
                    })
                    .await
                    .unwrap(),
            );
            request.execution = match request.execution {
                engine::Execution::Start(graphs) | engine::Execution::Resume(graphs) => {
                    engine::Execution::Resume(graphs)
                }
            };
        }
        seen.lock().unwrap().clear();
        let outcome = engine::run(request).await.unwrap();
        assert_eq!(outcome.status, RunStatus::Success, "{outcome:?}");
        let calls = seen.lock().unwrap();
        assert!(calls.len() >= 2, "both native stages call the model");
        assert!(
            calls
                .iter()
                .all(|header| header.as_deref() == Some(&run_id))
        );
    }
    seen.lock().unwrap().clear();
    client
        .complete(
            fabro_llm::Request::builder()
                .model(format!("openai/{OPENAI_MODEL}"))
                .user("outside a run")
                .build()
                .unwrap(),
        )
        .await
        .unwrap();
    assert_eq!(*seen.lock().unwrap(), vec![None]);
}