add and use command-router to simplify the way signal-gateway processes admin commands

This commit is contained in:
Chris Beck
2025-12-08 18:39:11 -07:00
parent 84e2d57688
commit 996ed1ed0f
4 changed files with 282 additions and 154 deletions
+24 -5
View File
@@ -5,7 +5,7 @@
use conf::{Conf, Subcommands};
use hyper::service::service_fn;
use hyper_util::{rt::TokioIo, server::conn::auto};
use signal_gateway::{Gateway, GatewayConfig};
use signal_gateway::{CommandRouter, Gateway, GatewayConfig, Handling};
use std::{env, fs, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
use tokio::net::TcpListener;
use tokio_util::sync::CancellationToken;
@@ -110,10 +110,29 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let token = CancellationToken::new();
let message_handler = config.admin_handler.map(|cmd| match cmd {
AdminHandlerCommand::AdminHttp(config) => config.into_handler(),
});
let gateway = Gateway::new(config.gateway, token.clone(), message_handler).await;
// Build the command router
let mut router_builder = CommandRouter::builder()
.route("--help", Handling::Help)
.route("-h", Handling::Help);
// Add custom handler for "#" prefix if configured
if let Some(admin_handler) = config.admin_handler {
let handler = match admin_handler {
AdminHandlerCommand::AdminHttp(config) => config.into_handler(),
};
router_builder = router_builder.route("#", Handling::Custom(handler));
}
// Add gateway commands for "/" prefix
router_builder = router_builder.route("/", Handling::GatewayCommand);
// Add Claude as default handler if configured
if config.gateway.claude.is_some() {
router_builder = router_builder.route("", Handling::Claude);
}
let command_router = router_builder.build();
let gateway = Gateway::new(config.gateway, token.clone(), command_router).await;
let listener = TcpListener::bind(config.http_listen_addr).await.unwrap();
info!("Listening for http on {}", config.http_listen_addr);
@@ -0,0 +1,190 @@
//! Command routing for Signal admin messages.
use crate::message_handler::MessageHandler;
use std::fmt::Write;
/// How a command prefix should be handled.
pub enum Handling {
/// Show the help/routing table.
Help,
/// Parse as a gateway command (e.g., /log, /query, /alerts).
GatewayCommand,
/// Send to Claude AI.
Claude,
/// Delegate to a custom message handler.
Custom(Box<dyn MessageHandler>),
}
/// Routes incoming messages to appropriate handlers based on prefix matching.
///
/// The router contains an ordered list of (prefix, handling) pairs. When routing,
/// it finds the first prefix that matches the start of the message.
pub struct CommandRouter {
routes: Vec<(String, Handling)>,
}
impl Default for CommandRouter {
fn default() -> Self {
Self { routes: Vec::new() }
}
}
impl CommandRouter {
/// Create a builder for constructing a CommandRouter.
pub fn builder() -> CommandRouterBuilder {
CommandRouterBuilder::default()
}
/// Route a message to the appropriate handler.
///
/// Returns the handling and the message with the prefix stripped for the first
/// prefix that matches the start of the message, or None if no prefix matches.
pub fn route<'a>(&self, message: &'a str) -> Option<(&'a str, &Handling)> {
self.routes.iter().find_map(|(prefix, handling)| {
message
.strip_prefix(prefix.as_str())
.map(|stripped| (stripped, handling))
})
}
/// Format the routing table as a help string.
pub fn help(&self) -> String {
let mut text = String::from("Command routing:\n");
for (prefix, handling) in &self.routes {
let prefix_display = if prefix.is_empty() {
"(default)"
} else {
prefix
};
let handling_display = match handling {
Handling::Help => "Show this help",
Handling::GatewayCommand => "Gateway commands (/log, /query, /alerts, etc.)",
Handling::Claude => "Claude AI",
Handling::Custom(_) => "Custom handler",
};
writeln!(&mut text, " {prefix_display:12} → {handling_display}").unwrap();
}
text
}
}
/// Builder for constructing a CommandRouter.
#[derive(Default)]
pub struct CommandRouterBuilder {
routes: Vec<(String, Handling)>,
}
impl CommandRouterBuilder {
/// Add a route mapping a prefix to a handling.
///
/// Routes are matched in the order they are added.
pub fn route(mut self, prefix: impl Into<String>, handling: Handling) -> Self {
self.routes.push((prefix.into(), handling));
self
}
/// Build the CommandRouter.
///
/// # Panics
///
/// Panics if any route is completely inaccessible because an earlier prefix
/// is a prefix of it.
pub fn build(self) -> CommandRouter {
// Check for inaccessible routes
for (i, (prefix_i, _)) in self.routes.iter().enumerate() {
for (j, (prefix_j, _)) in self.routes.iter().enumerate().skip(i + 1) {
if prefix_j.starts_with(prefix_i) {
panic!(
"Route '{}' at index {} is inaccessible because '{}' at index {} is a prefix of it",
prefix_j, j, prefix_i, i
);
}
}
}
CommandRouter {
routes: self.routes,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_route_exact_match() {
let router = CommandRouter::builder()
.route("/help", Handling::Help)
.route("/", Handling::GatewayCommand)
.build();
assert!(matches!(router.route("/help"), Some(("", Handling::Help))));
assert!(matches!(
router.route("/log"),
Some(("log", Handling::GatewayCommand))
));
}
#[test]
fn test_route_prefix_match() {
let router = CommandRouter::builder()
.route("#", Handling::Help)
.route("", Handling::Claude)
.build();
assert!(matches!(
router.route("#anything"),
Some(("anything", Handling::Help))
));
assert!(matches!(
router.route("hello"),
Some(("hello", Handling::Claude))
));
}
#[test]
fn test_route_no_match() {
let router = CommandRouter::builder()
.route("/", Handling::GatewayCommand)
.build();
assert!(router.route("hello").is_none());
}
#[test]
fn test_route_order_matters() {
let router = CommandRouter::builder()
.route("--help", Handling::Help)
.route("-h", Handling::Help)
.route("/", Handling::GatewayCommand)
.route("", Handling::Claude)
.build();
assert!(matches!(router.route("--help"), Some(("", Handling::Help))));
assert!(matches!(router.route("-h"), Some(("", Handling::Help))));
assert!(matches!(
router.route("/log"),
Some(("log", Handling::GatewayCommand))
));
assert!(matches!(
router.route("hello"),
Some(("hello", Handling::Claude))
));
}
#[test]
#[should_panic(expected = "inaccessible")]
fn test_inaccessible_route_panics() {
CommandRouter::builder()
.route("/", Handling::GatewayCommand)
.route("/help", Handling::Help) // This is inaccessible because "/" comes first
.build();
}
#[test]
fn test_empty_router() {
let router = CommandRouter::default();
assert!(router.route("anything").is_none());
}
}
+67 -148
View File
@@ -6,9 +6,7 @@ use crate::{
alertmanager::AlertPost,
claude::{ClaudeApi, ClaudeConfig, SentBy, Tool, ToolExecutor},
log_message::{LogMessage, Origin},
message_handler::{
AdminMessage, AdminMessageResponse, Context, MessageHandler, MessageHandlerResult,
},
message_handler::{AdminMessage, AdminMessageResponse, Context, MessageHandlerResult},
prometheus::{Prometheus, PrometheusConfig},
signal_jsonrpc::{
Envelope, MessageTarget, RpcClient, RpcClientError, SignalMessage, connect_tcp,
@@ -36,6 +34,9 @@ use tracing::{debug, error, info, warn};
mod signal_trust_set;
pub use signal_trust_set::SignalTrustSet;
mod command_router;
pub use command_router::{CommandRouter, CommandRouterBuilder, Handling};
mod log_buffer;
mod log_handler;
use log_handler::{LogHandler, LogHandlerConfig};
@@ -137,15 +138,8 @@ enum GatewayCommand {
/// Show current alerts from prometheus
#[conf(name = "alerts", alias = "ALERTS")]
Alerts,
/// Ask Claude AI a question
#[conf(name = "c", alias = "C")]
Claude {
/// The prompt to send to Claude
#[conf(repeat, pos)]
prompt: Vec<String>,
},
/// Stop current Claude request
#[conf(name = "cs", alias = "CS")]
#[conf(name = "stop", alias = "STOP")]
ClaudeStop,
/// Compact Claude's message history
#[conf(name = "compact", alias = "COMPACT")]
@@ -161,21 +155,6 @@ fn parse_gateway_command(s: &str) -> Result<GatewayCommand, String> {
return Err("Empty command".to_string());
}
// Special case for /c command: take everything after "c " as a single prompt
// This avoids issues with dashes being interpreted as flags
if s.eq_ignore_ascii_case("c") {
return Err("Empty prompt".to_string());
}
if let Some(prompt) = s.strip_prefix("c ").or_else(|| s.strip_prefix("C ")) {
let prompt = prompt.trim();
if prompt.is_empty() {
return Err("Empty prompt".to_string());
}
return Ok(GatewayCommand::Claude {
prompt: vec![prompt.to_string()],
});
}
// Parse using Conf, treating the input as command line arguments
// Prepend a dummy program name since Conf expects argv[0]
let args = std::iter::once("signal-gateway")
@@ -251,8 +230,8 @@ pub struct Gateway {
prometheus: Option<Prometheus>,
/// Log handler for processing log messages from all origins.
log_handler: LogHandler,
/// Handler for admin messages that don't start with `/`
message_handler: Option<Box<dyn MessageHandler>>,
/// Command router for dispatching admin messages.
command_router: CommandRouter,
/// Claude API client for AI-powered responses.
/// Initialized after Arc creation so it can hold a weak reference back to Gateway.
claude: OnceLock<Box<ClaudeApi>>,
@@ -263,7 +242,7 @@ impl Gateway {
pub async fn new(
config: GatewayConfig,
token: CancellationToken,
message_handler: Option<Box<dyn MessageHandler>>,
command_router: CommandRouter,
) -> Arc<Self> {
let (signal_alert_mq_tx, signal_alert_mq_rx) = unbounded_channel();
@@ -285,7 +264,7 @@ impl Gateway {
token,
prometheus,
log_handler,
message_handler,
command_router,
claude: OnceLock::new(),
});
@@ -456,6 +435,11 @@ impl Gateway {
})
);
// If no route matched, don't send a response
let Some(resp) = resp else {
continue;
};
let resp = resp.unwrap_or_else(|(code, msg)| {
let text = format!("{code}: {msg}");
error!("Message handler error: {text}");
@@ -489,53 +473,66 @@ impl Gateway {
}
}
// Returns Err in case of a timeout or handler error
// Returns Ok when success or error text is generated
async fn handle_signal_admin_message(&self, msg: &Envelope) -> MessageHandlerResult {
// Returns None if no route matched (don't send a response)
// Returns Some(Err) in case of a handler error
// Returns Some(Ok) when success or error text is generated
async fn handle_signal_admin_message(&self, msg: &Envelope) -> Option<MessageHandlerResult> {
let data = msg.data_message.as_ref().unwrap();
// Admin messages starting with / are handled by gateway
// Other messages are passed to the configured message handler
if data.message.starts_with("/") {
// Parse the command using conf
let cmd = parse_gateway_command(&data.message).map_err(|err| (400u16, err.into()))?;
let (stripped_message, handling) = self.command_router.route(&data.message)?;
// Record this as a system command in Claude's history, unless it's a Claude
// prompt command (which will be recorded when we call request())
let is_claude = matches!(cmd, GatewayCommand::Claude { .. });
if !is_claude {
match handling {
Handling::Help => {
// Record as system command
if let Some(claude) = self.claude.get() {
claude.record_message(
SentBy::UserToSystem,
&data.message,
data.timestamp,
);
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
}
Some(Ok(AdminMessageResponse::new(self.command_router.help())))
}
Handling::GatewayCommand => {
// Parse the command using conf (stripped message has "/" prefix removed)
let cmd = match parse_gateway_command(stripped_message) {
Ok(cmd) => cmd,
Err(err) => return Some(Err((400u16, err.into()))),
};
// Record this as a system command in Claude's history
if let Some(claude) = self.claude.get() {
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
}
let resp = match self.handle_gateway_command(cmd).await {
Ok(resp) => resp,
Err(err) => return Some(Err(err)),
};
Some(Ok(resp))
}
Handling::Claude => {
// Send directly to Claude (use stripped message)
let Some(claude) = self.claude.get() else {
return Some(Err((501u16, "Claude is not configured".into())));
};
match claude.request(stripped_message, data.timestamp).await {
Ok(response) => Some(Ok(AdminMessageResponse::new(response).from_claude())),
Err(err) => Some(Err((500, err.to_string().into()))),
}
}
Handling::Custom(handler) => {
// Record as system message
if let Some(claude) = self.claude.get() {
claude.record_message(SentBy::UserToSystem, &data.message, data.timestamp);
}
let resp = self.handle_gateway_command(cmd, data.timestamp).await?;
Ok(if is_claude { resp.from_claude() } else { resp })
} else if let Some(handler) = &self.message_handler {
// Record this as a system message in Claude's history (not directed at Claude)
if let Some(claude) = self.claude.get() {
claude.record_message(
SentBy::UserToSystem,
&data.message,
data.timestamp,
);
// Pass stripped message to custom handler
let admin_msg = AdminMessage {
message: stripped_message.to_owned(),
timestamp: data.timestamp,
sender_uuid: msg.source_uuid.clone(),
group_id: data.group_info.as_ref().map(|g| g.group_id.clone()),
};
Some(handler.handle_verified_signal_message(admin_msg, &GatewayContext).await)
}
let msg = AdminMessage {
message: data.message.clone(),
timestamp: data.timestamp,
sender_uuid: msg.source_uuid.clone(),
group_id: data.group_info.as_ref().map(|g| g.group_id.clone()),
};
handler
.handle_verified_signal_message(msg, &GatewayContext)
.await
} else {
Err((501u16, "No message handler configured".into()))
}
}
@@ -601,11 +598,7 @@ impl Gateway {
}
}
async fn handle_gateway_command(
&self,
cmd: GatewayCommand,
ts_ms: u64,
) -> MessageHandlerResult {
async fn handle_gateway_command(&self, cmd: GatewayCommand) -> MessageHandlerResult {
match cmd {
GatewayCommand::Log { filter } => {
let text = self.log_handler.format_logs(filter.as_deref()).await;
@@ -745,18 +738,6 @@ impl Gateway {
Err(err) => Err((500, err)),
}
}
GatewayCommand::Claude { prompt } => {
let claude = self
.claude
.get()
.ok_or_else(|| (501u16, "claude was not configured".into()))?;
let prompt_text = prompt.join(" ");
match claude.request(&prompt_text, ts_ms).await {
Ok(response) => Ok(AdminMessageResponse::new(response)),
Err(err) => Err((500, err.to_string().into())),
}
}
GatewayCommand::ClaudeStop => {
let claude = self
.claude
@@ -1017,66 +998,4 @@ mod tests {
assert!(parse_gateway_command("/").is_err());
assert!(parse_gateway_command("").is_err());
}
#[test]
fn test_parse_claude_command() {
// Simple text - now captured as single string
let cmd = parse_gateway_command("/c hello world").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["hello world"]);
} else {
panic!("Expected Claude command");
}
// Paragraph of text
let cmd = parse_gateway_command("/c This is a longer prompt with multiple words").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["This is a longer prompt with multiple words"]);
} else {
panic!("Expected Claude command");
}
// Text with double dash - now preserved
let cmd = parse_gateway_command("/c hello -- world").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["hello -- world"]);
} else {
panic!("Expected Claude command");
}
// Text with flag-like content - now works
let cmd = parse_gateway_command("/c what does -f mean").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["what does -f mean"]);
} else {
panic!("Expected Claude command");
}
let cmd = parse_gateway_command("/c what does --flag mean").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["what does --flag mean"]);
} else {
panic!("Expected Claude command");
}
// Text with dashes in words
let cmd = parse_gateway_command("/c explain self-documenting code").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["explain self-documenting code"]);
} else {
panic!("Expected Claude command");
}
// Uppercase /C works too
let cmd = parse_gateway_command("/C hello").unwrap();
if let GatewayCommand::Claude { prompt } = cmd {
assert_eq!(prompt, vec!["hello"]);
} else {
panic!("Expected Claude command");
}
// Empty prompt should fail
assert!(parse_gateway_command("/c ").is_err());
assert!(parse_gateway_command("/c").is_err()); // No space, falls through to conf parser
}
}
+1 -1
View File
@@ -19,7 +19,7 @@ pub(crate) mod rate_limiter;
pub(crate) mod signal_jsonrpc;
pub(crate) mod transports;
pub use gateway::{Gateway, GatewayConfig};
pub use gateway::{CommandRouter, CommandRouterBuilder, Gateway, GatewayConfig, Handling};
pub use log_message::{Level, LogFilter, LogMessage, LogMessageBuilder};
pub use message_handler::{
AdminMessage, AdminMessageResponse, Context, MessageHandler, MessageHandlerResult,