diff --git a/Cargo.lock b/Cargo.lock index 8053c9987..d5f5b2402 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6383,6 +6383,7 @@ dependencies = [ name = "yoi-workspace-server" version = "0.1.0" dependencies = [ + "agen", "async-trait", "axum", "chrono", @@ -6410,6 +6411,7 @@ dependencies = [ "tokio", "tokio-tungstenite 0.29.0", "toml", + "tools", "tower", "tracing", "tracing-subscriber", diff --git a/crates/tools/tests/fixtures/view-image/valid.gif b/crates/tools/tests/fixtures/view-image/valid.gif new file mode 100644 index 000000000..f1eda11f8 Binary files /dev/null and b/crates/tools/tests/fixtures/view-image/valid.gif differ diff --git a/crates/tools/tests/fixtures/view-image/valid.jpg b/crates/tools/tests/fixtures/view-image/valid.jpg new file mode 100644 index 000000000..60191034e Binary files /dev/null and b/crates/tools/tests/fixtures/view-image/valid.jpg differ diff --git a/crates/tools/tests/fixtures/view-image/valid.png b/crates/tools/tests/fixtures/view-image/valid.png new file mode 100644 index 000000000..5d5869660 Binary files /dev/null and b/crates/tools/tests/fixtures/view-image/valid.png differ diff --git a/crates/tools/tests/fixtures/view-image/valid.webp b/crates/tools/tests/fixtures/view-image/valid.webp new file mode 100644 index 000000000..0c9cd5469 Binary files /dev/null and b/crates/tools/tests/fixtures/view-image/valid.webp differ diff --git a/crates/tools/tests/integration.rs b/crates/tools/tests/integration.rs index eb4ca3c7f..6cf261ccb 100644 --- a/crates/tools/tests/integration.rs +++ b/crates/tools/tests/integration.rs @@ -18,7 +18,9 @@ use tools::{ Tracker, core_builtin_tools, routed_builtin_tools, routed_view_image_tool, view_image_tool, }; use workdir::{ - LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirSessionHandle, WorkdirSessionRouter, + LocalWorkdirSession, WorkdirAttachmentAlias, WorkdirPath, WorkdirSessionHandle, + WorkdirSessionRouter, WorkdirToolBroker, WorkdirToolScope, WorkdirToolScopePermission, + WorkdirToolScopeRule, }; fn scope_with_spill(workspace: &Path, spill: &Path) -> Scope { @@ -94,6 +96,33 @@ fn setup_routed() -> ( (left, right, spill, router, Registry::new(definitions)) } +fn valid_gif_with_len(target_len: usize) -> Vec { + let fixture = include_bytes!("fixtures/view-image/valid.gif"); + assert_eq!( + fixture.last(), + Some(&0x3b), + "GIF fixture must end in trailer" + ); + let mut bytes = fixture[..fixture.len() - 1].to_vec(); + bytes.extend_from_slice(&[0x21, 0xfe]); + let mut remaining = target_len + .checked_sub(bytes.len() + 2) + .expect("target must leave room for comment terminator and GIF trailer"); + while remaining > 0 { + let mut block = remaining.min(256); + if remaining - block == 1 { + block -= 1; + } + assert!(block >= 2); + bytes.push((block - 1) as u8); + bytes.extend(std::iter::repeat_n(b'x', block - 1)); + remaining -= block; + } + bytes.extend_from_slice(&[0, 0x3b]); + assert_eq!(bytes.len(), target_len); + bytes +} + async fn call(tool: &Arc, input: serde_json::Value) -> agen::tool::ToolOutput { tool.execute(&input.to_string(), Default::default()) .await @@ -546,7 +575,9 @@ async fn external_view_image_uses_a_protocol_bounded_request() { truncate_first_binary_response: false, mutate_after_first_binary_response: None, }); - let png = b"\x89PNG\r\n\x1a\nexternal\0image\nbody\xff"; + let png = include_bytes!("fixtures/view-image/valid.png"); + assert!(png.contains(&0)); + assert!(png.iter().any(|byte| *byte >= 0x80)); std::fs::write(dir.path().join("image.png"), png).unwrap(); let (_meta, tool) = view_image_tool(session)(); @@ -562,7 +593,7 @@ async fn external_view_image_rejects_truncated_and_cross_version_chunks() { let spill = TempDir::new().unwrap(); let image_path = dir.path().join("image.png"); let scope = scope_with_spill(dir.path(), spill.path()); - std::fs::write(&image_path, b"\x89PNG\r\n\x1a\npartial").unwrap(); + std::fs::write(&image_path, include_bytes!("fixtures/view-image/valid.png")).unwrap(); let truncated: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { inner: LocalWorkdirSession::new(scope.clone(), dir.path().to_path_buf()), truncate_first_binary_response: true, @@ -577,8 +608,7 @@ async fn external_view_image_rejects_truncated_and_cross_version_chunks() { "{error}" ); - let mut large = vec![b'x'; 1024 * 1024 + 32]; - large[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); + let large = valid_gif_with_len(1024 * 1024 + 32); std::fs::write(&image_path, large).unwrap(); let changing: WorkdirSessionHandle = Arc::new(ExternalValidatingSession { inner: LocalWorkdirSession::new(scope, dir.path().to_path_buf()), @@ -601,22 +631,98 @@ async fn view_image_enforces_the_ten_mib_boundary_without_partial_attachments() let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); let (_meta, tool) = view_image_tool(session)(); - let mut png = vec![0xff; LIMIT]; - png[..8].copy_from_slice(b"\x89PNG\r\n\x1a\n"); - std::fs::write(dir.path().join("boundary.png"), &png).unwrap(); + let mut image = valid_gif_with_len(LIMIT); + std::fs::write(dir.path().join("boundary.gif"), &image).unwrap(); - let output = call(&tool, json!({ "path": "boundary.png" })).await; - let agen::tool::Attachment::Image(image) = &output.attachments[0]; - assert_eq!(image.data(), png); + let output = call(&tool, json!({ "path": "boundary.gif" })).await; + let agen::tool::Attachment::Image(attachment) = &output.attachments[0]; + assert_eq!(attachment.mime_type(), "image/gif"); + assert_eq!(attachment.data(), image); - png.push(0); - std::fs::write(dir.path().join("boundary.png"), png).unwrap(); - let error = call_err(&tool, json!({ "path": "boundary.png" })) + image.push(0); + std::fs::write(dir.path().join("boundary.gif"), image).unwrap(); + let error = call_err(&tool, json!({ "path": "boundary.gif" })) .await .to_string(); assert!(error.contains("10485760-byte limit"), "{error}"); } +#[tokio::test] +async fn view_image_attaches_valid_supported_image_fixtures() { + let dir = TempDir::new().unwrap(); + let spill = TempDir::new().unwrap(); + let scope = scope_with_spill(dir.path(), spill.path()); + let session: WorkdirSessionHandle = + Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); + let (_meta, tool) = view_image_tool(session)(); + let fixtures: [(&str, &[u8], &str); 4] = [ + ( + "valid.png", + include_bytes!("fixtures/view-image/valid.png"), + "image/png", + ), + ( + "valid.jpg", + include_bytes!("fixtures/view-image/valid.jpg"), + "image/jpeg", + ), + ( + "valid.gif", + include_bytes!("fixtures/view-image/valid.gif"), + "image/gif", + ), + ( + "valid.webp", + include_bytes!("fixtures/view-image/valid.webp"), + "image/webp", + ), + ]; + + for (name, fixture, mime) in fixtures { + std::fs::write(dir.path().join(name), fixture).unwrap(); + let output = call(&tool, json!({ "path": name })).await; + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.mime_type(), mime); + assert_eq!(image.data(), fixture); + } +} + +#[tokio::test] +async fn delegated_view_image_preserves_caller_relative_path_identity() { + let dir = TempDir::new().unwrap(); + std::fs::create_dir(dir.path().join("docs")).unwrap(); + let fixture = include_bytes!("fixtures/view-image/valid.png"); + std::fs::write(dir.path().join("docs/image.png"), fixture).unwrap(); + let source: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new( + Scope::writable(dir.path()).unwrap(), + dir.path().to_path_buf(), + )); + let broker = WorkdirToolBroker::new(source); + let lease = broker + .scope(WorkdirToolScope { + rules: vec![WorkdirToolScopeRule { + target: WorkdirPath::new("docs").unwrap(), + permission: WorkdirToolScopePermission::Read, + recursive: true, + symlink_policy: Default::default(), + }], + cwd: WorkdirPath::new("docs").unwrap(), + command: false, + }) + .await + .unwrap(); + let (_meta, tool) = view_image_tool(lease.tool_session())(); + + let output = call(&tool, json!({ "path": "image.png" })).await; + assert!(output.summary.contains("image.png")); + assert!(!output.summary.contains("docs/image.png")); + let agen::tool::Attachment::Image(image) = &output.attachments[0]; + assert_eq!(image.data(), fixture); + + lease.close().await.unwrap(); + broker.close().await.unwrap(); +} + #[tokio::test] async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let dir = TempDir::new().unwrap(); @@ -624,7 +730,7 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { let scope = scope_with_spill(dir.path(), spill.path()); let session: WorkdirSessionHandle = Arc::new(LocalWorkdirSession::new(scope, dir.path().to_path_buf())); - let png = b"\x89PNG\r\n\x1a\nprivate-image-body"; + let png = include_bytes!("fixtures/view-image/valid.png"); std::fs::write(dir.path().join("image.png"), png).unwrap(); let definition = view_image_tool(session); let (_meta, tool) = definition(); @@ -635,7 +741,6 @@ async fn view_image_reads_scoped_bytes_into_durable_tool_detail() { assert_eq!(image.mime_type(), "image/png"); assert_eq!(image.data(), png); let serialized = serde_json::to_string(&output).unwrap(); - assert!(!serialized.contains("private-image-body")); assert!(serialized.contains("attachments")); let restored: agen::tool::ToolOutput = serde_json::from_str(&serialized).unwrap(); let agen::tool::Attachment::Image(restored_image) = &restored.attachments[0]; diff --git a/crates/workdir/src/scope.rs b/crates/workdir/src/scope.rs index efe773d9f..a58e0b01d 100644 --- a/crates/workdir/src/scope.rs +++ b/crates/workdir/src/scope.rs @@ -1016,12 +1016,15 @@ impl WorkdirSession for ScopedWorkdirSession { &self, mut request: ReadBytesRequest, ) -> Result { + let caller_path = request.path.clone(); let path = self .resolve_operation_path(&request.path, WorkdirToolScopePermission::Read) .await?; self.ensure_read(&path, WorkdirSessionCapability::Read)?; request.path = path; - self.source.read_bytes(request).await + let mut result = self.source.read_bytes(request).await?; + result.path = caller_path; + Ok(result) } async fn write(&self, mut request: WriteRequest) -> Result { @@ -1913,14 +1916,9 @@ mod tests { child.read(read("readme.md")).await.unwrap().bytes, b"visible" ); - assert_eq!( - child - .read_bytes(read_bytes("readme.md")) - .await - .unwrap() - .bytes, - b"visible" - ); + let binary = child.read_bytes(read_bytes("readme.md")).await.unwrap(); + assert_eq!(binary.path, fs_path("readme.md")); + assert_eq!(binary.bytes, b"visible"); assert!(matches!( child.write(write("new.md", "no")).await, Err(WorkdirError::Denied(_)) diff --git a/crates/workspace-server/Cargo.toml b/crates/workspace-server/Cargo.toml index 2438e67c8..1c50732f0 100644 --- a/crates/workspace-server/Cargo.toml +++ b/crates/workspace-server/Cargo.toml @@ -50,7 +50,9 @@ zeroize.workspace = true webauthn-rs = { workspace = true } [dev-dependencies] +agen.workspace = true tempfile.workspace = true +tools.workspace = true 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"] } diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 35f37069b..5aad53234 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -5,7 +5,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex, RwLock, Weak}; use async_trait::async_trait; -use axum::extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}; +use axum::extract::ws::{CloseFrame, Message as WsMessage, WebSocket, WebSocketUpgrade}; use axum::extract::{DefaultBodyLimit, Extension, Path as AxumPath, Query, Request, State}; #[cfg(test)] use axum::http::header::{CACHE_CONTROL, ETAG, IF_NONE_MATCH}; @@ -17472,41 +17472,116 @@ async fn send_external_workdir_server_frame( ) } +async fn reject_external_workdir_provider(socket: &mut WebSocket, code: u16, mut reason: String) { + const MAX_CLOSE_REASON_BYTES: usize = 123; + if reason.len() > MAX_CLOSE_REASON_BYTES { + reason.truncate(MAX_CLOSE_REASON_BYTES); + } + let _ = socket + .send(WsMessage::Close(Some(CloseFrame { + code, + reason: reason.into(), + }))) + .await; +} + +fn external_workdir_registration_parse_error(text: &str) -> String { + let actual_version = serde_json::from_str::(text) + .ok() + .and_then(|value| value.get("version")?.as_u64()); + match actual_version { + Some(actual) + if actual != workdir::external::EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION as u64 => + { + format!( + "unsupported External Workdir protocol version {actual}; Backend requires {}", + workdir::external::EXTERNAL_WORKDIR_PROVIDER_PROTOCOL_VERSION + ) + } + _ => "invalid External Workdir provider registration".to_string(), + } +} + async fn serve_external_workdir_provider( api: WorkspaceApi, initial_grant: ExternalWorkdirGrantRecord, mut socket: WebSocket, ) { - let registration = + let registration_text = match tokio::time::timeout(std::time::Duration::from_secs(10), socket.next()).await { - Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => { - serde_json::from_str::(&text).ok() + Ok(Some(Ok(WsMessage::Text(text)))) if text.len() <= 64 * 1024 => text, + _ => { + reject_external_workdir_provider( + &mut socket, + 1008, + "External Workdir provider registration was missing or oversized".to_string(), + ) + .await; + return; } - _ => None, }; - let Some(ExternalWorkdirProviderFrame { + let registration_frame = + match serde_json::from_str::(®istration_text) { + Ok(frame) => frame, + Err(_) => { + let reason = external_workdir_registration_parse_error(®istration_text); + reject_external_workdir_provider(&mut socket, 1002, reason).await; + return; + } + }; + let ExternalWorkdirProviderFrame { message: ExternalWorkdirProviderMessage::Register { registration }, .. - }) = registration + } = registration_frame else { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1008, + "expected External Workdir provider registration".to_string(), + ) + .await; return; }; let expected_capabilities = match external_workdir_capabilities(&initial_grant.permissions) { Ok(capabilities) => capabilities, Err(_) => { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1011, + "Backend could not validate External Workdir capabilities".to_string(), + ) + .await; return; } }; + if registration.read_limits != workdir::BoundedReadLimits::EXTERNAL_DEFAULT { + let required = workdir::BoundedReadLimits::EXTERNAL_DEFAULT; + reject_external_workdir_provider( + &mut socket, + 1008, + format!( + "provider read limits source={} response={}; Backend requires source={} response={}", + registration.read_limits.max_source_bytes, + registration.read_limits.max_response_bytes, + required.max_source_bytes, + required.max_response_bytes, + ), + ) + .await; + return; + } if registration.grant_id.as_str() != initial_grant.grant_id || registration.workdir_id.as_str() != initial_grant.workdir_id || registration.provider_instance_id.as_str() != initial_grant.provider_instance_id || registration.capabilities != expected_capabilities - || registration.read_limits != workdir::BoundedReadLimits::EXTERNAL_DEFAULT || registration.generation == 0 { - let _ = socket.close().await; + reject_external_workdir_provider( + &mut socket, + 1008, + "External Workdir provider registration does not match the grant".to_string(), + ) + .await; return; } let now = Utc::now().to_rfc3339_opts(SecondsFormat::Nanos, true); @@ -26473,6 +26548,26 @@ mod tests { WorkspaceRuntimeBinding, }; + fn valid_gif_with_len(target_len: usize) -> Vec { + let fixture = include_bytes!("../../tools/tests/fixtures/view-image/valid.gif"); + assert_eq!(fixture.last(), Some(&0x3b)); + let mut bytes = fixture[..fixture.len() - 1].to_vec(); + bytes.extend_from_slice(&[0x21, 0xfe]); + let mut remaining = target_len.checked_sub(bytes.len() + 2).unwrap(); + while remaining > 0 { + let mut block = remaining.min(256); + if remaining - block == 1 { + block -= 1; + } + bytes.push((block - 1) as u8); + bytes.extend(std::iter::repeat_n(b'x', block - 1)); + remaining -= block; + } + bytes.extend_from_slice(&[0, 0x3b]); + assert_eq!(bytes.len(), target_len); + bytes + } + #[test] fn external_permission_levels_map_to_existing_bits_and_reject_legacy_partial_grants() { let bits = |capabilities| { @@ -28018,6 +28113,50 @@ mod tests { }) }; + let (mut old_provider, _) = connect_async(&url).await.unwrap(); + let old_registration = serde_json::to_string(®ister(1)).unwrap().replacen( + r#""version":2"#, + r#""version":1"#, + 1, + ); + old_provider + .send(Message::Text(old_registration.into())) + .await + .unwrap(); + let rejected = old_provider.next().await.unwrap().unwrap(); + let Message::Close(Some(frame)) = rejected else { + panic!("old provider must receive a bounded compatibility close reason"); + }; + assert!( + frame + .reason + .contains("unsupported External Workdir protocol version 1") + ); + assert!(frame.reason.contains("requires 2")); + + let (mut smaller_provider, _) = connect_async(&url).await.unwrap(); + let mut smaller_registration = serde_json::to_value(register(1)).unwrap(); + smaller_registration["message"]["registration"]["read_limits"] = serde_json::json!({ + "max_source_bytes": 32 * 1024 * 1024, + "max_response_bytes": 512 * 1024, + }); + smaller_provider + .send(Message::Text( + serde_json::to_string(&smaller_registration).unwrap().into(), + )) + .await + .unwrap(); + let rejected = smaller_provider.next().await.unwrap().unwrap(); + let Message::Close(Some(frame)) = rejected else { + panic!("smaller provider limits must receive the required effective limits"); + }; + assert!(frame.reason.contains("response=524288"), "{}", frame.reason); + assert!( + frame.reason.contains("response=1048576"), + "{}", + frame.reason + ); + let (mut first, _) = connect_async(&url).await.unwrap(); first .send(Message::Text( @@ -28128,8 +28267,8 @@ mod tests { ); let provider_root = tempfile::tempdir().unwrap(); fs::write(provider_root.path().join("events.jsonl"), b"one\ntwo\n").unwrap(); - let image_bytes = b"\x89PNG\r\n\x1a\nexternal\0image\n\xff"; - fs::write(provider_root.path().join("image.png"), image_bytes).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), @@ -28199,17 +28338,72 @@ mod tests { let read_result = read.await.unwrap().unwrap(); assert_eq!(read_result.bytes, b"one\ntwo\n"); - let read_bytes = { + let view_image = { let broker_session = broker_session.clone(); tokio::spawn(async move { - broker_session - .read_bytes(workdir::ReadBytesRequest { - path: workdir::WorkdirPath::new("image.png").unwrap(), - offset: 0, - max_bytes: 1024, - expected_hash: None, - }) - .await + let (_meta, tool) = tools::view_image_tool(broker_session)(); + tool.execute( + 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::(&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 broker_session = broker_session.clone(); + tokio::spawn(async move { + let (_meta, tool) = tools::view_image_tool(broker_session)(); + tool.execute( + r#"{"path":"missing.gif"}"#, + agen::tool::ToolExecutionContext::default(), + ) + .await }) }; let (operation_id, operation) = loop { @@ -28228,21 +28422,22 @@ mod tests { operation.as_inner(), WorkdirSessionOperation::ReadBytes(_) )); - let result = workdir::http::dispatch_workdir_session_operation( + let workdir_error = workdir::http::dispatch_workdir_session_operation( &local_session, operation.into_inner(), ) .await - .unwrap() - .try_into() - .unwrap(); + .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::Completed { result }, + outcome: ExternalWorkdirOperationOutcome::Failed { error }, }, )) .unwrap() @@ -28250,9 +28445,8 @@ mod tests { )) .await .unwrap(); - let read_bytes_result = read_bytes.await.unwrap().unwrap(); - assert_eq!(read_bytes_result.bytes, image_bytes); - assert!(read_bytes_result.eof); + let error = missing_image.await.unwrap().unwrap_err().to_string(); + assert!(error.contains("not found"), "{error}"); let edit = { let broker_session = broker_session.clone(); diff --git a/crates/yoi/src/workdir_share.rs b/crates/yoi/src/workdir_share.rs index a2d3bc6b2..d6dd9d187 100644 --- a/crates/yoi/src/workdir_share.rs +++ b/crates/yoi/src/workdir_share.rs @@ -37,6 +37,39 @@ pub(crate) struct WorkdirShareOptions { const PROVIDER_OPERATION_TIMEOUT: Duration = Duration::from_secs(30); +fn validate_provider_registration(message: Message, generation: u64) -> Result<(), String> { + let registered = match message { + Message::Text(text) => text, + Message::Close(Some(frame)) => { + let reason = frame.reason.trim(); + return Err(if reason.is_empty() { + "Backend rejected the External Workdir provider registration".to_string() + } else { + format!("Backend rejected the External Workdir provider registration: {reason}") + }); + } + Message::Close(None) => { + return Err("Backend rejected the External Workdir provider registration".to_string()); + } + _ => { + return Err( + "Backend returned an invalid External Workdir registration frame".to_string(), + ); + } + }; + let frame = serde_json::from_str::(®istered) + .map_err(|error| format!("Backend returned an invalid External Workdir frame: {error}"))?; + if matches!( + frame.message, + ExternalWorkdirServerMessage::Registered { generation: registered_generation, .. } + if registered_generation == generation + ) { + Ok(()) + } else { + Err("Backend rejected the External Workdir provider registration".to_string()) + } +} + fn provider_operation_timeout(operation: &workdir::http::WorkdirSessionOperation) -> Duration { if matches!( operation, @@ -672,18 +705,7 @@ async fn serve_provider_connection( .map_err(|error| format!("External Workdir provider registration failed: {error}"))? } }; - let Message::Text(registered) = registered else { - return Err("Backend returned an invalid External Workdir registration frame".to_string()); - }; - let frame = serde_json::from_str::(®istered) - .map_err(|error| format!("Backend returned an invalid External Workdir frame: {error}"))?; - if !matches!( - frame.message, - ExternalWorkdirServerMessage::Registered { generation: registered_generation, .. } - if registered_generation == generation - ) { - return Err("Backend rejected the External Workdir provider registration".to_string()); - } + validate_provider_registration(registered, generation)?; println!("Connection: online (generation {generation})"); let (completion_sender, mut completions) = @@ -1047,6 +1069,23 @@ mod tests { } } + #[test] + fn provider_registration_surfaces_backend_close_reason() { + use tokio_tungstenite::tungstenite::protocol::CloseFrame; + use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; + + let error = validate_provider_registration( + Message::Close(Some(CloseFrame { + code: CloseCode::Policy, + reason: "provider read limits source=1024 response=512; Backend requires source=2048 response=1024".into(), + })), + 1, + ) + .unwrap_err(); + assert!(error.contains("provider read limits"), "{error}"); + assert!(error.contains("Backend requires"), "{error}"); + } + #[test] fn interactive_share_defaults_to_first_workspace_read_only_and_unlimited() { let workspaces = [workspace("workspace-a", "Alpha")];