fix: exercise production CLI image provider

Cargo.lock

@@ -6423,6 +6423,7 @@ dependencies = [
  "workdir",
  "worker",
  "worker-runtime",
+ "yoi",
  "zeroize",
 ]
 

crates/workspace-server/Cargo.toml

@@ -53,6 +53,7 @@ webauthn-rs = { workspace = true }
 agen.workspace = true
 tempfile.workspace = true
 tools.workspace = true
+yoi = { path = "../yoi" }
 tower = { workspace = true, features = ["util"] }
 tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
 webauthn-authenticator-rs = { version = "0.5.5", default-features = false, features = ["softpasskey"] }

crates/workspace-server/src/server.rs

@@ -28044,7 +28044,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn external_provider_websocket_routes_files_and_commands_then_reconnects_and_revokes() {
+    async fn external_provider_websocket_runs_production_cli_pump_for_images_and_revoke() {
         let dir = tempfile::tempdir().unwrap();
         init_clean_git_workspace(dir.path());
         let api = test_api(dir.path()).await;
@@ -28184,20 +28184,49 @@ mod tests {
             tokio::time::sleep(std::time::Duration::from_millis(10)).await;
         }
 
-        let (mut second, _) = connect_async(&url).await.unwrap();
+        let provider_root = tempfile::tempdir().unwrap();
-        second
+        fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap();
-            .send(Message::Text(
+        let image_bytes = valid_gif_with_len(1024 * 1024 + 257);
-                serde_json::to_string(&register(2)).unwrap().into(),
+        fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap();
-            ))
+        let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap();
+        let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned(
+            workdir::Workdir::new(&grant.working_directory_id),
+            pinned_provider_root,
+            workdir::BoundedReadLimits::EXTERNAL_DEFAULT,
+            workdir::WorkdirSessionCapabilities::ALL,
+        )
+        .unwrap();
+        let (second, _) = connect_async(&url).await.unwrap();
+        let registration = register(2);
+        let provider = tokio::spawn(async move {
+            yoi::workdir_share::serve_external_workdir_provider_socket(
+                second,
+                registration,
+                &local_session,
+                2,
+                std::future::pending::<std::result::Result<(), String>>(),
+            )
             .await
-            .unwrap();
+        });
-        let registered = second.next().await.unwrap().unwrap().into_text().unwrap();
+        for _ in 0..50 {
-        assert!(matches!(
+            if api
-            serde_json::from_str::<ExternalWorkdirServerFrame>(&registered)
+                .store
+                .get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id)
                 .unwrap()
-                .message,
+                .is_some_and(|stored| stored.status == "online")
-            ExternalWorkdirServerMessage::Registered { generation: 2, .. }
+            {
-        ));
+                break;
+            }
+            tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+        }
+        assert_eq!(
+            api.store
+                .get_external_workdir_grant(TEST_WORKSPACE_ID, &grant.grant_id)
+                .unwrap()
+                .unwrap()
+                .status,
+            "online"
+        );
 
         api.store
             .upsert_workdir_registry(&WorkdirRegistryRecord {
@@ -28265,18 +28294,6 @@ mod tests {
             external_capabilities,
             workdir::WorkdirSessionCapabilities::ALL
         );
-        let provider_root = tempfile::tempdir().unwrap();
-        fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap();
-        let image_bytes = valid_gif_with_len(1024 * 1024 + 257);
-        fs::write(provider_root.path().join("image.gif"), &image_bytes).unwrap();
-        let pinned_provider_root = workdir::ExternalWorkdirRoot::pin(provider_root.path()).unwrap();
-        let local_session = workdir::LocalWorkdirSession::external_with_capabilities_pinned(
-            workdir::Workdir::new(&grant.working_directory_id),
-            pinned_provider_root,
-            workdir::BoundedReadLimits::EXTERNAL_DEFAULT,
-            workdir::WorkdirSessionCapabilities::ALL,
-        )
-        .unwrap();
         let session_lock = current_worker_session_lock(&api, &worker);
         let session_guard = session_lock.lock().await;
         let broker_session =
@@ -28288,272 +28305,60 @@ mod tests {
             workdir::WorkdirSessionCapabilities::ALL
         );
         drop(session_guard);
-        let read = {
+        let read_result = broker_session
-            let broker_session = broker_session.clone();
+            .read(workdir::ReadRequest {
-            tokio::spawn(async move {
+                path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
-                broker_session
+                offset: 0,
-                    .read(workdir::ReadRequest {
+                limit: 10,
-                        path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
+                max_bytes: 1024,
-                        offset: 0,
-                        limit: 10,
-                        max_bytes: 1024,
-                    })
-                    .await
             })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
             .unwrap();
-        let read_result = read.await.unwrap().unwrap();
         assert_eq!(read_result.bytes, b"one\ntwo\n");
 
-        let view_image = {
+        let (_meta, view_image) = tools::view_image_tool(broker_session.clone())();
-            let broker_session = broker_session.clone();
+        let output = view_image
-            tokio::spawn(async move {
+            .execute(
-                let (_meta, tool) = tools::view_image_tool(broker_session)();
+                r#"{"path":"image.gif"}"#,
-                tool.execute(
+                agen::tool::ToolExecutionContext::default(),
-                    r#"{"path":"image.gif"}"#,
-                    agen::tool::ToolExecutionContext::default(),
-                )
-                .await
-            })
-        };
-        for expected_offset in [0_u64, 1024 * 1024] {
-            let (operation_id, operation) = loop {
-                let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-                let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-                if let ExternalWorkdirServerMessage::Operation {
-                    generation: 2,
-                    operation_id,
-                    operation,
-                } = frame.message
-                {
-                    break (operation_id, operation);
-                }
-            };
-            let WorkdirSessionOperation::ReadBytes(request) = operation.as_inner() else {
-                panic!("ViewImage must use the bounded binary operation");
-            };
-            assert_eq!(request.offset, expected_offset);
-            assert!(request.max_bytes <= 1024 * 1024);
-            let result = workdir::http::dispatch_workdir_session_operation(
-                &local_session,
-                operation.into_inner(),
             )
             .await
-            .unwrap()
-            .try_into()
             .unwrap();
-            second
-                .send(Message::Text(
-                    serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                        ExternalWorkdirProviderMessage::OperationResult {
-                            generation: 2,
-                            operation_id,
-                            outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                        },
-                    ))
-                    .unwrap()
-                    .into(),
-                ))
-                .await
-                .unwrap();
-        }
-        let output = view_image.await.unwrap().unwrap();
         let agen::tool::Attachment::Image(image) = &output.attachments[0];
         assert_eq!(image.mime_type(), "image/gif");
         assert_eq!(image.data(), image_bytes);
 
-        let missing_image = {
+        let error = view_image
-            let broker_session = broker_session.clone();
+            .execute(
-            tokio::spawn(async move {
+                r#"{"path":"missing.gif"}"#,
-                let (_meta, tool) = tools::view_image_tool(broker_session)();
+                agen::tool::ToolExecutionContext::default(),
-                tool.execute(
+            )
-                    r#"{"path":"missing.gif"}"#,
-                    agen::tool::ToolExecutionContext::default(),
-                )
-                .await
-            })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::ReadBytes(_)
-        ));
-        let workdir_error = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap_err();
-        let error = workdir::external::ExternalWorkdirOperationError::from_transport_error(
-            workdir::http::WorkdirTransportError::from_workdir_error(&workdir_error),
-        );
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Failed { error },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
-            .unwrap();
+            .unwrap_err()
-        let error = missing_image.await.unwrap().unwrap_err().to_string();
+            .to_string();
         assert!(error.contains("not found"), "{error}");
 
-        let edit = {
+        let edit = broker_session
-            let broker_session = broker_session.clone();
+            .edit(workdir::EditRequest {
-            let expected_hash = read_result.content_hash;
+                path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
-            tokio::spawn(async move {
+                old_string: "two".to_string(),
-                broker_session
+                new_string: "changed".to_string(),
-                    .edit(workdir::EditRequest {
+                replace_all: false,
-                        path: workdir::WorkdirPath::new("events.jsonl").unwrap(),
+                expected_hash: read_result.content_hash,
-                        old_string: "two".to_string(),
-                        new_string: "changed".to_string(),
-                        replace_all: false,
-                        expected_hash,
-                    })
-                    .await
             })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::Edit(_)
-        ));
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
             .unwrap();
-        assert_eq!(edit.await.unwrap().unwrap().replacements, 1);
+        assert_eq!(edit.replacements, 1);
 
-        let write = {
+        let write = broker_session
-            let broker_session = broker_session.clone();
+            .write(workdir::WriteRequest {
-            tokio::spawn(async move {
+                path: workdir::WorkdirPath::new("created.txt").unwrap(),
-                broker_session
+                content: b"created".to_vec(),
-                    .write(workdir::WriteRequest {
+                expected_hash: None,
-                        path: workdir::WorkdirPath::new("created.txt").unwrap(),
-                        content: b"created".to_vec(),
-                        expected_hash: None,
-                    })
-                    .await
             })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::Write(_)
-        ));
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
             .unwrap();
-        assert!(write.await.unwrap().unwrap().created);
+        assert!(write.created);
         assert_eq!(
             fs::read_to_string(provider_root.path().join("events.jsonl")).unwrap(),
             "one\nchanged\n"
@@ -28562,64 +28367,17 @@ mod tests {
             fs::read_to_string(provider_root.path().join("created.txt")).unwrap(),
             "created"
         );
-        let command_start = {
+        let command_handle = broker_session
-            let broker_session = broker_session.clone();
+            .start_command(workdir::CommandRequest {
-            tokio::spawn(async move {
+                command: "sleep 30 & child=$!; echo $child > child.pid; wait $child".to_string(),
-                broker_session
+                timeout_secs: 30,
-                    .start_command(workdir::CommandRequest {
+                output_limit: 1024,
-                        command: "sleep 30 & child=$!; echo $child > child.pid; wait $child"
+                cwd: workdir::WorkdirPath::root(),
-                            .to_string(),
+                spill_dir: Some("/worker/private/spill".into()),
-                        timeout_secs: 30,
+                tool_call_id: Some("call-e2e".to_string()),
-                        output_limit: 1024,
-                        cwd: workdir::WorkdirPath::root(),
-                        spill_dir: Some("/worker/private/spill".into()),
-                        tool_call_id: Some("call-e2e".to_string()),
-                    })
-                    .await
             })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::CommandStart(workdir::CommandRequest {
-                spill_dir: None,
-                ..
-            })
-        ));
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
             .unwrap();
-        let command_handle = command_start.await.unwrap().unwrap();
 
         for _ in 0..100 {
             if provider_root.path().join("child.pid").is_file() {
@@ -28629,162 +28387,25 @@ mod tests {
         }
         assert!(provider_root.path().join("child.pid").is_file());
 
-        let command_output = {
+        let command_output = broker_session
-            let broker_session = broker_session.clone();
+            .command_output(workdir::CommandOutputRequest {
-            let handle = command_handle.clone();
+                handle: command_handle,
-            tokio::spawn(async move {
+                cursor: 0,
-                broker_session
+                limit: 1024,
-                    .command_output(workdir::CommandOutputRequest {
-                        handle,
-                        cursor: 0,
-                        limit: 1024,
-                        wait: false,
-                    })
-                    .await
-            })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest {
                 wait: false,
-                ..
             })
-        ));
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
             .unwrap();
-        let command_output = command_output.await.unwrap().unwrap();
         assert_eq!(command_output.status, workdir::CommandStatus::Running);
         assert_eq!(command_output.output_path, None);
 
-        let revoker = {
+        let Json(revoked) = scoped_revoke_external_workdir_grant(
-            let api = api.clone();
+            State(api.clone()),
-            let grant_id = grant.grant_id.clone();
+            Extension(test_browser_request_actor()),
-            tokio::spawn(async move {
+            AxumPath((TEST_WORKSPACE_ID.to_string(), grant.grant_id.clone())),
-                scoped_revoke_external_workdir_grant(
-                    State(api),
-                    Extension(test_browser_request_actor()),
-                    AxumPath((TEST_WORKSPACE_ID.to_string(), grant_id)),
-                )
-                .await
-            })
-        };
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::CommandCancel(_)
-        ));
-        let result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
-        )
-        .await
-        .unwrap()
-        .try_into()
-        .unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
-            .await
-            .unwrap();
-
-        let (operation_id, operation) = loop {
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if let ExternalWorkdirServerMessage::Operation {
-                generation: 2,
-                operation_id,
-                operation,
-            } = frame.message
-            {
-                break (operation_id, operation);
-            }
-        };
-        assert!(matches!(
-            operation.as_inner(),
-            WorkdirSessionOperation::CommandOutput(workdir::CommandOutputRequest {
-                wait: true,
-                ..
-            })
-        ));
-        let provider_result = workdir::http::dispatch_workdir_session_operation(
-            &local_session,
-            operation.into_inner(),
         )
         .await
         .unwrap();
-        assert!(matches!(
-            provider_result,
-            WorkdirSessionOperationResult::CommandOutput(ref output)
-                if output.status != workdir::CommandStatus::Running
-        ));
-        let result = provider_result.try_into().unwrap();
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::OperationResult {
-                        generation: 2,
-                        operation_id,
-                        outcome: ExternalWorkdirOperationOutcome::Completed { result },
-                    },
-                ))
-                .unwrap()
-                .into(),
-            ))
-            .await
-            .unwrap();
-        let Json(revoked) = revoker.await.unwrap().unwrap();
         assert_eq!(revoked.status, "revoked");
         let active_links = api
             .store
@@ -28796,27 +28417,12 @@ mod tests {
         assert_eq!(active_links.len(), 1);
         assert_eq!(active_links[0].alias, "repository");
         assert_eq!(active_links[0].workdir_id, "analysis-repository");
-        let revoke = loop {
+        let provider_end = tokio::time::timeout(std::time::Duration::from_secs(2), provider)
-            let text = second.next().await.unwrap().unwrap().into_text().unwrap();
-            let frame = serde_json::from_str::<ExternalWorkdirServerFrame>(&text).unwrap();
-            if matches!(frame.message, ExternalWorkdirServerMessage::Revoke { .. }) {
-                break frame;
-            }
-        };
-        assert!(matches!(
-            revoke.message,
-            ExternalWorkdirServerMessage::Revoke { generation: 2, .. }
-        ));
-        second
-            .send(Message::Text(
-                serde_json::to_string(&ExternalWorkdirProviderFrame::current(
-                    ExternalWorkdirProviderMessage::RevokeAcknowledged { generation: 2 },
-                ))
-                .unwrap()
-                .into(),
-            ))
             .await
+            .expect("production CLI provider pump did not acknowledge revocation")
+            .unwrap()
             .unwrap();
+        assert_eq!(provider_end, yoi::workdir_share::ProviderEnd::Revoked);
         server.abort();
     }
 

crates/yoi/src/lib.rs

@@ -0,0 +1 @@
+pub mod workdir_share;

crates/yoi/src/main.rs

@@ -4,10 +4,11 @@ mod objective_cli;
 mod plugin_cli;
 mod session_cli;
 mod ticket_cli;
-mod workdir_share;
 mod worker_cleanup_cli;
 mod workspace_bootstrap;
 
+use yoi::workdir_share;
+
 use std::collections::BTreeMap;
 use std::fmt;
 use std::fs;

crates/yoi/src/workdir_share.rs

@@ -25,7 +25,7 @@ use workdir::{
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub(crate) struct WorkdirShareOptions {
+pub struct WorkdirShareOptions {
     pub path: PathBuf,
     pub workspace_id: Option<String>,
     pub backend_url: String,
@@ -82,7 +82,7 @@ fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum ProviderEnd {
+pub enum ProviderEnd {
     Disconnected,
     Revoked,
     Interrupted,
@@ -221,7 +221,7 @@ struct ResolvedWorkdirShareOptions {
     permissions: server_api::ExternalWorkdirPermissions,
 }
 
-pub(crate) fn parse_ttl(value: &str) -> Result<Duration, String> {
+pub fn parse_ttl(value: &str) -> Result<Duration, String> {
     let (number, multiplier) = if let Some(value) = value.strip_suffix('h') {
         (value, 60 * 60)
     } else if let Some(value) = value.strip_suffix('m') {
@@ -457,7 +457,7 @@ async fn resolve_share_options(
     configure_interactively(options, &workspaces, &mut input, &mut output)
 }
 
-pub(crate) async fn run(options: WorkdirShareOptions) -> Result<(), String> {
+pub async fn run(options: WorkdirShareOptions) -> Result<(), String> {
     let Some(options) = resolve_share_options(options).await? else {
         eprintln!("External Workdir share cancelled; no grant was created.");
         return Ok(());
@@ -669,7 +669,7 @@ async fn serve_provider_connection(
         .max_message_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES))
         .max_frame_size(Some(workdir::external::MAX_EXTERNAL_SERVER_FRAME_BYTES));
     let connection = connect_async_with_config(request, Some(websocket_config), false);
-    let (mut socket, _) = tokio::select! {
+    let (socket, _) = tokio::select! {
         signal = tokio::signal::ctrl_c() => {
             signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?;
             revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
@@ -691,11 +691,40 @@ async fn serve_provider_connection(
                 read_limits: BoundedReadLimits::EXTERNAL_DEFAULT,
             },
         });
+    let end =
+        serve_external_workdir_provider_socket(socket, registration, session, generation, async {
+            tokio::signal::ctrl_c()
+                .await
+                .map_err(|error| format!("failed to wait for Ctrl-C: {error}"))
+        })
+        .await?;
+    if end == ProviderEnd::Interrupted {
+        revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
+    }
+    Ok(end)
+}
+
+/// Run the production External Workdir provider protocol over an established WebSocket.
+///
+/// The CLI uses this pump after opening its authenticated Backend connection. Keeping the
+/// protocol pump separate lets composed tests exercise the same registration, operation,
+/// cancellation, and result-forwarding implementation without spawning a second CLI process.
+pub async fn serve_external_workdir_provider_socket<S, I>(
+    mut socket: tokio_tungstenite::WebSocketStream<S>,
+    registration: ExternalWorkdirProviderFrame,
+    session: &LocalWorkdirSession,
+    generation: u64,
+    interrupt: I,
+) -> Result<ProviderEnd, String>
+where
+    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
+    I: std::future::Future<Output = Result<(), String>>,
+{
+    tokio::pin!(interrupt);
     send_provider_frame(&mut socket, &registration).await?;
     let registered = tokio::select! {
-        signal = tokio::signal::ctrl_c() => {
+        signal = &mut interrupt => {
-            signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?;
+            signal?;
-            revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
             return Ok(ProviderEnd::Interrupted);
         }
         registered = tokio::time::timeout(Duration::from_secs(10), socket.next()) => {
@@ -713,10 +742,9 @@ async fn serve_provider_connection(
     let mut operations = HashMap::<String, OperationTask>::new();
     loop {
         tokio::select! {
-            signal = tokio::signal::ctrl_c() => {
+            signal = &mut interrupt => {
-                signal.map_err(|error| format!("failed to wait for Ctrl-C: {error}"))?;
+                signal?;
                 stop_provider_operations(session, &mut operations, &mut completions).await?;
-                revoke(server_api_client, &options.workspace_id, &grant.grant_id).await?;
                 return Ok(ProviderEnd::Interrupted);
             }
             completion = completions.recv() => {