From 71eb830086d918ddf6d76f63d882eb15e186807a Mon Sep 17 00:00:00 2001 From: Seongjae Date: Thu, 17 Sep 2026 20:45:16 +0900 Subject: [PATCH] Treat partial UDS writes as TransportAmbiguous. write_frame I/O after poll_write may have put bytes on the socket, so exec must keep the workspace lease. Encode and !alive stay BeforeDispatch. --runner is a ValueEnum so unknown values are rejected instead of silently using in-process. Co-authored-by: Cursor --- crates/runner/src/uds.rs | 156 ++++++++++++++++++++++++++++++++++-- crates/server/src/config.rs | 35 +++++++- crates/server/src/main.rs | 47 ++++++----- docs/ko/operations.md | 1 + docs/operations.md | 3 +- docs/translations.json | 4 +- 6 files changed, 213 insertions(+), 33 deletions(-) diff --git a/crates/runner/src/uds.rs b/crates/runner/src/uds.rs index c41295a..2e209ed 100644 --- a/crates/runner/src/uds.rs +++ b/crates/runner/src/uds.rs @@ -5,14 +5,17 @@ //! per worker. Replay is connection-local. use std::collections::HashMap; +use std::io; use std::path::Path; +use std::pin::Pin; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; +use std::task::{Context, Poll}; use std::time::Duration; use codespace_domain::{ErrorBody, ErrorCode, FindResult, ProcessId, ReadResult}; use codespace_policy::Workspace; -use tokio::io::AsyncWriteExt; +use tokio::io::{AsyncWrite, AsyncWriteExt}; use tokio::net::UnixStream; use tokio::sync::{oneshot, Mutex}; use tokio::time::timeout; @@ -147,12 +150,16 @@ impl UdsRunner { .insert(request_id.clone(), tx); } let envelope = WireEnvelope::request(request_id.clone(), op); - { + let write_result = { let mut writer = self.shared.writer.lock().await; - if let Err(err) = write_frame(&mut *writer, &envelope).await { - self.shared.pending.lock().await.remove(&request_id); - return Err(RunnerError::before_dispatch(err.to_string())); + write_envelope(&mut *writer, &envelope).await + }; + if let Err(err) = write_result { + self.shared.pending.lock().await.remove(&request_id); + if matches!(err, RunnerError::TransportAmbiguous { .. }) { + self.close().await; } + return Err(err); } match timeout(RUNNER_CALL_DEADLINE, rx).await { Ok(Ok(Ok(response))) => decode_response(response), @@ -215,6 +222,56 @@ async fn fail_pending(shared: &Shared, err: RunnerError) { } } +/// Records whether `write_frame` reached socket I/O. Encode errors happen +/// before `poll_write`; a later `write_all` / `flush` failure is ambiguous. +struct WriteProbe { + inner: W, + started: bool, +} + +impl AsyncWrite for WriteProbe { + fn poll_write( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + self.started = true; + Pin::new(&mut self.inner).poll_write(cx, buf) + } + + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.inner).poll_flush(cx) + } + + fn poll_shutdown( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + Pin::new(&mut self.inner).poll_shutdown(cx) + } +} + +fn write_failure(err: io::Error, started: bool) -> RunnerError { + if started { + RunnerError::ambiguous(err.to_string()) + } else { + RunnerError::before_dispatch(err.to_string()) + } +} + +async fn write_envelope( + writer: &mut W, + envelope: &WireEnvelope, +) -> Result<(), RunnerError> { + let mut probe = WriteProbe { + inner: writer, + started: false, + }; + write_frame(&mut probe, envelope) + .await + .map_err(|err| write_failure(err, probe.started)) +} + fn decode_response(response: WireEnvelope) -> Result { if response.ok == Some(true) { response.result.ok_or_else(|| { @@ -362,3 +419,92 @@ impl Runner for UdsRunner { } } } + +#[cfg(test)] +mod tests { + use super::*; + + struct FailOnWrite; + + impl AsyncWrite for FailOnWrite { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + _buf: &[u8], + ) -> Poll> { + Poll::Ready(Err(io::Error::other("write failed"))) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + + struct FailOnFlush; + + impl AsyncWrite for FailOnFlush { + fn poll_write( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + buf: &[u8], + ) -> Poll> { + Poll::Ready(Ok(buf.len())) + } + + fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Err(io::Error::other("flush failed"))) + } + + fn poll_shutdown( + self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + Poll::Ready(Ok(())) + } + } + + fn hello_envelope() -> WireEnvelope { + WireEnvelope::request("rrpc-1".into(), RunnerOp::Hello) + } + + #[test] + fn encode_or_unstarted_write_is_before_dispatch() { + let err = write_failure(io::Error::other("encode failed"), false); + assert!( + matches!(err, RunnerError::TransportBeforeDispatch { .. }), + "{err:?}" + ); + } + + #[tokio::test] + async fn first_poll_write_error_is_ambiguous() { + let mut writer = FailOnWrite; + let err = write_envelope(&mut writer, &hello_envelope()) + .await + .expect_err("write should fail"); + // `exec_command` keeps the workspace lease only on this variant. + assert!( + matches!(err, RunnerError::TransportAmbiguous { .. }), + "{err:?}" + ); + } + + #[tokio::test] + async fn flush_error_after_write_is_ambiguous() { + let mut writer = FailOnFlush; + let err = write_envelope(&mut writer, &hello_envelope()) + .await + .expect_err("flush should fail"); + assert!( + matches!(err, RunnerError::TransportAmbiguous { .. }), + "{err:?}" + ); + } +} diff --git a/crates/server/src/config.rs b/crates/server/src/config.rs index 8389993..0b03035 100644 --- a/crates/server/src/config.rs +++ b/crates/server/src/config.rs @@ -9,6 +9,13 @@ pub enum TransportMode { Http, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)] +pub enum RunnerMode { + #[value(name = "in-process")] + InProcess, + Uds, +} + #[derive(Debug, Clone, Parser)] #[command( name = "codespace-mcp", @@ -42,9 +49,9 @@ pub struct Cli { #[arg(long, env = "CODESPACE_OPERATIONS_DB")] pub operations_db: Option, - /// Runner backend. `in-process` is the default host supervisor. `uds` uses UdsRunner. - #[arg(long, env = "CODESPACE_RUNNER", default_value = "in-process")] - pub runner: String, + /// Runner backend. Allowed values: `in-process` (default host supervisor) or `uds`. + #[arg(long, env = "CODESPACE_RUNNER", value_enum, default_value_t = RunnerMode::InProcess)] + pub runner: RunnerMode, /// Unix socket for `CODESPACE_RUNNER=uds` when connecting to an /// already-running worker. Spawn path uses `--runner-dir` instead @@ -91,3 +98,25 @@ pub struct HttpConfig { pub port: u16, pub bearer_token: Option, } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runner_mode_rejects_unknown_value() { + let err = Cli::try_parse_from(["codespace-mcp", "--runner", "udss"]) + .expect_err("unknown runner value"); + let message = err.to_string(); + assert!(message.contains("udss"), "{message}"); + } + + #[test] + fn runner_mode_accepts_uds_and_in_process() { + let uds = Cli::try_parse_from(["codespace-mcp", "--runner", "uds"]).expect("uds"); + assert_eq!(uds.runner, RunnerMode::Uds); + let host = + Cli::try_parse_from(["codespace-mcp", "--runner", "in-process"]).expect("in-process"); + assert_eq!(host.runner, RunnerMode::InProcess); + } +} diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 4311ff1..ea4b384 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -4,7 +4,7 @@ use anyhow::Result; use clap::Parser; use codespace_policy::Registry; use codespace_runner::{RuntimeBackend, ShellRelease}; -use codespace_server::config::{Cli, TransportMode}; +use codespace_server::config::{Cli, RunnerMode, TransportMode}; use codespace_server::http::serve_http; use codespace_server::logging; use codespace_server::runtime::RuntimeProcess; @@ -44,30 +44,33 @@ async fn start_runner(cli: &Cli, store: Arc) -> Result { store_for_lease.release_process(process_id); }) }; - if cli.runner != "uds" { - return Ok(StartedRunner { + match cli.runner { + RunnerMode::InProcess => Ok(StartedRunner { runner: RuntimeBackend::in_process(on_release), _process: None, - }); - } - if let Some(bin) = &cli.runtime_bin { - let (process, runner) = - RuntimeProcess::spawn(bin, cli.runner_dir.as_deref(), on_release, store).await?; - return Ok(StartedRunner { - runner: RuntimeBackend::Uds(runner), - _process: Some(process), - }); + }), + RunnerMode::Uds => { + if let Some(bin) = &cli.runtime_bin { + let (process, runner) = + RuntimeProcess::spawn(bin, cli.runner_dir.as_deref(), on_release, store) + .await?; + return Ok(StartedRunner { + runner: RuntimeBackend::Uds(runner), + _process: Some(process), + }); + } + let socket = cli.runner_socket.as_ref().ok_or_else(|| { + anyhow::anyhow!( + "uds runner requires --runtime-bin (spawn) or --runner-socket (connect to an existing worker)" + ) + })?; + let runner = RuntimeProcess::connect_existing(socket, on_release, store).await?; + Ok(StartedRunner { + runner: RuntimeBackend::Uds(runner), + _process: None, + }) + } } - let socket = cli.runner_socket.as_ref().ok_or_else(|| { - anyhow::anyhow!( - "uds runner requires --runtime-bin (spawn) or --runner-socket (connect to an existing worker)" - ) - })?; - let runner = RuntimeProcess::connect_existing(socket, on_release, store).await?; - Ok(StartedRunner { - runner: RuntimeBackend::Uds(runner), - _process: None, - }) } fn load_registry(cli: &Cli) -> Result { diff --git a/docs/ko/operations.md b/docs/ko/operations.md index bedb584..5e0f8fc 100644 --- a/docs/ko/operations.md +++ b/docs/ko/operations.md @@ -111,6 +111,7 @@ export CODESPACE_OPERATIONS_DB="$PWD/data/operations.sqlite" 선택적 러너 워커(여전히 호스트 exec이며 Linux 격리가 아님). UDS는 1:1입니다. 게이트웨이가 `RuntimeProcess`(자식, 비공개 0700 디렉터리, `$dir/runner.sock`)를 소유합니다. 재연결은 없습니다. +`--runner` / `CODESPACE_RUNNER`의 허용 값은 `in-process`와 `uds`뿐입니다. `--runner-dir` / `CODESPACE_RUNNER_DIR`은 그 unique leaf의 부모가 될 수 있습니다. `/`, `/tmp`, `/var/tmp`, `$HOME`을 디렉터리 자체로 주면 거절합니다. `--runner-socket`은 이미 떠 있는 워커에 연결할 때만 쓰며 diff --git a/docs/operations.md b/docs/operations.md index 26206db..d90a5e7 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -109,7 +109,8 @@ survive restart. Opt-in runner worker (still host exec, not Linux isolation). UDS is 1:1: the gateway owns `RuntimeProcess` (child, private 0700 directory, -`$dir/runner.sock`). There is no reconnect. `--runner-dir` / +`$dir/runner.sock`). There is no reconnect. Allowed `--runner` / +`CODESPACE_RUNNER` values are `in-process` and `uds` only. `--runner-dir` / `CODESPACE_RUNNER_DIR` may name a parent for that unique leaf; `/`, `/tmp`, `/var/tmp`, and `$HOME` are rejected as the directory itself. `--runner-socket` is only for connecting to an already-running worker diff --git a/docs/translations.json b/docs/translations.json index 85f460b..487b510 100644 --- a/docs/translations.json +++ b/docs/translations.json @@ -91,8 +91,8 @@ "워크스페이스-레지스트리", "이-문서가-검증하지-않는-것" ], - "source_sha256": "ea0ce5511b84242e0988ce4eda7db0e91acd110c8d0e3b99d4e23e2ef58fae94", - "translation_sha256": "e83905f5420b58679530037f7aa56574af242e7b9cefced658a127bcf992c5e0" + "source_sha256": "1b5fe5472ed386cf78fcaf86569cd4f8d1b341eff80c1018fdbbbf70113b9625", + "translation_sha256": "b9458a5fdef87b8346f16c97a4436a9885a9337622cd0b16b31d35e8fa17b6a9" }, { "id": "chatgpt-connector",