diff --git a/src/crates/execution/tool-execution/src/exec_command.rs b/src/crates/execution/tool-execution/src/exec_command.rs index 63229e5195..19b5474b16 100644 --- a/src/crates/execution/tool-execution/src/exec_command.rs +++ b/src/crates/execution/tool-execution/src/exec_command.rs @@ -940,6 +940,12 @@ fn completion_status_lines(data: &Value) -> Vec { status_lines.push(format!( "Process is still running. session_id: {session_id}" )); + } else if completion_status == Some("exited") { + // The process finished but the transport never reported a status. Say so + // instead of leaving the caller with "Process status unavailable", and + // never invent a code — an invented failure reads as a real one. + status_lines + .push("Process exited, but no exit code was reported by the transport.".to_string()); } status_lines @@ -1021,6 +1027,33 @@ mod tests { assert!(rendered.contains("\n\n")); } + #[test] + fn command_response_says_an_exited_process_had_no_reported_exit_code() { + let data = json!({ + "wall_time_seconds": 0.1, + "output": "workspace output\n", + "tty": false, + "session_id": null, + "exit_code": null, + "completion": { + "status": "exited", + "source": "process" + } + }); + + let rendered = render_exec_command_response_for_assistant(&data); + + assert!(rendered.contains("Process exited, but no exit code was reported")); + assert!( + !rendered.contains("Process status unavailable."), + "a completed process is not an unknown state" + ); + assert!( + !rendered.contains("code -1"), + "an unknown status must never be rendered as a failing exit code" + ); + } + #[test] fn write_stdin_response_reports_external_interrupt() { let data = json!({ diff --git a/src/crates/services/services-integrations/src/remote_ssh/manager.rs b/src/crates/services/services-integrations/src/remote_ssh/manager.rs index 6376f36459..276eb2c7d9 100644 --- a/src/crates/services/services-integrations/src/remote_ssh/manager.rs +++ b/src/crates/services/services-integrations/src/remote_ssh/manager.rs @@ -3324,6 +3324,13 @@ impl SSHConnectionManager { } Some(russh::ChannelMsg::ExitSignal { signal_name, .. }) => { interrupted = interrupted || matches!(signal_name, Sig::INT | Sig::TERM); + // A signal death is still a resolved status. Without this the + // result fell through to the unknown `-1` below. + if exit_status.is_none() { + exit_status = crate::remote_ssh::transport::ssh_exit_code_for_signal( + &signal_name, + ); + } log::debug!( "Remote exec exit signal received: signal={:?}, stdout_len={}, stderr_len={}, duration_ms={}, command_preview={}", signal_name, diff --git a/src/crates/services/services-integrations/src/remote_ssh/remote_exec.rs b/src/crates/services/services-integrations/src/remote_ssh/remote_exec.rs index 4131ea9016..35bc0614ba 100644 --- a/src/crates/services/services-integrations/src/remote_ssh/remote_exec.rs +++ b/src/crates/services/services-integrations/src/remote_ssh/remote_exec.rs @@ -4,6 +4,7 @@ //! workspaces while keeping tool-owned command sessions separate from UI //! terminal sessions. +use crate::remote_ssh::transport::{ssh_exit_code_for_signal, SSH_EXIT_STATUS_AFTER_EOF_GRACE}; use crate::remote_ssh::SSHConnectionManager; use anyhow::{anyhow, Context}; use rand::Rng; @@ -820,6 +821,7 @@ async fn workspace_pipe_owner( let (mut stdin, mut stdout, mut stderr, control, completion) = transport.into_parts(); let mut completion_task = tokio::spawn(completion.wait()); let mut exit_code = None; + let mut process_completed = false; let mut control_state: Option = None; let mut stdout_closed = false; let mut stderr_closed = false; @@ -827,7 +829,7 @@ async fn workspace_pipe_owner( let mut stderr_buffer = vec![0u8; 16 * 1024]; loop { - if exit_code.is_some() && stdout_closed && stderr_closed { + if process_completed && stdout_closed && stderr_closed { break; } if let Some(state) = control_state { @@ -910,11 +912,12 @@ async fn workspace_pipe_owner( } } - completed = &mut completion_task, if exit_code.is_none() => { + // An unknown status stays unknown. Reporting a synthetic `-1` here + // used to make a command that ran fine look like it failed, and the + // model cannot tell that apart from a real non-zero exit. + completed = &mut completion_task, if !process_completed => { exit_code = completed.ok().and_then(|exit| exit.exit_code); - if exit_code.is_none() { - exit_code = Some(-1); - } + process_completed = true; } _ = tokio::time::sleep(wait_budget), if control_state.is_some() => {} @@ -935,14 +938,21 @@ async fn remote_pty_owner( ) { let mut exit_code = None; let mut close_after_control_at: Option = None; + let mut exit_status_deadline: Option = None; loop { if close_after_control_at.is_some_and(|deadline| Instant::now() >= deadline) { let _ = channel.close().await; break; } + if exit_status_deadline.is_some_and(|deadline| Instant::now() >= deadline) { + break; + } let wait_budget = close_after_control_at + .into_iter() + .chain(exit_status_deadline) + .min() .map(|deadline| deadline.saturating_duration_since(Instant::now())) .filter(|duration| !duration.is_zero()) .unwrap_or_else(|| Duration::from_millis(100)); @@ -982,21 +992,35 @@ async fn remote_pty_owner( } Some(ChannelMsg::ExitStatus { exit_status }) => { exit_code = Some(exit_status as i32); + if exit_status_deadline.is_some() { + break; + } + } + Some(ChannelMsg::ExitSignal { ref signal_name, .. }) => { + if exit_code.is_none() { + exit_code = ssh_exit_code_for_signal(signal_name); + } + if exit_status_deadline.is_some() && exit_code.is_some() { + break; + } } - Some(ChannelMsg::ExitSignal { signal_name, .. }) => { - exit_code = Some(match signal_name { - Sig::INT => 130, - Sig::KILL => 137, - Sig::TERM => 143, - _ => -1, + // See `run_ssh_channel`: EOF may precede the exit status, so + // keep the channel open long enough to collect it. + Some(ChannelMsg::Eof) => { + if exit_code.is_some() { + break; + } + exit_status_deadline.get_or_insert_with(|| { + Instant::now() + SSH_EXIT_STATUS_AFTER_EOF_GRACE }); } - Some(ChannelMsg::Eof) | Some(ChannelMsg::Close) | None => break, + Some(ChannelMsg::Close) | None => break, Some(_) => {} } } - _ = tokio::time::sleep(wait_budget), if close_after_control_at.is_some() => {} + _ = tokio::time::sleep(wait_budget), + if close_after_control_at.is_some() || exit_status_deadline.is_some() => {} } } @@ -1354,9 +1378,59 @@ fn new_chunk_id() -> String { #[cfg(test)] mod tests { use super::{ - decode_utf8_stream, new_session_id, HeadTailText, OutputStream, PendingUtf8Streams, + decode_utf8_stream, new_session_id, workspace_pipe_owner, HeadTailText, OutputState, + OutputStream, PendingUtf8Streams, }; + use crate::remote_ssh::transport::WorkspaceStdio; use std::collections::HashMap; + use std::sync::Arc; + use tokio::sync::mpsc; + use tokio::time::Duration; + + #[cfg(unix)] + async fn pipe_owner_exit_code(script: &str) -> Option { + let transport = + WorkspaceStdio::spawn_local_process("sh", &["-lc".to_string(), script.to_string()]) + .expect("local workspace process should start"); + let output = Arc::new(OutputState::new(None)); + let (_command_tx, command_rx) = mpsc::channel(8); + tokio::spawn(workspace_pipe_owner( + transport, + command_rx, + Arc::clone(&output), + )); + + tokio::time::timeout(Duration::from_secs(10), output.wait_closed()) + .await + .expect("pipe owner should close the output state") + } + + #[tokio::test] + #[cfg(unix)] + async fn pipe_owner_reports_successful_process_exit_code() { + assert_eq!(pipe_owner_exit_code("df -h >/dev/null 2>&1").await, Some(0)); + } + + #[tokio::test] + #[cfg(unix)] + async fn pipe_owner_reports_failing_process_exit_code() { + assert_eq!(pipe_owner_exit_code("exit 3").await, Some(3)); + } + + #[tokio::test] + #[cfg(unix)] + async fn pipe_owner_reports_exit_code_after_large_output() { + assert_eq!( + pipe_owner_exit_code("head -c 400000 /dev/zero | tr '\\0' 'a'").await, + Some(0) + ); + } + + #[tokio::test] + #[cfg(unix)] + async fn pipe_owner_reports_signal_death_as_conventional_status() { + assert_eq!(pipe_owner_exit_code("kill -TERM $$").await, Some(143)); + } #[test] fn remote_exec_session_ids_match_local_test_baseline() { diff --git a/src/crates/services/services-integrations/src/remote_ssh/transport.rs b/src/crates/services/services-integrations/src/remote_ssh/transport.rs index 75308bab05..6567a01604 100644 --- a/src/crates/services/services-integrations/src/remote_ssh/transport.rs +++ b/src/crates/services/services-integrations/src/remote_ssh/transport.rs @@ -16,12 +16,69 @@ use std::process::Stdio; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::task::{Context as TaskContext, Poll}; +use std::time::Duration; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, DuplexStream, ReadBuf}; use tokio::sync::{mpsc, watch}; use tokio_util::sync::CancellationToken; const WORKSPACE_STDIO_BUFFER_SIZE: usize = 256 * 1024; +/// How long to keep an SSH channel open after `SSH_MSG_CHANNEL_EOF` while the +/// exit status is still missing. +/// +/// EOF only says the peer will send no more data. RFC 4254 §6.10 leaves the +/// ordering of the `exit-status` request free, and OpenSSH in practice flushes +/// EOF from its channel loop before it reaps the child and reports the status. +/// Closing on EOF therefore threw away the exit code of nearly every +/// short-lived command. Waiting for `SSH_MSG_CHANNEL_CLOSE` costs nothing in +/// the normal case because it follows immediately; the grace only bounds +/// servers that go quiet without closing. +pub(crate) const SSH_EXIT_STATUS_AFTER_EOF_GRACE: Duration = Duration::from_secs(5); + +/// Map an SSH `exit-signal` to the conventional `128 + signal` wait status. +/// +/// Returns `None` for signals with no portable number so callers can report an +/// unknown status instead of inventing a misleading exit code. +#[cfg(feature = "remote-ssh-concrete")] +pub(crate) fn ssh_exit_code_for_signal(signal: &Sig) -> Option { + let number = match signal { + Sig::HUP => 1, + Sig::INT => 2, + Sig::QUIT => 3, + Sig::ILL => 4, + Sig::ABRT => 6, + Sig::FPE => 8, + Sig::KILL => 9, + Sig::SEGV => 11, + Sig::PIPE => 13, + Sig::ALRM => 14, + Sig::TERM => 15, + Sig::USR1 => 10, + Sig::Custom(_) => return None, + }; + Some(128 + number) +} + +/// Map a locally waited child status to an exit code. +/// +/// `ExitStatus::code()` is `None` when a process dies from a signal, which is +/// exactly how interrupt and kill end a supervised workspace command. Report +/// the conventional `128 + signal` status for those instead of losing it. +fn local_process_exit_code(status: std::process::ExitStatus) -> Option { + #[cfg(unix)] + { + use std::os::unix::process::ExitStatusExt; + + status + .code() + .or_else(|| status.signal().map(|signal| 128 + signal)) + } + #[cfg(not(unix))] + { + status.code() + } +} + pub type WorkspaceReader = Pin>; pub type WorkspaceWriter = Pin>; pub(crate) type WorkspaceSignalHook = Arc< @@ -342,8 +399,17 @@ async fn run_ssh_channel(mut channel: Channel, mut pipes: WorkspacePipeOwne let mut stdin_buffer = vec![0u8; 16 * 1024]; let mut stdin_closed = false; let mut exit_code = None; + // Set once the peer sends EOF while the exit status is still missing, so a + // server that never follows up cannot hold the channel open forever. + let mut exit_status_deadline: Option = None; loop { + let wait_budget = exit_status_deadline + .map(|deadline| deadline.saturating_duration_since(tokio::time::Instant::now())); + if wait_budget.is_some_and(|budget| budget.is_zero()) { + break; + } + tokio::select! { biased; @@ -389,16 +455,37 @@ async fn run_ssh_channel(mut channel: Channel, mut pipes: WorkspacePipeOwne } Some(ChannelMsg::ExitStatus { exit_status }) => { exit_code = Some(exit_status as i32); + // EOF already arrived, so the status was the last thing + // worth waiting for. Do not linger for CHANNEL_CLOSE. + if exit_status_deadline.is_some() { + break; + } } - Some(ChannelMsg::ExitSignal { signal_name, .. }) => { - exit_code = Some(match signal_name { - Sig::INT => 130, - Sig::KILL => 137, - Sig::TERM => 143, - _ => -1, - }); + Some(ChannelMsg::ExitSignal { ref signal_name, .. }) => { + // A server sends either exit-status or exit-signal. Keep + // whichever arrived first rather than letting an + // unmappable signal erase a known code. + if exit_code.is_none() { + exit_code = ssh_exit_code_for_signal(signal_name); + } + if exit_status_deadline.is_some() && exit_code.is_some() { + break; + } } - Some(ChannelMsg::Eof) | Some(ChannelMsg::Close) | None => break, + // EOF is not the end of the channel: the exit status is + // still allowed to follow, and OpenSSH usually sends it + // afterwards. Keep draining until CLOSE, or until the grace + // window expires, so the status is not silently dropped. + Some(ChannelMsg::Eof) => { + if exit_code.is_some() { + break; + } + exit_status_deadline + .get_or_insert_with(|| { + tokio::time::Instant::now() + SSH_EXIT_STATUS_AFTER_EOF_GRACE + }); + } + Some(ChannelMsg::Close) | None => break, Some(_) => {} } } @@ -409,6 +496,10 @@ async fn run_ssh_channel(mut channel: Channel, mut pipes: WorkspacePipeOwne exit_code.get_or_insert(137); break; } + + _ = tokio::time::sleep(wait_budget.unwrap_or_default()), if wait_budget.is_some() => { + break; + } } } @@ -446,8 +537,11 @@ async fn run_local_process( let exit_code = loop { tokio::select! { status = child.wait() => { - break status.ok().and_then(|status| status.code()); + break status.ok().and_then(local_process_exit_code); } + // We are the ones ending the process here, so a signal death says + // nothing the caller does not already know. Prefer a status the + // child chose for itself and otherwise report the requested intent. signal = pipes.control_rx.recv() => { let fallback = match signal { Some(WorkspaceProcessSignal::Interrupt) => 130, @@ -472,10 +566,242 @@ async fn run_local_process( .send(Some(WorkspaceProcessExit { exit_code })); } +/// In-process SSH server that lets the channel-owner loop be tested against the +/// message orderings real servers use, without needing a live host. +#[cfg(test)] +mod ssh_channel_tests { + use super::*; + use russh::server::{Auth, Msg as ServerMsg, Server as _, Session}; + use russh::{Channel as RusshChannel, ChannelId, CryptoVec}; + use std::net::SocketAddr; + use tokio::net::TcpListener; + + /// What the fake server sends once the client asks it to run something. + #[derive(Clone, Copy, PartialEq, Eq)] + enum ExitReport { + /// EOF first, exit status afterwards — what OpenSSH does for a + /// short-lived command, because its channel loop flushes EOF before it + /// reaps the child and reports the status. + EofBeforeExitStatus, + /// Exit status first, then EOF. + ExitStatusBeforeEof, + /// EOF, then a signal death, then close. + EofBeforeExitSignal, + /// EOF and close with no status at all. + NoStatus, + } + + #[derive(Clone)] + struct TestServer { + report: ExitReport, + } + + impl russh::server::Server for TestServer { + type Handler = Self; + + fn new_client(&mut self, _peer: Option) -> Self { + self.clone() + } + } + + #[async_trait::async_trait] + impl russh::server::Handler for TestServer { + type Error = russh::Error; + + async fn auth_none(&mut self, _user: &str) -> Result { + Ok(Auth::Accept) + } + + async fn auth_password( + &mut self, + _user: &str, + _password: &str, + ) -> Result { + Ok(Auth::Accept) + } + + async fn channel_open_session( + &mut self, + _channel: RusshChannel, + _session: &mut Session, + ) -> Result { + Ok(true) + } + + async fn exec_request( + &mut self, + channel: ChannelId, + _data: &[u8], + session: &mut Session, + ) -> Result<(), Self::Error> { + let handle = session.handle(); + let report = self.report; + tokio::spawn(async move { + let _ = handle + .data(channel, CryptoVec::from_slice(b"workspace output\n")) + .await; + let settle = || tokio::time::sleep(Duration::from_millis(30)); + match report { + ExitReport::EofBeforeExitStatus => { + let _ = handle.eof(channel).await; + settle().await; + let _ = handle.exit_status_request(channel, 7).await; + let _ = handle.close(channel).await; + } + ExitReport::ExitStatusBeforeEof => { + let _ = handle.exit_status_request(channel, 7).await; + settle().await; + let _ = handle.eof(channel).await; + let _ = handle.close(channel).await; + } + ExitReport::EofBeforeExitSignal => { + let _ = handle.eof(channel).await; + settle().await; + let _ = handle + .exit_signal_request(channel, russh::Sig::TERM, false, String::new(), String::new()) + .await; + let _ = handle.close(channel).await; + } + ExitReport::NoStatus => { + let _ = handle.eof(channel).await; + let _ = handle.close(channel).await; + } + } + }); + Ok(()) + } + } + + struct TestClient; + + #[async_trait::async_trait] + impl russh::client::Handler for TestClient { + type Error = russh::Error; + + async fn check_server_key( + &mut self, + _key: &russh_keys::key::PublicKey, + ) -> Result { + Ok(true) + } + } + + async fn workspace_exit_for(report: ExitReport) -> WorkspaceProcessExit { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("test SSH listener should bind"); + let address = listener + .local_addr() + .expect("test SSH listener should report its address"); + let server_config = Arc::new(russh::server::Config { + keys: vec![ + russh_keys::key::KeyPair::generate_ed25519().expect("test host key should generate"), + ], + ..Default::default() + }); + tokio::spawn(async move { + let mut server = TestServer { report }; + let _ = server.run_on_socket(server_config, &listener).await; + }); + + let client_config = Arc::new(russh::client::Config::default()); + let mut handle = russh::client::connect(client_config, address, TestClient) + .await + .expect("test client should connect"); + assert!( + handle + .authenticate_password("tester", "tester") + .await + .expect("test authentication should complete"), + "test server should accept the password" + ); + let channel = handle + .channel_open_session() + .await + .expect("test channel should open"); + channel + .exec(true, "df -h") + .await + .expect("test exec should start"); + + let transport = WorkspaceStdio::from_ssh_channel(channel); + let (_stdin, mut stdout, _stderr, _control, completion) = transport.into_parts(); + let mut stdout_bytes = Vec::new(); + let _ = stdout.read_to_end(&mut stdout_bytes).await; + assert_eq!(stdout_bytes, b"workspace output\n"); + + tokio::time::timeout(Duration::from_secs(20), completion.wait()) + .await + .expect("channel owner should report completion") + } + + #[tokio::test] + async fn exit_status_sent_after_eof_is_still_reported() { + let exit = workspace_exit_for(ExitReport::EofBeforeExitStatus).await; + + assert_eq!( + exit.exit_code, + Some(7), + "EOF does not end an SSH channel; the exit status may follow it" + ); + } + + #[tokio::test] + async fn exit_status_sent_before_eof_is_reported() { + let exit = workspace_exit_for(ExitReport::ExitStatusBeforeEof).await; + + assert_eq!(exit.exit_code, Some(7)); + } + + #[tokio::test] + async fn exit_signal_sent_after_eof_maps_to_a_conventional_status() { + let exit = workspace_exit_for(ExitReport::EofBeforeExitSignal).await; + + assert_eq!(exit.exit_code, Some(143)); + } + + #[tokio::test] + async fn missing_exit_status_stays_unknown() { + let exit = workspace_exit_for(ExitReport::NoStatus).await; + + assert_eq!( + exit.exit_code, None, + "an unreported status must not be turned into a synthetic failure code" + ); + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn ssh_exit_signals_map_to_conventional_wait_statuses() { + assert_eq!(ssh_exit_code_for_signal(&Sig::INT), Some(130)); + assert_eq!(ssh_exit_code_for_signal(&Sig::KILL), Some(137)); + assert_eq!(ssh_exit_code_for_signal(&Sig::TERM), Some(143)); + assert_eq!( + ssh_exit_code_for_signal(&Sig::Custom("WEIRD".to_string())), + None, + "an unmappable signal must stay unknown rather than become -1" + ); + } + + #[tokio::test] + #[cfg(unix)] + async fn local_process_signal_death_reports_a_conventional_status() { + let transport = + WorkspaceStdio::spawn_local_process("sh", &["-lc".to_string(), "kill -TERM $$".to_string()]) + .unwrap(); + let (_stdin, _stdout, _stderr, _control, completion) = transport.into_parts(); + + let exit = tokio::time::timeout(Duration::from_secs(5), completion.wait()) + .await + .expect("signal death should complete the supervised process"); + + assert_eq!(exit.exit_code, Some(143)); + } + #[tokio::test] #[cfg(unix)] async fn local_process_round_trips_stdin_stdout_and_exit_status() { diff --git a/src/crates/services/terminal/src/exec.rs b/src/crates/services/terminal/src/exec.rs index 956db3853b..0b2af05a99 100644 --- a/src/crates/services/terminal/src/exec.rs +++ b/src/crates/services/terminal/src/exec.rs @@ -1233,7 +1233,7 @@ async fn spawn_pipe_process(request: &ExecCommandRequest) -> TerminalResult>(1); #[cfg(unix)] tokio::spawn(async move { - let code = child.wait().await.ok().and_then(|status| status.code()); + let code = child.wait().await.ok().and_then(local_pipe_exit_code); let _ = child_exit_tx.send(code).await; }); #[cfg(unix)] @@ -1242,6 +1242,14 @@ async fn spawn_pipe_process(request: &ExecCommandRequest) -> TerminalResult = None; + // `request_control` takes the terminator, so the control sender is + // dropped as soon as the first interrupt or kill is queued. Once that + // happens `recv()` resolves to `None` instantly and forever, so the + // branch has to be disarmed: leaving it armed starves the reader-done + // branch, the loop never reaches its exit condition, and the session is + // left open with no exit code (while re-signalling a dead process group + // in a hot loop). + let mut control_closed = false; loop { if let Some(state) = control_state { @@ -1269,11 +1277,13 @@ async fn spawn_pipe_process(request: &ExecCommandRequest) -> TerminalResult { - control_state = request_unix_pipe_control( - pipe_pgid, - action.unwrap_or(ExecControlAction::Kill), - ); + action = control_rx.recv(), if !control_closed => { + match action { + Some(action) => { + control_state = request_unix_pipe_control(pipe_pgid, action); + } + None => control_closed = true, + } } done = reader_done_rx.recv(), if remaining_readers > 0 => { @@ -1332,6 +1342,19 @@ async fn spawn_pipe_process(request: &ExecCommandRequest) -> TerminalResult Option { + use std::os::unix::process::ExitStatusExt; + + status.code().or_else(|| status.signal().map(|s| 128 + s)) +} + #[cfg(unix)] fn configure_pipe_process_group(command: &mut Command) { unsafe { @@ -1749,6 +1772,32 @@ mod tests { assert!(response.output.contains("bitfun_exec_test")); } + #[cfg(unix)] + #[tokio::test] + async fn pipe_exec_reports_signal_death_as_a_conventional_status() { + let manager = ExecProcessManager::default(); + let response = manager + .exec_command(ExecCommandRequest { + argv: shell_argv("kill -TERM $$"), + cwd: std::env::current_dir().expect("current dir"), + env: HashMap::new(), + tty: false, + yield_time_ms: Some(5_000), + max_output_chars: Some(10_000), + lifecycle_tx: None, + output_capture_tx: None, + }) + .await + .expect("exec command should run"); + + assert!(response.session_id.is_none()); + assert_eq!( + response.exit_code, + Some(143), + "a signal death is a known status, not an unknown one" + ); + } + #[tokio::test] async fn delayed_poll_returns_unread_output_after_process_exit() { let manager = ExecProcessManager::default(); diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.test.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.test.ts index 1d8d481649..5082f7cc2a 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.test.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.test.ts @@ -7,15 +7,23 @@ import { saveDialogTurnToDisk, } from './PersistenceModule'; -const saveSessionTurn = vi.fn(); -const saveSessionMetadata = vi.fn(); -const loadSessionMetadata = vi.fn(); +// Vitest hoists `vi.mock` factories above ordinary module-scope declarations, +// so a plain `const` referenced inside the factory is still in its temporal +// dead zone when the factory runs. `vi.hoisted` hoists the value itself to +// the same point, ahead of `vi.mock`, so the factory can see it. +const { mockSaveSessionTurn, mockSaveSessionMetadata, mockLoadSessionMetadata } = vi.hoisted( + () => ({ + mockSaveSessionTurn: vi.fn(), + mockSaveSessionMetadata: vi.fn(), + mockLoadSessionMetadata: vi.fn(), + }) +); vi.mock('@/infrastructure/api/service-api/SessionAPI', () => ({ sessionAPI: { - saveSessionTurn, - saveSessionMetadata, - loadSessionMetadata, + saveSessionTurn: mockSaveSessionTurn, + saveSessionMetadata: mockSaveSessionMetadata, + loadSessionMetadata: mockLoadSessionMetadata, }, })); @@ -84,9 +92,9 @@ async function flushMicrotasks(): Promise { describe('PersistenceModule', () => { beforeEach(() => { vi.useFakeTimers(); - saveSessionTurn.mockResolvedValue(undefined); - saveSessionMetadata.mockResolvedValue(undefined); - loadSessionMetadata.mockResolvedValue(null); + mockSaveSessionTurn.mockResolvedValue(undefined); + mockSaveSessionMetadata.mockResolvedValue(undefined); + mockLoadSessionMetadata.mockResolvedValue(null); }); afterEach(() => { @@ -273,14 +281,14 @@ describe('PersistenceModule', () => { immediateSaveDialogTurn(context, SESSION_ID, TURN_ID); await flushMicrotasks(); - expect(saveSessionTurn).not.toHaveBeenCalled(); + expect(mockSaveSessionTurn).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(499); - expect(saveSessionTurn).not.toHaveBeenCalled(); + expect(mockSaveSessionTurn).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(1); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(1); }); it('checkpoints continuous streamed output without waiting for a quiet period', async () => { @@ -291,16 +299,16 @@ describe('PersistenceModule', () => { await vi.advanceTimersByTimeAsync(1000); debouncedSaveDialogTurn(context, SESSION_ID, TURN_ID, 2000); await vi.advanceTimersByTimeAsync(999); - expect(saveSessionTurn).not.toHaveBeenCalled(); + expect(mockSaveSessionTurn).not.toHaveBeenCalled(); await vi.advanceTimersByTimeAsync(1); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(1); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(1); debouncedSaveDialogTurn(context, SESSION_ID, TURN_ID, 2000); await vi.advanceTimersByTimeAsync(2000); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(2); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(2); }); it('flushes terminal turn saves immediately', async () => { @@ -311,7 +319,7 @@ describe('PersistenceModule', () => { await vi.advanceTimersByTimeAsync(0); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(1); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(1); expect(context.saveDebouncers.size).toBe(0); }); @@ -325,11 +333,11 @@ describe('PersistenceModule', () => { await saveDialogTurnToDisk(context, SESSION_ID, TURN_ID); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(1); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(1); expect(context.saveDebouncers.size).toBe(0); await vi.advanceTimersByTimeAsync(500); await flushMicrotasks(); - expect(saveSessionTurn).toHaveBeenCalledTimes(1); + expect(mockSaveSessionTurn).toHaveBeenCalledTimes(1); }); });