From 3a90ac2f1bc7ef2349819225afc0c38dc57abe14 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gon=C3=A7alo=20Carvalho?= Date: Wed, 27 May 2026 06:09:08 +0100 Subject: [PATCH 1/2] Harden bridge HID access with broker --- .../src/bin/agent-notify-hid-broker.rs | 40 +++ crates/agent-notify-bridge/src/display.rs | 265 ++++++++++++++++++ crates/agent-notify-bridge/src/hid_broker.rs | 171 +++++++++++ crates/agent-notify-bridge/src/ipc.rs | 114 ++++++++ crates/agent-notify-bridge/src/lib.rs | 11 + crates/agent-notify-bridge/src/main.rs | 15 +- crates/agent-notify-bridge/src/uhk.rs | 27 +- crates/agent-notify-bridge/src/worker.rs | 68 +++-- 8 files changed, 647 insertions(+), 64 deletions(-) create mode 100644 crates/agent-notify-bridge/src/bin/agent-notify-hid-broker.rs create mode 100644 crates/agent-notify-bridge/src/display.rs create mode 100644 crates/agent-notify-bridge/src/hid_broker.rs create mode 100644 crates/agent-notify-bridge/src/ipc.rs create mode 100644 crates/agent-notify-bridge/src/lib.rs diff --git a/crates/agent-notify-bridge/src/bin/agent-notify-hid-broker.rs b/crates/agent-notify-bridge/src/bin/agent-notify-hid-broker.rs new file mode 100644 index 0000000..08e20f6 --- /dev/null +++ b/crates/agent-notify-bridge/src/bin/agent-notify-hid-broker.rs @@ -0,0 +1,40 @@ +use agent_notify_bridge::hid_broker::{MockHidBackend, RealHidBackend, run_stdio}; +use clap::Parser; +use std::io::{BufReader, BufWriter}; +use tracing_subscriber::{EnvFilter, fmt}; + +#[derive(Debug, Parser)] +struct Args { + /// Run the broker over stdin/stdout. This is the default transport used by + /// agent-notify-bridge; stdout is reserved for IPC responses. + #[arg(long, default_value_t = true)] + stdio: bool, + /// Log generated UHK commands instead of touching HID hardware. + #[arg(long)] + mock_hid: bool, +} + +fn main() -> anyhow::Result<()> { + fmt() + .with_env_filter(EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into())) + .with_writer(std::io::stderr) + .init(); + + let args = Args::parse(); + if !args.stdio { + anyhow::bail!("only --stdio transport is supported"); + } + + let stdin = std::io::stdin(); + let stdout = std::io::stdout(); + let mut reader = BufReader::new(stdin.lock()); + let mut writer = BufWriter::new(stdout.lock()); + + if args.mock_hid { + let mut backend = MockHidBackend::default(); + run_stdio(&mut backend, &mut reader, &mut writer) + } else { + let mut backend = RealHidBackend; + run_stdio(&mut backend, &mut reader, &mut writer) + } +} diff --git a/crates/agent-notify-bridge/src/display.rs b/crates/agent-notify-bridge/src/display.rs new file mode 100644 index 0000000..b91820f --- /dev/null +++ b/crates/agent-notify-bridge/src/display.rs @@ -0,0 +1,265 @@ +#[cfg(windows)] +use crate::ipc::{HidBrokerRequest, HidBrokerResponse, read_message, write_message}; +use crate::uhk; +use agent_notify_core::{AgentEvent, clear_macro_command, macro_command_for_event}; +#[cfg(windows)] +use anyhow::{Context, bail}; +#[cfg(windows)] +use std::io::{BufReader, BufWriter}; +#[cfg(any(test, windows))] +use std::path::{Path, PathBuf}; +#[cfg(windows)] +use std::process::{Child, ChildStdin, ChildStdout, Command, Stdio}; + +pub struct DisplayAdapter { + inner: DisplayAdapterInner, +} + +enum DisplayAdapterInner { + Mock(MockDisplayAdapter), + Platform(PlatformDisplayAdapter), +} + +impl DisplayAdapter { + pub fn new(mock: bool) -> Self { + let inner = if mock { + DisplayAdapterInner::Mock(MockDisplayAdapter) + } else { + DisplayAdapterInner::Platform(PlatformDisplayAdapter::new()) + }; + Self { inner } + } + + pub fn keyboard_present(&mut self) -> bool { + match &mut self.inner { + DisplayAdapterInner::Mock(display) => display.keyboard_present(), + DisplayAdapterInner::Platform(display) => display.keyboard_present(), + } + } + + pub fn display_event(&mut self, event: &AgentEvent) -> anyhow::Result { + match &mut self.inner { + DisplayAdapterInner::Mock(display) => display.display_event(event), + DisplayAdapterInner::Platform(display) => display.display_event(event), + } + } + + pub fn clear(&mut self, reason: &str) -> anyhow::Result<()> { + match &mut self.inner { + DisplayAdapterInner::Mock(display) => display.clear(reason), + DisplayAdapterInner::Platform(display) => display.clear(reason), + } + } +} + +struct MockDisplayAdapter; + +impl MockDisplayAdapter { + fn keyboard_present(&mut self) -> bool { + true + } + + fn display_event(&mut self, event: &AgentEvent) -> anyhow::Result { + let command = macro_command_for_event(event)?; + tracing::info!(%command, "mock UHK display"); + Ok(command) + } + + fn clear(&mut self, reason: &str) -> anyhow::Result<()> { + let command = clear_macro_command(); + tracing::info!(%command, %reason, "mock UHK display clear"); + Ok(()) + } +} + +#[cfg(windows)] +struct PlatformDisplayAdapter { + broker: Option, +} + +#[cfg(windows)] +impl PlatformDisplayAdapter { + fn new() -> Self { + Self { broker: None } + } + + fn keyboard_present(&mut self) -> bool { + match self.request(HidBrokerRequest::ProbeKeyboard) { + Ok(HidBrokerResponse::KeyboardPresent { present }) => present, + Ok(response) => { + tracing::warn!(?response, "unexpected HID broker keyboard probe response"); + false + } + Err(err) => { + tracing::warn!(?err, "failed to probe keyboard through HID broker"); + false + } + } + } + + fn display_event(&mut self, event: &AgentEvent) -> anyhow::Result { + match self.request(HidBrokerRequest::SetDisplay { + event: event.clone(), + })? { + HidBrokerResponse::Ok { + display: Some(display), + } => Ok(display), + HidBrokerResponse::Ok { display: None } => { + bail!("HID broker did not return displayed command") + } + HidBrokerResponse::Error { code, message } => { + bail!("HID broker rejected display request ({code}): {message}") + } + response => bail!("unexpected HID broker display response: {response:?}"), + } + } + + fn clear(&mut self, reason: &str) -> anyhow::Result<()> { + match self.request(HidBrokerRequest::Clear { + reason: reason.to_string(), + })? { + HidBrokerResponse::Ok { .. } => Ok(()), + HidBrokerResponse::Error { code, message } => { + bail!("HID broker rejected clear request ({code}): {message}") + } + response => bail!("unexpected HID broker clear response: {response:?}"), + } + } + + fn request(&mut self, request: HidBrokerRequest) -> anyhow::Result { + match self.request_once(request.clone()) { + Ok(response) => Ok(response), + Err(first_err) => { + self.broker = None; + tracing::warn!(?first_err, "restarting HID broker after IPC failure"); + self.request_once(request) + } + } + } + + fn request_once(&mut self, request: HidBrokerRequest) -> anyhow::Result { + if self.broker.is_none() { + self.broker = Some(HidBrokerClient::spawn().context("failed to start HID broker")?); + } + self.broker + .as_mut() + .expect("broker is initialized") + .request(request) + } +} + +#[cfg(not(windows))] +struct PlatformDisplayAdapter; + +#[cfg(not(windows))] +impl PlatformDisplayAdapter { + fn new() -> Self { + Self + } + + fn keyboard_present(&mut self) -> bool { + uhk::keyboard_present() + } + + fn display_event(&mut self, event: &AgentEvent) -> anyhow::Result { + let command = macro_command_for_event(event)?; + uhk::display_macro_command(&command)?; + Ok(command) + } + + fn clear(&mut self, _reason: &str) -> anyhow::Result<()> { + uhk::display_macro_command(clear_macro_command()) + } +} + +#[cfg(windows)] +struct HidBrokerClient { + child: Child, + stdin: BufWriter, + stdout: BufReader, +} + +#[cfg(windows)] +impl HidBrokerClient { + fn spawn() -> anyhow::Result { + let current_exe = std::env::current_exe().context("failed to locate bridge executable")?; + let broker_path = broker_path_for_current_exe(¤t_exe); + let mut child = Command::new(&broker_path) + .arg("--stdio") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::inherit()) + .spawn() + .with_context(|| format!("failed to spawn {}", broker_path.display()))?; + let stdin = child.stdin.take().context("HID broker stdin missing")?; + let stdout = child.stdout.take().context("HID broker stdout missing")?; + Ok(Self { + child, + stdin: BufWriter::new(stdin), + stdout: BufReader::new(stdout), + }) + } + + fn request(&mut self, request: HidBrokerRequest) -> anyhow::Result { + write_message(&mut self.stdin, &request)?; + let response = read_message(&mut self.stdout)?.context("HID broker closed its stdout")?; + Ok(response) + } +} + +#[cfg(windows)] +impl Drop for HidBrokerClient { + fn drop(&mut self) { + let _ = write_message(&mut self.stdin, &HidBrokerRequest::Shutdown); + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +#[cfg(any(test, windows))] +fn broker_path_for_current_exe(current_exe: &Path) -> PathBuf { + current_exe.with_file_name(format!( + "agent-notify-hid-broker{}", + std::env::consts::EXE_SUFFIX + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_notify_core::{AgentEventInput, AgentState}; + + fn sample_event() -> AgentEvent { + AgentEventInput { + agent: "codex".to_string(), + host: "workstation".to_string(), + repo: Some("agent-notify".to_string()), + state: AgentState::Done, + summary: Some("complete".to_string()), + priority: None, + ttl_seconds: Some(60), + run_id: None, + } + .into_event() + .unwrap() + } + + #[test] + fn mock_display_generates_macro_without_broker() { + let mut display = DisplayAdapter::new(true); + assert!(display.keyboard_present()); + let command = display.display_event(&sample_event()).unwrap(); + assert!(command.starts_with("notify ")); + display.clear("test").unwrap(); + } + + #[test] + fn broker_path_is_sibling_binary() { + let current = Path::new(r"C:\tools\agent-notify-bridge.exe"); + let broker = broker_path_for_current_exe(current); + assert!(broker.ends_with(format!( + "agent-notify-hid-broker{}", + std::env::consts::EXE_SUFFIX + ))); + } +} diff --git a/crates/agent-notify-bridge/src/hid_broker.rs b/crates/agent-notify-bridge/src/hid_broker.rs new file mode 100644 index 0000000..17f4b91 --- /dev/null +++ b/crates/agent-notify-bridge/src/hid_broker.rs @@ -0,0 +1,171 @@ +use crate::ipc::{HidBrokerRequest, HidBrokerResponse, read_message, write_message}; +use crate::uhk; +use agent_notify_core::{clear_macro_command, macro_command_for_event}; +use std::io::{BufRead, Write}; +use tracing::info; + +pub trait HidBackend { + fn keyboard_present(&mut self) -> bool; + fn display_macro_command(&mut self, command: &str) -> anyhow::Result<()>; +} + +#[derive(Debug, Default)] +pub struct RealHidBackend; + +impl HidBackend for RealHidBackend { + fn keyboard_present(&mut self) -> bool { + uhk::keyboard_present() + } + + fn display_macro_command(&mut self, command: &str) -> anyhow::Result<()> { + uhk::display_macro_command(command) + } +} + +#[derive(Debug, Default)] +pub struct MockHidBackend { + pub commands: Vec, +} + +impl HidBackend for MockHidBackend { + fn keyboard_present(&mut self) -> bool { + true + } + + fn display_macro_command(&mut self, command: &str) -> anyhow::Result<()> { + info!(%command, "mock UHK display"); + self.commands.push(command.to_string()); + Ok(()) + } +} + +pub fn handle_request(backend: &mut B, request: HidBrokerRequest) -> HidBrokerResponse +where + B: HidBackend, +{ + match request { + HidBrokerRequest::ProbeKeyboard => HidBrokerResponse::KeyboardPresent { + present: backend.keyboard_present(), + }, + HidBrokerRequest::SetDisplay { event } => match macro_command_for_event(&event) { + Ok(command) => match backend.display_macro_command(&command) { + Ok(()) => HidBrokerResponse::Ok { + display: Some(command), + }, + Err(err) => error_response("hid_write_failed", err), + }, + Err(err) => error_response("invalid_event", err), + }, + HidBrokerRequest::Clear { reason } => { + let command = clear_macro_command(); + match backend.display_macro_command(command) { + Ok(()) => { + info!(%reason, "cleared UHK display"); + HidBrokerResponse::Ok { display: None } + } + Err(err) => error_response("hid_clear_failed", err), + } + } + HidBrokerRequest::Shutdown => HidBrokerResponse::Ok { display: None }, + } +} + +pub fn run_stdio(backend: &mut B, reader: &mut R, writer: &mut W) -> anyhow::Result<()> +where + B: HidBackend, + R: BufRead, + W: Write, +{ + while let Some(request) = read_message::<_, HidBrokerRequest>(reader)? { + let shutdown = matches!(request, HidBrokerRequest::Shutdown); + let response = handle_request(backend, request); + write_message(writer, &response)?; + if shutdown { + break; + } + } + Ok(()) +} + +fn error_response(code: &'static str, err: impl std::fmt::Display) -> HidBrokerResponse { + HidBrokerResponse::Error { + code: code.to_string(), + message: err.to_string(), + } +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_notify_core::{AgentEventInput, AgentState}; + use std::io::{BufReader, Cursor}; + + fn sample_event() -> agent_notify_core::AgentEvent { + AgentEventInput { + agent: "codex".to_string(), + host: "workstation".to_string(), + repo: Some("agent-notify".to_string()), + state: AgentState::WaitingInput, + summary: Some("waiting for input".to_string()), + priority: None, + ttl_seconds: Some(60), + run_id: None, + } + .into_event() + .unwrap() + } + + #[test] + fn set_display_generates_macro_inside_broker() { + let mut backend = MockHidBackend::default(); + let response = handle_request( + &mut backend, + HidBrokerRequest::SetDisplay { + event: sample_event(), + }, + ); + + assert_eq!( + response, + HidBrokerResponse::Ok { + display: backend.commands.first().cloned() + } + ); + assert_eq!(backend.commands.len(), 1); + assert!(backend.commands[0].starts_with("setLedTxt 0 notification")); + } + + #[test] + fn clear_generates_known_clear_macro_inside_broker() { + let mut backend = MockHidBackend::default(); + let response = handle_request( + &mut backend, + HidBrokerRequest::Clear { + reason: "test".to_string(), + }, + ); + + assert_eq!(response, HidBrokerResponse::Ok { display: None }); + assert_eq!(backend.commands, vec![clear_macro_command().to_string()]); + } + + #[test] + fn stdio_loop_processes_requests_until_shutdown() { + let requests = [HidBrokerRequest::ProbeKeyboard, HidBrokerRequest::Shutdown]; + let mut input = Vec::new(); + for request in requests { + write_message(&mut input, &request).unwrap(); + } + + let mut backend = MockHidBackend::default(); + let mut reader = BufReader::new(Cursor::new(input)); + let mut output = Vec::new(); + run_stdio(&mut backend, &mut reader, &mut output).unwrap(); + + let mut reader = BufReader::new(Cursor::new(output)); + let first: HidBrokerResponse = read_message(&mut reader).unwrap().unwrap(); + let second: HidBrokerResponse = read_message(&mut reader).unwrap().unwrap(); + assert_eq!(first, HidBrokerResponse::KeyboardPresent { present: true }); + assert_eq!(second, HidBrokerResponse::Ok { display: None }); + } +} diff --git a/crates/agent-notify-bridge/src/ipc.rs b/crates/agent-notify-bridge/src/ipc.rs new file mode 100644 index 0000000..a5ef987 --- /dev/null +++ b/crates/agent-notify-bridge/src/ipc.rs @@ -0,0 +1,114 @@ +use agent_notify_core::AgentEvent; +use anyhow::{Context, bail}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use std::io::{BufRead, Write}; + +pub const MAX_IPC_MESSAGE_BYTES: usize = 8192; + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum HidBrokerRequest { + ProbeKeyboard, + SetDisplay { event: AgentEvent }, + Clear { reason: String }, + Shutdown, +} + +#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum HidBrokerResponse { + Ok { display: Option }, + KeyboardPresent { present: bool }, + Error { code: String, message: String }, +} + +pub fn write_message(writer: &mut W, message: &T) -> anyhow::Result<()> +where + W: Write, + T: Serialize, +{ + let mut encoded = serde_json::to_vec(message).context("failed to encode IPC message")?; + if encoded.len() > MAX_IPC_MESSAGE_BYTES { + bail!("IPC message exceeds {MAX_IPC_MESSAGE_BYTES} bytes"); + } + encoded.push(b'\n'); + writer + .write_all(&encoded) + .context("failed to write IPC message")?; + writer.flush().context("failed to flush IPC message")?; + Ok(()) +} + +pub fn read_message(reader: &mut R) -> anyhow::Result> +where + R: BufRead, + T: DeserializeOwned, +{ + let mut encoded = Vec::new(); + let size = reader + .read_until(b'\n', &mut encoded) + .context("failed to read IPC message")?; + if size == 0 { + return Ok(None); + } + if encoded.len() > MAX_IPC_MESSAGE_BYTES { + bail!("IPC message exceeds {MAX_IPC_MESSAGE_BYTES} bytes"); + } + if encoded.last() == Some(&b'\n') { + encoded.pop(); + } + let message = serde_json::from_slice(&encoded).context("failed to decode IPC message")?; + Ok(Some(message)) +} + +#[cfg(test)] +mod tests { + use super::*; + use agent_notify_core::{AgentEventInput, AgentState}; + use std::io::{BufReader, Cursor}; + + fn sample_event() -> AgentEvent { + AgentEventInput { + agent: "codex".to_string(), + host: "workstation".to_string(), + repo: Some("agent-notify".to_string()), + state: AgentState::WaitingInput, + summary: Some("waiting for input".to_string()), + priority: None, + ttl_seconds: Some(60), + run_id: None, + } + .into_event() + .unwrap() + } + + #[test] + fn request_round_trips_as_json_line() { + let request = HidBrokerRequest::SetDisplay { + event: sample_event(), + }; + let mut encoded = Vec::new(); + write_message(&mut encoded, &request).unwrap(); + + let mut reader = BufReader::new(Cursor::new(encoded)); + let decoded: HidBrokerRequest = read_message(&mut reader).unwrap().unwrap(); + assert!(matches!(decoded, HidBrokerRequest::SetDisplay { .. })); + } + + #[test] + fn unknown_raw_macro_request_is_rejected() { + let raw = br#"{"type":"display_macro","command":"notify \"hello\""} +"#; + let mut reader = BufReader::new(Cursor::new(raw)); + let err = read_message::<_, HidBrokerRequest>(&mut reader).unwrap_err(); + assert!(err.to_string().contains("failed to decode IPC message")); + } + + #[test] + fn oversized_messages_are_rejected() { + let raw = vec![b'a'; MAX_IPC_MESSAGE_BYTES + 1]; + let mut reader = BufReader::new(Cursor::new(raw)); + let err = read_message::<_, HidBrokerRequest>(&mut reader).unwrap_err(); + assert!(err.to_string().contains("exceeds")); + } +} diff --git a/crates/agent-notify-bridge/src/lib.rs b/crates/agent-notify-bridge/src/lib.rs new file mode 100644 index 0000000..1bf039b --- /dev/null +++ b/crates/agent-notify-bridge/src/lib.rs @@ -0,0 +1,11 @@ +pub mod display; +pub mod hid_broker; +#[cfg(windows)] +pub mod icons; +pub mod ipc; +pub mod settings; +#[cfg(windows)] +pub mod tray; +pub mod uhk; +pub mod url; +pub mod worker; diff --git a/crates/agent-notify-bridge/src/main.rs b/crates/agent-notify-bridge/src/main.rs index 48fdb0b..bebe3df 100644 --- a/crates/agent-notify-bridge/src/main.rs +++ b/crates/agent-notify-bridge/src/main.rs @@ -1,14 +1,5 @@ -#[cfg(windows)] -mod icons; -mod settings; -#[cfg(windows)] -mod tray; -mod uhk; -mod url; -mod worker; - +use agent_notify_bridge::settings::{load_config, validate}; use clap::Parser; -use settings::{load_config, validate}; use tracing_subscriber::{EnvFilter, fmt}; #[derive(Debug, Parser)] @@ -46,11 +37,11 @@ fn main() -> anyhow::Result<()> { #[cfg(windows)] { - tray::run_windows_tray(config) + agent_notify_bridge::tray::run_windows_tray(config) } #[cfg(not(windows))] { - worker::run_console(config) + agent_notify_bridge::worker::run_console(config) } } diff --git a/crates/agent-notify-bridge/src/uhk.rs b/crates/agent-notify-bridge/src/uhk.rs index 2f61f7d..b222619 100644 --- a/crates/agent-notify-bridge/src/uhk.rs +++ b/crates/agent-notify-bridge/src/uhk.rs @@ -3,31 +3,12 @@ use agent_notify_core::uhk_exec_macro_report; #[cfg(windows)] use anyhow::Context; -pub struct DisplayAdapter { - mock: bool, +pub fn keyboard_present() -> bool { + platform::keyboard_present() } -impl DisplayAdapter { - pub fn new(mock: bool) -> Self { - Self { mock } - } - - pub fn keyboard_present(&self) -> bool { - if self.mock { - return true; - } - - platform::keyboard_present() - } - - pub fn display_macro_command(&self, command: &str) -> anyhow::Result<()> { - if self.mock { - tracing::info!(%command, "mock UHK display"); - return Ok(()); - } - - platform::display_macro_command(command) - } +pub fn display_macro_command(command: &str) -> anyhow::Result<()> { + platform::display_macro_command(command) } #[cfg(windows)] diff --git a/crates/agent-notify-bridge/src/worker.rs b/crates/agent-notify-bridge/src/worker.rs index 4fb4fc7..cc7ec75 100644 --- a/crates/agent-notify-bridge/src/worker.rs +++ b/crates/agent-notify-bridge/src/worker.rs @@ -1,9 +1,9 @@ +use crate::display::DisplayAdapter; use crate::settings::BridgeConfig; -use crate::uhk::DisplayAdapter; use crate::url::{redacted_url, websocket_url}; use agent_notify_core::{ - AgentEventInput, AgentState, BridgeClientMessage, BridgeServerMessage, clear_macro_command, - local_hostname as detect_local_hostname, macro_command_for_event, + AgentEvent, AgentEventInput, AgentState, BridgeClientMessage, BridgeServerMessage, + local_hostname as detect_local_hostname, }; use anyhow::Context; use futures_util::{SinkExt, StreamExt}; @@ -89,6 +89,12 @@ impl BridgeRuntimeState { } } +impl Default for BridgeRuntimeState { + fn default() -> Self { + Self::new() + } +} + #[derive(Debug, Eq, PartialEq)] enum BridgeExit { Disconnected, @@ -183,13 +189,13 @@ async fn bridge_session( icon_sink: &IconSink, ) -> anyhow::Result { let url = websocket_url(&config.server_url, &config.token)?; - let display = DisplayAdapter::new(config.mock_display); + let mut display = DisplayAdapter::new(config.mock_display); let (mut ws, _) = tokio_tungstenite::connect_async(&url) .await .context("failed to connect websocket")?; info!(url = %redacted_url(&url), "connected to agent-notify server"); - send_status(&mut ws, config, &display, state, None).await?; + send_status(&mut ws, config, &mut display, state, None).await?; ws.send(Message::Text( serde_json::to_string(&BridgeClientMessage::RequestLatest)?.into(), )) @@ -209,7 +215,7 @@ async fn bridge_session( loop { tokio::select! { _ = heartbeat.tick() => { - send_status(&mut ws, config, &display, state, last_display.clone()).await?; + send_status(&mut ws, config, &mut display, state, last_display.clone()).await?; } command = commands.recv() => { match command { @@ -218,26 +224,26 @@ async fn bridge_session( Some(BridgeCommand::SetPaused(paused)) => { state.paused.store(paused, Ordering::Relaxed); icon_sink(effective_icon(current_state, paused)); - send_status(&mut ws, config, &display, state, last_display.clone()).await?; + send_status(&mut ws, config, &mut display, state, last_display.clone()).await?; } Some(BridgeCommand::Test) => { - let command = test_macro_command(config)?; - display.display_macro_command(&command)?; - last_display = Some(command); + let event = test_event(config)?; + let display_text = display.display_event(&event)?; + last_display = Some(display_text); // The test notification is a synthetic Done event. current_state = Some(AgentState::Done); icon_sink(effective_icon(current_state, state.paused.load(Ordering::Relaxed))); - send_status(&mut ws, config, &display, state, last_display.clone()).await?; + send_status(&mut ws, config, &mut display, state, last_display.clone()).await?; } Some(BridgeCommand::Dismiss) => { ws.send(Message::Text( serde_json::to_string(&BridgeClientMessage::DismissLatest)?.into(), )) .await?; - clear_display(&display, &mut last_display, "tray"); + clear_display(&mut display, &mut last_display, "tray"); current_state = None; icon_sink(effective_icon(current_state, state.paused.load(Ordering::Relaxed))); - send_status(&mut ws, config, &display, state, last_display.clone()).await?; + send_status(&mut ws, config, &mut display, state, last_display.clone()).await?; } None => return Ok(BridgeExit::Quit), } @@ -254,18 +260,23 @@ async fn bridge_session( if state.paused.load(Ordering::Relaxed) { continue; } - let command = macro_command_for_event(&event)?; - if let Err(err) = display.display_macro_command(&command) { - warn!(?err, %command, "failed to update UHK display"); + let display_text = match display.display_event(&event) { + Ok(display_text) => display_text, + Err(err) => { + warn!(?err, "failed to update UHK display"); + continue; + } + }; + if display_text.is_empty() { continue; } - last_display = Some(command); + last_display = Some(display_text); current_state = Some(event.state); icon_sink(effective_icon(current_state, false)); } BridgeServerMessage::Clear { reason } => { info!(%reason, "clear requested"); - clear_display(&display, &mut last_display, &reason); + clear_display(&mut display, &mut last_display, &reason); current_state = None; icon_sink(effective_icon( current_state, @@ -285,8 +296,8 @@ async fn bridge_session( /// Build the test notification through the same path real events take, so any /// macro-formatting regression surfaces from the tray Test action too. -fn test_macro_command(config: &BridgeConfig) -> anyhow::Result { - let event = AgentEventInput { +fn test_event(config: &BridgeConfig) -> anyhow::Result { + Ok(AgentEventInput { agent: "agent-notify".to_string(), host: config.hostname.clone().unwrap_or_else(local_hostname), repo: None, @@ -296,14 +307,12 @@ fn test_macro_command(config: &BridgeConfig) -> anyhow::Result { ttl_seconds: Some(60), run_id: None, } - .into_event()?; - Ok(macro_command_for_event(&event)?) + .into_event()?) } -fn clear_display(display: &DisplayAdapter, last_display: &mut Option, reason: &str) { - let command = clear_macro_command(); - if let Err(err) = display.display_macro_command(command) { - warn!(?err, %command, %reason, "failed to clear UHK display"); +fn clear_display(display: &mut DisplayAdapter, last_display: &mut Option, reason: &str) { + if let Err(err) = display.clear(reason) { + warn!(?err, %reason, "failed to clear UHK display"); return; } @@ -315,7 +324,7 @@ async fn send_status( tokio_tungstenite::MaybeTlsStream, >, config: &BridgeConfig, - display: &DisplayAdapter, + display: &mut DisplayAdapter, state: &BridgeRuntimeState, last_display: Option, ) -> anyhow::Result<()> { @@ -390,14 +399,15 @@ mod tests { } #[test] - fn test_macro_command_fits_uhk_payload() { + fn test_event_fits_uhk_payload() { let config = BridgeConfig { server_url: "http://127.0.0.1:8787".to_string(), token: "change-me".to_string(), hostname: Some("workstation".to_string()), mock_display: true, }; - let command = test_macro_command(&config).unwrap(); + let event = test_event(&config).unwrap(); + let command = agent_notify_core::macro_command_for_event(&event).unwrap(); assert!(command.len() <= agent_notify_core::UHK_MAX_MACRO_COMMAND_BYTES); } } From d2bc0581e59ac964d731c95ca540a279b417a98c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Gon=C3=A7alo=20Carvalho?= Date: Wed, 27 May 2026 19:21:21 +0100 Subject: [PATCH 2/2] Address bridge broker review comments --- crates/agent-notify-bridge/Cargo.toml | 1 + crates/agent-notify-bridge/src/ipc.rs | 23 ++++++++++++++++++++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/crates/agent-notify-bridge/Cargo.toml b/crates/agent-notify-bridge/Cargo.toml index 6e3a37d..1ecdd83 100644 --- a/crates/agent-notify-bridge/Cargo.toml +++ b/crates/agent-notify-bridge/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition.workspace = true license.workspace = true repository.workspace = true +default-run = "agent-notify-bridge" [dependencies] agent-notify-core.workspace = true diff --git a/crates/agent-notify-bridge/src/ipc.rs b/crates/agent-notify-bridge/src/ipc.rs index a5ef987..af90cfd 100644 --- a/crates/agent-notify-bridge/src/ipc.rs +++ b/crates/agent-notify-bridge/src/ipc.rs @@ -51,12 +51,12 @@ where if size == 0 { return Ok(None); } - if encoded.len() > MAX_IPC_MESSAGE_BYTES { - bail!("IPC message exceeds {MAX_IPC_MESSAGE_BYTES} bytes"); - } if encoded.last() == Some(&b'\n') { encoded.pop(); } + if encoded.len() > MAX_IPC_MESSAGE_BYTES { + bail!("IPC message exceeds {MAX_IPC_MESSAGE_BYTES} bytes"); + } let message = serde_json::from_slice(&encoded).context("failed to decode IPC message")?; Ok(Some(message)) } @@ -111,4 +111,21 @@ mod tests { let err = read_message::<_, HidBrokerRequest>(&mut reader).unwrap_err(); assert!(err.to_string().contains("exceeds")); } + + #[test] + fn max_size_message_with_newline_is_accepted_by_reader() { + let fixed_len = r#"{"type":"error","code":"x","message":""#.len() + r#""}"#.len(); + let message = "x".repeat(MAX_IPC_MESSAGE_BYTES - fixed_len); + let response = HidBrokerResponse::Error { + code: "x".to_string(), + message, + }; + let mut encoded = Vec::new(); + write_message(&mut encoded, &response).unwrap(); + assert_eq!(encoded.len(), MAX_IPC_MESSAGE_BYTES + 1); + + let mut reader = BufReader::new(Cursor::new(encoded)); + let decoded: HidBrokerResponse = read_message(&mut reader).unwrap().unwrap(); + assert!(matches!(decoded, HidBrokerResponse::Error { .. })); + } }