Rust: add race-safe canvas availability waiter

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<String>,
+        /// 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<String>,
+    },
+
     /// 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 \

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<DiscoveredCanvas, Error> {
+        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

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<String>,
+    /// 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<String>, 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<String>) -> 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 —