bitrouter 1.0.0-alpha.2

//! Management command logic — the implementations behind the CLI subcommands.
//! Kept in the lib so it is testable and so a future desktop front-end can
//! reuse it without dragging in the binary.

use anyhow::{Context, Result};

use bitrouter_sdk::caller::CallerContext;
use bitrouter_sdk::config::{Config, ConfigRoutingTable};
use bitrouter_sdk::language_model::{RoutingPrefs, RoutingTable};

use crate::auth::{NewApiKey, db as auth_db, generate};
use crate::daemon::RouteHop;

/// The starter `bitrouter.yaml` written by `bitrouter init`. Mirrors
/// the zero-config in-memory default so a user who runs `init` and
/// then deletes nothing gets the same behaviour as running bitrouter
/// with no config at all — only the file is here to be edited.
pub const STARTER_CONFIG: &str = r#"# bitrouter.yaml — generated by `bitrouter init`.
#
# This file mirrors the zero-config in-memory default. Delete any
# section you don't want to override, or add custom providers /
# guardrails / policies / MCP servers as needed.

server:
  # Bind localhost only by default — combined with `skip_auth: true`,
  # exposing `0.0.0.0` would let any host on the network use bitrouter
  # without a credential check. Flip to `0.0.0.0:4356` only when you've
  # also set `skip_auth: false` and minted virtual keys with
  # `bitrouter key sign`.
  listen: "127.0.0.1:4356"
  # skip_auth=true admits credential-less local requests (cc-switch /
  # litellm style). Set to false and create keys with `bitrouter key
  # sign` for a multi-tenant deployment.
  skip_auth: true

database:
  # Relative paths are interpreted against the bitrouter home (the
  # directory containing this file) — the daemon chdirs there on
  # startup, so `./bitrouter.db` always lands alongside the config.
  url: "sqlite://./bitrouter.db"

# Built-in providers auto-enable when their env var is set. An entry
# with `inherit_defaults: true` (the workspace default) plus an empty
# body is enough — `api_base`, `api_protocol`, `api_key`, and the
# routable models all flow from the built-in catalog. Comment a line
# out (or set `active: false`) to disable a provider you don't want.
providers:
  openai: {}        # uses OPENAI_API_KEY
  anthropic: {}     # uses ANTHROPIC_API_KEY
  google: {}        # uses GEMINI_API_KEY
  openrouter: {}    # uses OPENROUTER_API_KEY
  # github-copilot needs `bitrouter login github-copilot` first:
  # github-copilot: {}
  # A provider can hold multiple accounts — e.g. two subscriptions to
  # the same upstream. `account_strategy` is `failover` (try the first,
  # drop to the next on a retryable / out-of-credits error — default)
  # or `balance` (spread load evenly across accounts):
  # opencode-go: { account_strategy: failover, accounts: [{ api_key: "${OPENCODE_GO_KEY_A}", label: primary }, { api_key: "${OPENCODE_GO_KEY_B}", label: backup }] }

# Upstream MCP (Model Context Protocol) servers — `bitrouter tools
# list/status` reads this, and `POST /mcp/<name>` proxies JSON-RPC
# against the entry. Spec:
# https://modelcontextprotocol.io/specification/2025-06-18
# mcp_servers:
#   ctx7:
#     name: ctx7
#     transport:
#       type: http
#       url: https://mcp.context7.com/mcp
#       headers:
#         Authorization: "Bearer ${CTX7_TOKEN}"
#   git:
#     name: git
#     transport:
#       type: stdio
#       command: uvx
#       args: ["mcp-server-git"]

# Upstream ACP (Agent Client Protocol) agents — invoked from an editor
# via `bitrouter agent-proxy <id>`. Spec:
# https://agentclientprotocol.com/protocol/schema
# agents:
#   claude:
#     name: claude
#     transport:
#       type: stdio
#       command: npx
#       args: ["-y", "@zed-industries/claude-code-acp@latest"]
#   codex:
#     name: codex
#     transport:
#       type: stdio
#       command: npx
#       args: ["-y", "@zed-industries/codex-acp@latest"]

inherit_defaults: true
"#;

