sausage/host/src/harness/credentials.rs

//! Credential request snapshots and host-only model proxy configuration.
use crate::cred::{store::Cred, types::Subject};
use deno_core::{OpState, op2};
use deno_error::JsErrorBox;
use serde_json::{Value, json};
use std::{cell::RefCell, rc::Rc, sync::Arc};

pub(super) struct Notifications {
    filed: tokio::sync::Mutex<tokio::sync::broadcast::Receiver<String>>,
}
impl Notifications {
    pub(super) fn new(cred: &Cred) -> Rc<Self> {
        Rc::new(Self {
            filed: tokio::sync::Mutex::new(cred.subscribe_filed()),
        })
    }
}
#[op2]
pub(super) async fn op_credential_next(state: Rc<RefCell<OpState>>) {
    let notifications = state.borrow().borrow::<Rc<Notifications>>().clone();
    let mut filed = notifications.filed.lock().await;
    let _ = tokio::time::timeout(std::time::Duration::from_secs(1), filed.recv()).await;
}

#[op2]
#[serde]
pub(super) fn op_credential_requests(state: &mut OpState) -> Result<serde_json::Value, JsErrorBox> {
    let cred = state.borrow::<Arc<Cred>>();
    let work = || -> anyhow::Result<Value> {
        cred.expire_requests()?;
        // TODO(T27): prune retained records once a day after the library adds Cred::prune.
        let mut rows = Vec::new();
        for request in cred.request_list()? {
            let Subject::Conversation(conversation) = request.subject else {
                continue;
            };
            let names = request
                .payload
                .get("items")
                .and_then(Value::as_array)
                .map(|items| {
                    items
                        .iter()
                        .filter_map(Value::as_str)
                        .map(|id| {
                            cred.item_get(id)
                                .map(|item| format!("`{}`", item.name))
                                .unwrap_or_else(|_| "the credential".into())
                        })
                        .collect::<Vec<_>>()
                        .join(", ")
                })
                .unwrap_or_else(|| "the credential change".into());
            let message = match request.state.as_str() {
                "approved" => format!("Your request for {names} was approved. Retry the command."),
                "denied" => format!("Your request for {names} was denied."),
                "expired" => format!(
                    "Your request for {names} expired. File a new request if you still need it."
                ),
                _ => String::new(),
            };
            rows.push(json!({"id":request.id,"conversation":conversation,"state":request.state,"message":message}));
        }
        Ok(json!(rows))
    };
    work().map_err(|_| JsErrorBox::generic("Cannot read credential requests"))
}

pub(super) fn proxy(cred: &Cred) -> anyhow::Result<Option<Value>> {
    use crate::cred::{
        items::{Field, Item},
        types::{Caller, FieldType},
    };
    let caller = Caller {
        pid: 0,
        program: "host".into(),
        conversations: vec![],
    };
    let item = cred
        .item_list()?
        .into_iter()
        .find(|item| item.name == "llm-proxy");
    let item = match item {
        Some(item) => item,
        None => {
            let (Ok(base), Ok(key)) = (
                std::env::var("PI_PROXY_BASE_URL"),
                std::env::var("PI_PROXY_API_KEY"),
            ) else {
                return Ok(None);
            };
            let mut store = cred.access()?;
            let master = store.master.clone();
            let tx = store
                .db
                .transaction_with_behavior(rusqlite::TransactionBehavior::Immediate)?;
            let mut item = crate::cred::items::insert_new(
                &tx,
                Item {
                    id: String::new(),
                    name: "llm-proxy".into(),
                    version: 0,
                    free_use: false,
                    rules: vec![],
                    validate_js: None,
                    transform_js: "export default (ctx) => ctx.request".into(),
                    fields: vec![
                        Field {
                            field_id: String::new(),
                            name: "base_url".into(),
                            public: true,
                            value: Some(base),
                            needs_value: false,
                            field_type: FieldType::Text,
                            key_type: None,
                        },
                        Field {
                            field_id: String::new(),
                            name: "api_key".into(),
                            public: false,
                            value: None,
                            needs_value: true,
                            field_type: FieldType::Text,
                            key_type: None,
                        },
                    ],
                },
                &caller,
                Some("host bootstrap"),
            )?;
            let field = &item
                .fields
                .iter()
                .find(|f| f.name == "api_key")
                .unwrap()
                .field_id;
            let sealed =
                crate::cred::seal::seal(&master, &format!("{}|{field}", item.id), key.as_bytes())?;
            tx.execute(
                "INSERT INTO fields VALUES(?1,?2,?3)",
                rusqlite::params![item.id, field, sealed],
            )?;
            item.fields
                .iter_mut()
                .find(|f| f.name == "api_key")
                .unwrap()
                .needs_value = false;
            crate::cred::items::save(&tx, &item)?;
            tx.commit()?;
            item
        }
    };
    Ok(Some(cred.credential(&item, true)?))
}

#[op2]
#[serde]
pub(super) fn op_model_proxy(state: &mut OpState) -> Result<serde_json::Value, JsErrorBox> {
    proxy(state.borrow::<Arc<Cred>>())
        .map(|v| v.unwrap_or(Value::Null))
        .map_err(|_| JsErrorBox::generic("Cannot read llm-proxy credential"))
}