ai-crew-sync 0.5.3

use ai_crew_sync::{MIGRATOR, admin, client, serve, webhooks};
use anyhow::Context;
use clap::{Args, Parser, Subcommand};
use sqlx::postgres::PgPoolOptions;
use uuid::Uuid;

#[derive(Parser)]
#[command(
    name = "ai-crew-sync",
    version,
    about = "MCP coordination bus for a team of AI coding agents, backed by Postgres"
)]
struct Cli {
    /// Postgres connection string.
    #[arg(long, env = "DATABASE_URL", global = true)]
    database_url: Option<String>,

    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Apply pending database migrations and exit.
    Migrate,
    /// Run the MCP server.
    Serve(ServeArgs),
    /// Manage teams.
    #[command(subcommand)]
    Team(TeamCmd),
    /// Manage agents (one per teammate's coding agent).
    #[command(subcommand)]
    Agent(AgentCmd),
    /// Manage bearer tokens.
    #[command(subcommand)]
    Token(TokenCmd),
    /// Manage outgoing webhooks (Slack/Discord/generic).
    #[command(subcommand)]
    Webhook(WebhookCmd),
    /// Talk to a running bus from the console, as an agent. Everything the MCP
    /// tools can do: send/read messages, claim tasks, notes, presence.
    Client(client::ClientArgs),
    /// Print a ready-to-paste .mcp.json snippet.
    McpConfig {
        /// Public URL of the /mcp endpoint.
        #[arg(long, default_value = "http://localhost:8787/mcp")]
        url: String,
        /// The token issued to this agent.
        #[arg(long)]
        token: String,
    },
}

#[derive(Args)]
struct ServeArgs {
    /// Address to bind.
    #[arg(long, env = "BUS_BIND", default_value = "0.0.0.0:8787")]
    bind: String,

    /// Comma-separated hostnames accepted in the Host header. Use "*" to accept
    /// any host, which is fine behind a proxy that already validates it.
    #[arg(
        long,
        env = "BUS_ALLOWED_HOSTS",
        default_value = "localhost,127.0.0.1,0.0.0.0,[::1]"
    )]
    allowed_hosts: String,

    /// Comma-separated browser origins to accept. Empty disables the check.
    #[arg(long, env = "BUS_ALLOWED_ORIGINS", default_value = "")]
    allowed_origins: String,

    /// Run migrations on startup.
    #[arg(long, env = "BUS_AUTO_MIGRATE", default_value_t = true)]
    auto_migrate: bool,

    /// Largest MCP request body accepted, in bytes. Rejected with 413 before
    /// the JSON is parsed.
    #[arg(long, env = "BUS_MAX_REQUEST_BYTES",
          default_value_t = serve::DEFAULT_MAX_REQUEST_BYTES)]
    max_request_bytes: usize,

    /// Requests per minute per token, enforced in-process. 0 disables it.
    /// With several replicas the effective ceiling is per replica — put a
    /// hard global limit in the reverse proxy.
    #[arg(long, env = "BUS_RATE_LIMIT_PER_MINUTE",
          default_value_t = serve::DEFAULT_RATE_LIMIT_PER_MINUTE)]
    rate_limit_per_minute: u32,

    /// Signs the dashboard's read-only session cookies. Set the same value on
    /// every replica so a session works across all of them; when unset a
    /// random key is generated at startup, so sessions end at restart.
    #[arg(long, env = "BUS_DASHBOARD_SECRET")]
    dashboard_secret: Option<String>,
}

#[derive(Subcommand)]
enum TeamCmd {
    /// Set or clear a team's attachment storage quota.
    Quota {
        #[arg(long)]
        team: String,
        /// Total attachment bytes allowed. Omit to clear (unlimited).
        #[arg(long)]
        bytes: Option<i64>,
    },
    /// Report what a team is storing (counts and bytes; never content).
    Usage {
        #[arg(long)]
        team: String,
    },
    /// Trim history older than N days. Dry run unless --apply is passed.
    Prune {
        #[arg(long)]
        team: String,
        #[arg(long, default_value_t = 90)]
        older_than_days: i64,
        /// Actually delete. Without this the command only reports.
        #[arg(long)]
        apply: bool,
    },
    Create {
        #[arg(long)]
        slug: String,
        #[arg(long)]
        name: Option<String>,
    },
    List,
}

#[derive(Subcommand)]
enum AgentCmd {
    Add {
        #[arg(long)]
        team: String,
        #[arg(long)]
        name: String,
        #[arg(long)]
        display_name: Option<String>,
        /// Also mint a token for the new agent.
        #[arg(long, default_value_t = true)]
        with_token: bool,
    },
    List {
        #[arg(long)]
        team: String,
    },
    Disable {
        #[arg(long)]
        team: String,
        #[arg(long)]
        name: String,
    },
}

#[derive(Subcommand)]
enum WebhookCmd {
    Add {
        #[arg(long)]
        team: String,
        /// Destination URL (Slack/Discord webhook URL, or any JSON endpoint).
        #[arg(long)]
        url: String,
        /// Payload format: slack, discord or generic.
        #[arg(long, default_value = "slack")]
        kind: String,
        /// Comma-separated event kinds: message,task,lock,note.
        #[arg(long, default_value = "message,task")]
        events: String,
        /// Only forward messages from this channel (messages only).
        #[arg(long)]
        channel: Option<String>,
    },
    List {
        #[arg(long)]
        team: String,
    },
    Remove {
        #[arg(long)]
        id: Uuid,
    },
}

