From caac7a3a573e94587a4aa8fed024facb70830b34 Mon Sep 17 00:00:00 2001 From: sh14y Date: Fri, 10 Apr 2026 22:25:34 +0800 Subject: [PATCH 1/5] feat: add !model command for runtime model switching via session/set_model - Add ModelInfo struct and model state fields to AcpConnection - Parse available models from session/new response - Add session_set_model() and resolve_model_alias() methods - Add !model Discord command to list and switch models - Support aliases: opus, sonnet, haiku, auto Kiro ACP session/set_model verified on Kiro CLI v1.29.3 --- src/acp/connection.rs | 90 +++++++++++++++++++++++++++++++++++++++++++ src/discord.rs | 64 ++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+) diff --git a/src/acp/connection.rs b/src/acp/connection.rs index 9fc5a1f1..775d035b 100644 --- a/src/acp/connection.rs +++ b/src/acp/connection.rs @@ -20,6 +20,13 @@ fn expand_env(val: &str) -> String { } use tokio::time::Instant; +#[derive(Debug, Clone)] +pub struct ModelInfo { + pub model_id: String, + pub name: String, + pub description: String, +} + pub struct AcpConnection { _proc: Child, stdin: Arc>, @@ -29,6 +36,8 @@ pub struct AcpConnection { pub acp_session_id: Option, pub last_active: Instant, pub session_reset: bool, + pub current_model: String, + pub available_models: Vec, _reader_handle: JoinHandle<()>, } @@ -163,6 +172,8 @@ impl AcpConnection { acp_session_id: None, last_active: Instant::now(), session_reset: false, + current_model: "auto".to_string(), + available_models: Vec::new(), _reader_handle: reader_handle, }) } @@ -239,9 +250,88 @@ impl AcpConnection { info!(session_id = %session_id, "session created"); self.acp_session_id = Some(session_id.clone()); + + if let Some(models) = resp.result.as_ref().and_then(|r| r.get("models")) { + if let Some(current) = models.get("currentModelId").and_then(|v| v.as_str()) { + self.current_model = current.to_string(); + } + if let Some(arr) = models.get("availableModels").and_then(|v| v.as_array()) { + self.available_models = arr + .iter() + .filter_map(|m| { + let model_id = m.get("modelId")?.as_str()?.to_string(); + let name = m + .get("name") + .and_then(|v| v.as_str()) + .unwrap_or(&model_id) + .to_string(); + let description = m + .get("description") + .and_then(|v| v.as_str()) + .unwrap_or("") + .to_string(); + if description.contains("[Deprecated]") + || description.contains("[Internal]") + { + return None; + } + Some(ModelInfo { + model_id, + name, + description, + }) + }) + .collect(); + info!( + count = self.available_models.len(), + current = %self.current_model, + "parsed available models" + ); + } + } + Ok(session_id) } + pub async fn session_set_model(&mut self, model_id: &str) -> Result<()> { + let session_id = self + .acp_session_id + .as_ref() + .ok_or_else(|| anyhow!("no active session"))? + .clone(); + self.send_request( + "session/set_model", + Some(json!({ + "sessionId": session_id, + "modelId": model_id, + })), + ) + .await?; + self.current_model = model_id.to_string(); + Ok(()) + } + + pub fn resolve_model_alias(&self, input: &str) -> Option { + let lower = input.to_lowercase(); + if self.available_models.iter().any(|m| m.model_id == lower) { + return Some(lower); + } + let candidate = match lower.as_str() { + "auto" => "auto", + "opus" => "claude-opus-4.6", + "sonnet" => "claude-sonnet-4.6", + "haiku" => "claude-haiku-4.5", + other => other, + }; + if candidate == "auto" + || self.available_models.iter().any(|m| m.model_id == candidate) + { + Some(candidate.to_string()) + } else { + None + } + } + /// Send a prompt and return a receiver for streaming notifications. /// The final message on the channel will have id set (the prompt response). pub async fn session_prompt( diff --git a/src/discord.rs b/src/discord.rs index 5b4bb8b0..5180f43c 100644 --- a/src/discord.rs +++ b/src/discord.rs @@ -82,6 +82,70 @@ impl EventHandler for Handler { return; } + // Handle !model command (intercept before normal prompt flow) + if prompt.starts_with("!model") { + let arg = prompt[6..].trim().to_string(); + let thread_key = msg.channel_id.get().to_string(); + + // Ensure session exists so we have available_models populated + if let Err(e) = self.pool.get_or_create(&thread_key).await { + let _ = msg.channel_id.say(&ctx.http, format!("⚠️ Failed to start agent: {e}")).await; + return; + } + + let reply = self + .pool + .with_connection(&thread_key, |conn| { + let arg = arg.clone(); + Box::pin(async move { + if arg.is_empty() { + // List models + if conn.available_models.is_empty() { + return Ok::( + "_(no models reported by agent — backend may not support model switching)_".to_string(), + ); + } + let mut out = String::from("**Available models:**\n"); + for m in &conn.available_models { + let marker = if m.model_id == conn.current_model { "**▶**" } else { "•" }; + out.push_str(&format!( + "{} `{}` — {}\n", + marker, m.model_id, m.name + )); + } + out.push_str(&format!("\nCurrent: `{}`\n", conn.current_model)); + out.push_str("\nAliases: `opus`, `sonnet`, `haiku`, `auto`"); + Ok(out) + } else { + match conn.resolve_model_alias(&arg) { + Some(model_id) => { + match conn.session_set_model(&model_id).await { + Ok(()) => Ok(format!("✅ Switched to `{model_id}`")), + Err(e) => Ok(format!("⚠️ Failed to set model: {e}")), + } + } + None => { + let mut out = format!("❌ Unknown model: `{arg}`\n\n**Available:**\n"); + for m in &conn.available_models { + out.push_str(&format!("• `{}`\n", m.model_id)); + } + out.push_str("\nAliases: `opus`, `sonnet`, `haiku`, `auto`"); + Ok(out) + } + } + } + }) + }) + .await; + + let text = match reply { + Ok(t) => t, + Err(e) => format!("⚠️ {e}"), + }; + let _ = msg.channel_id.say(&ctx.http, text).await; + return; + } + // Inject structured sender context so the downstream CLI can identify who sent the message let display_name = msg.member.as_ref() .and_then(|m| m.nick.as_ref()) From 9151206a5f021759997a4dda5ee6c6647323eec0 Mon Sep 17 00:00:00 2001 From: sh14y Date: Fri, 10 Apr 2026 23:56:44 +0800 Subject: [PATCH 2/5] fix: self-bootstrap kiro creds and config.toml inside the binary MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Move the entrypoint.sh logic into main.rs so openab is independent of any external shell wrapper. Before loading config, it now: 1. Restores ~/.local/share/kiro-cli/data.sqlite3 from KIRO_CRED_B64 if missing (idempotent — skips when file already exists, e.g. from a mounted volume). 2. Generates the config file from DISCORD_BOT_TOKEN / DISCORD_CHANNEL_ID env vars when the path doesn't exist. Secrets are written as literal ${VAR} placeholders so they get expanded by load_config() at read time and never land on disk in plaintext. Why: Zeabur deployments were intermittently running images built without the entrypoint shim (e.g. when the source branch was switched mid-flight), producing "failed to read /etc/openab/config.toml" with no entrypoint logs. With the bootstrap inside the binary, the deploy is robust to any image build path that happens to omit our shell wrapper. The Dockerfile entrypoint.sh remains as a belt-and-suspenders no-op when present. --- Cargo.lock | 1 + Cargo.toml | 1 + src/main.rs | 121 +++++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 121 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 7fe18255..cd348477 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -764,6 +764,7 @@ name = "openab" version = "0.1.0" dependencies = [ "anyhow", + "base64", "rand 0.8.5", "regex", "serde", diff --git a/Cargo.toml b/Cargo.toml index edfdf870..49448663 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -15,3 +15,4 @@ uuid = { version = "1", features = ["v4"] } regex = "1" anyhow = "1" rand = "0.8" +base64 = "0.22" diff --git a/src/main.rs b/src/main.rs index 05bbfd84..ba0419e3 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,11 +4,12 @@ mod discord; mod format; mod reactions; +use base64::Engine; use serenity::prelude::*; use std::collections::HashSet; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::sync::Arc; -use tracing::info; +use tracing::{info, warn}; #[tokio::main] async fn main() -> anyhow::Result<()> { @@ -24,6 +25,15 @@ async fn main() -> anyhow::Result<()> { .map(PathBuf::from) .unwrap_or_else(|| PathBuf::from("config.toml")); + // Self-bootstrap: ensure Kiro credentials and config.toml exist before loading. + // This makes openab independent of any external entrypoint shim — useful when + // deployed to platforms (e.g. Zeabur) that may build the image without our + // entrypoint wrapper. Idempotent: only acts when the target file is missing. + bootstrap_kiro_credentials(); + if let Err(e) = bootstrap_config(&config_path) { + warn!(error = %e, path = %config_path.display(), "config bootstrap from env vars failed"); + } + let cfg = config::load_config(&config_path)?; info!( agent_cmd = %cfg.agent.command, @@ -84,6 +94,113 @@ async fn main() -> anyhow::Result<()> { Ok(()) } +/// Restore Kiro CLI credentials from KIRO_CRED_B64 if the target file is missing. +/// No-op when the env var is unset or the file already exists (e.g. mounted via volume). +fn bootstrap_kiro_credentials() { + let b64 = match std::env::var("KIRO_CRED_B64") { + Ok(v) if !v.is_empty() => v, + _ => { + info!("[bootstrap] KIRO_CRED_B64 not set, skipping kiro-cli credential restore"); + return; + } + }; + let home = match std::env::var("HOME") { + Ok(h) => PathBuf::from(h), + Err(_) => { + warn!("[bootstrap] HOME not set, cannot restore kiro-cli credentials"); + return; + } + }; + let target_dir = home.join(".local/share/kiro-cli"); + let target_file = target_dir.join("data.sqlite3"); + if target_file.exists() && std::fs::metadata(&target_file).map(|m| m.len() > 0).unwrap_or(false) { + info!(path = %target_file.display(), "[bootstrap] kiro-cli credentials already present, skipping restore"); + return; + } + if let Err(e) = std::fs::create_dir_all(&target_dir) { + warn!(error = %e, "[bootstrap] failed to create kiro-cli data dir"); + return; + } + let bytes = match base64::engine::general_purpose::STANDARD.decode(b64.trim()) { + Ok(b) => b, + Err(e) => { + warn!(error = %e, "[bootstrap] KIRO_CRED_B64 is not valid base64"); + return; + } + }; + if let Err(e) = std::fs::write(&target_file, &bytes) { + warn!(error = %e, "[bootstrap] failed to write kiro-cli credentials"); + return; + } + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(&target_file, std::fs::Permissions::from_mode(0o600)); + } + info!(path = %target_file.display(), bytes = bytes.len(), "[bootstrap] restored kiro-cli credentials from KIRO_CRED_B64"); +} + +/// Generate config.toml from environment variables if it doesn't exist. +/// Secrets are written as literal `${VAR}` placeholders — they are expanded by +/// `config::load_config()` at read time, so the bot token never lands on disk. +fn bootstrap_config(config_path: &Path) -> anyhow::Result<()> { + if config_path.exists() { + info!(path = %config_path.display(), "[bootstrap] config exists, skipping generation"); + return Ok(()); + } + if std::env::var("DISCORD_BOT_TOKEN").is_err() { + info!("[bootstrap] DISCORD_BOT_TOKEN not set, skipping config generation"); + return Ok(()); + } + let channel = std::env::var("DISCORD_CHANNEL_ID").unwrap_or_default(); + let template = format!( + r#"[discord] +bot_token = "${{DISCORD_BOT_TOKEN}}" +allowed_channels = ["{channel}"] + +[agent] +command = "kiro-cli" +args = ["acp", "--trust-all-tools"] +working_dir = "/home/agent" + +[pool] +max_sessions = 10 +session_ttl_hours = 24 + +[reactions] +enabled = true +remove_after_reply = false + +[reactions.emojis] +queued = "👀" +thinking = "🤔" +tool = "🔥" +coding = "👨‍💻" +web = "⚡" +done = "🆗" +error = "😱" + +[reactions.timing] +debounce_ms = 700 +stall_soft_ms = 10000 +stall_hard_ms = 30000 +done_hold_ms = 1500 +error_hold_ms = 2500 +"# + ); + if let Some(parent) = config_path.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::write(config_path, template)?; + #[cfg(unix)] + { + use std::os::unix::fs::PermissionsExt; + let _ = std::fs::set_permissions(config_path, std::fs::Permissions::from_mode(0o600)); + } + info!(path = %config_path.display(), "[bootstrap] generated config.toml from environment variables"); + Ok(()) +} + fn parse_id_set(raw: &[String], label: &str) -> anyhow::Result> { let set: HashSet = raw .iter() From bdc886df0dcf83d3b7eed986c6bd9c43d199966d Mon Sep 17 00:00:00 2001 From: sh14y Date: Sat, 11 Apr 2026 09:58:19 +0800 Subject: [PATCH 3/5] feat: replace !model text command with /model slash command + autocomplete MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Discord's slash commands give a much better UX than text commands: - Typing /m surfaces /model in Discord's native command picker - Autocomplete on the model option shows the live list of available models (and aliases) with arrow-key selection - Discord renders the response as an interaction reply instead of cluttering the channel as a regular bot message Implementation notes: - SessionPool gains a cached_models snapshot updated on every session_new(). Autocomplete must respond within Discord's 3-second deadline, which rules out spawning a fresh kiro-cli (~10s cold start). - main.rs adds a background warmup task that creates one session at startup so the cache is populated before the first user invocation. - /model is registered as a guild command on ready() for instant propagation (global commands take up to 1 hour). - The interaction handler defers the response on the set path because cold session creation can exceed Discord's 3s window for non-deferred replies. - README documents the applications.commands scope requirement; without it the slash command will not appear in Discord clients. The !model text command is removed entirely — /model is its replacement, not an addition. --- README.md | 15 ++- src/acp/pool.rs | 23 +++- src/discord.rs | 310 +++++++++++++++++++++++++++++++++++++----------- src/main.rs | 16 +++ 4 files changed, 295 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index 4cb03ab2..52b4408e 100644 --- a/README.md +++ b/README.md @@ -36,9 +36,11 @@ In short: 1. Go to https://discord.com/developers/applications and create an application 2. Bot tab → enable **Message Content Intent** -3. OAuth2 → URL Generator → scope: `bot` → permissions: Send Messages, Send Messages in Threads, Create Public Threads, Read Message History, Add Reactions, Manage Messages +3. OAuth2 → URL Generator → scopes: **`bot` + `applications.commands`** → permissions: Send Messages, Send Messages in Threads, Create Public Threads, Read Message History, Add Reactions, Manage Messages 4. Invite the bot to your server using the generated URL +> **Note:** `applications.commands` scope is required for the `/model` slash command. If you previously invited the bot with only `bot` scope, re-invite it using the new URL — the bot account stays the same, no data is lost. + ### 2. Configure ```bash @@ -82,6 +84,17 @@ In your Discord channel: The bot creates a thread. After that, just type in the thread — no @mention needed. +### Slash Commands + +| Command | Description | +|---------|-------------| +| `/model` | Show the current model and list all available models | +| `/model ` | Switch to a specific model. Supports aliases (`auto`, `opus`, `sonnet`, `haiku`) and full model ids — Discord's autocomplete shows the full list as you type. | + +The model picker uses Discord's native autocomplete: type `/m` and Discord will suggest `/model`; tab into the field and use ↑↓ to pick from the live list of available models reported by your agent backend. + +> Slash commands require the `applications.commands` OAuth scope (see [Create a Discord Bot](#1-create-a-discord-bot) above). If `/model` does not appear in your Discord client, re-invite the bot with the correct scope. + ## Pluggable Agent Backends Supports Kiro CLI, Claude Code, Codex, Gemini, and any ACP-compatible CLI. diff --git a/src/acp/pool.rs b/src/acp/pool.rs index a2c8a06c..4f7a6636 100644 --- a/src/acp/pool.rs +++ b/src/acp/pool.rs @@ -1,4 +1,4 @@ -use crate::acp::connection::AcpConnection; +use crate::acp::connection::{AcpConnection, ModelInfo}; use crate::config::AgentConfig; use anyhow::{anyhow, Result}; use std::collections::HashMap; @@ -10,6 +10,11 @@ pub struct SessionPool { connections: RwLock>, config: AgentConfig, max_sessions: usize, + /// Snapshot of available models from the most recent session creation. + /// Populated on first session_new() so slash command autocomplete can serve + /// suggestions without spawning a fresh agent (which takes ~10s). + cached_models: RwLock>, + cached_current_model: RwLock, } impl SessionPool { @@ -18,9 +23,19 @@ impl SessionPool { connections: RwLock::new(HashMap::new()), config, max_sessions, + cached_models: RwLock::new(Vec::new()), + cached_current_model: RwLock::new("auto".to_string()), } } + pub async fn cached_models(&self) -> Vec { + self.cached_models.read().await.clone() + } + + pub async fn cached_current_model(&self) -> String { + self.cached_current_model.read().await.clone() + } + pub async fn get_or_create(&self, thread_id: &str) -> Result<()> { // Check if alive connection exists { @@ -59,6 +74,12 @@ impl SessionPool { conn.initialize().await?; conn.session_new(&self.config.working_dir).await?; + // Refresh model cache snapshot for slash command autocomplete. + if !conn.available_models.is_empty() { + *self.cached_models.write().await = conn.available_models.clone(); + } + *self.cached_current_model.write().await = conn.current_model.clone(); + let is_rebuild = conns.contains_key(thread_id); if is_rebuild { conn.session_reset = true; diff --git a/src/discord.rs b/src/discord.rs index 5180f43c..77e8d27b 100644 --- a/src/discord.rs +++ b/src/discord.rs @@ -3,6 +3,13 @@ use crate::config::ReactionsConfig; use crate::format; use crate::reactions::StatusReactionController; use serenity::async_trait; +use serenity::builder::{ + AutocompleteChoice, CreateAutocompleteResponse, CreateCommand, CreateCommandOption, + CreateInteractionResponse, CreateInteractionResponseMessage, EditInteractionResponse, +}; +use serenity::model::application::{ + CommandDataOptionValue, CommandInteraction, CommandOptionType, Interaction, +}; use serenity::model::channel::{Message, ReactionType}; use serenity::model::gateway::Ready; use serenity::model::id::{ChannelId, MessageId}; @@ -10,7 +17,15 @@ use serenity::prelude::*; use std::collections::HashSet; use std::sync::Arc; use tokio::sync::watch; -use tracing::{error, info}; +use tracing::{error, info, warn}; + +/// Alias shortcuts shown in autocomplete and resolved by AcpConnection::resolve_model_alias. +const MODEL_ALIASES: &[(&str, &str)] = &[ + ("auto", "auto"), + ("opus", "claude-opus-4.6"), + ("sonnet", "claude-sonnet-4.6"), + ("haiku", "claude-haiku-4.5"), +]; pub struct Handler { pub pool: Arc, @@ -82,70 +97,6 @@ impl EventHandler for Handler { return; } - // Handle !model command (intercept before normal prompt flow) - if prompt.starts_with("!model") { - let arg = prompt[6..].trim().to_string(); - let thread_key = msg.channel_id.get().to_string(); - - // Ensure session exists so we have available_models populated - if let Err(e) = self.pool.get_or_create(&thread_key).await { - let _ = msg.channel_id.say(&ctx.http, format!("⚠️ Failed to start agent: {e}")).await; - return; - } - - let reply = self - .pool - .with_connection(&thread_key, |conn| { - let arg = arg.clone(); - Box::pin(async move { - if arg.is_empty() { - // List models - if conn.available_models.is_empty() { - return Ok::( - "_(no models reported by agent — backend may not support model switching)_".to_string(), - ); - } - let mut out = String::from("**Available models:**\n"); - for m in &conn.available_models { - let marker = if m.model_id == conn.current_model { "**▶**" } else { "•" }; - out.push_str(&format!( - "{} `{}` — {}\n", - marker, m.model_id, m.name - )); - } - out.push_str(&format!("\nCurrent: `{}`\n", conn.current_model)); - out.push_str("\nAliases: `opus`, `sonnet`, `haiku`, `auto`"); - Ok(out) - } else { - match conn.resolve_model_alias(&arg) { - Some(model_id) => { - match conn.session_set_model(&model_id).await { - Ok(()) => Ok(format!("✅ Switched to `{model_id}`")), - Err(e) => Ok(format!("⚠️ Failed to set model: {e}")), - } - } - None => { - let mut out = format!("❌ Unknown model: `{arg}`\n\n**Available:**\n"); - for m in &conn.available_models { - out.push_str(&format!("• `{}`\n", m.model_id)); - } - out.push_str("\nAliases: `opus`, `sonnet`, `haiku`, `auto`"); - Ok(out) - } - } - } - }) - }) - .await; - - let text = match reply { - Ok(t) => t, - Err(e) => format!("⚠️ {e}"), - }; - let _ = msg.channel_id.say(&ctx.http, text).await; - return; - } - // Inject structured sender context so the downstream CLI can identify who sent the message let display_name = msg.member.as_ref() .and_then(|m| m.nick.as_ref()) @@ -243,8 +194,233 @@ impl EventHandler for Handler { } } - async fn ready(&self, _ctx: Context, ready: Ready) { - info!(user = %ready.user.name, "discord bot connected"); + async fn ready(&self, ctx: Context, ready: Ready) { + info!(user = %ready.user.name, guilds = ready.guilds.len(), "discord bot connected"); + + // Register /model as a guild command in every guild we're in. + // Guild commands appear instantly (vs. global commands which can take + // up to 1 hour to propagate). + let cmd = CreateCommand::new("model") + .description("Switch or query the AI model used by this bot") + .add_option( + CreateCommandOption::new( + CommandOptionType::String, + "model", + "Model id or alias (auto / opus / sonnet / haiku) — leave empty to view current", + ) + .required(false) + .set_autocomplete(true), + ); + + for guild in &ready.guilds { + match guild.id.set_commands(&ctx.http, vec![cmd.clone()]).await { + Ok(cmds) => info!(guild_id = %guild.id, count = cmds.len(), "registered slash commands"), + Err(e) => error!(guild_id = %guild.id, error = %e, "failed to register slash commands"), + } + } + } + + async fn interaction_create(&self, ctx: Context, interaction: Interaction) { + match interaction { + Interaction::Command(cmd) if cmd.data.name == "model" => { + self.handle_model_command(&ctx, &cmd).await; + } + Interaction::Autocomplete(ac) if ac.data.name == "model" => { + self.handle_model_autocomplete(&ctx, &ac).await; + } + _ => {} + } + } +} + +impl Handler { + /// Resolve `partial` (the user's typing in the autocomplete field) to up + /// to 25 suggestions, drawing from cached aliases + cached model ids. + async fn handle_model_autocomplete(&self, ctx: &Context, ac: &CommandInteraction) { + let partial = ac + .data + .options + .first() + .and_then(|o| match &o.value { + CommandDataOptionValue::Autocomplete { value, .. } => Some(value.as_str()), + CommandDataOptionValue::String(s) => Some(s.as_str()), + _ => None, + }) + .unwrap_or("") + .to_lowercase(); + + let models = self.pool.cached_models().await; + let current = self.pool.cached_current_model().await; + + let mut choices: Vec = Vec::new(); + + // Aliases first — cheap shortcuts most users will reach for + for (alias, target) in MODEL_ALIASES { + if !partial.is_empty() && !alias.starts_with(&partial) { + continue; + } + // Only surface an alias if its target is actually available + // (or if it's "auto", which is always valid). + if *target != "auto" && !models.iter().any(|m| m.model_id == *target) { + continue; + } + let label = if *target == "auto" { + "auto (smart routing)".to_string() + } else { + format!("{alias} → {target}") + }; + choices.push(AutocompleteChoice::new(label, (*alias).to_string())); + if choices.len() >= 25 { + break; + } + } + + // Then real model ids + for m in &models { + if choices.len() >= 25 { + break; + } + if !partial.is_empty() && !m.model_id.to_lowercase().contains(&partial) { + continue; + } + let marker = if m.model_id == current { " (current)" } else { "" }; + let label = format!("{}{marker}", m.model_id); + choices.push(AutocompleteChoice::new(label, m.model_id.clone())); + } + + let response = CreateInteractionResponse::Autocomplete( + CreateAutocompleteResponse::new().set_choices(choices), + ); + if let Err(e) = ac.create_response(&ctx.http, response).await { + warn!(error = %e, "failed to send autocomplete response"); + } + } + + /// Handle the actual /model command submission. + async fn handle_model_command(&self, ctx: &Context, cmd: &CommandInteraction) { + // Allowlist: channel + let channel_id = cmd.channel_id.get(); + let in_allowed_channel = + self.allowed_channels.is_empty() || self.allowed_channels.contains(&channel_id); + + let in_thread = if !in_allowed_channel { + match cmd.channel_id.to_channel(&ctx.http).await { + Ok(serenity::model::channel::Channel::Guild(gc)) => gc + .parent_id + .map_or(false, |pid| self.allowed_channels.contains(&pid.get())), + _ => false, + } + } else { + false + }; + + if !in_allowed_channel && !in_thread { + let _ = cmd + .create_response( + &ctx.http, + CreateInteractionResponse::Message( + CreateInteractionResponseMessage::new() + .content("⚠️ This channel is not allowlisted.") + .ephemeral(true), + ), + ) + .await; + return; + } + + // Allowlist: user + if !self.allowed_users.is_empty() && !self.allowed_users.contains(&cmd.user.id.get()) { + let _ = cmd + .create_response( + &ctx.http, + CreateInteractionResponse::Message( + CreateInteractionResponseMessage::new() + .content("🚫 You are not authorized to use this command.") + .ephemeral(true), + ), + ) + .await; + return; + } + + // Extract the model option (None → list current) + let arg = cmd.data.options.first().and_then(|o| match &o.value { + CommandDataOptionValue::String(s) => { + let trimmed = s.trim(); + if trimmed.is_empty() { + None + } else { + Some(trimmed.to_string()) + } + } + _ => None, + }); + + // Defer the response — session spawn can take up to ~10s on cold pool + if let Err(e) = cmd.defer(&ctx.http).await { + error!(error = %e, "failed to defer /model response"); + return; + } + + let thread_key = cmd.channel_id.get().to_string(); + if let Err(e) = self.pool.get_or_create(&thread_key).await { + let _ = cmd + .edit_response( + &ctx.http, + EditInteractionResponse::new() + .content(format!("⚠️ Failed to start agent: {e}")), + ) + .await; + return; + } + + let reply = self + .pool + .with_connection(&thread_key, |conn| { + let arg = arg.clone(); + Box::pin(async move { + match arg { + None => { + if conn.available_models.is_empty() { + return Ok::( + "_(no models reported by agent — backend may not support model switching)_" + .to_string(), + ); + } + let mut out = String::from("**Available models:**\n"); + for m in &conn.available_models { + let marker = if m.model_id == conn.current_model { + "**▶**" + } else { + "•" + }; + out.push_str(&format!( + "{} `{}` — {}\n", + marker, m.model_id, m.name + )); + } + out.push_str(&format!("\nCurrent: `{}`", conn.current_model)); + Ok(out) + } + Some(input) => match conn.resolve_model_alias(&input) { + Some(model_id) => match conn.session_set_model(&model_id).await { + Ok(()) => Ok(format!("✅ Switched to `{model_id}`")), + Err(e) => Ok(format!("⚠️ Failed to set model: {e}")), + }, + None => Ok(format!("❌ Unknown model: `{input}`")), + }, + } + }) + }) + .await; + + let text = match reply { + Ok(t) => t, + Err(e) => format!("⚠️ {e}"), + }; + let _ = cmd + .edit_response(&ctx.http, EditInteractionResponse::new().content(text)) + .await; } } diff --git a/src/main.rs b/src/main.rs index ba0419e3..38d3766c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -66,6 +66,22 @@ async fn main() -> anyhow::Result<()> { .event_handler(handler) .await?; + // Warmup: spawn a background session so the model cache is populated + // before the first /model autocomplete fires. Without this, the first user + // to open the autocomplete picker would see an empty list (since spawning + // an agent takes ~10s — far over Discord's 3s autocomplete deadline). + let warmup_pool = pool.clone(); + tokio::spawn(async move { + info!("[warmup] preloading model cache"); + match warmup_pool.get_or_create("__warmup__").await { + Ok(()) => { + let count = warmup_pool.cached_models().await.len(); + info!(count, "[warmup] model cache populated"); + } + Err(e) => warn!(error = %e, "[warmup] failed to preload model cache"), + } + }); + // Spawn cleanup task let cleanup_pool = pool.clone(); let cleanup_handle = tokio::spawn(async move { From 90a5dd52cd749896d86ea7c4338e2546552d0613 Mon Sep 17 00:00:00 2001 From: sh14y Date: Wed, 15 Apr 2026 22:25:08 +0800 Subject: [PATCH 4/5] fix: enforce channel+user allowlist on /model slash command and autocomplete - Add interaction_allowlist_ok() helper (channel, thread-parent, user checks) - Apply to handle_model_command: reject with ephemeral reply if not allowed - Apply to handle_model_autocomplete: return empty choices if not allowed Addresses review feedback from thepagent on PR #186 Co-Authored-By: Claude Opus 4.6 --- src/discord.rs | 85 ++++++++++++++++++++++++++++++-------------------- 1 file changed, 52 insertions(+), 33 deletions(-) diff --git a/src/discord.rs b/src/discord.rs index 77e8d27b..69315cb7 100644 --- a/src/discord.rs +++ b/src/discord.rs @@ -234,9 +234,56 @@ impl EventHandler for Handler { } impl Handler { + /// Shared allowlist check for slash command + autocomplete paths. + /// Returns true when the interaction should be processed. + async fn interaction_allowlist_ok( + &self, + ctx: &Context, + channel_id: ChannelId, + user_id: u64, + ) -> bool { + let cid = channel_id.get(); + let in_allowed_channel = + self.allowed_channels.is_empty() || self.allowed_channels.contains(&cid); + + let in_thread = if !in_allowed_channel { + match channel_id.to_channel(&ctx.http).await { + Ok(serenity::model::channel::Channel::Guild(gc)) => gc + .parent_id + .map_or(false, |pid| self.allowed_channels.contains(&pid.get())), + _ => false, + } + } else { + false + }; + + if !in_allowed_channel && !in_thread { + return false; + } + + if !self.allowed_users.is_empty() && !self.allowed_users.contains(&user_id) { + return false; + } + + true + } + /// Resolve `partial` (the user's typing in the autocomplete field) to up /// to 25 suggestions, drawing from cached aliases + cached model ids. async fn handle_model_autocomplete(&self, ctx: &Context, ac: &CommandInteraction) { + // Silently return empty choices when the caller is not allowlisted — + // no point revealing available models to denied users. + if !self + .interaction_allowlist_ok(ctx, ac.channel_id, ac.user.id.get()) + .await + { + let empty = CreateInteractionResponse::Autocomplete( + CreateAutocompleteResponse::new().set_choices(Vec::new()), + ); + let _ = ac.create_response(&ctx.http, empty).await; + return; + } + let partial = ac .data .options @@ -298,44 +345,16 @@ impl Handler { /// Handle the actual /model command submission. async fn handle_model_command(&self, ctx: &Context, cmd: &CommandInteraction) { - // Allowlist: channel - let channel_id = cmd.channel_id.get(); - let in_allowed_channel = - self.allowed_channels.is_empty() || self.allowed_channels.contains(&channel_id); - - let in_thread = if !in_allowed_channel { - match cmd.channel_id.to_channel(&ctx.http).await { - Ok(serenity::model::channel::Channel::Guild(gc)) => gc - .parent_id - .map_or(false, |pid| self.allowed_channels.contains(&pid.get())), - _ => false, - } - } else { - false - }; - - if !in_allowed_channel && !in_thread { - let _ = cmd - .create_response( - &ctx.http, - CreateInteractionResponse::Message( - CreateInteractionResponseMessage::new() - .content("⚠️ This channel is not allowlisted.") - .ephemeral(true), - ), - ) - .await; - return; - } - - // Allowlist: user - if !self.allowed_users.is_empty() && !self.allowed_users.contains(&cmd.user.id.get()) { + if !self + .interaction_allowlist_ok(ctx, cmd.channel_id, cmd.user.id.get()) + .await + { let _ = cmd .create_response( &ctx.http, CreateInteractionResponse::Message( CreateInteractionResponseMessage::new() - .content("🚫 You are not authorized to use this command.") + .content("🚫 You are not authorized to use this command here.") .ephemeral(true), ), ) From 583f6d318db6a1a20b288caf6916f87bde56202c Mon Sep 17 00:00:00 2001 From: sh14y Date: Wed, 15 Apr 2026 22:39:45 +0800 Subject: [PATCH 5/5] test: add unit tests for allowlist_decision Extract pure allowlist logic into ctx-free allowlist_decision() so it can be unit-tested without a Discord runtime. interaction_allowlist_ok() becomes a thin wrapper that resolves the thread parent via HTTP and delegates. Covers: empty lists, channel in/out, thread parent in/out, user in/out, and the "thread-allowed + user-denied" combination. Co-Authored-By: Claude Opus 4.6 --- src/discord.rs | 124 ++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 107 insertions(+), 17 deletions(-) diff --git a/src/discord.rs b/src/discord.rs index 69315cb7..8c42f6e0 100644 --- a/src/discord.rs +++ b/src/discord.rs @@ -233,6 +233,29 @@ impl EventHandler for Handler { } } +/// Pure allowlist decision function — no Discord runtime required, fully unit-testable. +/// `parent_channel_id` is the resolved parent of a thread (None if not a thread or +/// lookup failed). Empty allowlists mean "allow everything". +fn allowlist_decision( + allowed_channels: &HashSet, + allowed_users: &HashSet, + channel_id: u64, + parent_channel_id: Option, + user_id: u64, +) -> bool { + let in_allowed_channel = allowed_channels.is_empty() || allowed_channels.contains(&channel_id); + let in_thread = !in_allowed_channel + && parent_channel_id.map_or(false, |pid| allowed_channels.contains(&pid)); + + if !in_allowed_channel && !in_thread { + return false; + } + if !allowed_users.is_empty() && !allowed_users.contains(&user_id) { + return false; + } + true +} + impl Handler { /// Shared allowlist check for slash command + autocomplete paths. /// Returns true when the interaction should be processed. @@ -243,29 +266,26 @@ impl Handler { user_id: u64, ) -> bool { let cid = channel_id.get(); - let in_allowed_channel = - self.allowed_channels.is_empty() || self.allowed_channels.contains(&cid); - let in_thread = if !in_allowed_channel { + // Only resolve the parent channel when the direct channel isn't already + // allowlisted — avoids an HTTP round-trip on the happy path and keeps + // autocomplete inside Discord's 3-second deadline. + let parent_id = if !self.allowed_channels.is_empty() && !self.allowed_channels.contains(&cid) { match channel_id.to_channel(&ctx.http).await { - Ok(serenity::model::channel::Channel::Guild(gc)) => gc - .parent_id - .map_or(false, |pid| self.allowed_channels.contains(&pid.get())), - _ => false, + Ok(serenity::model::channel::Channel::Guild(gc)) => gc.parent_id.map(|p| p.get()), + _ => None, } } else { - false + None }; - if !in_allowed_channel && !in_thread { - return false; - } - - if !self.allowed_users.is_empty() && !self.allowed_users.contains(&user_id) { - return false; - } - - true + allowlist_decision( + &self.allowed_channels, + &self.allowed_users, + cid, + parent_id, + user_id, + ) } /// Resolve `partial` (the user's typing in the autocomplete field) to up @@ -637,3 +657,73 @@ async fn get_or_create_thread(ctx: &Context, msg: &Message, prompt: &str) -> any Ok(thread.id.get()) } + +#[cfg(test)] +mod tests { + use super::*; + + fn set(ids: &[u64]) -> HashSet { + ids.iter().copied().collect() + } + + #[test] + fn empty_allowlists_allow_everything() { + assert!(allowlist_decision(&set(&[]), &set(&[]), 123, None, 999)); + assert!(allowlist_decision(&set(&[]), &set(&[]), 456, Some(789), 42)); + } + + #[test] + fn channel_in_allowlist_is_allowed() { + let ch = set(&[100, 200]); + assert!(allowlist_decision(&ch, &set(&[]), 100, None, 1)); + assert!(allowlist_decision(&ch, &set(&[]), 200, None, 1)); + } + + #[test] + fn channel_not_in_allowlist_is_denied() { + let ch = set(&[100]); + assert!(!allowlist_decision(&ch, &set(&[]), 999, None, 1)); + } + + #[test] + fn thread_parent_in_allowlist_is_allowed() { + let ch = set(&[100]); + // Direct channel 555 not allowlisted, but parent 100 is. + assert!(allowlist_decision(&ch, &set(&[]), 555, Some(100), 1)); + } + + #[test] + fn thread_parent_not_in_allowlist_is_denied() { + let ch = set(&[100]); + assert!(!allowlist_decision(&ch, &set(&[]), 555, Some(222), 1)); + } + + #[test] + fn user_not_in_allowlist_is_denied_even_if_channel_ok() { + let ch = set(&[100]); + let users = set(&[42]); + assert!(!allowlist_decision(&ch, &users, 100, None, 999)); + } + + #[test] + fn user_in_allowlist_with_channel_ok_is_allowed() { + let ch = set(&[100]); + let users = set(&[42]); + assert!(allowlist_decision(&ch, &users, 100, None, 42)); + } + + #[test] + fn empty_channel_allowlist_still_enforces_user_allowlist() { + let users = set(&[42]); + assert!(allowlist_decision(&set(&[]), &users, 999, None, 42)); + assert!(!allowlist_decision(&set(&[]), &users, 999, None, 7)); + } + + #[test] + fn thread_allowed_channel_plus_denied_user() { + let ch = set(&[100]); + let users = set(&[42]); + // Thread under allowed parent, but user not in allowlist. + assert!(!allowlist_decision(&ch, &users, 555, Some(100), 7)); + } +}