From 04a9e31abd6ad8158541cb3274fa4a65d4fba7d3 Mon Sep 17 00:00:00 2001 From: Matt Ellis Date: Thu, 6 Aug 2026 21:59:38 -0700 Subject: [PATCH] Add race-safe canvas availability waiter Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- rust/src/errors.rs | 48 ++++++++ rust/src/session.rs | 109 +++++++++++++++++- rust/src/types.rs | 29 +++++ rust/tests/session_test.rs | 226 ++++++++++++++++++++++++++++++++++++- 4 files changed, 407 insertions(+), 5 deletions(-) diff --git a/rust/src/errors.rs b/rust/src/errors.rs index 6e05bbfae1..a5265f837c 100644 --- a/rust/src/errors.rs +++ b/rust/src/errors.rs @@ -131,6 +131,24 @@ pub enum SessionErrorKind { /// The session event loop exited before a pending `send_and_wait` completed. EventLoopClosed, + /// A requested canvas did not become available before the wait timeout. + CanvasWaitTimeout { + /// Provider-local canvas identifier. + canvas_id: String, + /// Owning provider identifier, when the wait was provider-specific. + extension_id: Option, + /// Configured wait timeout. + timeout: Duration, + }, + + /// The session shut down while waiting for a canvas. + CanvasWaitSessionClosed { + /// Provider-local canvas identifier. + canvas_id: String, + /// Owning provider identifier, when the wait was provider-specific. + extension_id: Option, + }, + /// Elicitation is not supported by the host. /// Check `session.capabilities().ui.elicitation` before calling UI methods. ElicitationNotSupported, @@ -166,6 +184,36 @@ impl fmt::Display for SessionErrorKind { SessionErrorKind::EventLoopClosed => { write!(f, "event loop closed before session reached idle") } + SessionErrorKind::CanvasWaitTimeout { + canvas_id, + extension_id, + timeout, + } => { + if let Some(extension_id) = extension_id { + write!( + f, + "timed out after {timeout:?} waiting for canvas {extension_id}:{canvas_id}" + ) + } else { + write!( + f, + "timed out after {timeout:?} waiting for canvas {canvas_id}" + ) + } + } + SessionErrorKind::CanvasWaitSessionClosed { + canvas_id, + extension_id, + } => { + if let Some(extension_id) = extension_id { + write!( + f, + "session closed while waiting for canvas {extension_id}:{canvas_id}" + ) + } else { + write!(f, "session closed while waiting for canvas {canvas_id}") + } + } SessionErrorKind::ElicitationNotSupported => write!( f, "elicitation not supported by host \ diff --git a/rust/src/session.rs b/rust/src/session.rs index d505541a50..794c5b773b 100644 --- a/rust/src/session.rs +++ b/rust/src/session.rs @@ -12,8 +12,8 @@ use tracing::{Instrument, warn}; use crate::canvas::CanvasHandler; use crate::generated::api_types::{ - LogRequest, ModelSwitchToRequest, OpenCanvasInstance, RegisterEventInterestParams, - ToolsGetCurrentMetadataResult, rpc_methods, + DiscoveredCanvas, LogRequest, ModelSwitchToRequest, OpenCanvasInstance, + RegisterEventInterestParams, ToolsGetCurrentMetadataResult, rpc_methods, }; use crate::generated::session_events::{ CommandExecuteData, ElicitationRequestedData, ExternalToolRequestedData, McpOauthRequiredData, @@ -27,6 +27,7 @@ use crate::handler::{ use crate::hooks::SessionHooks; use crate::provider_token::BearerTokenProvider; use crate::session_fs::SessionFsProvider; +use crate::subscription::RecvErrorKind; use crate::trace_context::inject_trace_context; use crate::transforms::SystemMessageTransform; use crate::types::{ @@ -35,7 +36,7 @@ use crate::types::{ PermissionRequestData, RequestId, ResumeSessionConfig, ResumeSessionResult, SectionOverride, SessionCapabilities, SessionConfig, SessionEvent, SessionId, SetModelOptions, SystemMessageConfig, ToolInvocation, ToolResult, ToolResultExpanded, TraceContext, - UiInputOptions, ensure_attachment_display_names, + UiInputOptions, WaitForCanvasOptions, ensure_attachment_display_names, }; use crate::{ Client, Error, ErrorKind, JsonRpcResponse, SessionErrorKind, SessionEventNotification, @@ -292,6 +293,108 @@ impl Session { crate::subscription::EventSubscription::new(self.event_tx.subscribe()) } + /// Wait for a canvas declaration to become available in this session. + /// + /// The waiter subscribes before reading the current registry, so a + /// declaration registered concurrently with the initial `canvas.list` + /// request cannot be missed. Subscription lag triggers another registry + /// read rather than failing the wait. + /// + /// This waits only for the requested canvas. It does not indicate that + /// extension initialization or the complete canvas registry has settled. + /// + /// # Cancel safety + /// + /// **Cancel-safe.** Dropping this future cancels the wait without changing + /// session or canvas state. + pub async fn wait_for_canvas( + &self, + options: WaitForCanvasOptions, + ) -> Result { + let WaitForCanvasOptions { + canvas_id, + extension_id, + timeout, + } = options; + + let wait = async { + let mut events = self.subscribe(); + + loop { + let canvas_rpc = self.rpc().canvas(); + let canvases = tokio::select! { + biased; + _ = self.shutdown.cancelled() => { + return Err(ErrorKind::Session( + SessionErrorKind::CanvasWaitSessionClosed { + canvas_id: canvas_id.clone(), + extension_id: extension_id.clone(), + }, + ) + .into()); + } + result = canvas_rpc.list() => result?.canvases, + }; + if let Some(canvas) = canvases.into_iter().find(|canvas| { + canvas.canvas_id == canvas_id + && extension_id + .as_ref() + .is_none_or(|id| canvas.extension_id == *id) + }) { + return Ok(canvas); + } + + loop { + tokio::select! { + biased; + _ = self.shutdown.cancelled() => { + return Err(ErrorKind::Session( + SessionErrorKind::CanvasWaitSessionClosed { + canvas_id: canvas_id.clone(), + extension_id: extension_id.clone(), + }, + ) + .into()); + } + result = events.recv() => { + match result { + Ok(event) + if event.parsed_type() + == SessionEventType::SessionCanvasRegistryChanged => + { + break; + } + Ok(_) => {} + Err(error) => match error.kind() { + RecvErrorKind::Lagged(_) => break, + RecvErrorKind::Closed => { + return Err(ErrorKind::Session( + SessionErrorKind::CanvasWaitSessionClosed { + canvas_id: canvas_id.clone(), + extension_id: extension_id.clone(), + }, + ) + .into()); + } + }, + } + } + } + } + } + }; + + match tokio::time::timeout(timeout, wait).await { + Ok(result) => result, + Err(_) => Err(ErrorKind::Session(SessionErrorKind::CanvasWaitTimeout { + canvas_id, + extension_id, + timeout, + }) + .into()), + } + } + /// The underlying Client (for advanced use cases). pub fn client(&self) -> &Client { &self.client diff --git a/rust/src/types.rs b/rust/src/types.rs index 37d3b248bf..b59362cb1b 100644 --- a/rust/src/types.rs +++ b/rust/src/types.rs @@ -4487,6 +4487,35 @@ impl SetModelOptions { } } +/// Options for [`Session::wait_for_canvas`](crate::session::Session::wait_for_canvas). +#[derive(Debug, Clone)] +#[non_exhaustive] +pub struct WaitForCanvasOptions { + /// Provider-local canvas identifier to wait for. + pub canvas_id: String, + /// Owning provider identifier used to disambiguate canvases with the same ID. + pub extension_id: Option, + /// Maximum time to wait for the canvas to become available. + pub timeout: Duration, +} + +impl WaitForCanvasOptions { + /// Wait for a canvas with the given provider-local ID. + pub fn new(canvas_id: impl Into, timeout: Duration) -> Self { + Self { + canvas_id: canvas_id.into(), + extension_id: None, + timeout, + } + } + + /// Require the canvas to belong to the given provider. + pub fn with_extension_id(mut self, extension_id: impl Into) -> Self { + self.extension_id = Some(extension_id.into()); + self + } +} + /// Response from the top-level `ping` RPC. /// /// The `protocol_version` field is the most commonly-inspected piece — diff --git a/rust/tests/session_test.rs b/rust/tests/session_test.rs index 41cae3e950..9f6b669e5c 100644 --- a/rust/tests/session_test.rs +++ b/rust/tests/session_test.rs @@ -23,9 +23,11 @@ use github_copilot_sdk::types::{ CanvasProviderIdentity, CloudSessionOptions, CloudSessionRepository, CommandContext, CommandDefinition, CommandHandler, DeliveryMode, ElicitationRequest, ElicitationResult, ExitPlanModeData, ExtensionInfo, MessageOptions, RequestId, SessionConfig, SessionId, - SetModelOptions, Tool, ToolInvocation, ToolResult, + SetModelOptions, Tool, ToolInvocation, ToolResult, WaitForCanvasOptions, +}; +use github_copilot_sdk::{ + Client, ContextTier, ErrorKind, ProtocolErrorKind, SessionErrorKind, tool, }; -use github_copilot_sdk::{Client, ContextTier, ErrorKind, ProtocolErrorKind, tool}; use serde_json::Value; use tokio::io::{AsyncWrite, AsyncWriteExt, duplex}; use tokio::time::timeout; @@ -77,6 +79,15 @@ fn test_canvas_handler() -> Arc { Arc::new(TestCanvasHandler) } +fn discovered_canvas(canvas_id: &str, extension_id: &str) -> Value { + serde_json::json!({ + "canvasId": canvas_id, + "extensionId": extension_id, + "displayName": "Test Canvas", + "description": "Test canvas description" + }) +} + async fn write_framed(writer: &mut (impl AsyncWrite + Unpin), body: &[u8]) { let header = format!("Content-Length: {}\r\n\r\n", body.len()); writer.write_all(header.as_bytes()).await.unwrap(); @@ -3392,6 +3403,217 @@ async fn resume_session_sends_canvas_fields_and_captures_open_canvases() { assert_eq!(caps.ui.unwrap().canvases, Some(true)); } +#[tokio::test] +async fn wait_for_canvas_returns_matching_existing_canvas() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas( + WaitForCanvasOptions::new("counter", TIMEOUT) + .with_extension_id("project:counter"), + ) + .await + } + }); + + let request = server.read_request().await; + assert_eq!(request["method"], "session.canvas.list"); + server + .respond( + &request, + serde_json::json!({ + "canvases": [ + discovered_canvas("counter", "user:counter"), + discovered_canvas("counter", "project:counter") + ] + }), + ) + .await; + + let canvas = timeout(TIMEOUT, waiter).await.unwrap().unwrap().unwrap(); + assert_eq!(canvas.canvas_id, "counter"); + assert_eq!(canvas.extension_id, "project:counter"); +} + +#[tokio::test] +async fn wait_for_canvas_handles_registration_during_initial_list() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas(WaitForCanvasOptions::new("counter", TIMEOUT)) + .await + } + }); + + let initial_list = server.read_request().await; + assert_eq!(initial_list["method"], "session.canvas.list"); + server + .send_event( + "session.canvas.registry_changed", + serde_json::json!({ + "canvases": [discovered_canvas("counter", "project:counter")] + }), + ) + .await; + server + .respond(&initial_list, serde_json::json!({ "canvases": [] })) + .await; + + let refreshed_list = server.read_request().await; + assert_eq!(refreshed_list["method"], "session.canvas.list"); + server + .respond( + &refreshed_list, + serde_json::json!({ + "canvases": [discovered_canvas("counter", "project:counter")] + }), + ) + .await; + + let canvas = timeout(TIMEOUT, waiter).await.unwrap().unwrap().unwrap(); + assert_eq!(canvas.extension_id, "project:counter"); +} + +#[tokio::test] +async fn wait_for_canvas_resyncs_after_subscription_lag() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas(WaitForCanvasOptions::new("counter", TIMEOUT)) + .await + } + }); + + let initial_list = server.read_request().await; + for sequence in 0..600 { + server + .send_event("test.event", serde_json::json!({ "sequence": sequence })) + .await; + } + tokio::time::sleep(Duration::from_millis(50)).await; + server + .respond(&initial_list, serde_json::json!({ "canvases": [] })) + .await; + + let refreshed_list = timeout(TIMEOUT, server.read_request()).await.unwrap(); + assert_eq!(refreshed_list["method"], "session.canvas.list"); + server + .respond( + &refreshed_list, + serde_json::json!({ + "canvases": [discovered_canvas("counter", "project:counter")] + }), + ) + .await; + + let canvas = timeout(TIMEOUT, waiter).await.unwrap().unwrap().unwrap(); + assert_eq!(canvas.canvas_id, "counter"); +} + +#[tokio::test] +async fn wait_for_canvas_times_out_without_sending_a_turn() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas(WaitForCanvasOptions::new( + "missing", + Duration::from_millis(50), + )) + .await + } + }); + + let request = server.read_request().await; + assert_eq!(request["method"], "session.canvas.list"); + server + .respond(&request, serde_json::json!({ "canvases": [] })) + .await; + + let error = timeout(TIMEOUT, waiter) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(matches!( + error.kind(), + ErrorKind::Session(SessionErrorKind::CanvasWaitTimeout { canvas_id, .. }) + if canvas_id == "missing" + )); +} + +#[tokio::test] +async fn wait_for_canvas_stops_when_session_closes() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas(WaitForCanvasOptions::new("missing", TIMEOUT)) + .await + } + }); + + let request = server.read_request().await; + assert_eq!(request["method"], "session.canvas.list"); + server + .respond(&request, serde_json::json!({ "canvases": [] })) + .await; + session.stop_event_loop().await; + + let error = timeout(TIMEOUT, waiter) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(matches!( + error.kind(), + ErrorKind::Session(SessionErrorKind::CanvasWaitSessionClosed { canvas_id, .. }) + if canvas_id == "missing" + )); +} + +#[tokio::test] +async fn wait_for_canvas_stops_when_session_closes_during_list() { + let (session, mut server) = create_session_pair().await; + let session = Arc::new(session); + let waiter = tokio::spawn({ + let session = session.clone(); + async move { + session + .wait_for_canvas(WaitForCanvasOptions::new("missing", TIMEOUT)) + .await + } + }); + + let request = server.read_request().await; + assert_eq!(request["method"], "session.canvas.list"); + session.stop_event_loop().await; + + let error = timeout(TIMEOUT, waiter) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(matches!( + error.kind(), + ErrorKind::Session(SessionErrorKind::CanvasWaitSessionClosed { canvas_id, .. }) + if canvas_id == "missing" + )); +} + #[tokio::test] async fn session_canvas_opened_updates_open_canvas_snapshots() { let (session, mut server) = create_session_pair().await;