#[derive(Subcommand)]
enum TokenCmd {
    Issue {
        #[arg(long)]
        team: String,
        #[arg(long)]
        agent: String,
        #[arg(long)]
        label: Option<String>,
    },
    List {
        #[arg(long)]
        team: String,
    },
    Revoke {
        #[arg(long)]
        id: Uuid,
    },
}

fn split_csv(s: &str) -> Vec<String> {
    if s.trim() == "*" {
        return Vec::new();
    }
    s.split(',')
        .map(|v| v.trim().to_owned())
        .filter(|v| !v.is_empty())
        .collect()
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let _ = dotenvy::dotenv();
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "ai_crew_sync=info,tower_http=info,warn".into()),
        )
        .init();

    let cli = Cli::parse();

    // Commands that talk to the bus over HTTP (or to nothing at all) do not
    // need a database connection.
    match cli.command {
        Command::McpConfig { url, token } => {
            admin::print_mcp_config(&url, &token);
            return Ok(());
        }
        Command::Client(args) => return client::run(args).await,
        _ => {}
    }

    let url = cli
        .database_url
        .clone()
        .context("DATABASE_URL is not set (pass --database-url or set the env var)")?;
    let pool = PgPoolOptions::new()
        .max_connections(20)
        .acquire_timeout(std::time::Duration::from_secs(10))
        .connect(&url)
        .await
        .context("could not connect to Postgres")?;

    dispatch(cli.command, pool).await
}

/// Every command that needs a database, once the pool exists.
///
/// Kept out of `main` so process setup — logging, `.env`, argument parsing,
/// the two commands that need no database — reads as its own short function
/// rather than as a preamble to a 100-line match.
async fn dispatch(command: Command, pool: sqlx::PgPool) -> anyhow::Result<()> {
    match command {
        // `main` routes these before opening a pool; they cannot arrive here.
        Command::McpConfig { .. } | Command::Client(_) => unreachable!("handled in main"),

        Command::Migrate => {
            MIGRATOR.run(&pool).await?;
            println!("migrations applied");
        }

        Command::Serve(args) => {
            if args.auto_migrate {
                MIGRATOR.run(&pool).await?;
                tracing::info!("migrations applied");
            }
            let allowed_hosts = split_csv(&args.allowed_hosts);
            if allowed_hosts.is_empty() {
                tracing::warn!(
                    "Host header validation is disabled; make sure a proxy in front of this \
                     service validates it"
                );
            }
            let dashboard_secret = match args.dashboard_secret {
                Some(s) if !s.trim().is_empty() => s.into_bytes(),
                _ => {
                    tracing::warn!(
                        "BUS_DASHBOARD_SECRET is unset: dashboard sessions end at restart \
                         and are not shared between replicas"
                    );
                    ai_crew_sync::auth::generate_token().into_bytes()
                }
            };
            serve::run(
                pool,
                serve::ServeOptions {
                    bind: args.bind,
                    allowed_hosts,
                    allowed_origins: split_csv(&args.allowed_origins),
                    max_request_bytes: args.max_request_bytes,
                    rate_limit_per_minute: args.rate_limit_per_minute,
                    dashboard_secret,
                },
            )
            .await?;
        }

        Command::Team(cmd) => match cmd {
            TeamCmd::Create { slug, name } => admin::team_create(&pool, &slug, name).await?,
            TeamCmd::List => admin::team_list(&pool).await?,
            TeamCmd::Quota { team, bytes } => admin::team_quota(&pool, &team, bytes).await?,
            TeamCmd::Usage { team } => admin::team_usage(&pool, &team).await?,
            TeamCmd::Prune {
                team,
                older_than_days,
                apply,
            } => admin::team_prune(&pool, &team, older_than_days, apply).await?,
        },

        Command::Agent(cmd) => match cmd {
            AgentCmd::Add {
                team,
                name,
                display_name,
                with_token,
            } => admin::agent_add(&pool, &team, &name, display_name, with_token).await?,
            AgentCmd::List { team } => admin::agent_list(&pool, &team).await?,
            AgentCmd::Disable { team, name } => admin::agent_disable(&pool, &team, &name).await?,
        },

        Command::Token(cmd) => match cmd {
            TokenCmd::Issue { team, agent, label } => {
                admin::token_issue(&pool, &team, &agent, label).await?
            }
            TokenCmd::List { team } => admin::token_list(&pool, &team).await?,
            TokenCmd::Revoke { id } => admin::token_revoke(&pool, id).await?,
        },

        Command::Webhook(cmd) => match cmd {
            WebhookCmd::Add {
                team,
                url,
                kind,
                events,
                channel,
            } => webhooks::webhook_add(&pool, &team, &url, &kind, &events, channel).await?,
            WebhookCmd::List { team } => webhooks::webhook_list(&pool, &team).await?,
            WebhookCmd::Remove { id } => webhooks::webhook_remove(&pool, id).await?,
        },
    }

    Ok(())
}