/// `bitrouter init` — write a starter config to `path` (refuses to overwrite).
pub async fn init(path: &std::path::Path) -> Result<()> {
    if tokio::fs::try_exists(path).await.unwrap_or(false) {
        anyhow::bail!("{} already exists — refusing to overwrite", path.display());
    }
    tokio::fs::write(path, STARTER_CONFIG)
        .await
        .with_context(|| format!("writing {}", path.display()))?;
    Ok(())
}

/// The outcome of `bitrouter key generate` — the plaintext secret is shown to
/// the user **once** and never persisted (only its SHA-256 hash is stored).
pub struct GeneratedKey {
    /// The api key id.
    pub id: String,
    /// The plaintext `brvk_…` secret — display once, then discard.
    pub secret: String,
}

/// `bitrouter key sign` — mint a `brvk_` virtual key for `user_id`, persist only
/// its hash, and return the plaintext once. v1 does **not** sign a JWT (004
/// §3.0); the command keeps the `sign` name but creates a DB-backed virtual key.
pub async fn key_sign(
    db_url: &str,
    user_id: &str,
    policy_id: Option<&str>,
) -> Result<GeneratedKey> {
    let db = crate::db::connect(db_url)
        .await
        .with_context(|| format!("connecting to database {db_url}"))?;
    crate::db::run_migrations(&db)
        .await
        .context("running database migrations")?;
    auth_db::upsert_user(&db, user_id)
        .await
        .context("creating user row")?;

    let key = generate();
    let id = format!("brvk_id_{}", &key.hash[..16]);
    auth_db::insert_api_key(
        &db,
        &NewApiKey {
            id: id.clone(),
            key_hash: key.hash.clone(),
            user_id: user_id.to_string(),
            spend_limit_micro_usd: None,
            rpm_limit: None,
            policy_id: policy_id.map(|s| s.to_string()),
        },
    )
    .await
    .context("inserting api key")?;

    Ok(GeneratedKey {
        id,
        secret: key.secret,
    })
}

/// `bitrouter models [--provider <id>]` — list routable models, optionally
/// filtered to those a given provider declares.
pub async fn list_models(
    config: &Config,
    provider_filter: Option<&str>,
) -> Result<Vec<(String, Vec<String>)>> {
    let mut resolved = config.clone();
    bitrouter_providers::apply_builtin_defaults(&mut resolved);
    // Best-effort `/models` discovery for any provider with
    // `auto_discover: true`. Zero-config built-ins set the flag, so
    // this turns `bitrouter models` from "empty" into a useful catalog
    // when the user has the API key set. Failures log and leave the
    // provider with no models — no error bubbles to the user.
    bitrouter_sdk::config::discover_models(&mut resolved).await;
    let table = ConfigRoutingTable::from_config(resolved);
    Ok(table
        .list_models()
        .into_iter()
        .filter(|m| match provider_filter {
            Some(p) => m.providers.iter().any(|x| x == p),
            None => true,
        })
        .map(|m| (m.id, m.providers))
        .collect())
}

/// `bitrouter route <model>` — resolve a model name through the routing table,
/// **standalone** (no running daemon needed). Returns the fallback chain.
pub async fn resolve_route(config: &Config, model: &str) -> Result<Vec<RouteHop>> {
    let mut resolved = config.clone();
    bitrouter_providers::apply_builtin_defaults(&mut resolved);
    let table = ConfigRoutingTable::from_config(resolved);
    let chain = table
        .route_chain(model, &RoutingPrefs::default(), &CallerContext::local())
        .await
        .with_context(|| format!("resolving model '{model}'"))?;
    Ok(chain
        .into_iter()
        .map(|t| RouteHop {
            provider: t.provider_name,
            service_id: t.service_id,
            api_protocol: format!("{:?}", t.api_protocol).to_lowercase(),
        })
        .collect())
}

/// One provider entry for `bitrouter providers list`.
pub struct ProviderListing {
    /// Provider id.
    pub id: String,
    /// Upstream API base.
    pub api_base: String,
    /// Whether the provider is active / routable.
    pub active: bool,
    /// Number of declared models.
    pub model_count: usize,
}

/// `bitrouter providers list` — list every configured provider.
/// Applies built-in catalog defaults first so a zero-config or
/// minimally-declared provider (`openai: {}`) shows its real `api_base`
/// and `active` state instead of empty placeholders.
pub fn list_providers(config: &Config) -> Vec<ProviderListing> {
    let mut resolved = config.clone();
    bitrouter_providers::apply_builtin_defaults(&mut resolved);
    let mut providers: Vec<ProviderListing> = resolved
        .providers
        .iter()
        .map(|(id, p)| ProviderListing {
            id: id.clone(),
            api_base: p.api_base.clone(),
            active: p.active,
            model_count: p.models.len(),
        })
        .collect();
    providers.sort_by(|a, b| a.id.cmp(&b.id));
    providers
}

