use std::sync::Arc;
use serde::Deserialize;
use serde_json::{Value, json};
use tokio::{
sync::{Mutex, broadcast, mpsc},
time::Duration,
};
use tokio_util::sync::CancellationToken;
use crate::{
domain::{
codec::decode_base64,
errors::AgentResult,
policy::CommandPolicy,
protocol::{AgentMessage, BackendMessage},
},
infrastructure::{
directory_browser::DirectoryBrowser, host::build_host_hello, pi_history::PiHistoryStore,
pi_models::PiModels,
},
operational::{
connection_routes::{
AuthenticatedRequestOwner, ProcessContext, RouteAction, process_message,
},
pi_rpc_manager::PiRpcManager,
session_manager::{SessionManager, SessionOutput},
skill_service::SkillService,
},
presentation::{
connection_helpers::{
chat_response_error, error_message, hostname, invalid, now, spawn_background_request,
spawn_skill_request,
},
connection_output::{
OutputMessage, OutputSender, OutputTasks, output_channel, spawn_output_task,
},
ui_channel::{UI_CLIENT_QUEUE_CAPACITY, UiDirectChatRequest, UiHub, UiInboundMessage},
ui_handshake::UiPairingService,
},
};
const LOCAL_AGENT_ID: &str = "local";
const CLIENT_APPROVAL_RECONCILIATION_INTERVAL: Duration = Duration::from_secs(5);
const OMITTED_TRANSCRIPT_RAW: &str = "[transcript payload omitted]";
const MAX_PROTOCOL_RAW_COMPATIBILITY_BYTES: usize = 256;
struct UiServerState {
manager: Arc<Mutex<SessionManager>>,
pi_manager: Option<Arc<PiRpcManager>>,
policy: CommandPolicy,
hub: Arc<UiHub>,
output_tx: OutputSender,
output_tasks: Mutex<OutputTasks>,
skill_service: SkillService,
response_tx: mpsc::UnboundedSender<AgentMessage>,
history_operations: Arc<Mutex<()>>,
connected_at: String,
history_store: Option<PiHistoryStore>,
directory_browser: Option<DirectoryBrowser>,
pi_models: Option<PiModels>,
pairing_service: Arc<UiPairingService>,
}
pub(crate) struct UiServerResources {
pub manager: Arc<Mutex<SessionManager>>,
pub pi_manager: Option<Arc<PiRpcManager>>,
pub history_store: Option<PiHistoryStore>,
pub directory_browser: Option<DirectoryBrowser>,
pub pi_models: Option<PiModels>,
pub skill_service: SkillService,
pub history_operations: Arc<Mutex<()>>,
pub pairing_service: Arc<UiPairingService>,
}
#[derive(Clone)]
pub(crate) struct UiApplication {
state: Arc<UiServerState>,
}
impl UiApplication {
pub(crate) fn hub(&self) -> Arc<UiHub> {
self.state.hub.clone()
}
#[cfg(test)]
pub(crate) fn pairing_service(&self) -> Arc<UiPairingService> {
self.state.pairing_service.clone()
}
}
pub(crate) struct RunningUiApplication {
application: UiApplication,
pub(crate) task: tokio::task::JoinHandle<AgentResult<()>>,
}
impl RunningUiApplication {
pub(crate) fn application(&self) -> UiApplication {
self.application.clone()
}
#[cfg(test)]
pub(crate) fn hub(&self) -> Arc<UiHub> {
self.application.hub()
}
#[cfg(test)]
pub(crate) fn pairing_service(&self) -> Arc<UiPairingService> {
self.application.pairing_service()
}
}
#[derive(Debug, Deserialize)]
#[serde(tag = "kind")]
enum UiMessage {
#[serde(rename = "agent.send")]
AgentSend {
#[serde(rename = "agentId")]
agent_id: Option<String>,
message: BackendMessage,
},
#[serde(rename = "select.agent")]
SelectAgent {
#[serde(rename = "agentId")]
agent_id: String,
},
}
pub(crate) async fn start_application(
output_buffer_bytes: usize,
resources: UiServerResources,
cancellation: CancellationToken,
) -> AgentResult<RunningUiApplication> {
start_application_with_options(
output_buffer_bytes,
resources,
cancellation,
Duration::ZERO,
None,
CLIENT_APPROVAL_RECONCILIATION_INTERVAL,
)
.await
}
#[cfg(test)]
pub(crate) async fn start_application_with_reconciliation_interval_for_test(
output_buffer_bytes: usize,
resources: UiServerResources,
cancellation: CancellationToken,
reconciliation_interval: Duration,
) -> AgentResult<RunningUiApplication> {
start_application_with_options(
output_buffer_bytes,
resources,
cancellation,
Duration::ZERO,
None,
reconciliation_interval,
)
.await
}
async fn start_application_with_options(
output_buffer_bytes: usize,
resources: UiServerResources,
cancellation: CancellationToken,
chat_event_delay: Duration,
chat_event_pause: Option<(usize, Duration)>,
reconciliation_interval: Duration,
) -> AgentResult<RunningUiApplication> {
let UiServerResources {
manager,
pi_manager,
history_store,
directory_browser,
pi_models,
skill_service,
history_operations,
pairing_service,
} = resources;
let (output_tx, output_rx) = output_channel(output_buffer_bytes);
let (response_tx, response_rx) = mpsc::unbounded_channel();
let (ui_incoming_tx, ui_incoming_rx) = mpsc::channel(UI_CLIENT_QUEUE_CAPACITY);
let (owner_expired_tx, owner_expired_rx) = mpsc::unbounded_channel();
let chat_events = pi_manager.as_ref().map(|manager| manager.subscribe());
let policy = manager.lock().await.policy_snapshot();
let connected_at = now();
let initial_snapshot = snapshot_value(directory_browser.as_ref(), &connected_at);
let hub = UiHub::new_with_owner_expiry(initial_snapshot, ui_incoming_tx, owner_expired_tx);
let state = Arc::new(UiServerState {
manager,
pi_manager,
policy,
hub,
output_tx,
output_tasks: Mutex::new(OutputTasks::default()),
skill_service,
response_tx,
history_operations,
connected_at,
history_store,
directory_browser,
pi_models,
pairing_service,
});
let application = UiApplication {
state: state.clone(),
};
let task = tokio::spawn(async move {
tokio::select! {
_ = cancellation.cancelled() => Ok(()),
() = broadcast_outputs(state.clone(), output_rx) => Ok(()),
() = broadcast_response_results(state.clone(), response_rx) => Ok(()),
() = broadcast_chat_events(state.clone(), chat_events, chat_event_delay, chat_event_pause) => Ok(()),
() = process_ui_messages(state.clone(), ui_incoming_rx) => Ok(()),
() = expire_ui_owners(state.clone(), owner_expired_rx) => Ok(()),
() = reconcile_approved_clients(state.clone(), cancellation.clone(), reconciliation_interval) => Ok(()),
}
});
Ok(RunningUiApplication { application, task })
}
async fn expire_ui_owners(
state: Arc<UiServerState>,
mut owners: mpsc::UnboundedReceiver<AuthenticatedRequestOwner>,
) {
while let Some(owner) = owners.recv().await {
if let Some(manager) = state.pi_manager.as_ref() {
manager.expire_transcript_owner(&owner).await;
}
}
}
async fn reconcile_approved_clients(
state: Arc<UiServerState>,
cancellation: CancellationToken,
reconciliation_interval: Duration,
) {
let minimum_interval = Duration::from_millis(1);
let mut interval = tokio::time::interval(reconciliation_interval.max(minimum_interval));
interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
loop {
tokio::select! {
_ = cancellation.cancelled() => return,
_ = interval.tick() => {
match state.pairing_service.current_approved_clients() {
Ok(approved_clients) => {
state.hub.disconnect_revoked_clients(&approved_clients).await;
}
Err(error) => {
tracing::warn!(error = %error, "could not reconcile revoked UI clients");
}
}
}
}
}
}
async fn process_ui_messages(
state: Arc<UiServerState>,
mut incoming: mpsc::Receiver<UiInboundMessage>,
) {
while let Some(message) = incoming.recv().await {
if let Err(error) =
handle_ui_text(&message.text, message.client_id, message.source, &state).await
{
tracing::debug!(client_id = message.client_id, error = %error, "ui channel message failed");
}
}
}
async fn handle_ui_text(
text: &str,
client_id: u64,
source: crate::pairing::SourceIdentity,
state: &Arc<UiServerState>,
) -> AgentResult<()> {
let payload = match serde_json::from_str::<UiMessage>(text) {
Ok(payload) => payload,
Err(err) => {
send_ui_error(state, client_id, format!("Invalid UI JSON: {err}")).await;
return Ok(());
}
};
match payload {
UiMessage::SelectAgent { agent_id } => {
if agent_id != LOCAL_AGENT_ID {
send_ui_error(state, client_id, format!("Unknown agent: {agent_id}")).await;
return Ok(());
}
broadcast_value(state, snapshot(state), None).await;
}
UiMessage::AgentSend { agent_id, message } => {
if let Some(agent_id) = agent_id.as_deref()
&& agent_id != LOCAL_AGENT_ID
{
send_ui_error(state, client_id, format!("Unknown agent: {agent_id}")).await;
return Ok(());
}
let Some(liveness) = state.hub.owner_liveness(client_id).await else {
return Ok(());
};
process_ui_agent_message(
state,
message,
client_id,
AuthenticatedRequestOwner::ui_with_liveness(source, client_id, liveness),
)
.await?;
}
}
Ok(())
}
async fn process_ui_agent_message(
state: &Arc<UiServerState>,
message: BackendMessage,
client_id: u64,
request_owner: AuthenticatedRequestOwner,
) -> AgentResult<()> {
let session_id = message.session_id().map(str::to_string);
let request_id = message.request_id().map(str::to_string);
let is_chat_command = message.is_chat_command();
let routed_message = message.clone();
let context = ProcessContext {
manager: &state.manager,
pi_manager: state.pi_manager.as_ref(),
policy: &state.policy,
history_operations: &state.history_operations,
history_store: state.history_store.as_ref(),
directory_browser: state.directory_browser.as_ref(),
pi_models: state.pi_models.as_ref(),
request_owner,
};
match process_message(routed_message, context).await {
Ok(RouteAction::Direct { response, output }) => {
broadcast_backend_message(state, &message).await?;
if let Some(output) = output {
track_output_task(state, output).await;
}
if let Some(response) = response {
broadcast_agent_message(state, &response.message).await?;
}
}
Ok(RouteAction::Skill(request)) => {
broadcast_backend_message(state, &message).await?;
spawn_skill_request(
state.skill_service.clone(),
request,
state.response_tx.clone(),
);
}
Ok(RouteAction::Chat(request)) => {
send_ui_backend_message(state, client_id, &message).await?;
match state.hub.begin_direct_chat_request(client_id).await {
UiDirectChatRequest::Ready { closed, _permit } => {
spawn_direct_ui_chat_request(
state.clone(),
client_id,
request,
closed,
_permit,
);
}
UiDirectChatRequest::Busy => {
let error = crate::domain::errors::AgentError::new(
crate::domain::errors::ErrorCode::SnapshotBusy,
"too many pending chat requests for this UI client",
);
let response = chat_response_error(
session_id.as_deref(),
request_id.as_deref().unwrap_or_default(),
&error,
);
let _ = send_ui_agent_message(state, client_id, &response).await;
}
UiDirectChatRequest::Closed => {}
}
}
Ok(RouteAction::Background(request)) => {
broadcast_backend_message(state, &message).await?;
spawn_background_request(
request,
state.response_tx.clone(),
state.history_operations.clone(),
);
}
Err(err) => {
let error = if is_chat_command {
chat_response_error(
session_id.as_deref(),
request_id.as_deref().unwrap_or_default(),
&err,
)
} else {
error_message(session_id.as_deref(), request_id.as_deref(), &err)
};
if is_chat_command {
let _ = send_ui_agent_message(state, client_id, &error).await;
} else {
broadcast_agent_message(state, &error).await?;
}
}
}
Ok(())
}
fn spawn_direct_ui_chat_request(
state: Arc<UiServerState>,
client_id: u64,
request: crate::operational::connection_routes::ChatRequest,
_closed: CancellationToken,
_permit: tokio::sync::OwnedSemaphorePermit,
) {
let session_id = request.session_id().to_owned();
let request_id = request.request_id().to_owned();
tokio::spawn(async move {
let response = match request.execute().await {
Ok(Some(message)) => message,
Ok(None) => return,
Err(error) => chat_response_error(Some(&session_id), &request_id, &error),
};
let _ = send_ui_agent_message(&state, client_id, &response).await;
});
}
async fn broadcast_chat_events(
state: Arc<UiServerState>,
receiver: Option<broadcast::Receiver<crate::domain::chat::SequencedChatEvent>>,
start_delay: Duration,
pause_after_deliveries: Option<(usize, Duration)>,
) {
let Some(mut receiver) = receiver else {
std::future::pending::<()>().await;
return;
};
if !start_delay.is_zero() {
tokio::time::sleep(start_delay).await;
}
let mut deliveries = 0_usize;
loop {
match receiver.recv().await {
Ok(record) => {
if broadcast_agent_message(&state, &AgentMessage::from_chat_event(record))
.await
.is_err()
{
return;
}
deliveries = deliveries.saturating_add(1);
if let Some((pause_after, pause)) = pause_after_deliveries
&& deliveries == pause_after
{
tokio::time::sleep(pause).await;
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
let Some(manager) = state.pi_manager.as_ref() else {
continue;
};
let session_ids = state.hub.resync_session_ids(manager.active_session_ids());
for session_id in session_ids {
let message = AgentMessage::ChatResyncRequired {
event_sequence: None,
session_id,
reason: format!(
"local UI missed {skipped} chat events; request a snapshot"
),
};
if broadcast_agent_message(&state, &message).await.is_err() {
return;
}
}
}
Err(broadcast::error::RecvError::Closed) => return,
}
}
}
async fn broadcast_response_results(
state: Arc<UiServerState>,
mut results: mpsc::UnboundedReceiver<AgentMessage>,
) {
while let Some(message) = results.recv().await {
if broadcast_agent_message(&state, &message).await.is_err() {
break;
}
}
}
async fn track_output_task(state: &Arc<UiServerState>, output: SessionOutput) {
let task = spawn_output_task(output, state.output_tx.clone());
state.output_tasks.lock().await.push(task);
}
async fn broadcast_outputs(
state: Arc<UiServerState>,
mut output_rx: crate::presentation::connection_output::OutputReceiver,
) {
while let Some(output) = output_rx.recv().await {
if broadcast_output(&state, output).await.is_err() {
break;
}
}
}
async fn broadcast_output(state: &Arc<UiServerState>, output: OutputMessage) -> AgentResult<()> {
let exit_session_id = output.exit_session_id().map(str::to_string);
broadcast_agent_message(state, &output.message).await?;
output.mark_sent();
if let Some(session_id) = exit_session_id {
let _ = state.manager.lock().await.remove_session(&session_id);
}
Ok(())
}
async fn broadcast_backend_message(
state: &Arc<UiServerState>,
message: &BackendMessage,
) -> AgentResult<()> {
let raw = serde_json::to_string(message).map_err(|err| invalid(err.to_string()))?;
let value = serde_json::to_value(message).map_err(|err| invalid(err.to_string()))?;
let _ = broadcast_value(state, protocol_event("ui->agent", value, raw, None), None).await;
Ok(())
}
async fn broadcast_agent_message(
state: &Arc<UiServerState>,
message: &AgentMessage,
) -> AgentResult<()> {
let raw = serde_json::to_string(message).map_err(|err| invalid(err.to_string()))?;
let value = serde_json::to_value(message).map_err(|err| invalid(err.to_string()))?;
let decoded = decoded_terminal_output(message);
let _ = broadcast_value(
state,
protocol_event("agent->ui", value, raw, decoded),
Some(message),
)
.await;
Ok(())
}
async fn send_ui_backend_message(
state: &Arc<UiServerState>,
client_id: u64,
message: &BackendMessage,
) -> AgentResult<bool> {
let raw = serde_json::to_string(message).map_err(|err| invalid(err.to_string()))?;
let value = serde_json::to_value(message).map_err(|err| invalid(err.to_string()))?;
Ok(state
.hub
.send_value(client_id, protocol_event("ui->agent", value, raw, None))
.await)
}
async fn send_ui_agent_message(
state: &Arc<UiServerState>,
client_id: u64,
message: &AgentMessage,
) -> AgentResult<bool> {
let raw = serde_json::to_string(message).map_err(|err| invalid(err.to_string()))?;
let value = serde_json::to_value(message).map_err(|err| invalid(err.to_string()))?;
let delivered = state
.hub
.send_value(
client_id,
protocol_event("agent->ui", value, raw, decoded_terminal_output(message)),
)
.await;
if delivered {
state.hub.record_delivered_chat_message(message);
}
Ok(delivered)
}
fn protocol_event(
direction: &str,
message: Value,
raw: String,
decoded_output: Option<String>,
) -> Value {
let raw = if serialized_raw_compatibility_copy_is_large(&raw) {
OMITTED_TRANSCRIPT_RAW.to_owned()
} else {
raw
};
let mut value = json!({
"kind": "protocol",
"direction": direction,
"agentId": LOCAL_AGENT_ID,
"message": message,
"raw": raw,
"timestamp": now(),
});
if let Some(decoded_output) = decoded_output {
value["decodedOutput"] = Value::String(decoded_output);
}
value
}
fn serialized_raw_compatibility_copy_is_large(raw: &str) -> bool {
serde_json::to_vec(raw)
.map(|encoded| encoded.len() > MAX_PROTOCOL_RAW_COMPATIBILITY_BYTES)
.unwrap_or(true)
}
fn decoded_terminal_output(message: &AgentMessage) -> Option<String> {
let AgentMessage::TerminalOutput { data_base64, .. } = message else {
return None;
};
let bytes = decode_base64(data_base64).ok()?;
Some(String::from_utf8_lossy(&bytes).into_owned())
}
#[cfg(test)]
#[allow(clippy::items_after_test_module)]
mod protocol_event_tests {
use super::*;
use crate::{
domain::errors::{AgentError, ErrorCode},
presentation::connection_helpers::chat_response_error,
};
use regy_ui_wire::AUTHENTICATED_MAX_BYTES;
#[test]
fn protocol_event_elides_any_large_serialized_raw_compatibility_copy() {
let message = AgentMessage::ChatExited {
event_sequence: None,
session_id: "chat".into(),
message: "x".repeat(3 * 1024 * 1024),
};
let raw = serde_json::to_string(&message).unwrap();
let event = protocol_event(
"agent->ui",
serde_json::to_value(&message).unwrap(),
raw.clone(),
None,
);
assert_eq!(event["raw"], OMITTED_TRANSCRIPT_RAW);
assert!(serde_json::to_vec(&event).unwrap().len() <= AUTHENTICATED_MAX_BYTES);
let small_message = AgentMessage::ChatExited {
event_sequence: None,
session_id: "chat".into(),
message: "small".into(),
};
let small_raw = serde_json::to_string(&small_message).unwrap();
let small_event = protocol_event(
"agent->ui",
serde_json::to_value(&small_message).unwrap(),
small_raw.clone(),
None,
);
assert_eq!(small_event["raw"], small_raw);
}
#[test]
fn truncated_pi_failure_serializes_as_a_typed_outer_frame_within_the_wire_cap() {
let message = chat_response_error(
Some("chat"),
"request",
&AgentError::new(ErrorCode::HistoryReadFailed, "x".repeat(2 * 1024 * 1024)),
);
let raw = serde_json::to_string(&message).unwrap();
let event = protocol_event(
"agent->ui",
serde_json::to_value(&message).unwrap(),
raw,
None,
);
assert_eq!(event["message"]["type"], "chat.response.error");
assert_eq!(event["message"]["code"], "HISTORY_READ_FAILED");
assert!(
event["message"]["message"]
.as_str()
.unwrap()
.ends_with(" [truncated]")
);
assert_eq!(event["raw"], OMITTED_TRANSCRIPT_RAW);
assert!(serde_json::to_vec(&event).unwrap().len() <= AUTHENTICATED_MAX_BYTES);
}
#[test]
fn successful_extension_response_serializes_within_the_authenticated_outer_cap() {
let message = AgentMessage::ChatExtensionResponded {
event_sequence: None,
session_id: "s".repeat(128),
request_id: "r".repeat(128),
extension_request_id: "e".repeat(128),
};
let raw = serde_json::to_string(&message).unwrap();
let event = protocol_event(
"agent->ui",
serde_json::to_value(&message).unwrap(),
raw,
None,
);
assert_eq!(event["message"]["type"], "chat.extension.responded");
assert!(serde_json::to_vec(&event).unwrap().len() <= AUTHENTICATED_MAX_BYTES);
}
}
fn snapshot(state: &UiServerState) -> Value {
snapshot_value(state.directory_browser.as_ref(), &state.connected_at)
}
fn snapshot_value(directory_browser: Option<&DirectoryBrowser>, connected_at: &str) -> Value {
let home_dir =
directory_browser.map(|browser| browser.home_dir().to_string_lossy().into_owned());
let username = whoami::username().unwrap_or_else(|_| "unknown".to_string());
let hello = build_host_hello(LOCAL_AGENT_ID, hostname(), username, home_dir);
json!({
"kind": "snapshot",
"activeAgentId": LOCAL_AGENT_ID,
"agents": [{
"id": LOCAL_AGENT_ID,
"connectedAt": connected_at,
"hello": hello,
}],
})
}
async fn broadcast_value(
state: &Arc<UiServerState>,
value: Value,
delivered_chat_message: Option<&AgentMessage>,
) -> bool {
let delivered = state.hub.broadcast_value(&value).await;
if delivered && let Some(message) = delivered_chat_message {
state.hub.record_delivered_chat_message(message);
}
delivered
}
async fn send_ui_error(state: &UiServerState, client_id: u64, message: String) {
state
.hub
.send_value(
client_id,
json!({
"kind": "error",
"message": message,
"timestamp": now(),
}),
)
.await;
}