sausage/host/tests/durable.rs

//! Credential decisions wake an idle faux-model conversation through durable tasks.
use sausage::{
    cred::{
        items::Item,
        requests::{GrantRequest, Overrides},
        store::Cred,
        types::Caller,
        webauthn::Decision,
    },
    harness::{run_module, run_module_with_credentials},
    world::World,
};
use std::sync::Arc;

#[tokio::test(flavor = "current_thread")]
async fn restart_child() -> anyhow::Result<()> {
    let Ok(dir) = std::env::var("DURABLE_RESTART_DIR") else {
        return Ok(());
    };
    let dir = std::path::Path::new(&dir);
    let source = format!(
        "globalThis.expectedState = 'approved';\n{}",
        include_str!("../js/durable-test-bundle.js")
    );
    run_module(
        World::open(dir)?,
        sausage::process::exec(World::open(dir)?)?,
        &source,
    )
    .await
}

#[test]
fn sigkill_preserves_request_wait() -> anyhow::Result<()> {
    use std::{
        process::Command,
        time::{Duration, Instant},
    };
    let dir = std::env::temp_dir().join(format!("sausage-durable-restart-{}", std::process::id()));
    let world = World::open(&dir)?;
    let cred = Cred::open(&world.host_dir)?;
    let admin = Caller {
        pid: 0,
        program: "test".into(),
        conversations: vec![],
    };
    let item = cred.create_item(
        Item {
            id: String::new(),
            name: "github".into(),
            version: 0,
            free_use: false,
            rules: vec![],
            validate_js: None,
            transform_js: "export default (ctx) => ctx.request".into(),
            fields: vec![],
        },
        &admin,
        cred.unlocked()?,
    )?;
    let mut child = Command::new(std::env::current_exe()?)
        .args(["--exact", "restart_child", "--nocapture"])
        .env("DURABLE_RESTART_DIR", &dir)
        .spawn()?;
    let deadline = Instant::now() + Duration::from_secs(30);
    let conversation = loop {
        if let Ok(id) = std::fs::read_to_string(dir.join("world/conversation"))
            && !id.is_empty()
        {
            break id;
        }
        if Instant::now() > deadline {
            child.kill()?;
            child.wait()?;
            anyhow::bail!("Conversation did not start");
        }
        std::thread::yield_now();
    };
    let caller = Caller {
        conversations: vec![conversation],
        ..admin
    };
    let id = cred.file_request(
        &caller,
        "restart",
        GrantRequest {
            items: vec![item.id],
            duration_secs: None,
            max_uses: None,
            validator: None,
        },
    )?;
    let db = rusqlite::Connection::open(dir.join("host/harness.sqlite"))?;
    loop {
        let count: i64 = db.query_row(
            "SELECT count(*) FROM tasks WHERE json_extract(kind, '$')='sausage.credential.wait'",
            [],
            |r| r.get(0),
        )?;
        if count == 1 {
            break;
        }
        if Instant::now() > deadline {
            child.kill()?;
            child.wait()?;
            anyhow::bail!("Durable wait did not start");
        }
        std::thread::yield_now();
    }
    child.kill()?;
    child.wait()?;
    cred.approve(
        &id,
        Decision::Approve,
        Overrides::default(),
        cred.unlocked()?,
    )?;
    for _ in 0..2 {
        let mut resumed = Command::new(std::env::current_exe()?)
            .args(["--exact", "restart_child", "--nocapture"])
            .env("DURABLE_RESTART_DIR", &dir)
            .spawn()?;
        let deadline = Instant::now() + Duration::from_secs(30);
        loop {
            if let Some(status) = resumed.try_wait()? {
                assert!(status.success());
                break;
            }
            if Instant::now() > deadline {
                resumed.kill()?;
                resumed.wait()?;
                anyhow::bail!("Restart did not finish");
            }
            std::thread::yield_now();
        }
    }
    drop(db);
    std::fs::remove_dir_all(dir)?;
    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn proxy_child() -> anyhow::Result<()> {
    let Ok(dir) = std::env::var("DURABLE_PROXY_DIR") else {
        return Ok(());
    };
    let dir = std::path::Path::new(&dir);
    let source = format!(
        "globalThis.proxyOnly = true; const proxy = Deno.core.ops.op_model_proxy(); if (proxy.base_url !== 'https://proxy.invalid' || proxy.api_key !== 'faux-private-token') throw new Error('Wrong proxy item');\n{}",
        include_str!("../js/durable-test-bundle.js")
    );
    run_module(
        World::open(dir)?,
        sausage::process::exec(World::open(dir)?)?,
        &source,
    )
    .await?;
    let error = run_module(
        World::open(dir)?,
        sausage::process::exec(World::open(dir)?)?,
        "throw new Error(Deno.core.ops.op_model_proxy().api_key)",
    )
    .await
    .unwrap_err();
    assert!(!error.to_string().contains("faux-private-token"));
    assert!(error.to_string().contains("[redacted]"));
    Ok(())
}

#[test]
fn proxy_bootstrap_once() -> anyhow::Result<()> {
    let dir = std::env::temp_dir().join(format!("sausage-durable-proxy-{}", std::process::id()));
    for start in 0..2 {
        let status = std::process::Command::new(std::env::current_exe()?)
            .args(["--exact", "proxy_child", "--nocapture"])
            .env("DURABLE_PROXY_DIR", &dir)
            .env(
                "PI_PROXY_BASE_URL",
                if start == 0 {
                    "https://proxy.invalid"
                } else {
                    "https://ignored.invalid"
                },
            )
            .env(
                "PI_PROXY_API_KEY",
                if start == 0 {
                    "faux-private-token"
                } else {
                    "ignored-token"
                },
            )
            .status()?;
        assert!(status.success());
    }
    let cred = Cred::open(&dir.join("host"))?;
    assert_eq!(
        cred.item_list()?
            .iter()
            .filter(|i| i.name == "llm-proxy")
            .count(),
        1
    );
    for file in [
        "host/cred.sqlite",
        "host/cred.sqlite-wal",
        "host/harness.sqlite",
        "host/harness.sqlite-wal",
    ] {
        if let Ok(bytes) = std::fs::read(dir.join(file)) {
            assert!(
                !bytes
                    .windows(b"faux-private-token".len())
                    .any(|w| w == b"faux-private-token")
            );
        }
    }
    assert!(!dir.join("world/conversation").exists());
    drop(cred);
    std::fs::remove_dir_all(dir)?;
    Ok(())
}

#[tokio::test(flavor = "current_thread")]
async fn decisions_wake_idle_conversation() -> anyhow::Result<()> {
    for state in ["approved", "denied", "expired"] {
        let dir =
            std::env::temp_dir().join(format!("sausage-durable-{}-{state}", std::process::id()));
        let world = World::open(&dir)?;
        let cred = Arc::new(Cred::open(&world.host_dir)?);
        let admin = Caller {
            pid: 0,
            program: "test".into(),
            conversations: vec![],
        };
        let item = cred.create_item(
            Item {
                id: String::new(),
                name: "github".into(),
                version: 0,
                free_use: false,
                rules: vec![],
                validate_js: None,
                transform_js: "export default (ctx) => ctx.request".into(),
                fields: vec![],
            },
            &admin,
            cred.unlocked()?,
        )?;
        cred.file_request(
            &admin,
            "program",
            GrantRequest {
                items: vec![item.id.clone()],
                duration_secs: None,
                max_uses: None,
                validator: None,
            },
        )?;
        let source = format!(
            "globalThis.expectedState = {state:?};\n{}",
            include_str!("../js/durable-test-bundle.js")
        );
        let decide = async {
            let conversation = loop {
                if let Ok(id) = std::fs::read_to_string(dir.join("world/conversation"))
                    && !id.is_empty()
                {
                    break id;
                }
                tokio::task::yield_now().await;
            };
            let caller = Caller {
                conversations: vec![conversation],
                ..admin
            };
            let id = cred.file_request(
                &caller,
                "test",
                GrantRequest {
                    items: vec![item.id],
                    duration_secs: None,
                    max_uses: None,
                    validator: None,
                },
            )?;
            let harness_db = rusqlite::Connection::open(dir.join("host/harness.sqlite"))?;
            loop {
                let count: i64 = harness_db.query_row(
                    "SELECT count(*) FROM tasks WHERE json_extract(kind, '$')='sausage.credential.wait'",
                    [],
                    |r| r.get(0),
                )?;
                if count == 1 {
                    break;
                }
                tokio::task::yield_now().await;
            }
            match state {
                "approved" => cred.approve(
                    &id,
                    Decision::Approve,
                    Overrides::default(),
                    cred.unlocked()?,
                )?,
                "denied" => cred.deny(&id, Decision::Deny, cred.unlocked()?)?,
                _ => {
                    let db = rusqlite::Connection::open(dir.join("host/cred.sqlite"))?;
                    db.execute("UPDATE requests SET expires=0 WHERE id=?", [&id])?;
                }
            }
            Ok::<_, anyhow::Error>(())
        };
        let exec = sausage::process::exec(World::open(&dir)?)?;
        tokio::time::timeout(std::time::Duration::from_secs(20), async {
            let (a, b) = tokio::join!(
                run_module_with_credentials(world, Arc::clone(&exec), cred.clone(), &source),
                decide
            );
            a?;
            b?;
            Ok::<_, anyhow::Error>(())
        })
        .await??;
        let db = rusqlite::Connection::open(dir.join("host/harness.sqlite"))?;
        let tasks: i64 = db.query_row(
            "SELECT count(*) FROM tasks WHERE json_extract(kind, '$')='sausage.credential.wait'",
            [],
            |r| r.get(0),
        )?;
        assert_eq!(tasks, 1, "Program requests must create no durable task");
        drop(db);
        std::fs::remove_dir_all(dir)?;
    }
    Ok(())
}