/// Reject a policy id that would escape its directory or hide as a dotfile.
/// Allowed: ASCII letters / digits / `_` / `-` / `.` (not as the leading char).
fn validate_policy_id(id: &str) -> Result<()> {
    if id.is_empty() {
        anyhow::bail!("policy id must not be empty");
    }
    if id.starts_with('.') {
        anyhow::bail!("policy id must not start with '.' (no dotfiles / `..`)");
    }
    if id.chars().any(|c| {
        c == '/'
            || c == '\\'
            || c == '\0'
            || c.is_whitespace()
            || !(c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '.')
    }) {
        anyhow::bail!(
            "policy id '{id}' contains invalid characters (allowed: letters, digits, `_`, `-`, `.`)"
        );
    }
    Ok(())
}

/// `bitrouter policy create` — write a starter policy file into `policy_dir`.
/// Refuses to overwrite an existing policy of the same id. The `id` must be a
/// plain filename (no path separators, no `..`, no leading dot) so the joined
/// path cannot escape `policy_dir`.
pub async fn create_policy(policy_dir: &std::path::Path, id: &str) -> Result<std::path::PathBuf> {
    validate_policy_id(id)?;
    tokio::fs::create_dir_all(policy_dir)
        .await
        .with_context(|| format!("creating policy dir {}", policy_dir.display()))?;
    let path = policy_dir.join(format!("{id}.yaml"));
    if tokio::fs::try_exists(&path).await.unwrap_or(false) {
        anyhow::bail!("policy '{id}' already exists at {}", path.display());
    }
    let body = format!(
        "# policy '{id}' — edit to taste, then bind it to a key:\n\
         #   bitrouter key sign --user <id> --policy {id}\n\
         id: {id}\n\
         # allowed_models: [gpt-5, claude-sonnet-4-6]\n\
         # denied_models: []\n\
         # max_spend_micro_usd: 1000000\n\
         # max_requests_per_minute: 60\n\
         # allowed_tools: [search]\n"
    );
    tokio::fs::write(&path, body)
        .await
        .with_context(|| format!("writing {}", path.display()))?;
    Ok(path)
}

/// `bitrouter login <provider>` — run the provider's OAuth flow + persist a
/// token. Today only `github-copilot` is wired; other oauth providers can
/// register via [`bitrouter_providers::builtin`].
pub async fn login_provider(provider_id: &str) -> Result<()> {
    use bitrouter_providers::{AuthScheme, builtin};

    let entry = builtin::find(provider_id).ok_or_else(|| {
        anyhow::anyhow!(
            "unknown provider '{provider_id}' — try `bitrouter providers list` for the catalog"
        )
    })?;
    let (client_id, scope, device_url, token_url) = match &entry.auth {
        AuthScheme::Oauth { params, .. } => {
            let client_id = params
                .get("client_id")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "provider '{provider_id}': auth.params.client_id is missing in the TOML"
                    )
                })?
                .to_string();
            let scope = params
                .get("scope")
                .and_then(|v| v.as_str())
                .map(str::to_string);
            let device_url = params
                .get("device_authorization_endpoint")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "provider '{provider_id}': auth.params.device_authorization_endpoint is missing"
                    )
                })?
                .to_string();
            let token_url = params
                .get("token_endpoint")
                .and_then(|v| v.as_str())
                .ok_or_else(|| {
                    anyhow::anyhow!(
                        "provider '{provider_id}': auth.params.token_endpoint is missing"
                    )
                })?
                .to_string();
            (client_id, scope, device_url, token_url)
        }
        other => anyhow::bail!(
            "provider '{provider_id}' is not OAuth (auth scheme: {:?})",
            std::mem::discriminant(other)
        ),
    };

    let flow = bitrouter_providers::oauth::DeviceCodeFlow::new(
        bitrouter_providers::oauth::DeviceCodeParams {
            client_id,
            scope,
            device_authorization_endpoint: device_url,
            token_endpoint: token_url,
        },
    )
    .with_context(|| format!("building OAuth flow for {provider_id}"))?;

    let device_code = flow
        .request_device_code()
        .await
        .with_context(|| format!("requesting OAuth device code for {provider_id}"))?;

    eprintln!();
    eprintln!("  OAuth device authorization for {}", entry.display_name);
    eprintln!("  ──────────────────────────────────────────");
    if let Some(complete) = &device_code.verification_uri_complete {
        eprintln!("  Open:  {complete}");
    } else {
        eprintln!("  Open:  {}", device_code.verification_uri);
    }
    eprintln!("  Code:  {}", device_code.user_code);
    eprintln!();
    eprintln!("  Waiting for authorization…");

    let mut interval = std::time::Duration::from_secs(device_code.interval.max(1));
    let token = loop {
        tokio::time::sleep(interval).await;
        match flow.poll_once(&device_code.device_code).await? {
            bitrouter_providers::oauth::device_code::PollOutcome::Token(t) => break t,
            bitrouter_providers::oauth::device_code::PollOutcome::Pending => continue,
            bitrouter_providers::oauth::device_code::PollOutcome::SlowDown => {
                // RFC 8628 §3.5: server requested back-off — add 5s.
                interval += std::time::Duration::from_secs(5);
            }
        }
    };

    let mut store = bitrouter_providers::oauth::TokenStore::default_path()
        .context("opening OAuth token store")?;
    store
        .set(provider_id, token)
        .with_context(|| format!("persisting OAuth token for {provider_id}"))?;

    eprintln!("  ✓ Authorized. Token saved to {}", store.path().display());
    eprintln!();
    Ok(())
}

