Treat partial UDS writes as TransportAmbiguous

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<W> {
+    inner: W,
+    started: bool,
+}
+
+impl<W: AsyncWrite + Unpin> AsyncWrite for WriteProbe<W> {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<Result<usize, io::Error>> {
+        self.started = true;
+        Pin::new(&mut self.inner).poll_write(cx, buf)
+    }
+
+    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        Pin::new(&mut self.inner).poll_flush(cx)
+    }
+
+    fn poll_shutdown(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<(), io::Error>> {
+        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<W: AsyncWrite + Unpin>(
+    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<RunnerOpResult, RunnerError> {
     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<Result<usize, io::Error>> {
+            Poll::Ready(Err(io::Error::other("write failed")))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+            Poll::Ready(Ok(()))
+        }
+
+        fn poll_shutdown(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), io::Error>> {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    struct FailOnFlush;
+
+    impl AsyncWrite for FailOnFlush {
+        fn poll_write(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            buf: &[u8],
+        ) -> Poll<Result<usize, io::Error>> {
+            Poll::Ready(Ok(buf.len()))
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+            Poll::Ready(Err(io::Error::other("flush failed")))
+        }
+
+        fn poll_shutdown(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+        ) -> Poll<Result<(), io::Error>> {
+            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:?}"
+        );
+    }
+}

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<std::path::PathBuf>,
 
-    /// 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<String>,
 }
+
+#[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);
+    }
+}

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<Store>) -> Result<StartedRunner> {
             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<Registry> {

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`은 이미 떠 있는 워커에 연결할 때만 쓰며

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

docs/translations.json

@@ -91,8 +91,8 @@
         "워크스페이스-레지스트리",
         "이-문서가-검증하지-않는-것"
       ],
-      "source_sha256": "ea0ce5511b84242e0988ce4eda7db0e91acd110c8d0e3b99d4e23e2ef58fae94",
-      "translation_sha256": "e83905f5420b58679530037f7aa56574af242e7b9cefced658a127bcf992c5e0"
+      "source_sha256": "1b5fe5472ed386cf78fcaf86569cd4f8d1b341eff80c1018fdbbbf70113b9625",
+      "translation_sha256": "b9458a5fdef87b8346f16c97a4436a9885a9337622cd0b16b31d35e8fa17b6a9"
     },
     {
       "id": "chatgpt-connector",