diff --git a/docs/en/developer-guide/task-runtime-state-machine.md b/docs/en/developer-guide/task-runtime-state-machine.md index 965d90632f..0357c18070 100644 --- a/docs/en/developer-guide/task-runtime-state-machine.md +++ b/docs/en/developer-guide/task-runtime-state-machine.md @@ -97,6 +97,21 @@ When a task becomes terminal: `chat:done` only means the current message stream ended. It does not prove the task lifecycle is terminal. If queued messages are still blocked after `chat:done`, `useChatStreamHandlers` calls the current task machine with `requestRuntimeCheck('queued-message-blocked')` so the frontend converges back to the server state. +### Cancellation And State Convergence + +Runtime task cancellation must converge two state surfaces: + +- Runtime task state exposed by the executor, such as `running` and `status` from `runtime.tasks.list`. Sidebars, close guards, and task lists consume this state. +- Message stream state inside the current pane, such as assistant messages that are still `assistant:streaming`. The composer stop button, send lock, and local streaming UI consume this state. + +The cancellation flow must follow these rules: + +- After receiving `runtime.tasks.cancel`, the executor must interrupt the active turn and wait for the app-server or child process to actually exit. If it does not exit before the timeout, return `cancel_timeout`; the frontend must not treat that as stopped. +- Runners that can spawn child processes, such as the Codex app-server, must be managed and terminated as a process group so background workers cannot keep running after the parent process exits. +- After the frontend receives an accepted cancellation result, it must refresh the runtime work list so sidebars and close guards observe `running: false`. +- The pane session must locally finalize any still-streaming assistant messages as cancelled because cancellation may not emit a normal `chat:done` or response completed event. +- Updating only runtime work leaves the composer looking busy. Updating only message streams leaves the sidebar looking busy. Both state surfaces must be updated in the same stop path. + ## Change Guidelines When adding or changing task runtime UI: diff --git a/docs/zh/developer-guide/task-runtime-state-machine.md b/docs/zh/developer-guide/task-runtime-state-machine.md index 9bf1673195..fbe7369b7f 100644 --- a/docs/zh/developer-guide/task-runtime-state-machine.md +++ b/docs/zh/developer-guide/task-runtime-state-machine.md @@ -97,6 +97,21 @@ WebSocket 连接恢复也由状态机按这个顺序统一处理：外部触发 `chat:done` 只表示当前消息流结束，不等价于任务生命周期终态。若 `chat:done` 后运行态仍阻塞 queued message，`useChatStreamHandlers` 会触发当前任务状态机的 `requestRuntimeCheck('queued-message-blocked')`，让前端重新对齐服务端状态。 +### 取消与状态收敛 + +运行时任务取消必须同时收敛两类状态： + +- executor 暴露的运行时任务状态，例如 `runtime.tasks.list` 返回的 `running` 和 `status`，用于侧边栏运行标记、关闭保护和任务列表。 +- 当前 pane 内的消息流状态，例如仍处于 `assistant:streaming` 的消息，用于输入框停止按钮、发送锁和本地流式 UI。 + +取消流程必须遵循以下规则： + +- executor 收到 `runtime.tasks.cancel` 后必须中断当前 turn，并等待 app-server 或子进程实际退出；如果超过超时时间仍未退出，应返回 `cancel_timeout`，前端不能把它当作已停止。 +- Codex app-server 这类会派生子进程的运行器必须按进程组管理和终止，避免只停止父进程而后台 worker 继续执行。 +- 前端收到已接受的取消结果后，必须刷新 runtime work 列表，让侧边栏和关闭保护看到 `running: false`。 +- pane session 必须把本地仍在 streaming 的 assistant 消息固化为 cancelled，因为取消不一定会产生正常的 `chat:done` 或 response completed 事件。 +- 只更新 runtime work 会导致输入框仍显示运行中；只更新消息流会导致侧边栏仍显示运行中。两类状态必须在同一个停止路径里同时更新。 + ## 修改指南 新增或修改任务运行态 UI 时： diff --git a/executor/src/agents/codex.rs b/executor/src/agents/codex.rs index e348473164..64f4141665 100644 --- a/executor/src/agents/codex.rs +++ b/executor/src/agents/codex.rs @@ -19,7 +19,7 @@ use serde_json::{json, Value}; use tokio::{ io::{AsyncBufReadExt, AsyncWriteExt, BufReader}, process::{Child, ChildStdin, ChildStdout, Command}, - sync::{mpsc, Mutex}, + sync::{mpsc, oneshot, Mutex}, time::timeout, }; @@ -38,6 +38,7 @@ use super::{model_id, prompt_text}; const DEFAULT_CODEX_RPC_TIMEOUT_SECONDS: u64 = 300; const DEFAULT_PROVIDER_ID: &str = "wecode-openai"; +pub const CODEX_APP_SERVER_TURN_CANCELLED: &str = "codex app-server turn cancelled"; const DEFAULT_PROVIDER_NAME: &str = "wecode openai"; const DEFAULT_REASONING_EFFORT: &str = "medium"; const DEFAULT_NO_PROXY: &str = "localhost,127.0.0.1,::1,host.docker.internal"; @@ -210,14 +211,13 @@ struct CodexAppServerProcess { impl CodexAppServerProcess { async fn shutdown(mut self) { - let _ = self.child.start_kill(); - let _ = self.child.wait().await; + terminate_codex_app_server_child(&mut self.child).await; } } impl Drop for CodexAppServerProcess { fn drop(&mut self) { - let _ = self.child.start_kill(); + signal_codex_app_server_child(&mut self.child); } } @@ -254,8 +254,7 @@ async fn start_persistent_codex_app_server(binary: &str) -> Result Ok(CodexAppServerProcess { child, rpc }), Err(error) => { - let _ = child.start_kill(); - let _ = child.wait().await; + terminate_codex_app_server_child(&mut child).await; Err(error) } } @@ -267,6 +266,25 @@ pub async fn run_codex_app_server_turn( resume_thread_id: Option, initial_thread_name: Option, notifications: Option, +) -> Result { + run_codex_app_server_turn_with_cancel( + binary, + request, + resume_thread_id, + initial_thread_name, + notifications, + None, + ) + .await +} + +pub async fn run_codex_app_server_turn_with_cancel( + binary: &str, + request: ExecutionRequest, + resume_thread_id: Option, + initial_thread_name: Option, + notifications: Option, + mut cancellation: Option>, ) -> Result { let prepared = prepare_codex_execution_request(request); let launch_config = build_codex_launch_config(&prepared.request); @@ -390,9 +408,15 @@ pub async fn run_codex_app_server_turn( ), ) .await?; - let outcome = rpc - .read_turn(turn_request_id, &mut state, notifications) - .await?; + let outcome = if let Some(cancellation) = cancellation.as_mut() { + tokio::select! { + outcome = rpc.read_turn(turn_request_id, &mut state, notifications) => outcome?, + _ = cancellation => return Err(CODEX_APP_SERVER_TURN_CANCELLED.to_owned()), + } + } else { + rpc.read_turn(turn_request_id, &mut state, notifications) + .await? + }; turn_fields.push(("outcome", codex_outcome_name(&outcome).to_owned())); if let ExecutionOutcome::Failed { message } = &outcome { turn_fields.push(("error", message.clone())); @@ -403,8 +427,7 @@ pub async fn run_codex_app_server_turn( } .await; - let _ = child.start_kill(); - let _ = child.wait().await; + terminate_codex_app_server_child(&mut child).await; if let Err(error) = &result { let mut failed_fields = fields.clone(); failed_fields.push(("error", error.clone())); @@ -445,7 +468,9 @@ fn spawn_codex_app_server( env::var("PATH").ok().as_deref().unwrap_or_default(), ), ); + configure_codex_app_server_process_group(&mut command); command + .kill_on_drop(true) .stdin(Stdio::piped()) .stdout(Stdio::piped()) .stderr(Stdio::inherit()) @@ -453,6 +478,44 @@ fn spawn_codex_app_server( .map_err(|error| format!("failed to start codex app-server: {error}")) } +#[cfg(unix)] +fn configure_codex_app_server_process_group(command: &mut Command) { + unsafe { + command.pre_exec(|| { + if libc::setpgid(0, 0) == 0 { + Ok(()) + } else { + Err(std::io::Error::last_os_error()) + } + }); + } +} + +#[cfg(not(unix))] +fn configure_codex_app_server_process_group(_command: &mut Command) {} + +async fn terminate_codex_app_server_child(child: &mut Child) { + signal_codex_app_server_child(child); + let _ = child.wait().await; +} + +fn signal_codex_app_server_child(child: &mut Child) { + if matches!(child.try_wait(), Ok(Some(_))) { + return; + } + + #[cfg(unix)] + if let Some(process_group_id) = child.id() { + unsafe { + let _ = libc::kill(-(process_group_id as libc::pid_t), libc::SIGTERM); + let _ = libc::kill(-(process_group_id as libc::pid_t), libc::SIGKILL); + } + return; + } + + let _ = child.start_kill(); +} + async fn with_rpc_timeout( operation: &str, timeout_seconds: u64, diff --git a/executor/src/agents/mod.rs b/executor/src/agents/mod.rs index a3b2f6f2db..83ba50d102 100644 --- a/executor/src/agents/mod.rs +++ b/executor/src/agents/mod.rs @@ -33,8 +33,9 @@ use claude_code::{ }; pub use claude_options::{extract_claude_options, ClaudeOptions}; pub use codex::{ - run_codex_app_server_turn, CodexAppServerClient, CodexAppServerEngine, CodexAppServerTurn, - CodexCancellationState, CodexNotificationSender, CodexTurnInterrupter, + run_codex_app_server_turn, run_codex_app_server_turn_with_cancel, CodexAppServerClient, + CodexAppServerEngine, CodexAppServerTurn, CodexCancellationState, CodexNotificationSender, + CodexTurnInterrupter, CODEX_APP_SERVER_TURN_CANCELLED, }; pub use dify::{build_dify_config, saved_dify_task_id, DifyEngine}; pub use image_validator::ImageValidatorEngine; diff --git a/executor/src/runtime_work/handler.rs b/executor/src/runtime_work/handler.rs index baaab67236..d830ac2817 100644 --- a/executor/src/runtime_work/handler.rs +++ b/executor/src/runtime_work/handler.rs @@ -3,7 +3,7 @@ // SPDX-License-Identifier: Apache-2.0 use std::{ - collections::HashSet, + collections::{HashMap, HashSet}, future::Future, path::Path, pin::Pin, @@ -12,10 +12,13 @@ use std::{ }; use serde_json::{json, Map, Value}; -use tokio::sync::{broadcast, mpsc}; +use tokio::sync::{broadcast, mpsc, oneshot}; use crate::{ - agents::{run_codex_app_server_turn, CodexAppServerClient, CodexNotificationSender}, + agents::{ + run_codex_app_server_turn_with_cancel, CodexAppServerClient, CodexNotificationSender, + CODEX_APP_SERVER_TURN_CANCELLED, + }, local::app_ipc::{AppIpcError, RuntimeWorkHandler}, logging::log_executor_event, protocol::ExecutionRequest, @@ -65,11 +68,17 @@ pub struct RuntimeWorkRpcHandler { codex_app_server: CodexAppServerClient, event_tx: Option>, active_local_tasks: Arc>>, + active_turn_cancellations: Arc>>, store: RuntimeWorkStore, transcript_cache: TranscriptCache, thread_list_cache: CodexThreadListCache, } +struct ActiveTurnCancellation { + cancel: oneshot::Sender<()>, + stopped: oneshot::Receiver<()>, +} + impl RuntimeWorkRpcHandler { pub fn new(device_id: impl Into, codex_binary: impl Into) -> Self { let codex_binary = codex_binary.into(); @@ -79,6 +88,7 @@ impl RuntimeWorkRpcHandler { codex_app_server: CodexAppServerClient::new(codex_binary), event_tx: None, active_local_tasks: Arc::new(Mutex::new(HashSet::new())), + active_turn_cancellations: Arc::new(Mutex::new(HashMap::new())), store: RuntimeWorkStore::from_env(), transcript_cache: TranscriptCache::default(), thread_list_cache: CodexThreadListCache::default(), @@ -684,6 +694,16 @@ impl RuntimeWorkRpcHandler { link.updated_at = now_ms(); }) .or_else(|| self.local_task_link(&local_task_id)); + if !self.abort_active_turn(&local_task_id).await { + return Ok(json!({ + "success": false, + "accepted": false, + "localTaskId": local_task_id, + "runtime": "codex", + "error": "runtime task did not stop within timeout", + "code": "cancel_timeout", + })); + } Ok(match link { Some(link) => task_action_success(&link), @@ -896,39 +916,59 @@ impl RuntimeWorkRpcHandler { log_executor_event("runtime work turn spawning", &fields); self.mark_active_local_task(&local_task_id); + let (cancel_tx, cancel_rx) = oneshot::channel(); + let (stopped_tx, stopped_rx) = oneshot::channel(); + self.set_active_turn_cancellation( + local_task_id.clone(), + ActiveTurnCancellation { + cancel: cancel_tx, + stopped: stopped_rx, + }, + ); let handler = self.clone(); - tokio::spawn(async move { + let turn_local_task_id = local_task_id.clone(); + let turn_handle = tokio::spawn(async move { emit_response_event( &handler.event_tx, &handler.device_id, "response.created", - &local_task_id, + &turn_local_task_id, &request, json!({"response": {"status": "in_progress"}}), ); let (notifications, mapper_task) = - handler.spawn_notification_mapper(local_task_id.clone(), request.clone()); - let result = run_codex_app_server_turn( + handler.spawn_notification_mapper(turn_local_task_id.clone(), request.clone()); + let result = run_codex_app_server_turn_with_cancel( &handler.codex_binary, request.clone(), resume_thread_id, initial_thread_name, notifications, + Some(cancel_rx), ) .await; if let Err(error) = mapper_task.await { log_executor_event( "runtime work notification mapper failed", &[ - ("local_task_id", local_task_id.clone()), + ("local_task_id", turn_local_task_id.clone()), ("error", error.to_string()), ], ); } - handler.handle_turn_result(&local_task_id, &request, result); + if matches!(result.as_ref(), Err(error) if error == CODEX_APP_SERVER_TURN_CANCELLED) { + handler.clear_active_turn_cancellation(&turn_local_task_id); + handler.unmark_active_local_task(&turn_local_task_id); + let _ = stopped_tx.send(()); + return; + } + + handler.handle_turn_result(&turn_local_task_id, &request, result); + let _ = stopped_tx.send(()); }); + drop(turn_handle); } fn handle_turn_result( @@ -1391,6 +1431,43 @@ impl RuntimeWorkRpcHandler { .remove(local_task_id); } + fn set_active_turn_cancellation(&self, local_task_id: String, control: ActiveTurnCancellation) { + if let Some(previous) = self + .active_turn_cancellations + .lock() + .expect("active turn cancellation map lock should not be poisoned") + .insert(local_task_id, control) + { + let _ = previous.cancel.send(()); + } + } + + fn clear_active_turn_cancellation(&self, local_task_id: &str) { + self.active_turn_cancellations + .lock() + .expect("active turn cancellation map lock should not be poisoned") + .remove(local_task_id); + } + + async fn abort_active_turn(&self, local_task_id: &str) -> bool { + let control = { + self.active_turn_cancellations + .lock() + .expect("active turn cancellation map lock should not be poisoned") + .remove(local_task_id) + }; + if let Some(control) = control { + let _ = control.cancel.send(()); + let stopped = + tokio::time::timeout(std::time::Duration::from_secs(10), control.stopped).await; + if stopped.is_err() { + return false; + } + } + self.unmark_active_local_task(local_task_id); + true + } + fn is_active_local_task(&self, local_task_id: &str) -> bool { self.active_local_tasks .lock() @@ -1399,6 +1476,7 @@ impl RuntimeWorkRpcHandler { } fn finish_local_task(&self, local_task_id: &str, thread_id: Option, status: &str) { + self.clear_active_turn_cancellation(local_task_id); let invalidate_thread_id = thread_id.clone(); self.store.update_task(local_task_id, |link| { if thread_id.is_some() { diff --git a/executor/src/runtime_work/store.rs b/executor/src/runtime_work/store.rs index 3ab7970870..2a692e489d 100644 --- a/executor/src/runtime_work/store.rs +++ b/executor/src/runtime_work/store.rs @@ -123,10 +123,10 @@ impl RuntimeWorkStore { }); if let Ok(payload) = payload { let temp_path = temporary_index_path(&self.index_path); - if fs::write(&temp_path, payload).is_ok() { - if fs::rename(&temp_path, &self.index_path).is_err() { - let _ = fs::remove_file(temp_path); - } + if fs::write(&temp_path, payload).is_ok() + && fs::rename(&temp_path, &self.index_path).is_err() + { + let _ = fs::remove_file(temp_path); } } } diff --git a/executor/tests/app_runtime_work_contract.rs b/executor/tests/app_runtime_work_contract.rs index 7a2e5d778a..b7574f5fe3 100644 --- a/executor/tests/app_runtime_work_contract.rs +++ b/executor/tests/app_runtime_work_contract.rs @@ -1126,21 +1126,25 @@ fn write_fake_codex_running_tool(log_path: &Path) -> PathBuf { LOG_PATH='{}' while IFS= read -r line; do printf '%s\n' "$line" >> "$LOG_PATH" + request_id=$(printf '%s\n' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') case "$line" in *'"method":"initialize"'*) - printf '%s\n' '{{"id":1,"result":{{"protocolVersion":1}}}}' + printf '%s\n' '{{"id":'"$request_id"',"result":{{"protocolVersion":1}}}}' ;; *'"method":"initialized"'*) ;; *'"method":"thread/start"'*) - printf '%s\n' '{{"id":2,"result":{{"thread":{{"id":"thread-running-tool"}}}}}}' + printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-running-tool"}}}}}}' + ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' ;; *'"method":"thread/read"'*) - printf '%s\n' '{{"id":2,"result":{{"thread":{{"id":"thread-running-tool","cwd":"/tmp/project","name":"Running tool task","preview":"run tests","path":"/tmp/codex/thread-running-tool.jsonl","createdAt":1780000000,"updatedAt":1780000060,"status":"idle","turns":[]}}}}}}' + printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-running-tool","cwd":"/tmp/project","name":"Running tool task","preview":"run tests","path":"/tmp/codex/thread-running-tool.jsonl","createdAt":1780000000,"updatedAt":1780000060,"status":"idle","turns":[]}}}}}}' exit 0 ;; *'"method":"turn/start"'*) - printf '%s\n' '{{"id":3,"result":{{"turn":{{"id":"turn-running-tool","status":"inProgress"}}}}}}' + printf '%s\n' '{{"id":'"$request_id"',"result":{{"turn":{{"id":"turn-running-tool","status":"inProgress"}}}}}}' printf '%s\n' '{{"method":"item/reasoning/delta","params":{{"delta":"checking context"}}}}' printf '%s\n' '{{"method":"item/agentMessage/delta","params":{{"delta":"running tests","phase":"commentary"}}}}' printf '%s\n' '{{"method":"item/started","params":{{"item":{{"id":"cmd-1","type":"commandExecution","command":"cargo test","cwd":"/tmp/project","status":"inProgress"}}}}}}' diff --git a/executor/tests/app_runtime_work_send_contract.rs b/executor/tests/app_runtime_work_send_contract.rs index 5f880ec12b..c0e8125d1a 100644 --- a/executor/tests/app_runtime_work_send_contract.rs +++ b/executor/tests/app_runtime_work_send_contract.rs @@ -96,7 +96,7 @@ async fn runtime_tasks_send_accepts_address_content_source_and_attachments() { assert_eq!(created["accepted"], true); wait_for_thread_mapping(&handler, "local-task-1", "thread-1").await; wait_for_turn_count(&log_path, 1).await; - wait_for_response_completed(&mut events).await; + wait_for_response_event(&mut events, "response.completed", 2001).await; wait_until_task_idle(&handler, "local-task-1").await; drain_events(&mut events); @@ -354,7 +354,7 @@ async fn runtime_tasks_send_includes_local_text_attachment_content() { .expect("create should be accepted"); wait_for_thread_mapping(&handler, "local-task-text", "thread-1").await; wait_for_turn_count(&log_path, 1).await; - wait_for_response_completed(&mut events).await; + wait_for_response_event(&mut events, "response.completed", 2002).await; wait_until_task_idle(&handler, "local-task-text").await; drain_events(&mut events); @@ -510,6 +510,104 @@ async fn runtime_tasks_send_rejects_running_local_task_until_cancelled() { assert_eq!(cancelled["accepted"], true); } +#[tokio::test] +async fn runtime_tasks_cancel_terminates_running_codex_turn() { + let _lock = env_lock().await; + let _home = EnvGuard::set( + "WEGENT_EXECUTOR_HOME", + &temp_path("runtime-cancel-kill-home", "dir") + .display() + .to_string(), + ); + let _codex_home = EnvGuard::set( + "CODEX_HOME", + &temp_path("runtime-cancel-kill-codex-home", "dir") + .display() + .to_string(), + ); + let log_path = temp_path("runtime-cancel-kill-log", "jsonl"); + let fake_codex = write_fake_codex_traps_termination(&log_path); + let handler = RuntimeWorkRpcHandler::new("device-1", fake_codex.display().to_string()); + + handler + .handle_runtime_rpc(json!({ + "method": "runtime.tasks.create", + "payload": { + "localTaskId": "local-task-cancel", + "workspacePath": "/tmp/project", + "message": "first turn", + "modelId": "gpt-5.5" + } + })) + .await + .expect("create should be accepted"); + let pid = wait_for_logged_pid(&log_path, "turn-pid:").await; + + let cancelled = handler + .handle_runtime_rpc(json!({ + "method": "runtime.tasks.cancel", + "payload": { + "workspacePath": "/tmp/project", + "localTaskId": "local-task-cancel" + } + })) + .await + .expect("cancel should be accepted"); + + assert_eq!(cancelled["accepted"], true); + assert_process_exited(pid).await; +} + +#[tokio::test] +async fn runtime_tasks_cancel_terminates_codex_turn_process_group() { + let _lock = env_lock().await; + let _home = EnvGuard::set( + "WEGENT_EXECUTOR_HOME", + &temp_path("runtime-cancel-group-home", "dir") + .display() + .to_string(), + ); + let _codex_home = EnvGuard::set( + "CODEX_HOME", + &temp_path("runtime-cancel-group-codex-home", "dir") + .display() + .to_string(), + ); + let log_path = temp_path("runtime-cancel-group-log", "jsonl"); + let fake_codex = write_fake_codex_spawns_worker(&log_path); + let handler = RuntimeWorkRpcHandler::new("device-1", fake_codex.display().to_string()); + + handler + .handle_runtime_rpc(json!({ + "method": "runtime.tasks.create", + "payload": { + "localTaskId": "local-task-cancel-group", + "workspacePath": "/tmp/project", + "message": "first turn", + "modelId": "gpt-5.5" + } + })) + .await + .expect("create should be accepted"); + let parent_pid = wait_for_logged_pid(&log_path, "turn-pid:").await; + let worker_pid = wait_for_logged_pid(&log_path, "worker-pid:").await; + + let cancelled = handler + .handle_runtime_rpc(json!({ + "method": "runtime.tasks.cancel", + "payload": { + "workspacePath": "/tmp/project", + "localTaskId": "local-task-cancel-group" + } + })) + .await + .expect("cancel should be accepted"); + + assert_eq!(cancelled["accepted"], true); + assert_process_exited(parent_pid).await; + assert_process_exited(worker_pid).await; +} + #[tokio::test] async fn runtime_tasks_send_after_cancel_resumes_started_thread_not_local_task_id() { let _lock = env_lock().await; @@ -654,6 +752,9 @@ while IFS= read -r line; do *'"method":"thread/resume"'*) printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-1"}}}}}}' ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' + ;; *'"method":"turn/start"'*) progress_id='progress-1' case "$line" in @@ -701,6 +802,9 @@ while IFS= read -r line; do *'"method":"thread/resume"'*) printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-1"}}}}}}' ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' + ;; *'"method":"turn/start"'*) printf '%s\n' '{{"id":'"$request_id"',"result":{{"turn":{{"id":"turn-1","status":"inProgress"}}}}}}' case "$line" in @@ -743,6 +847,9 @@ while IFS= read -r line; do *'"method":"thread/start"'*) printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-1"}}}}}}' ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' + ;; *'"method":"turn/start"'*) printf '%s\n' '{{"id":'"$request_id"',"result":{{"turn":{{"id":"turn-1","status":"inProgress"}}}}}}' sleep 1 @@ -758,6 +865,89 @@ done path } +fn write_fake_codex_traps_termination(log_path: &Path) -> PathBuf { + let path = temp_path("fake-codex-cancel-kill", "sh"); + let _ = fs::remove_file(log_path); + let content = format!( + r#"#!/bin/sh +LOG_PATH='{}' +case "$*" in + *' -c '*|'-c '*) + printf 'turn-pid:%s\n' "$$" >> "$LOG_PATH" + ;; + *) + printf 'persistent-pid:%s\n' "$$" >> "$LOG_PATH" + ;; +esac +while IFS= read -r line; do + printf '%s\n' "$line" >> "$LOG_PATH" + request_id=$(printf '%s\n' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') + case "$line" in + *'"method":"initialize"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"protocolVersion":1}}}}' + ;; + *'"method":"initialized"'*) + ;; + *'"method":"thread/start"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-1"}}}}}}' + ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' + ;; + *'"method":"turn/start"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"turn":{{"id":"turn-1","status":"inProgress"}}}}}}' + while true; do sleep 1; done + ;; + esac +done +"#, + log_path.display() + ); + write_executable(&path, &content); + path +} + +fn write_fake_codex_spawns_worker(log_path: &Path) -> PathBuf { + let path = temp_path("fake-codex-cancel-group", "sh"); + let _ = fs::remove_file(log_path); + let content = format!( + r#"#!/bin/sh +LOG_PATH='{}' +case "$*" in + *' -c '*|'-c '*) + printf 'turn-pid:%s\n' "$$" >> "$LOG_PATH" + ;; +esac +while IFS= read -r line; do + printf '%s\n' "$line" >> "$LOG_PATH" + request_id=$(printf '%s\n' "$line" | sed -n 's/.*"id":\([0-9][0-9]*\).*/\1/p') + case "$line" in + *'"method":"initialize"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"protocolVersion":1}}}}' + ;; + *'"method":"initialized"'*) + ;; + *'"method":"thread/start"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"thread":{{"id":"thread-1"}}}}}}' + ;; + *'"method":"thread/name/set"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{}}}}' + ;; + *'"method":"turn/start"'*) + printf '%s\n' '{{"id":'"$request_id"',"result":{{"turn":{{"id":"turn-1","status":"inProgress"}}}}}}' + sh -c 'trap "" TERM; while true; do sleep 1; done' & + printf 'worker-pid:%s\n' "$!" >> "$LOG_PATH" + while true; do sleep 1; done + ;; + esac +done +"#, + log_path.display() + ); + write_executable(&path, &content); + path +} + fn write_executable(path: &Path, content: &str) { fs::write(path, content).unwrap(); #[cfg(unix)] @@ -897,6 +1087,44 @@ async fn wait_for_method_count(log_path: &Path, method: &str, expected: usize) { panic!("expected at least {expected} {method} calls"); } +async fn wait_for_logged_pid(log_path: &Path, prefix: &str) -> u32 { + let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(2); + loop { + let content = fs::read_to_string(log_path).unwrap_or_default(); + if let Some(pid) = content.lines().find_map(|line| { + line.strip_prefix(prefix) + .and_then(|value| value.trim().parse::().ok()) + }) { + return pid; + } + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for fake codex pid with prefix {prefix:?}; content:\n{content}" + ); + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } +} + +async fn assert_process_exited(pid: u32) { + for _ in 0..50 { + if !process_exists(pid) { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + panic!("cancel returned before fake codex process {pid} exited"); +} + +fn process_exists(pid: u32) -> bool { + std::process::Command::new("kill") + .arg("-0") + .arg(pid.to_string()) + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .is_ok_and(|status| status.success()) +} + async fn wait_until_task_running(handler: &RuntimeWorkRpcHandler, local_task_id: &str) { for _ in 0..50 { let listed = handler @@ -929,13 +1157,13 @@ async fn wait_until_task_idle(handler: &RuntimeWorkRpcHandler, local_task_id: &s })) .await .expect("list should succeed"); - let idle = listed["workspaces"] + let running = listed["workspaces"] .as_array() .into_iter() .flatten() .flat_map(|workspace| workspace["localTasks"].as_array().into_iter().flatten()) - .any(|task| task["localTaskId"] == local_task_id && task["running"] == false); - if idle { + .any(|task| task["localTaskId"] == local_task_id && task["running"] == true); + if !running { return; } tokio::time::sleep(std::time::Duration::from_millis(20)).await; @@ -961,13 +1189,6 @@ fn drain_events(events: &mut broadcast::Receiver) { while events.try_recv().is_ok() {} } -async fn wait_for_response_completed(events: &mut broadcast::Receiver) { - recv_events_until(events, |runtime_events| { - find_runtime_event(runtime_events, "response.completed", |_| true).is_some() - }) - .await; -} - async fn recv_events_until(events: &mut broadcast::Receiver, mut done: F) -> Vec where F: FnMut(&[Value]) -> bool, @@ -994,6 +1215,20 @@ where } } +async fn wait_for_response_event( + events: &mut broadcast::Receiver, + event_name: &str, + subtask_id: i64, +) { + recv_events_until(events, |received| { + find_runtime_event(received, event_name, |event| { + event["payload"]["subtask_id"] == subtask_id + }) + .is_some() + }) + .await; +} + fn find_runtime_event<'a, F>( events: &'a [Value], event_name: &str, diff --git a/packages/chat-core/src/workbench-message-reducer.test.ts b/packages/chat-core/src/workbench-message-reducer.test.ts index 6c825221b5..6ef9670de9 100644 --- a/packages/chat-core/src/workbench-message-reducer.test.ts +++ b/packages/chat-core/src/workbench-message-reducer.test.ts @@ -2,7 +2,7 @@ import { describe, expect, test } from 'vitest' import { normalizeWorkbenchBlockStatus, reduceWorkbenchMessages, - type WorkbenchMessage, + type WorkbenchMessage } from './index' describe('reduceWorkbenchMessages', () => { @@ -15,19 +15,19 @@ describe('reduceWorkbenchMessages', () => { role: 'user', content: 'hello', status: 'done', - createdAt: '2026-05-25T00:00:00.000Z', - }, + createdAt: '2026-05-25T00:00:00.000Z' + } }) const withStart = reduceWorkbenchMessages(withUser, { type: 'assistant_started', taskId: 1, turnId: 9, - shellType: 'ClaudeCode', + shellType: 'ClaudeCode' }) const withChunk = reduceWorkbenchMessages(withStart, { type: 'assistant_chunk', turnId: 9, - content: 'hi', + content: 'hi' }) expect(withChunk).toHaveLength(2) @@ -36,7 +36,7 @@ describe('reduceWorkbenchMessages', () => { role: 'assistant', content: 'hi', status: 'streaming', - shellType: 'ClaudeCode', + shellType: 'ClaudeCode' }) }) @@ -48,8 +48,8 @@ describe('reduceWorkbenchMessages', () => { turnId: 9, content: 'hello', status: 'done', - createdAt: '2026-05-25T00:00:00.000Z', - }, + createdAt: '2026-05-25T00:00:00.000Z' + } ] const withTool = reduceWorkbenchMessages(state, { @@ -61,31 +61,31 @@ describe('reduceWorkbenchMessages', () => { type: 'tool', toolName: 'bash', status: 'pending', - createdAt: 1770000000000, - }, + createdAt: 1770000000000 + } }) const withChunk = reduceWorkbenchMessages(withTool, { type: 'assistant_chunk', turnId: 9, - content: 'Hi', + content: 'Hi' }) const done = reduceWorkbenchMessages(withChunk, { type: 'assistant_done', turnId: 9, - content: 'Hi', + content: 'Hi' }) expect(done).toHaveLength(2) expect(done[0]).toMatchObject({ role: 'user', content: 'hello', - status: 'done', + status: 'done' }) expect(done[1]).toMatchObject({ role: 'assistant', content: 'Hi', status: 'done', - blocks: [{ type: 'tool', toolName: 'bash', status: 'done' }], + blocks: [{ type: 'tool', toolName: 'bash', status: 'done' }] }) }) @@ -94,13 +94,13 @@ describe('reduceWorkbenchMessages', () => { reduceWorkbenchMessages([], { type: 'assistant_started', taskId: 1, - turnId: 9, + turnId: 9 }), { type: 'assistant_chunk', turnId: 9, content: '', - reasoningChunk: 'Running a command', + reasoningChunk: 'Running a command' } ) @@ -114,13 +114,13 @@ describe('reduceWorkbenchMessages', () => { toolName: 'bash', toolInput: { command: 'pwd' }, status: 'pending', - createdAt: 1770000000000, - }, + createdAt: 1770000000000 + } }) expect(withTool[0].blocks).toMatchObject([ { type: 'thinking', status: 'done' }, - { type: 'tool', toolName: 'bash', status: 'pending' }, + { type: 'tool', toolName: 'bash', status: 'pending' } ]) }) @@ -129,12 +129,12 @@ describe('reduceWorkbenchMessages', () => { reduceWorkbenchMessages([], { type: 'assistant_started', taskId: 1, - turnId: 9, + turnId: 9 }), { type: 'assistant_chunk', turnId: 9, - content: 'Let me inspect the repository first.', + content: 'Let me inspect the repository first.' } ) @@ -148,8 +148,8 @@ describe('reduceWorkbenchMessages', () => { toolName: 'bash', toolInput: { command: 'ls' }, status: 'pending', - createdAt: 1770000000000, - }, + createdAt: 1770000000000 + } }) expect(withTool[0].content).toBe('') @@ -157,9 +157,9 @@ describe('reduceWorkbenchMessages', () => { { type: 'text', content: 'Let me inspect the repository first.', - status: 'done', + status: 'done' }, - { type: 'tool', toolName: 'bash', status: 'pending' }, + { type: 'tool', toolName: 'bash', status: 'pending' } ]) }) @@ -167,7 +167,7 @@ describe('reduceWorkbenchMessages', () => { const state = reduceWorkbenchMessages([], { type: 'assistant_started', taskId: 1, - turnId: 9, + turnId: 9 }) const done = reduceWorkbenchMessages(state, { @@ -181,7 +181,7 @@ describe('reduceWorkbenchMessages', () => { type: 'thinking', content: 'Drafting', status: 'streaming', - createdAt: 1770000000000, + createdAt: 1770000000000 }, { id: 'call_1', @@ -189,7 +189,7 @@ describe('reduceWorkbenchMessages', () => { type: 'tool', toolName: 'bash', status: 'pending', - createdAt: 1770000001000, + createdAt: 1770000001000 }, { id: 'text-real', @@ -197,30 +197,73 @@ describe('reduceWorkbenchMessages', () => { type: 'text', content: 'Final text', status: 'streaming', - createdAt: 1770000002000, - }, - ], + createdAt: 1770000002000 + } + ] }) expect(done[0].blocks).toMatchObject([ { id: 'thinking-real', type: 'thinking', status: 'done' }, { id: 'call_1', type: 'tool', status: 'done' }, - { id: 'text-real', type: 'text', status: 'done' }, + { id: 'text-real', type: 'text', status: 'done' } ]) }) + test('marks all streaming assistant messages cancelled when no turn is specified', () => { + const state = [ + { + id: 'user-1', + role: 'user' as const, + content: 'stop', + status: 'done' as const, + createdAt: '2026-05-25T00:00:00.000Z' + }, + { + id: 'assistant-old', + role: 'assistant' as const, + content: 'persisted streaming', + status: 'streaming' as const, + createdAt: '2026-05-25T00:00:00.000Z' + }, + { + id: 'assistant-new', + role: 'assistant' as const, + content: '', + status: 'streaming' as const, + turnId: 9, + createdAt: '2026-05-25T00:00:01.000Z' + } + ] + + const cancelled = reduceWorkbenchMessages(state, { + type: 'assistant_cancelled' + }) + + expect(cancelled[0]).toMatchObject({ role: 'user', status: 'done' }) + expect(cancelled[1]).toMatchObject({ + status: 'done', + runtimeStatus: 'cancelled', + stoppedNotice: true + }) + expect(cancelled[2]).toMatchObject({ + status: 'done', + runtimeStatus: 'cancelled', + stoppedNotice: true + }) + }) + test('clears stale stream error when an active block update arrives after disconnect', () => { const state = reduceWorkbenchMessages( reduceWorkbenchMessages([], { type: 'assistant_started', taskId: 1, - turnId: 9, + turnId: 9 }), { type: 'assistant_error', turnId: 9, error: 'Device disconnected', - errorType: 'container_error', + errorType: 'container_error' } ) @@ -233,8 +276,8 @@ describe('reduceWorkbenchMessages', () => { type: 'tool', toolName: 'bash', status: 'pending', - createdAt: 1770000000000, - }, + createdAt: 1770000000000 + } }) const resumed = reduceWorkbenchMessages(withTool, { type: 'block_updated', @@ -242,15 +285,20 @@ describe('reduceWorkbenchMessages', () => { blockId: 'call_1', updates: { status: 'streaming', - toolOutput: 'still running', - }, + toolOutput: 'still running' + } }) expect(resumed[0].status).toBe('streaming') expect(resumed[0].error).toBeUndefined() expect(resumed[0].errorType).toBeUndefined() expect(resumed[0].blocks).toMatchObject([ - { id: 'call_1', type: 'tool', status: 'streaming', toolOutput: 'still running' }, + { + id: 'call_1', + type: 'tool', + status: 'streaming', + toolOutput: 'still running' + } ]) }) @@ -259,13 +307,13 @@ describe('reduceWorkbenchMessages', () => { reduceWorkbenchMessages([], { type: 'assistant_started', taskId: 1, - turnId: 9, + turnId: 9 }), { type: 'assistant_error', turnId: 9, error: 'Codex CLI failed to resume thread: session not found', - errorType: 'execution_error', + errorType: 'execution_error' } ) @@ -273,13 +321,13 @@ describe('reduceWorkbenchMessages', () => { type: 'assistant_error', turnId: 9, error: 'Task failed with status: FAILED', - errorType: 'execution_error', + errorType: 'execution_error' }) expect(next[0]).toMatchObject({ status: 'failed', error: 'Codex CLI failed to resume thread: session not found', - errorType: 'execution_error', + errorType: 'execution_error' }) }) @@ -290,14 +338,13 @@ describe('reduceWorkbenchMessages', () => { role: 'assistant', content: 'hello', status: 'done', - createdAt: '2026-05-25T00:00:00.000Z', - }, + createdAt: '2026-05-25T00:00:00.000Z' + } ] - const next = reduceWorkbenchMessages( - state, - { type: 'unexpected' } as unknown as Parameters[1] - ) + const next = reduceWorkbenchMessages(state, { + type: 'unexpected' + } as unknown as Parameters[1]) expect(next).toBe(state) }) diff --git a/packages/chat-core/src/workbench-message-reducer.ts b/packages/chat-core/src/workbench-message-reducer.ts index fce437d3b6..0c3cdedbc8 100644 --- a/packages/chat-core/src/workbench-message-reducer.ts +++ b/packages/chat-core/src/workbench-message-reducer.ts @@ -4,6 +4,7 @@ export type WorkbenchMessageRole = 'user' | 'assistant' | 'system' export type WorkbenchMessageStatus = 'pending' | 'streaming' | 'done' | 'failed' +export type WorkbenchRuntimeMessageStatus = WorkbenchMessageStatus | 'cancelled' export type WorkbenchToolBlockStatus = | 'generating_arguments' | 'pending' @@ -40,8 +41,9 @@ export interface WorkbenchTextBlock extends BaseWorkbenchProcessingBlock { content: string } -export interface WorkbenchFileChangesBlock - extends BaseWorkbenchProcessingBlock { +export interface WorkbenchFileChangesBlock< + TFileChanges = unknown +> extends BaseWorkbenchProcessingBlock { type: 'file_changes' fileChanges: TFileChanges } @@ -66,6 +68,9 @@ export interface WorkbenchMessage blocks?: WorkbenchProcessingBlock[] fileChanges?: TFileChanges source?: MessageSource + runtimeStatus?: WorkbenchRuntimeMessageStatus | null + completedAt?: string | number | null + stoppedNotice?: boolean | null createdAt: string } @@ -79,7 +84,13 @@ type ProcessingBlockUpdate = { export type WorkbenchMessageAction = | { type: 'reset'; messages: WorkbenchMessage[] } | { type: 'user_added'; message: WorkbenchMessage } - | { type: 'assistant_started'; messageId?: string; taskId?: number; turnId?: number; shellType?: string } + | { + type: 'assistant_started' + messageId?: string + taskId?: number + turnId?: number + shellType?: string + } | { type: 'assistant_cached' messageId?: string @@ -104,12 +115,24 @@ export type WorkbenchMessageAction[] fileChanges?: TFileChanges } + | { + type: 'assistant_cancelled' + messageId?: string + turnId?: number + content?: string + } | { type: 'file_changes_updated' turnId: number fileChanges: TFileChanges } - | { type: 'assistant_error'; messageId?: string; turnId?: number; error: string; errorType?: string } + | { + type: 'assistant_error' + messageId?: string + turnId?: number + error: string + errorType?: string + } | { type: 'block_created' messageId?: string @@ -145,7 +168,7 @@ export function reduceWorkbenchMessages isAssistantMessageForAction(message, action))) { @@ -168,7 +191,7 @@ export function reduceWorkbenchMessages isAssistantMessageForAction(message, action))) { const message = createAssistantMessage({ messageId: action.messageId, turnId: action.turnId, - content: action.content, + content: action.content }) return [ ...state, { ...message, - blocks: getChunkBlocks(message, action), - }, + blocks: getChunkBlocks(message, action) + } ] } return state.map(message => @@ -204,7 +227,7 @@ export function reduceWorkbenchMessages @@ -229,10 +252,45 @@ export function reduceWorkbenchMessages isAssistantMessageForAction(message, action)) + ? (message: WorkbenchMessage) => + isAssistantMessageForAction(message, action) + : (message: WorkbenchMessage) => + message.role === 'assistant' && message.status === 'streaming' + if (!state.some(matches)) { + return [ + ...state, + createAssistantMessage({ + messageId: action.messageId, + turnId: action.turnId, + content: action.content ?? '', + status: 'done', + runtimeStatus: 'cancelled', + completedAt, + stoppedNotice: true + }) + ] + } + return state.map(message => + matches(message) + ? { + ...clearMessageError(message), + content: action.content ?? message.content, + status: 'done' as const, + runtimeStatus: 'cancelled' as const, + completedAt, + stoppedNotice: true, + blocks: finalizeProcessingBlocks(message.blocks, 'done') } : message ) + } case 'file_changes_updated': return state.map(message => isAssistantMessageForTurn(message, action.turnId) @@ -249,8 +307,8 @@ export function reduceWorkbenchMessages @@ -266,7 +324,7 @@ export function reduceWorkbenchMessages isAssistantMessageForAction(message, action))) { const message = createAssistantMessage({ messageId: action.messageId, - turnId: action.turnId, + turnId: action.turnId }) return [...state, createBlockCreatedMessage(message, action)] } @@ -290,9 +348,12 @@ export function reduceWorkbenchMessages block.id === action.blockId - ? ({ ...block, ...action.updates } as WorkbenchProcessingBlock) + ? ({ + ...block, + ...action.updates + } as WorkbenchProcessingBlock) : block - ), + ) } : message ) @@ -328,6 +389,9 @@ function createAssistantMessage({ fileChanges, error, errorType, + runtimeStatus, + completedAt, + stoppedNotice }: { messageId?: string taskId?: number @@ -339,6 +403,9 @@ function createAssistantMessage({ fileChanges?: TFileChanges error?: string errorType?: string + runtimeStatus?: WorkbenchRuntimeMessageStatus | null + completedAt?: string | number | null + stoppedNotice?: boolean | null }): WorkbenchMessage { return { id: messageId ?? `assistant-${turnId ?? Date.now()}`, @@ -352,7 +419,10 @@ function createAssistantMessage({ fileChanges, error, errorType, - createdAt: new Date().toISOString(), + runtimeStatus, + completedAt, + stoppedNotice, + createdAt: new Date().toISOString() } } @@ -373,16 +443,12 @@ function withActiveStreamState( return { ...clearMessageError(message), - status: 'streaming', + status: 'streaming' } } function isActiveBlockStatus(status?: WorkbenchToolBlockStatus): boolean { - return ( - status === 'generating_arguments' || - status === 'pending' || - status === 'streaming' - ) + return status === 'generating_arguments' || status === 'pending' || status === 'streaming' } function createBlockCreatedMessage( @@ -393,18 +459,19 @@ function createBlockCreatedMessage( const activeMessage = withActiveStreamState(message, isActiveBlockStatus(action.block.status)) return { ...activeMessage, - content: shouldMovePendingContentBeforeBlock(message, action.block) - ? '' - : message.content, + content: shouldMovePendingContentBeforeBlock(message, action.block) ? '' : message.content, blocks: mergeProcessingBlock( getBlocksBeforeIncomingBlock(message, turnId, action.block), action.block - ), + ) } } export function normalizeWorkbenchBlockStatus(status?: string): WorkbenchToolBlockStatus { - const normalizedStatus = status?.trim().toLowerCase().replace(/[\s-]+/g, '_') + const normalizedStatus = status + ?.trim() + .toLowerCase() + .replace(/[\s-]+/g, '_') switch (normalizedStatus) { case 'generating_arguments': @@ -455,7 +522,7 @@ function appendThinkingChunk( if (lastBlock?.type === 'thinking' && lastBlock.status === 'streaming') { nextBlocks[nextBlocks.length - 1] = { ...lastBlock, - content: lastBlock.content + chunk, + content: lastBlock.content + chunk } return nextBlocks } @@ -469,8 +536,8 @@ function appendThinkingChunk( type: 'thinking', content: chunk, status: 'streaming', - createdAt: Date.now(), - }, + createdAt: Date.now() + } ] } @@ -490,8 +557,8 @@ function getBlocksBeforeIncomingBlock( type: 'text', content: message.content, status: 'done', - createdAt: Date.now(), - }, + createdAt: Date.now() + } ] } @@ -510,10 +577,7 @@ function finalizeOpenNarrativeBlocks( blocks: WorkbenchProcessingBlock[] | undefined ): WorkbenchProcessingBlock[] { return (blocks ?? []).map(block => { - if ( - (block.type === 'thinking' || block.type === 'text') && - block.status === 'streaming' - ) { + if ((block.type === 'thinking' || block.type === 'text') && block.status === 'streaming') { return { ...block, status: 'done' as const } } @@ -534,7 +598,10 @@ function finalizeBlocks( return block } - return { ...block, status: finalStatus } as WorkbenchProcessingBlock + return { + ...block, + status: finalStatus + } as WorkbenchProcessingBlock }) } @@ -556,7 +623,7 @@ function mergeProcessingBlock( const nextBlocks = [...blocks] nextBlocks[index] = { ...nextBlocks[index], - ...incomingBlock, + ...incomingBlock } as WorkbenchProcessingBlock return nextBlocks } diff --git a/wework/src/components/layout/DesktopWorkbenchMain.tsx b/wework/src/components/layout/DesktopWorkbenchMain.tsx index 56c45b91af..67a2e9a580 100644 --- a/wework/src/components/layout/DesktopWorkbenchMain.tsx +++ b/wework/src/components/layout/DesktopWorkbenchMain.tsx @@ -127,7 +127,6 @@ const DesktopWorkbenchPane = memo(function DesktopWorkbenchPane({ projectChat, upgradeDevice, retryFailedMessage, - pauseCurrentResponse, loadTurnFileChangesDiff, revertTurnFileChanges, forkCurrentRuntimeTask, @@ -900,7 +899,7 @@ const DesktopWorkbenchPane = memo(function DesktopWorkbenchPane({ guidanceMessages={paneGuidanceMessages} codeComments={paneSession.codeCommentContexts} isStreaming={paneIsResponseStreaming} - onPause={() => void pauseCurrentResponse(paneMessages)} + onPause={() => void paneSession.pauseCurrentResponse()} onCancelQueuedMessage={paneSession.cancelQueuedMessage} onSendQueuedAsGuidance={paneSession.sendQueuedAsGuidance} onEditQueuedMessage={paneSession.editQueuedMessage} @@ -954,7 +953,7 @@ const DesktopWorkbenchPane = memo(function DesktopWorkbenchPane({ guidanceMessages={paneGuidanceMessages} codeComments={paneSession.codeCommentContexts} isStreaming={paneIsResponseStreaming} - onPause={() => void pauseCurrentResponse(paneMessages)} + onPause={() => void paneSession.pauseCurrentResponse()} onCancelQueuedMessage={paneSession.cancelQueuedMessage} onSendQueuedAsGuidance={paneSession.sendQueuedAsGuidance} onEditQueuedMessage={paneSession.editQueuedMessage} @@ -1045,7 +1044,7 @@ const DesktopWorkbenchPane = memo(function DesktopWorkbenchPane({ devices={devices} requiresStop={paneIsResponseStreaming} onOpenChange={setForkDialogOpen} - onStopCurrentResponse={() => pauseCurrentResponse(paneMessages)} + onStopCurrentResponse={() => paneSession.pauseCurrentResponse()} onPrepareDeviceWorkspace={prepareDeviceWorkspace} onDeleteDeviceWorkspace={deleteDeviceWorkspace} onGetDeviceHomeDirectory={getDeviceHomeDirectory} diff --git a/wework/src/components/layout/MobileWorkbenchLayout.tsx b/wework/src/components/layout/MobileWorkbenchLayout.tsx index 921614df0e..c3abd45353 100644 --- a/wework/src/components/layout/MobileWorkbenchLayout.tsx +++ b/wework/src/components/layout/MobileWorkbenchLayout.tsx @@ -64,7 +64,6 @@ const MobileWorkbenchPane = memo(function MobileWorkbenchPane({ projectChat, upgradeDevice, retryFailedMessage, - pauseCurrentResponse, loadTurnFileChangesDiff, revertTurnFileChanges, forkCurrentRuntimeTask, @@ -328,7 +327,7 @@ const MobileWorkbenchPane = memo(function MobileWorkbenchPane({ guidanceMessages={paneGuidanceMessages} codeComments={paneSession.codeCommentContexts} isStreaming={paneIsResponseStreaming} - onPause={() => void pauseCurrentResponse(paneMessages)} + onPause={() => void paneSession.pauseCurrentResponse()} onCancelQueuedMessage={paneSession.cancelQueuedMessage} onSendQueuedAsGuidance={paneSession.sendQueuedAsGuidance} onEditQueuedMessage={paneSession.editQueuedMessage} @@ -416,7 +415,7 @@ const MobileWorkbenchPane = memo(function MobileWorkbenchPane({ guidanceMessages={paneGuidanceMessages} codeComments={paneSession.codeCommentContexts} isStreaming={paneIsResponseStreaming} - onPause={() => void pauseCurrentResponse(paneMessages)} + onPause={() => void paneSession.pauseCurrentResponse()} onCancelQueuedMessage={paneSession.cancelQueuedMessage} onSendQueuedAsGuidance={paneSession.sendQueuedAsGuidance} onEditQueuedMessage={paneSession.editQueuedMessage} @@ -473,7 +472,7 @@ const MobileWorkbenchPane = memo(function MobileWorkbenchPane({ devices={state.devices} requiresStop={paneIsResponseStreaming} onOpenChange={setForkDialogOpen} - onStopCurrentResponse={() => pauseCurrentResponse(paneMessages)} + onStopCurrentResponse={() => paneSession.pauseCurrentResponse()} onPrepareDeviceWorkspace={onPrepareDeviceWorkspace} onDeleteDeviceWorkspace={onDeleteDeviceWorkspace} onGetDeviceHomeDirectory={onGetDeviceHomeDirectory} diff --git a/wework/src/components/layout/useWorkbenchPaneSession.ts b/wework/src/components/layout/useWorkbenchPaneSession.ts index e97ee2bc1d..238dde3f9a 100644 --- a/wework/src/components/layout/useWorkbenchPaneSession.ts +++ b/wework/src/components/layout/useWorkbenchPaneSession.ts @@ -563,7 +563,7 @@ export function useWorkbenchPaneSession({ currentRuntimeTask }: WorkbenchPaneSes if (activeAssistantMessage) { const action: RuntimePaneMessageAction = { - type: 'assistant_done', + type: 'assistant_cancelled', messageId: activeAssistantMessage.id, turnId: activeAssistantMessage.turnId, content: activeAssistantMessage.content, @@ -592,6 +592,17 @@ export function useWorkbenchPaneSession({ currentRuntimeTask }: WorkbenchPaneSes ] ) + const pauseCurrentResponse = useCallback(async () => { + if (!currentRuntimeTask || !activeAssistantMessage) return + + const cancelled = await cancelRuntimePaneTask(currentRuntimeTask) + if (!cancelled) return + + dispatchMessages({ + type: 'assistant_cancelled', + }) + }, [activeAssistantMessage, cancelRuntimePaneTask, currentRuntimeTask, dispatchMessages]) + const cancelGuidanceMessage = useCallback(() => undefined, []) return { @@ -617,6 +628,7 @@ export function useWorkbenchPaneSession({ currentRuntimeTask }: WorkbenchPaneSes sendQueuedAsGuidance, editQueuedMessage, cancelGuidanceMessage, + pauseCurrentResponse, } } diff --git a/wework/src/features/workbench/WorkbenchProvider.test.tsx b/wework/src/features/workbench/WorkbenchProvider.test.tsx index 5f1b7f84de..0b2ce97bcc 100644 --- a/wework/src/features/workbench/WorkbenchProvider.test.tsx +++ b/wework/src/features/workbench/WorkbenchProvider.test.tsx @@ -696,6 +696,9 @@ function RuntimeOpenProbe() { {paneSession.messages.map(message => message.content).join('|')} + + {paneSession.messages.map(message => `${message.role}:${message.status}`).join('|')} + {paneSession.transcriptLoading ? 'loading' : 'idle'} @@ -729,6 +732,9 @@ function RuntimeOpenProbe() { .join('|')} {workbench.state.error ?? ''} + + {workbench.currentRuntimeTaskRunning ? 'running' : 'idle'} + {fileChangesDiff} {fileChangesStatus} + { localTaskId: 'runtime-a', workspacePath: '/workspace/project-alpha', runtime: 'claude_code', - messages: [{ id: 'runtime-a:user:1', role: 'user', content: 'first message' }], + messages: [ + { id: 'runtime-a:user:1', role: 'user', content: 'first message' }, + { + id: 'runtime-a:assistant:1', + role: 'assistant', + content: 'working', + status: 'streaming', + }, + ], }), sendRuntimeMessage, cancelRuntimeTask, @@ -3816,6 +3833,151 @@ describe('WorkbenchProvider runtime tasks', () => { expect(screen.getByTestId('guidance-messages')).toHaveTextContent('') }) + test('refreshes runtime work after cancelling the active local task', async () => { + let streamHandlers: ChatStreamHandlers = {} + const subscribe = vi.fn((handlers: ChatStreamHandlers) => { + if (hasRuntimeStreamHandler(handlers)) streamHandlers = handlers + return vi.fn() + }) + const runningRuntimeWork = createRuntimeWork({ + projects: [ + { + project: { id: 7, name: 'Wegent' }, + deviceWorkspaces: [ + { + id: 22, + projectId: 7, + deviceId: 'device-1', + deviceName: 'Project Device', + deviceStatus: 'online', + workspacePath: '/workspace/project-alpha', + mapped: true, + available: true, + localTasks: [ + { + localTaskId: 'runtime-a', + workspacePath: '/workspace/project-alpha', + title: 'Runtime A', + runtime: 'claude_code', + running: true, + }, + ], + }, + ], + }, + ], + totalLocalTasks: 1, + }) + const idleRuntimeWork = createRuntimeWork({ + projects: [ + { + project: { id: 7, name: 'Wegent' }, + deviceWorkspaces: [ + { + id: 22, + projectId: 7, + deviceId: 'device-1', + deviceName: 'Project Device', + deviceStatus: 'online', + workspacePath: '/workspace/project-alpha', + mapped: true, + available: true, + localTasks: [ + { + localTaskId: 'runtime-a', + workspacePath: '/workspace/project-alpha', + title: 'Runtime A', + runtime: 'claude_code', + running: false, + status: 'cancelled', + }, + ], + }, + ], + }, + ], + totalLocalTasks: 1, + }) + let runtimeRunning = true + const listRuntimeWork = vi + .fn() + .mockImplementation(() => + Promise.resolve(runtimeRunning ? runningRuntimeWork : idleRuntimeWork) + ) + const cancelRuntimeTask = vi.fn().mockImplementation(() => { + runtimeRunning = false + return Promise.resolve({ + accepted: true, + localTaskId: 'runtime-a', + }) + }) + const runtimeWorkApi = createRuntimeWorkApiMock({ + listRuntimeWork, + getRuntimeTranscript: vi.fn().mockResolvedValue({ + localTaskId: 'runtime-a', + workspacePath: '/workspace/project-alpha', + runtime: 'claude_code', + messages: [ + { id: 'runtime-a:user:1', role: 'user', content: 'first message' }, + { + id: 'runtime-a:assistant:1', + role: 'assistant', + content: 'working', + status: 'streaming', + }, + ], + }), + cancelRuntimeTask, + }) + const services = createWorkbenchServices({ + runtimeWorkApi: runtimeWorkApi as WorkbenchServices['runtimeWorkApi'], + chatStream: { + subscribe, + } as unknown as WorkbenchServices['chatStream'], + }) + + renderWorkbench( + <> + + + , + services + ) + + await userEvent.click(await screen.findByText('open runtime a')) + await waitFor(() => + expect(screen.getByTestId('current-runtime-task-running')).toHaveTextContent('running') + ) + await act(async () => { + streamHandlers.onChatStart?.({ + message_id: 701, + task_id: 77, + subtask_id: 101, + shell_type: 'Chat', + device_id: 'device-1', + local_task_id: 'runtime-a', + }) + }) + const listCallsBeforeCancel = listRuntimeWork.mock.calls.length + await userEvent.click(screen.getByText('stop current response')) + + await waitFor(() => expect(cancelRuntimeTask).toHaveBeenCalledTimes(1)) + expect(cancelRuntimeTask).toHaveBeenCalledWith({ + deviceId: 'device-1', + workspacePath: '/workspace/project-alpha', + localTaskId: 'runtime-a', + }) + await waitFor(() => expect(listRuntimeWork).toHaveBeenCalledTimes(listCallsBeforeCancel + 1)) + await waitFor(() => + expect(screen.getByTestId('current-runtime-task-running')).toHaveTextContent('idle') + ) + await waitFor(() => + expect(screen.getByTestId('runtime-message-statuses')).not.toHaveTextContent( + 'assistant:streaming' + ) + ) + }) + test('pauses the active local task before sending queued guidance without DB task context', async () => { let streamHandlers: ChatStreamHandlers = {} const subscribe = vi.fn((handlers: ChatStreamHandlers) => { diff --git a/wework/src/features/workbench/useWorkbenchRuntimeMessaging.ts b/wework/src/features/workbench/useWorkbenchRuntimeMessaging.ts index b3de309a38..d7e899620f 100644 --- a/wework/src/features/workbench/useWorkbenchRuntimeMessaging.ts +++ b/wework/src/features/workbench/useWorkbenchRuntimeMessaging.ts @@ -143,9 +143,10 @@ export function useWorkbenchRuntimeMessaging({ }) return false } + await refreshWorkLists() return true }, - [dispatch, executorClient] + [dispatch, executorClient, refreshWorkLists] ) const buildSendPayload = useCallback( @@ -756,9 +757,11 @@ export function useWorkbenchRuntimeMessaging({ type: 'error_set', error: normalizeGuidanceError(ack.error ?? '取消当前回复失败'), }) + return } + await refreshWorkLists() }, - [dispatch, executorClient, state.currentRuntimeTask] + [dispatch, executorClient, refreshWorkLists, state.currentRuntimeTask] ) return {