/// `bitrouter logout <provider>` — drop the stored OAuth token, if any.
pub async fn logout_provider(provider_id: &str) -> Result<()> {
    let mut store = bitrouter_providers::oauth::TokenStore::default_path()
        .context("opening OAuth token store")?;
    match store
        .remove(provider_id)
        .with_context(|| format!("removing OAuth token for {provider_id}"))?
    {
        Some(_) => eprintln!("  ✓ Removed stored OAuth token for {provider_id}."),
        None => eprintln!("  No stored OAuth token for {provider_id}; nothing to remove."),
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn init_writes_skip_auth_true_config() {
        let dir = std::env::temp_dir().join(format!(
            "brinit-{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        tokio::fs::create_dir_all(&dir).await.unwrap();
        let path = dir.join("bitrouter.yaml");
        init(&path).await.unwrap();
        let written = tokio::fs::read_to_string(&path).await.unwrap();
        assert!(written.contains("skip_auth: true"));
        // and it parses back into a valid Config
        let cfg = bitrouter_sdk::config::parse_with(&written, |_| Some("x".to_string())).unwrap();
        assert!(cfg.server.skip_auth);
        // refuses to overwrite
        assert!(init(&path).await.is_err());
        let _ = tokio::fs::remove_dir_all(&dir).await;
    }

    #[tokio::test]
    async fn key_sign_produces_brvk_secret_and_persists_only_hash() {
        let key = key_sign("sqlite::memory:", "u1", Some("p1")).await.unwrap();
        assert!(key.secret.starts_with("brvk_"));
        assert!(key.id.starts_with("brvk_id_"));
    }

    /// Regression: the `mcp_servers:` and `agents:` blocks in
    /// `STARTER_CONFIG` are commented out by default, but if a user
    /// uncomments them the result must still parse against the current
    /// `Config` schema. This test strips the leading `# ` from those
    /// blocks and reparses; future schema drift fails CI here rather
    /// than fails the next user's `bitrouter init`.
    #[test]
    fn starter_config_parses_with_mcp_and_agents_examples_uncommented() {
        let uncommented: String = STARTER_CONFIG
            .lines()
            .map(|line| {
                // Uncomment any line inside the mcp_servers: / agents:
                // blocks. Those blocks are characterised by lines that
                // start with "# " followed by either the block-header
                // keyword or a key-value pair shaped like YAML.
                let trimmed = line.trim_start();
                if trimmed.starts_with("# mcp_servers:")
                    || trimmed.starts_with("# agents:")
                    || trimmed.starts_with("#   ")
                    || trimmed.starts_with("#     ")
                    || trimmed.starts_with("#       ")
                {
                    let leading_spaces = line.len() - trimmed.len();
                    let stripped = trimmed
                        .strip_prefix("# ")
                        .or_else(|| trimmed.strip_prefix("#"))
                        .unwrap_or(trimmed);
                    format!("{}{}", " ".repeat(leading_spaces), stripped)
                } else {
                    line.to_string()
                }
            })
            .collect::<Vec<_>>()
            .join("\n");
        let cfg = bitrouter_sdk::config::parse_with(&uncommented, |_| Some("x".to_string()))
            .expect("uncommented STARTER_CONFIG must parse against current schema");
        assert!(
            !cfg.mcp_servers.is_empty(),
            "mcp_servers block should land as parsed entries; got: {:?}",
            cfg.mcp_servers.keys().collect::<Vec<_>>()
        );
        assert!(
            !cfg.agents.is_empty(),
            "agents block should land as parsed entries; got: {:?}",
            cfg.agents.keys().collect::<Vec<_>>()
        );
    }

    fn sample_config() -> Config {
        let yaml = r#"
providers:
  openai:
    api_base: https://api.openai.com/v1
    api_key: k1
    models: [{ id: gpt-5 }, { id: shared }]
  anthropic:
    api_base: https://api.anthropic.com/v1
    api_key: k2
    models: [{ id: shared }]
"#;
        bitrouter_sdk::config::parse_with(yaml, |_| None).unwrap()
    }

    #[tokio::test]
    async fn list_models_filters_by_provider() {
        let cfg = sample_config();
        let all = list_models(&cfg, None).await.unwrap();
        assert_eq!(all.len(), 2); // gpt-5, shared
        let openai_only = list_models(&cfg, Some("openai")).await.unwrap();
        assert_eq!(openai_only.len(), 2); // gpt-5 + shared (openai declares both)
        let anthropic_only = list_models(&cfg, Some("anthropic")).await.unwrap();
        assert_eq!(anthropic_only.len(), 1); // only `shared`
        assert_eq!(anthropic_only[0].0, "shared");
    }

    #[tokio::test]
    async fn resolve_route_returns_the_cascade_chain() {
        let cfg = sample_config();
        // `shared` is declared by both providers → a 2-hop alphabetical chain
        let chain = resolve_route(&cfg, "shared").await.unwrap();
        assert_eq!(chain.len(), 2);
        assert_eq!(chain[0].provider, "anthropic");
        assert_eq!(chain[1].provider, "openai");
        // an unknown model is a clean error (no DEFAULT_PROVIDER fallback)
        assert!(resolve_route(&cfg, "no-such-model").await.is_err());
    }

    #[test]
    fn list_providers_sorts_and_summarises() {
        let providers = list_providers(&sample_config());
        assert_eq!(providers.len(), 2);
        assert_eq!(providers[0].id, "anthropic");
        assert_eq!(providers[1].id, "openai");
        assert_eq!(providers[1].model_count, 2);
        assert!(providers[1].active);
    }

    #[tokio::test]
    async fn create_policy_rejects_path_traversal_ids() {
        let dir = std::env::temp_dir().join(format!(
            "brpoltrav-{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        // Path separators, parent refs, dotfiles, control chars — all rejected
        // before any file is written.
        for bad in [
            "../etc/passwd",
            "foo/bar",
            "..",
            ".hidden",
            "",
            "a b",
            "x\0y",
        ] {
            let err = create_policy(&dir, bad).await.unwrap_err();
            let msg = format!("{err}");
            assert!(
                msg.contains("invalid characters")
                    || msg.contains("must not")
                    || msg.contains("empty"),
                "bad id '{bad}' should have been rejected; got: {msg}"
            );
        }
        // The dir was never created (because every call bailed before touching fs).
        // (We don't assert this strongly — a parallel test could create it — just clean up.)
        let _ = tokio::fs::remove_dir_all(&dir).await;
    }

    #[tokio::test]
    async fn create_policy_writes_a_starter_file() {
        let dir = std::env::temp_dir().join(format!(
            "brpol-{}",
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        let path = create_policy(&dir, "team-default").await.unwrap();
        let body = tokio::fs::read_to_string(&path).await.unwrap();
        assert!(body.contains("id: team-default"));
        // a re-create refuses to overwrite
        assert!(create_policy(&dir, "team-default").await.is_err());
        let _ = tokio::fs::remove_dir_all(&dir).await;
    }
}