add and use command-router to simplify the way signal-gateway processes admin commands
This commit is contained in:
@@ -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());
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user