blomios/IdeA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(web): gère les frames chat.attach/chat.send/chat.detach du structured chat — #231

Browse Source
This commit is contained in:
Blomios
2026-08-20 15:48:38 +02:00
parent 1f58b96f7e
commit 3db23321a7
3 changed files with 476 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

490
crates/web-server/src/lib.rs
View File

@ -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<u8>;
type ChatChunk = Value;
use backend::BackendCore;
const DEFAULT_LISTEN: &str = "127.0.0.1:17373";
@ -572,6 +576,7 @@ struct ServerState {
app: Arc<BackendCore>,
pairing_code: Mutex<Option<PairingCodeState>>,
ws_pty_bridge: Arc<OutputBridge<SessionId, PtyChunk>>,
ws_chat_bridge: Arc<OutputBridge<SessionId, ChatChunk>>,
active_connections: Arc<ActiveConnectionRegistry>,
_revocation_observer: Option<JoinHandle<()>>,
security_logger: Arc<dyn SecurityLogger>,
@ -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::<ServerFrame>(WS_OUTPUT_BUFFER);
let owned = Arc::new(Mutex::new(Vec::<(SessionId, u64)>::new()));
let owned = Arc::new(Mutex::new(Vec::<WsAttachment>::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<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) {
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<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> Result<(), ErrorDto> {
let request_value = frame
.payload
@ -1539,7 +1562,7 @@ async fn ws_launch_agent(
frame: &ClientFrame,
state: &Arc<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> Result<(), ErrorDto> {
let request_value = frame
.payload
@ -1556,7 +1579,7 @@ async fn send_launch_agent_attached(
frame: &ClientFrame,
state: &Arc<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
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<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> 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<ServerState>,
sid: SessionId,
sink: WsPtySink,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> 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<ServerState>) -> Result<(), ErrorD
fn ws_detach(
frame: &ClientFrame,
state: &Arc<ServerState>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> 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<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> 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<ServerState>,
tx: &mpsc::Sender<ServerFrame>,
) -> 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<ServerState>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> 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<ServerState>,
owned: &Arc<Mutex<Vec<(SessionId, u64)>>>,
owned: &Arc<Mutex<Vec<WsAttachment>>>,
) -> Result<(SessionId, Option<i32>), 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<i32>) -> Self {
Self {
kind: "terminal.status".to_owned(),
@ -5547,6 +5777,28 @@ impl OutputSink<PtyChunk> for WsPtySink {
}
}
struct WsChatSink {
session_id: SessionId,
tx: mpsc::Sender<ServerFrame>,
}
impl WsChatSink {
fn new(session_id: SessionId, tx: mpsc::Sender<ServerFrame>) -> Self {
Self { session_id, tx }
}
}
impl OutputSink<ChatChunk> 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<String>,
events: Vec<ReplyEvent>,
prompts: Mutex<Vec<String>>,
}
impl FakeChatSession {
fn new(id: SessionId, events: Vec<ReplyEvent>) -> 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<String> {
self.conversation_id.clone()
}
async fn send(&self, prompt: &str) -> Result<ReplyStream, AgentSessionError> {
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<dyn AgentSession>,
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<dyn AgentSession>,
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<dyn AgentSession>,
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();