diff --git a/Cargo.lock b/Cargo.lock index 26ee4aa1..247290f7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5303,6 +5303,7 @@ name = "web-server" version = "0.3.0" dependencies = [ "application", + "async-trait", "backend", "base64 0.22.1", "bytes", diff --git a/crates/web-server/Cargo.toml b/crates/web-server/Cargo.toml index ecc2515a..ce9528ba 100644 --- a/crates/web-server/Cargo.toml +++ b/crates/web-server/Cargo.toml @@ -18,6 +18,7 @@ path = "src/bin/idea-serve.rs" backend = { workspace = true } application = { workspace = true } domain = { workspace = true } +async-trait = { workspace = true } tokio = { workspace = true, features = ["io-std", "rt", "net"] } serde = { workspace = true } serde_json = { workspace = true } diff --git a/crates/web-server/src/lib.rs b/crates/web-server/src/lib.rs index 4157938b..c0278dcb 100644 --- a/crates/web-server/src/lib.rs +++ b/crates/web-server/src/lib.rs @@ -64,7 +64,10 @@ use application::{ UpdateProjectContextInput, UpdateProjectMcpToolPermissionsInput, UpdateProjectPermissionsInput, UpdateProjectSystemPermissionsInput, UpdateSkillInput, WriteToTerminalInput, }; -use domain::ports::PtyHandle; +use domain::ports::{ + PtyHandle, ReplyEvent, ReplyProgress, ReplyProgressKind, ReplyProgressSource, + ReplyProgressStage, +}; use domain::IssueActor; use domain::{ AuthenticatedDevice, DeviceId, DeviceName, Project, SessionId, SessionTokenHash, SkillRef, @@ -114,6 +117,7 @@ use backend::dto::{ }; use backend::events::DomainEventDto; type PtyChunk = Vec; +type ChatChunk = Value; use backend::BackendCore; const DEFAULT_LISTEN: &str = "127.0.0.1:17373"; @@ -572,6 +576,7 @@ struct ServerState { app: Arc, pairing_code: Mutex>, ws_pty_bridge: Arc>, + ws_chat_bridge: Arc>, active_connections: Arc, _revocation_observer: Option>, security_logger: Arc, @@ -612,6 +617,7 @@ impl ServerState { config, pairing_code: Mutex::new(None), ws_pty_bridge: Arc::new(OutputBridge::new()), + ws_chat_bridge: Arc::new(OutputBridge::new()), active_connections, _revocation_observer: revocation_observer, security_logger: Arc::new(StderrSecurityLogger), @@ -640,6 +646,7 @@ impl ServerState { config, pairing_code: Mutex::new(None), ws_pty_bridge: Arc::new(OutputBridge::new()), + ws_chat_bridge: Arc::new(OutputBridge::new()), active_connections, _revocation_observer: revocation_observer, security_logger, @@ -1349,7 +1356,7 @@ async fn run_ws_connection( ) -> Result<(), String> { let (mut reader, mut writer) = stream.into_split(); let (tx, mut rx) = mpsc::channel::(WS_OUTPUT_BUFFER); - let owned = Arc::new(Mutex::new(Vec::<(SessionId, u64)>::new())); + let owned = Arc::new(Mutex::new(Vec::::new())); let registration = state.active_connections.register(device.device_id); let mut shutdown = registration.shutdown; let event_relay_task = @@ -1421,8 +1428,15 @@ async fn run_ws_connection( } if let Ok(owned) = owned.lock() { - for (session, gen) in owned.iter() { - state.ws_pty_bridge.unregister_if(session, *gen); + for attachment in owned.iter() { + match *attachment { + WsAttachment::Pty { session, gen } => { + state.ws_pty_bridge.unregister_if(&session, gen); + } + WsAttachment::Chat { session, gen } => { + state.ws_chat_bridge.unregister_if(&session, gen); + } + } } } state @@ -1460,11 +1474,17 @@ fn spawn_ws_domain_event_relay( }) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum WsAttachment { + Pty { session: SessionId, gen: u64 }, + Chat { session: SessionId, gen: u64 }, +} + async fn handle_client_frame( frame: ClientFrame, state: &Arc, tx: &mpsc::Sender, - owned: &Arc>>, + owned: &Arc>>, ) { let result = match frame.kind.as_str() { "terminal.open" | "open_terminal" => ws_open_terminal(&frame, state, tx, owned).await, @@ -1473,6 +1493,9 @@ async fn handle_client_frame( "terminal.input" | "input" => ws_input(&frame, state), "terminal.resize" | "resize" => ws_resize(&frame, state), "terminal.detach" | "detach" => ws_detach(&frame, state, owned), + "chat.attach" => ws_attach_chat(&frame, state, tx, owned).await, + "chat.send" => ws_send_chat(&frame, state, tx).await, + "chat.detach" => ws_detach_chat(&frame, state, owned), "terminal.close" | "close" => match ws_close(&frame, state, owned).await { Ok((session_id, exit_code)) => { let _ = tx @@ -1508,7 +1531,7 @@ async fn ws_open_terminal( frame: &ClientFrame, state: &Arc, tx: &mpsc::Sender, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(), ErrorDto> { let request_value = frame .payload @@ -1539,7 +1562,7 @@ async fn ws_launch_agent( frame: &ClientFrame, state: &Arc, tx: &mpsc::Sender, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(), ErrorDto> { let request_value = frame .payload @@ -1556,7 +1579,7 @@ async fn send_launch_agent_attached( frame: &ClientFrame, state: &Arc, tx: &mpsc::Sender, - owned: &Arc>>, + owned: &Arc>>, output: application::LaunchAgentOutput, ) -> Result<(), ErrorDto> { if let Some(structured) = output.structured.as_ref() { @@ -1672,7 +1695,7 @@ async fn ws_attach_terminal( frame: &ClientFrame, state: &Arc, tx: &mpsc::Sender, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(), ErrorDto> { let session_id = payload_string(&frame.payload, "sessionId")?; let sid = parse_session_id(&session_id)?; @@ -1714,11 +1737,11 @@ fn attach_sink_and_pump( state: &Arc, sid: SessionId, sink: WsPtySink, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(), ErrorDto> { let gen = state.ws_pty_bridge.register(sid, Arc::new(sink)); if let Ok(mut owned) = owned.lock() { - owned.push((sid, gen)); + owned.push(WsAttachment::Pty { session: sid, gen }); } let handle = PtyHandle { session_id: sid }; let stream = state @@ -1775,22 +1798,218 @@ fn ws_resize(frame: &ClientFrame, state: &Arc) -> Result<(), ErrorD fn ws_detach( frame: &ClientFrame, state: &Arc, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(), ErrorDto> { let sid = parse_session_id(&payload_string(&frame.payload, "sessionId")?)?; if let Ok(mut owned) = owned.lock() { - if let Some(index) = owned.iter().rposition(|(session, _)| *session == sid) { - let (_, gen) = owned.remove(index); + if let Some(index) = owned.iter().rposition( + |attachment| matches!(attachment, WsAttachment::Pty { session, .. } if *session == sid), + ) { + let WsAttachment::Pty { gen, .. } = owned.remove(index) else { + unreachable!("rposition matched a PTY attachment"); + }; state.ws_pty_bridge.unregister_if(&sid, gen); } } Ok(()) } +async fn ws_attach_chat( + frame: &ClientFrame, + state: &Arc, + tx: &mpsc::Sender, + owned: &Arc>>, +) -> Result<(), ErrorDto> { + let sid = parse_session_id(&payload_string(&frame.payload, "sessionId")?)?; + let session = state + .app + .structured_sessions + .session(&sid) + .ok_or_else(|| ErrorDto { + code: "NOT_FOUND".to_owned(), + message: "structured chat session not found".to_owned(), + })?; + tx.send(ServerFrame::chat_attached( + &frame.id, + &application::StructuredSessionDescriptor { + session_id: sid, + agent_id: domain::AgentId::from_uuid(Uuid::nil()), + node_id: domain::NodeId::from_uuid(Uuid::nil()), + conversation_id: session.conversation_id(), + }, + None, + session.conversation_id().as_deref(), + )) + .await + .map_err(|_| ErrorDto { + code: "WS_CLOSED".to_owned(), + message: "websocket output closed".to_owned(), + })?; + + let gen = state + .ws_chat_bridge + .register(sid, Arc::new(WsChatSink::new(sid, tx.clone()))); + if let Ok(mut owned) = owned.lock() { + owned.push(WsAttachment::Chat { session: sid, gen }); + } + Ok(()) +} + +async fn ws_send_chat( + frame: &ClientFrame, + state: &Arc, + tx: &mpsc::Sender, +) -> Result<(), ErrorDto> { + let sid = parse_session_id(&payload_string(&frame.payload, "sessionId")?)?; + let prompt = payload_string(&frame.payload, "prompt")?; + let session = state + .app + .structured_sessions + .session(&sid) + .ok_or_else(|| ErrorDto { + code: "NOT_FOUND".to_owned(), + message: "structured chat session not found".to_owned(), + })?; + let stream = session.send(&prompt).await.map_err(|err| ErrorDto { + code: "CHAT_SEND_FAILED".to_owned(), + message: err.to_string(), + })?; + + let mut sent_ack = false; + for event in stream { + let chunk = reply_event_chunk(event); + if !sent_ack { + tx.send(ServerFrame::chat_output(Some(&frame.id), sid, chunk)) + .await + .map_err(|_| ErrorDto { + code: "WS_CLOSED".to_owned(), + message: "websocket output closed".to_owned(), + })?; + sent_ack = true; + } else { + tx.send(ServerFrame::chat_output(None, sid, chunk)) + .await + .map_err(|_| ErrorDto { + code: "WS_CLOSED".to_owned(), + message: "websocket output closed".to_owned(), + })?; + } + } + + if !sent_ack { + tx.send(ServerFrame::chat_output( + Some(&frame.id), + sid, + json!({ "kind": "final", "content": "" }), + )) + .await + .map_err(|_| ErrorDto { + code: "WS_CLOSED".to_owned(), + message: "websocket output closed".to_owned(), + })?; + } + + Ok(()) +} + +fn ws_detach_chat( + frame: &ClientFrame, + state: &Arc, + owned: &Arc>>, +) -> Result<(), ErrorDto> { + let sid = parse_session_id(&payload_string(&frame.payload, "sessionId")?)?; + if let Ok(mut owned) = owned.lock() { + if let Some(index) = owned + .iter() + .rposition(|attachment| matches!(attachment, WsAttachment::Chat { session, .. } if *session == sid)) + { + let WsAttachment::Chat { gen, .. } = owned.remove(index) else { + unreachable!("rposition matched a chat attachment"); + }; + state.ws_chat_bridge.unregister_if(&sid, gen); + } + } + Ok(()) +} + +fn reply_event_chunk(event: ReplyEvent) -> Value { + match event { + ReplyEvent::TextDelta { text } => json!({ "kind": "textDelta", "text": text }), + ReplyEvent::ToolActivity { label } => { + json!({ "kind": "toolActivity", "label": label }) + } + ReplyEvent::Progress { progress } => { + json!({ "kind": "progress", "progress": reply_progress_value(progress) }) + } + ReplyEvent::Announcement { text } => json!({ "kind": "textDelta", "text": text }), + ReplyEvent::Heartbeat => json!({ + "kind": "progress", + "progress": { + "source": "ideaLocal", + "kind": "other", + "stage": "info", + "label": "heartbeat" + } + }), + ReplyEvent::RateLimited { resets_at_ms } => json!({ + "kind": "progress", + "progress": { + "source": "providerNative", + "kind": "turn", + "stage": "info", + "label": "rate_limited", + "occurredAtMs": resets_at_ms + } + }), + ReplyEvent::Error { message } => json!({ "kind": "error", "message": message }), + ReplyEvent::Final { content } => json!({ "kind": "final", "content": content }), + } +} + +fn reply_progress_value(progress: ReplyProgress) -> Value { + json!({ + "progressId": progress.progress_id, + "source": reply_progress_source(progress.source), + "kind": reply_progress_kind(progress.kind), + "stage": reply_progress_stage(progress.stage), + "label": progress.label, + "text": progress.text, + "provider": progress.provider, + "nativeEvent": progress.native_event, + "toolName": progress.tool_name, + }) +} + +fn reply_progress_source(source: ReplyProgressSource) -> &'static str { + match source { + ReplyProgressSource::ProviderNative => "providerNative", + ReplyProgressSource::IdeaLocal => "ideaLocal", + } +} + +fn reply_progress_kind(kind: ReplyProgressKind) -> &'static str { + match kind { + ReplyProgressKind::Turn => "turn", + ReplyProgressKind::Message => "message", + ReplyProgressKind::Tool => "tool", + ReplyProgressKind::Mcp => "mcp", + ReplyProgressKind::Other => "other", + } +} + +fn reply_progress_stage(stage: ReplyProgressStage) -> &'static str { + match stage { + ReplyProgressStage::Started => "started", + ReplyProgressStage::Delta => "delta", + ReplyProgressStage::Completed => "completed", + ReplyProgressStage::Info => "info", + } +} + async fn ws_close( frame: &ClientFrame, state: &Arc, - owned: &Arc>>, + owned: &Arc>>, ) -> Result<(SessionId, Option), ErrorDto> { let sid = parse_session_id(&payload_string(&frame.payload, "sessionId")?)?; ws_detach(frame, state, owned)?; @@ -5472,6 +5691,17 @@ impl ServerFrame { } } + fn chat_output(request_id: Option<&str>, session_id: SessionId, chunk: Value) -> Self { + Self { + kind: "chat.output".to_owned(), + reply_to: request_id.map(ToOwned::to_owned), + payload: json!({ + "sessionId": session_id.to_string(), + "chunk": chunk, + }), + } + } + fn status(session_id: SessionId, status: &str, exit_code: Option) -> Self { Self { kind: "terminal.status".to_owned(), @@ -5547,6 +5777,28 @@ impl OutputSink for WsPtySink { } } +struct WsChatSink { + session_id: SessionId, + tx: mpsc::Sender, +} + +impl WsChatSink { + fn new(session_id: SessionId, tx: mpsc::Sender) -> Self { + Self { session_id, tx } + } +} + +impl OutputSink for WsChatSink { + fn send(&self, item: ChatChunk) -> Result<(), OutputSinkError> { + self.tx + .try_send(ServerFrame::chat_output(None, self.session_id, item)) + .map_err(|err| match err { + mpsc::error::TrySendError::Full(_) => OutputSinkError::Full, + mpsc::error::TrySendError::Closed(_) => OutputSinkError::Closed, + }) + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum WsOpcode { Text, @@ -5861,7 +6113,8 @@ mod tests { CreateAgentInput, CreateProjectInput, LaunchAgentOutput, SaveProfileInput, StructuredSessionDescriptor, }; - use domain::ports::EventBus; + use async_trait::async_trait; + use domain::ports::{AgentSession, AgentSessionError, EventBus, ReplyStream}; use domain::{ AgentId, AgentProfile, BackgroundTask, BackgroundTaskKind, BackgroundTaskWakePolicy, ContextInjection, NodeId, ProfileId, ProjectId, ProjectPath, PtySize, SessionKind, @@ -6563,6 +6816,47 @@ mod tests { } } + struct FakeChatSession { + id: SessionId, + conversation_id: Option, + events: Vec, + prompts: Mutex>, + } + + impl FakeChatSession { + fn new(id: SessionId, events: Vec) -> Self { + Self { + id, + conversation_id: Some("engine-chat-session".to_owned()), + events, + prompts: Mutex::new(Vec::new()), + } + } + } + + #[async_trait] + impl AgentSession for FakeChatSession { + fn id(&self) -> SessionId { + self.id + } + + fn conversation_id(&self) -> Option { + self.conversation_id.clone() + } + + async fn send(&self, prompt: &str) -> Result { + self.prompts + .lock() + .expect("prompt trace mutex poisoned") + .push(prompt.to_owned()); + Ok(Box::new(self.events.clone().into_iter())) + } + + async fn shutdown(&self) -> Result<(), AgentSessionError> { + Ok(()) + } + } + fn attached_session_id(frame: &ServerFrame) -> String { assert_eq!(frame.kind, "terminal.attached"); frame.payload["session"]["sessionId"] @@ -7125,6 +7419,170 @@ mod tests { assert_eq!(state.ws_pty_bridge.active_sessions(), 0); } + #[tokio::test] + async fn websocket_chat_attach_emits_chat_attached_and_registers_sink() { + let state = state(); + let session_id = SessionId::new_random(); + let agent_id = AgentId::from_uuid(Uuid::new_v4()); + let node_id = NodeId::from_uuid(Uuid::new_v4()); + let session = Arc::new(FakeChatSession::new( + session_id, + vec![ReplyEvent::Final { + content: "unused".to_owned(), + }], + )); + state.app.structured_sessions.insert( + Arc::clone(&session) as Arc, + agent_id, + node_id, + ); + let (tx, mut rx) = mpsc::channel(8); + let owned = Arc::new(Mutex::new(Vec::new())); + + handle_client_frame( + ws_frame( + "chat-attach", + "chat.attach", + json!({ "sessionId": session_id.to_string() }), + ), + &state, + &tx, + &owned, + ) + .await; + let attached = recv_server_frame(&mut rx).await; + + assert_eq!(attached.kind, "chat.attached"); + assert_eq!(attached.reply_to.as_deref(), Some("chat-attach")); + assert_eq!(attached.payload["sessionId"], session_id.to_string()); + assert_eq!( + attached.payload["session"]["sessionId"], + session_id.to_string() + ); + assert_eq!(attached.payload["cellKind"], "chat"); + assert_eq!(attached.payload["engineSessionId"], "engine-chat-session"); + assert_eq!(state.ws_chat_bridge.active_sessions(), 1); + } + + #[tokio::test] + async fn websocket_chat_send_acknowledges_first_chunk_and_streams_rest() { + let state = state(); + let session_id = SessionId::new_random(); + let agent_id = AgentId::from_uuid(Uuid::new_v4()); + let node_id = NodeId::from_uuid(Uuid::new_v4()); + let session = Arc::new(FakeChatSession::new( + session_id, + vec![ + ReplyEvent::TextDelta { + text: "Hello ".to_owned(), + }, + ReplyEvent::ToolActivity { + label: "read_file".to_owned(), + }, + ReplyEvent::Final { + content: "Hello back".to_owned(), + }, + ], + )); + state.app.structured_sessions.insert( + Arc::clone(&session) as Arc, + agent_id, + node_id, + ); + let (tx, mut rx) = mpsc::channel(16); + let owned = Arc::new(Mutex::new(Vec::new())); + + handle_client_frame( + ws_frame( + "chat-send", + "chat.send", + json!({ "sessionId": session_id.to_string(), "prompt": "Salut" }), + ), + &state, + &tx, + &owned, + ) + .await; + + let ack = recv_server_frame(&mut rx).await; + let tool = recv_server_frame(&mut rx).await; + let final_chunk = recv_server_frame(&mut rx).await; + + assert_eq!(ack.kind, "chat.output"); + assert_eq!(ack.reply_to.as_deref(), Some("chat-send")); + assert_eq!(ack.payload["sessionId"], session_id.to_string()); + assert_eq!( + ack.payload["chunk"], + json!({ "kind": "textDelta", "text": "Hello " }) + ); + assert_eq!(tool.kind, "chat.output"); + assert!(tool.reply_to.is_none()); + assert_eq!( + tool.payload["chunk"], + json!({ "kind": "toolActivity", "label": "read_file" }) + ); + assert_eq!(final_chunk.kind, "chat.output"); + assert!(final_chunk.reply_to.is_none()); + assert_eq!( + final_chunk.payload["chunk"], + json!({ "kind": "final", "content": "Hello back" }) + ); + assert_eq!( + session.prompts.lock().unwrap().as_slice(), + &["Salut".to_owned()] + ); + } + + #[tokio::test] + async fn websocket_chat_detach_unregisters_sink_without_killing_session() { + let state = state(); + let session_id = SessionId::new_random(); + let agent_id = AgentId::from_uuid(Uuid::new_v4()); + let node_id = NodeId::from_uuid(Uuid::new_v4()); + let session = Arc::new(FakeChatSession::new( + session_id, + vec![ReplyEvent::Final { + content: "still alive".to_owned(), + }], + )); + state.app.structured_sessions.insert( + Arc::clone(&session) as Arc, + agent_id, + node_id, + ); + let (tx, mut rx) = mpsc::channel(8); + let owned = Arc::new(Mutex::new(Vec::new())); + + handle_client_frame( + ws_frame( + "chat-attach", + "chat.attach", + json!({ "sessionId": session_id.to_string() }), + ), + &state, + &tx, + &owned, + ) + .await; + let _attached = recv_server_frame(&mut rx).await; + assert_eq!(state.ws_chat_bridge.active_sessions(), 1); + + handle_client_frame( + ws_frame( + "chat-detach", + "chat.detach", + json!({ "sessionId": session_id.to_string() }), + ), + &state, + &tx, + &owned, + ) + .await; + + assert_eq!(state.ws_chat_bridge.active_sessions(), 0); + assert!(state.app.structured_sessions.session(&session_id).is_some()); + } + #[tokio::test] async fn websocket_launch_agent_reattach_replays_scrollback_without_respawn() { let state = state();