From 78ead2cfdfc871571b878aa1c8a40c27aa066f0a Mon Sep 17 00:00:00 2001 From: stephen Date: Fri, 21 Aug 2026 13:04:07 -0700 Subject: [PATCH 1/3] refactor!: kill direct tokio sync coupling in hot paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bevy-prep PR 1: rig's hot paths no longer touch tokio primitives, so a tokio runtime is not required for streaming pause/resume, tool-server registration, or the copilot/chatgpt auth caches. - streaming: PauseControl drops its tokio::sync::watch channel for one Arc. poll_next integrates directly via the register-then-recheck AtomicWaker protocol, deleting the boxed resume_wait future; both #2258 H7 invariants (no busy re-wake, no lost resume race) are preserved and pinned by the existing tests. PauseControl is now Clone. - tool server: ToolServerHandle's tokio RwLock becomes std::sync::RwLock (no guard ever crossed an await; all eight sites are clone-under-lock or sync mutations). Poisoning is recovered via PoisonError::into_inner in one pair of private accessors. Registration-only methods de-async: add_tool, add_dynamic_tool, add_portable_dynamic_tool, append_toolset, remove_tool (MIGRATING.md entry included). rmcp.rs and discord_bot.rs keep tokio locks: the rmcp feature pulls tokio via the SDK regardless. - auth caches: the copilot/chatgpt Mutex<()> becomes async_lock::Mutex — the lock now wraps the state it serializes (token/key caches) instead of guarding code, and async-lock works on both native and wasm halves. - device-flow sleeps: tokio::time::sleep only compiled via feature unification (rig-core builds tokio without "time"); both call sites now use a new wasm_compat::sleep built on futures_timer, next to the existing timeout helper. Residual non-test tokio in rig-core is the openai realtime websocket (feature-gated, moves out in the transport-crate split) and test-only code. Verification: cargo check/clippy --workspace --all-features --all-targets clean; rig-core (1799) and rig-agent (598) tests pass; wasm32-unknown-unknown check (CI's command) passes. --- Cargo.lock | 1 + Cargo.toml | 1 + MIGRATING.md | 21 ++++ crates/rig-agent/src/agent/runner.rs | 6 +- crates/rig-agent/src/tool/rmcp.rs | 10 +- crates/rig-agent/src/tool/server.rs | 96 ++++++++++--------- crates/rig-core/Cargo.toml | 1 + .../src/providers/chatgpt/auth/mod.rs | 20 ++-- .../src/providers/chatgpt/auth/native.rs | 2 +- .../src/providers/copilot/auth/mod.rs | 23 +++-- .../src/providers/copilot/auth/native.rs | 2 +- crates/rig-core/src/streaming/mod.rs | 80 +++++++--------- crates/rig-core/src/wasm_compat.rs | 11 +++ .../providers/gemini/cassette/tool_server.rs | 6 +- tests/providers/gemini/tools_support.rs | 2 +- 15 files changed, 152 insertions(+), 130 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 39daa3dec9..83e5775924 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9925,6 +9925,7 @@ dependencies = [ "anyhow", "as-any", "assert_fs", + "async-lock", "async-stream", "base64 0.22.1", "bytes", diff --git a/Cargo.toml b/Cargo.toml index 2218be38ed..c29484d8ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -110,6 +110,7 @@ anyhow = "1" arrow-array = "58" as-any = "0.3" assert_fs = "1" +async-lock = "3" async-stream = "0.3" axum = "0.8" aws-config = { version = "1", default-features = false } diff --git a/MIGRATING.md b/MIGRATING.md index 21db55d4ce..7971f0834e 100644 --- a/MIGRATING.md +++ b/MIGRATING.md @@ -796,6 +796,27 @@ handed back a silently short list. ## 0.41 → next +### `ToolServerHandle` registration methods are now synchronous + +`add_tool`, `add_dynamic_tool`, `add_portable_dynamic_tool`, `append_toolset`, +and `remove_tool` only ever took a short registry lock that is never held +across an await; the lock is now a `std::sync::RwLock` and the methods are +plain `fn`. Drop the `.await`: + +```rust +// before +handle.add_tool(MyTool).await; + +// after +handle.add_tool(MyTool); +``` + +Execution and snapshot paths (`execute`, `get_tool_defs`) are unchanged and +remain async. This removes the last `tokio::sync` primitive from the +tool-server hot path; streaming pause/resume and the copilot/chatgpt auth +caches likewise moved off tokio primitives (no API change), so neither +rig-core nor rig-agent needs a tokio runtime for these paths. + ### Telemetry getters borrow: `ProviderResponseExt::get_response_id` / `get_response_model_name` return `Option<&str>` Both getters exist to hand a value to `tracing::Span::record`, which takes diff --git a/crates/rig-agent/src/agent/runner.rs b/crates/rig-agent/src/agent/runner.rs index 9f01fc6685..895bea320c 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -7913,7 +7913,7 @@ mod migrated_tests { request_started.notified().await; handle .add_tool(SecondGenerationTool(second_calls.clone())) - .await; + ; release_response.notify_one(); }; let (response, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async { @@ -7971,7 +7971,7 @@ mod migrated_tests { request_started.notified().await; handle .add_tool(SecondGenerationTool(second_calls.clone())) - .await; + ; release_response.notify_one(); }; let (final_output, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async { @@ -9232,7 +9232,7 @@ mod migrated_tests { _event: ModelTurnFinished<'_>, ) -> ModelTurnAction { if ctx.turn() == 1 { - self.handle.add_tool(FinalResultTool).await; + self.handle.add_tool(FinalResultTool); } ModelTurnAction::continue_run() diff --git a/crates/rig-agent/src/tool/rmcp.rs b/crates/rig-agent/src/tool/rmcp.rs index b6110b8114..aaf1b34cf2 100644 --- a/crates/rig-agent/src/tool/rmcp.rs +++ b/crates/rig-agent/src/tool/rmcp.rs @@ -702,10 +702,7 @@ impl McpClientHandler { tracing::debug!(refresh, "discarding stale initial MCP tool list"); return; } - managed.registrations = self - .tool_server_handle - .add_managed_erased_tools(tools) - .await; + managed.registrations = self.tool_server_handle.add_managed_erased_tools(tools); managed.committed_refresh = refresh; } @@ -718,8 +715,7 @@ impl McpClientHandler { let expected = managed.registrations.clone(); managed.registrations = self .tool_server_handle - .reconcile_managed_erased_tools(expected, tools) - .await; + .reconcile_managed_erased_tools(expected, tools); managed.committed_refresh = refresh; true } @@ -1793,7 +1789,7 @@ mod migrated_tests { handle .add_dynamic_tool(make_dynamic_tool("alpha", "Local alpha")) - .await; + ; server_control .set_tools(vec![make_tool("refresh_complete", "Refresh sentinel")]) .await; diff --git a/crates/rig-agent/src/tool/server.rs b/crates/rig-agent/src/tool/server.rs index bab669b7bd..a9d06255ae 100644 --- a/crates/rig-agent/src/tool/server.rs +++ b/crates/rig-agent/src/tool/server.rs @@ -4,7 +4,7 @@ use std::{collections::BTreeSet, sync::Arc}; use std::collections::HashMap; use indexmap::IndexMap; -use tokio::sync::RwLock; +use std::sync::{PoisonError, RwLock, RwLockReadGuard, RwLockWriteGuard}; #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] use crate::tool::ErasedTool; @@ -270,10 +270,23 @@ impl ToolServer { pub struct ToolServerHandle(Arc>); impl ToolServerHandle { + /// Shared registry state under the single poisoning policy: a panic + /// inside one of the short sync critical sections cannot leave the + /// registry logically torn, so a poisoned lock is recovered rather than + /// propagated. + fn state(&self) -> RwLockReadGuard<'_, ToolServerState> { + self.0.read().unwrap_or_else(PoisonError::into_inner) + } + + /// Exclusive registry state; same poisoning policy as [`Self::state`]. + fn state_mut(&self) -> RwLockWriteGuard<'_, ToolServerState> { + self.0.write().unwrap_or_else(PoisonError::into_inner) + } + /// Register through `add`, then drop any stale MCP managed-generation /// entry so the (re)registered name follows last-registration-wins. - async fn register(&self, add: impl FnOnce(&mut ToolSet) -> String) { - let mut state = self.0.write().await; + fn register(&self, add: impl FnOnce(&mut ToolSet) -> String) { + let mut state = self.state_mut(); let _name = add(&mut state.toolset); #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] state.managed_generations.remove(&_name); @@ -281,33 +294,31 @@ impl ToolServerHandle { /// Register a new static tool. Re-registering an existing name replaces /// the implementation (last wins) and keeps its position. - pub async fn add_tool(&self, tool: T) + pub fn add_tool(&self, tool: T) where T: Tool + 'static, { - self.register(|toolset| toolset.add_tool(tool)).await + self.register(|toolset| toolset.add_tool(tool)) } /// Register a runtime-defined static tool. - pub async fn add_dynamic_tool(&self, tool: DynamicTool) { + pub fn add_dynamic_tool(&self, tool: DynamicTool) { self.register(|toolset| toolset.add_dynamic_tool(tool)) - .await } /// Register a context-free dynamic tool through the classic adapter. - pub async fn add_portable_dynamic_tool(&self, tool: PortableDynamicTool) { + pub fn add_portable_dynamic_tool(&self, tool: PortableDynamicTool) { self.register(|toolset| toolset.add_portable_dynamic_tool(tool)) - .await } /// Atomically install the initial tools owned by one MCP handler. /// Initial connection retains the registry's last-registration-wins policy. #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] - pub(crate) async fn add_managed_erased_tools( + pub(crate) fn add_managed_erased_tools( &self, tools: Vec>, ) -> HashMap { - let mut state = self.0.write().await; + let mut state = self.state_mut(); let mut managed = HashMap::with_capacity(tools.len()); for tool in tools { @@ -336,12 +347,12 @@ impl ToolServerHandle { /// tool list. Existing names are changed only when their expected generation /// remains current; newer local or peer-handler registrations win. #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] - pub(crate) async fn reconcile_managed_erased_tools( + pub(crate) fn reconcile_managed_erased_tools( &self, mut expected: HashMap, tools: Vec>, ) -> HashMap { - let mut state = self.0.write().await; + let mut state = self.state_mut(); let mut refreshed = HashMap::with_capacity(tools.len()); let mut managed_order = Vec::with_capacity(tools.len()); let mut seen = std::collections::HashSet::with_capacity(tools.len()); @@ -421,8 +432,8 @@ impl ToolServerHandle { /// Merge an entire toolset into the server in registration order. /// Existing names are replaced (last wins) and keep their position. - pub async fn append_toolset(&self, toolset: ToolSet) { - let mut state = self.0.write().await; + pub fn append_toolset(&self, toolset: ToolSet) { + let mut state = self.state_mut(); #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] let names = toolset.tools.keys().cloned().collect::>(); state.toolset.add_tools(toolset); @@ -433,8 +444,8 @@ impl ToolServerHandle { } /// Remove a tool by name. - pub async fn remove_tool(&self, tool_name: &str) { - let mut state = self.0.write().await; + pub fn remove_tool(&self, tool_name: &str) { + let mut state = self.state_mut(); state.toolset.delete_tool(tool_name); #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] state.managed_generations.remove(tool_name); @@ -460,16 +471,16 @@ impl ToolServerHandle { /// Run `f` against the registry state, first retiring disconnected MCP /// tools (which needs a write lock) when that feature is compiled in. - async fn with_registry(&self, f: impl FnOnce(&ToolServerState) -> R) -> R { + fn with_registry(&self, f: impl FnOnce(&ToolServerState) -> R) -> R { #[cfg(all(feature = "rmcp", not(target_family = "wasm")))] { - let mut state = self.0.write().await; + let mut state = self.state_mut(); state.retire_disconnected_tools(); f(&state) } #[cfg(not(all(feature = "rmcp", not(target_family = "wasm"))))] { - let state = self.0.read().await; + let state = self.state(); f(&state) } } @@ -481,9 +492,7 @@ impl ToolServerHandle { args: &str, context: &ToolContext, ) -> ToolDispatch { - let tool = self - .with_registry(|state| state.toolset.get(tool_name).cloned()) - .await; + let tool = self.with_registry(|state| state.toolset.get(tool_name).cloned()); dispatch_tool(tool_name, args.to_string(), tool, context).await } @@ -507,7 +516,7 @@ impl ToolServerHandle { prompt: Option, ) -> Result { let retrieval_indexes = { - let state = self.0.read().await; + let state = self.state(); state.retrieval_indexes.clone() }; @@ -549,9 +558,7 @@ impl ToolServerHandle { Vec::new() }; - let tools = self - .with_registry(|state| snapshot_registered_tools(state, &dynamic_tool_ids)) - .await; + let tools = self.with_registry(|state| snapshot_registered_tools(state, &dynamic_tool_ids)); Ok(ToolRegistrySnapshot::new(tools)) } @@ -728,7 +735,7 @@ mod tests { let handle = server.run(); - handle.add_tool(MockAddTool).await; + handle.add_tool(MockAddTool); let res = handle.get_tool_defs(None).await.unwrap(); assert_eq!(res.len(), 1); @@ -740,7 +747,7 @@ mod tests { .unwrap(); assert_eq!(res, "7"); - handle.remove_tool("add").await; + handle.remove_tool("add"); let res = handle.get_tool_defs(None).await.unwrap(); assert_eq!(res.len(), 0); @@ -761,7 +768,7 @@ mod tests { description: "second schema", output: "second implementation", }) - .await; + ; assert_eq!(snapshot.definitions()[0].description, "first schema"); let dispatch = snapshot @@ -786,8 +793,8 @@ mod tests { pub async fn test_toolserver_append_toolset_matches_add_tool() { let mut via_add_tool = { let handle = ToolServer::new().run(); - handle.add_tool(MockAddTool).await; - handle.add_tool(MockSubtractTool).await; + handle.add_tool(MockAddTool); + handle.add_tool(MockSubtractTool); handle.get_tool_defs(None).await.unwrap() }; via_add_tool.sort_by(|a, b| a.name.cmp(&b.name)); @@ -797,7 +804,7 @@ mod tests { let mut toolset = ToolSet::default(); toolset.add_tool(MockAddTool); toolset.add_tool(MockSubtractTool); - handle.append_toolset(toolset).await; + handle.append_toolset(toolset); handle.get_tool_defs(None).await.unwrap() }; via_append_toolset.sort_by(|a, b| a.name.cmp(&b.name)); @@ -824,7 +831,7 @@ mod tests { #[tokio::test] pub async fn handle_add_tool_uses_canonical_static_name() { let handle = ToolServer::new().run(); - handle.add_tool(NamedTool::new()).await; + handle.add_tool(NamedTool::new()); let defs = handle.get_tool_defs(None).await.unwrap(); assert_eq!(defs.len(), 1); @@ -850,8 +857,8 @@ mod tests { #[tokio::test] pub async fn get_tool_defs_preserves_static_registration_order() { let handle = ToolServer::new().run(); - handle.add_tool(MockSubtractTool).await; - handle.add_tool(MockAddTool).await; + handle.add_tool(MockSubtractTool); + handle.add_tool(MockAddTool); let defs = handle.get_tool_defs(None).await.unwrap(); assert_eq!( @@ -906,11 +913,11 @@ mod tests { #[tokio::test] pub async fn duplicate_registration_advertises_one_definition() { let handle = ToolServer::new().tool(MockAddTool).run(); - handle.add_tool(MockAddTool).await; + handle.add_tool(MockAddTool); let mut toolset = ToolSet::default(); toolset.add_tool(MockAddTool); - handle.append_toolset(toolset).await; + handle.append_toolset(toolset); let defs = handle.get_tool_defs(None).await.unwrap(); assert_eq!( @@ -1029,15 +1036,10 @@ mod tests { // Wait until we are strictly inside `call()` started.notified().await; - // Try to write to the state (add a tool) while the tool call is mid-execution. - // If the read lock is incorrectly held across tool execution, this will deadlock. - let add_result = - tokio::time::timeout(Duration::from_secs(1), handle.add_tool(MockAddTool)).await; - - assert!( - add_result.is_ok(), - "Writing to ToolServer deadlocked! The read lock is being held across tool execution." - ); + // Write to the state (add a tool) while the tool call is mid-execution. + // If the read lock were incorrectly held across tool execution, this + // sync call would block forever and the test harness would time out. + handle.add_tool(MockAddTool); // Allow the background tool to finish and clean up allow_finish.notify_one(); diff --git a/crates/rig-core/Cargo.toml b/crates/rig-core/Cargo.toml index 823a232a20..ba584fa968 100644 --- a/crates/rig-core/Cargo.toml +++ b/crates/rig-core/Cargo.toml @@ -25,6 +25,7 @@ doctest = true [dependencies] as-any = { workspace = true } +async-lock = { workspace = true } async-stream = { workspace = true } base64 = { workspace = true } bytes = { workspace = true } diff --git a/crates/rig-core/src/providers/chatgpt/auth/mod.rs b/crates/rig-core/src/providers/chatgpt/auth/mod.rs index aadc89d532..3f16836c79 100644 --- a/crates/rig-core/src/providers/chatgpt/auth/mod.rs +++ b/crates/rig-core/src/providers/chatgpt/auth/mod.rs @@ -3,7 +3,7 @@ use std::fmt; use std::path::PathBuf; use std::sync::Arc; -use tokio::sync::Mutex; +use async_lock::Mutex; pub use crate::providers::internal::auth::{DeviceCodeHandler, DeviceCodePrompt}; @@ -38,15 +38,17 @@ impl fmt::Debug for AuthSource { #[derive(Clone)] pub struct Authenticator { source: AuthSource, - platform: platform::PlatformAuthenticator, - state_lock: Arc>, + /// The platform half owns the token/key caches (files plus their parsed + /// state); serializing access to it — rather than to a detached unit + /// lock — is what prevents concurrent refreshes from racing the cache. + platform: Arc>, } impl fmt::Debug for Authenticator { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Authenticator") .field("source", &self.source) - .field("platform", &self.platform) + .field("platform", &"") .finish() } } @@ -68,12 +70,11 @@ impl Authenticator { ) -> Self { Self { source, - platform: platform::PlatformAuthenticator::new( + platform: Arc::new(Mutex::new(platform::PlatformAuthenticator::new( auth_file, device_code_handler, allow_device_flow, - ), - state_lock: Arc::new(Mutex::new(())), + ))), } } @@ -86,10 +87,7 @@ impl Authenticator { access_token: access_token.clone(), account_id: account_id.clone(), }), - AuthSource::OAuth => { - let _guard = self.state_lock.lock().await; - self.platform.auth_context_oauth().await - } + AuthSource::OAuth => self.platform.lock().await.auth_context_oauth().await, } } } diff --git a/crates/rig-core/src/providers/chatgpt/auth/native.rs b/crates/rig-core/src/providers/chatgpt/auth/native.rs index d907394333..3f563f8acb 100644 --- a/crates/rig-core/src/providers/chatgpt/auth/native.rs +++ b/crates/rig-core/src/providers/chatgpt/auth/native.rs @@ -179,7 +179,7 @@ impl PlatformAuthenticator { let status = response.status(); if status.as_u16() == 403 || status.as_u16() == 404 { - tokio::time::sleep(std::time::Duration::from_secs(interval)).await; + crate::wasm_compat::sleep(std::time::Duration::from_secs(interval)).await; continue; } diff --git a/crates/rig-core/src/providers/copilot/auth/mod.rs b/crates/rig-core/src/providers/copilot/auth/mod.rs index d8a38ef455..5fe55eedb5 100644 --- a/crates/rig-core/src/providers/copilot/auth/mod.rs +++ b/crates/rig-core/src/providers/copilot/auth/mod.rs @@ -1,7 +1,7 @@ use std::fmt; use std::path::PathBuf; use std::sync::Arc; -use tokio::sync::Mutex; +use async_lock::Mutex; pub use crate::providers::internal::auth::{DeviceCodeHandler, DeviceCodePrompt}; @@ -35,15 +35,17 @@ impl fmt::Debug for AuthSource { #[derive(Clone)] pub struct Authenticator { source: AuthSource, - platform: platform::PlatformAuthenticator, - state_lock: Arc>, + /// The platform half owns the token/key caches (files plus their parsed + /// state); serializing access to it — rather than to a detached unit + /// lock — is what prevents concurrent refreshes from racing the cache. + platform: Arc>, } impl fmt::Debug for Authenticator { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Authenticator") .field("source", &self.source) - .field("platform", &self.platform) + .field("platform", &"") .finish() } } @@ -66,13 +68,12 @@ impl Authenticator { ) -> Self { Self { source, - platform: platform::PlatformAuthenticator::new( + platform: Arc::new(Mutex::new(platform::PlatformAuthenticator::new( access_token_file, api_key_file, device_code_handler, allow_device_flow, - ), - state_lock: Arc::new(Mutex::new(())), + ))), } } @@ -83,15 +84,13 @@ impl Authenticator { api_base: None, }), AuthSource::GitHubAccessToken(access_token) => { - let _guard = self.state_lock.lock().await; self.platform + .lock() + .await .auth_context_with_github_access_token(access_token) .await } - AuthSource::OAuth => { - let _guard = self.state_lock.lock().await; - self.platform.auth_context_oauth().await - } + AuthSource::OAuth => self.platform.lock().await.auth_context_oauth().await, } } } diff --git a/crates/rig-core/src/providers/copilot/auth/native.rs b/crates/rig-core/src/providers/copilot/auth/native.rs index 7d41e7bd3a..f0bfaf0ef4 100644 --- a/crates/rig-core/src/providers/copilot/auth/native.rs +++ b/crates/rig-core/src/providers/copilot/auth/native.rs @@ -224,7 +224,7 @@ impl PlatformAuthenticator { response.error.as_deref(), response.error_description.as_deref(), )?; - tokio::time::sleep(std::time::Duration::from_secs(interval)).await; + crate::wasm_compat::sleep(std::time::Duration::from_secs(interval)).await; } Err(AuthError::Message( diff --git a/crates/rig-core/src/streaming/mod.rs b/crates/rig-core/src/streaming/mod.rs index ecc39d5d96..646ed59daa 100644 --- a/crates/rig-core/src/streaming/mod.rs +++ b/crates/rig-core/src/streaming/mod.rs @@ -18,41 +18,54 @@ use futures::{Stream, StreamExt}; pub use identity::{MintKind, StreamPartId, SyntheticIds, WireId}; use parts::PartsAccumulator; use serde::{Deserialize, Serialize}; -use std::future::Future; use std::pin::Pin; -use std::sync::atomic::AtomicBool; +use futures::task::AtomicWaker; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use std::task::{Context, Poll}; -use tokio::sync::watch; + +/// Shared pause flag plus the parked consumer's waker. +/// +/// `AtomicWaker` holds a single waker, so this is correct only while one +/// task polls the stream — which `poll_next` taking `Pin<&mut Self>` +/// enforces. A design that shares one control across multiple streams must +/// switch to a multi-waiter primitive instead. +struct PauseState { + paused: AtomicBool, + waker: AtomicWaker, +} /// Control for pausing and resuming a streaming response +#[derive(Clone)] pub struct PauseControl { - pub(crate) paused_tx: watch::Sender, - pub(crate) paused_rx: watch::Receiver, + state: Arc, } impl PauseControl { /// Create a pause controller in the running state. pub fn new() -> Self { - let (paused_tx, paused_rx) = watch::channel(false); Self { - paused_tx, - paused_rx, + state: Arc::new(PauseState { + paused: AtomicBool::new(false), + waker: AtomicWaker::new(), + }), } } /// Pause polling of the public stream until [`PauseControl::resume`] is called. pub fn pause(&self) { - let _ = self.paused_tx.send(true); + self.state.paused.store(true, Ordering::Release); } /// Resume polling after a pause. pub fn resume(&self) { - let _ = self.paused_tx.send(false); + self.state.paused.store(false, Ordering::Release); + self.state.waker.wake(); } /// Returns whether the stream is currently paused. pub fn is_paused(&self) -> bool { - *self.paused_rx.borrow() + self.state.paused.load(Ordering::Acquire) } } @@ -888,16 +901,6 @@ where })) } -#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))] -/// Future a paused [`StreamingCompletionResponse`] parks on until resumed, on -/// native targets. -type ResumeWait = Pin + Send>>; - -#[cfg(all(target_arch = "wasm32", target_os = "unknown"))] -/// Future a paused [`StreamingCompletionResponse`] parks on until resumed, on -/// wasm targets. -type ResumeWait = Pin>>; - /// The response from a streaming completion request; /// message and response are populated at the end of the /// `inner` stream. @@ -923,9 +926,6 @@ pub struct StreamingCompletionResponse { /// stream — which `Stream` permits and combinators do — would otherwise /// replace a fully aggregated `choice` with empty text (#2258 H6). finished: bool, - /// Parked wait on the pause channel while [`PauseControl`] holds the - /// stream paused; `None` whenever the stream is running (#2258 H7). - resume_wait: Option, /// Rig-generated public correlators for reasoning parts, one per /// accumulation key: stable across a part's deltas, unique per run, and /// carrying nothing an accumulation key could leak. @@ -966,7 +966,6 @@ impl StreamingCompletionResponse { // empty text block the model had emitted. choice: Vec::new(), finished: false, - resume_wait: None, reasoning_correlators: std::collections::HashMap::new(), finished_reasoning_correlators: std::collections::HashMap::new(), response: None, @@ -1115,25 +1114,18 @@ impl Stream for StreamingCompletionResponse { } if stream.is_paused() { - // Park on the pause channel rather than re-waking immediately: a - // self-wake turns a pause into a busy poll loop that burns the - // executor for as long as the consumer stays paused (#2258 H7). - // `wait_for` evaluates the *current* value when it is first - // polled, so a resume racing this branch resolves it at once - // instead of parking forever on a notification already sent. - let wait = match stream.resume_wait.as_mut() { - Some(wait) => wait, - None => { - let mut paused_rx = stream.pause_control.paused_rx.clone(); - stream.resume_wait.insert(Box::pin(async move { - let _ = paused_rx.wait_for(|paused| !*paused).await; - })) - } - }; - if wait.as_mut().poll(cx).is_pending() { + // Park rather than re-waking immediately: a self-wake turns a + // pause into a busy poll loop that burns the executor for as long + // as the consumer stays paused (#2258 H7). Register-then-recheck + // is the `AtomicWaker` protocol that also closes the resume race: + // `resume` clears the flag before waking, and this poll registers + // its waker before re-reading the flag, so a resume racing this + // branch either sees the registered waker (and wakes the task) or + // is observed by the re-check below. + stream.pause_control.state.waker.register(cx.waker()); + if stream.is_paused() { return Poll::Pending; } - stream.resume_wait = None; } // Non-yielding events (`continue` arms: block bookkeeping, dropped @@ -2111,7 +2103,7 @@ mod tests { yield Ok(RawStreamingChoice::Message("hello".to_string())); }), ); - let resume = stream.pause_control.paused_tx.clone(); + let resume = stream.pause_control.clone(); stream.pause(); let mut task = tokio_test::task::spawn(stream); @@ -2124,7 +2116,7 @@ mod tests { "a paused stream must idle, not re-wake itself" ); - resume.send(false).expect("resume"); + resume.resume(); assert!(task.is_woken(), "resuming must wake the parked stream"); assert!(matches!( task.poll_next(), diff --git a/crates/rig-core/src/wasm_compat.rs b/crates/rig-core/src/wasm_compat.rs index 213cffd497..b55f09784c 100644 --- a/crates/rig-core/src/wasm_compat.rs +++ b/crates/rig-core/src/wasm_compat.rs @@ -115,6 +115,17 @@ where } } +/// Sleep for `duration`. +/// +/// A cross-platform (native + wasm) replacement for `tokio::time::sleep`, for +/// the same reason as [`timeout`]: rig's `tokio` dependency is built without +/// the `time` feature, and `tokio::time` does not function on wasm. Built on +/// [`futures_timer::Delay`], whose backend selection (background timer thread +/// natively, `setTimeout` on browser wasm) is documented on [`timeout`]. +pub async fn sleep(duration: std::time::Duration) { + futures_timer::Delay::new(duration).await; +} + #[macro_export] macro_rules! if_wasm { ($($tokens:tt)*) => { diff --git a/tests/providers/gemini/cassette/tool_server.rs b/tests/providers/gemini/cassette/tool_server.rs index 565e9e8960..2c5795c3c4 100644 --- a/tests/providers/gemini/cassette/tool_server.rs +++ b/tests/providers/gemini/cassette/tool_server.rs @@ -39,7 +39,7 @@ async fn add_tool_between_turns_appears_in_next_request() { .expect("first prompt should succeed with only the add tool"); assert_mentions_expected_number(&first, 42); - handle.add_tool(subtract).await; + handle.add_tool(subtract); let mut history = Vec::::new(); let second = agent @@ -88,7 +88,7 @@ async fn remove_tool_between_turns_drops_definition() { assert_mentions_expected_number(&first, 42); assert_eq!(add_counter.count(), 1, "add should execute on the first prompt"); - handle.remove_tool("subtract").await; + handle.remove_tool("subtract"); let mut history = Vec::::new(); let second = agent @@ -144,7 +144,7 @@ async fn shared_tool_server_handle_updates_all_agents() { .expect("the first agent should use the shared add tool"); assert_mentions_expected_number(&first, 42); - handle.add_tool(subtract).await; + handle.add_tool(subtract); let mut history = Vec::::new(); let second = second_agent diff --git a/tests/providers/gemini/tools_support.rs b/tests/providers/gemini/tools_support.rs index 80d4604681..00b75fa0cf 100644 --- a/tests/providers/gemini/tools_support.rs +++ b/tests/providers/gemini/tools_support.rs @@ -504,7 +504,7 @@ impl AgentHook for RemoveToolBeforeExecutionHook { event: ToolCallEvent<'_>, ) -> ToolCallAction { if event.tool_name == self.tool_name { - self.handle.remove_tool(self.tool_name).await; + self.handle.remove_tool(self.tool_name); } ToolCallAction::run() } From 52e55b90eab5caba9380316cafd85692c75a767d Mon Sep 17 00:00:00 2001 From: stephen Date: Fri, 21 Aug 2026 13:11:04 -0700 Subject: [PATCH 2/3] style: rustfmt --- crates/rig-agent/src/agent/runner.rs | 8 ++------ crates/rig-agent/src/tool/rmcp.rs | 4 +--- crates/rig-agent/src/tool/server.rs | 10 ++++------ crates/rig-core/src/providers/chatgpt/auth/mod.rs | 2 +- crates/rig-core/src/providers/copilot/auth/mod.rs | 2 +- crates/rig-core/src/streaming/mod.rs | 2 +- 6 files changed, 10 insertions(+), 18 deletions(-) diff --git a/crates/rig-agent/src/agent/runner.rs b/crates/rig-agent/src/agent/runner.rs index 895bea320c..dffc355691 100644 --- a/crates/rig-agent/src/agent/runner.rs +++ b/crates/rig-agent/src/agent/runner.rs @@ -7911,9 +7911,7 @@ mod migrated_tests { let run = runner.run(); let replace = async { request_started.notified().await; - handle - .add_tool(SecondGenerationTool(second_calls.clone())) - ; + handle.add_tool(SecondGenerationTool(second_calls.clone())); release_response.notify_one(); }; let (response, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async { @@ -7969,9 +7967,7 @@ mod migrated_tests { }; let replace = async { request_started.notified().await; - handle - .add_tool(SecondGenerationTool(second_calls.clone())) - ; + handle.add_tool(SecondGenerationTool(second_calls.clone())); release_response.notify_one(); }; let (final_output, ()) = tokio::time::timeout(std::time::Duration::from_secs(2), async { diff --git a/crates/rig-agent/src/tool/rmcp.rs b/crates/rig-agent/src/tool/rmcp.rs index aaf1b34cf2..2806669594 100644 --- a/crates/rig-agent/src/tool/rmcp.rs +++ b/crates/rig-agent/src/tool/rmcp.rs @@ -1787,9 +1787,7 @@ mod migrated_tests { let handle = ToolServer::new().run(); let (client, server_task) = connect(server, handle.clone()).await; - handle - .add_dynamic_tool(make_dynamic_tool("alpha", "Local alpha")) - ; + handle.add_dynamic_tool(make_dynamic_tool("alpha", "Local alpha")); server_control .set_tools(vec![make_tool("refresh_complete", "Refresh sentinel")]) .await; diff --git a/crates/rig-agent/src/tool/server.rs b/crates/rig-agent/src/tool/server.rs index a9d06255ae..551d621f4a 100644 --- a/crates/rig-agent/src/tool/server.rs +++ b/crates/rig-agent/src/tool/server.rs @@ -763,12 +763,10 @@ mod tests { .run(); let snapshot = handle.snapshot_tool_defs(None).await.unwrap(); - handle - .add_tool(ReplacementTool { - description: "second schema", - output: "second implementation", - }) - ; + handle.add_tool(ReplacementTool { + description: "second schema", + output: "second implementation", + }); assert_eq!(snapshot.definitions()[0].description, "first schema"); let dispatch = snapshot diff --git a/crates/rig-core/src/providers/chatgpt/auth/mod.rs b/crates/rig-core/src/providers/chatgpt/auth/mod.rs index 3f16836c79..34aac1c413 100644 --- a/crates/rig-core/src/providers/chatgpt/auth/mod.rs +++ b/crates/rig-core/src/providers/chatgpt/auth/mod.rs @@ -1,9 +1,9 @@ //! Shared ChatGPT authentication types and target-specific dispatch. +use async_lock::Mutex; use std::fmt; use std::path::PathBuf; use std::sync::Arc; -use async_lock::Mutex; pub use crate::providers::internal::auth::{DeviceCodeHandler, DeviceCodePrompt}; diff --git a/crates/rig-core/src/providers/copilot/auth/mod.rs b/crates/rig-core/src/providers/copilot/auth/mod.rs index 5fe55eedb5..3b142df717 100644 --- a/crates/rig-core/src/providers/copilot/auth/mod.rs +++ b/crates/rig-core/src/providers/copilot/auth/mod.rs @@ -1,7 +1,7 @@ +use async_lock::Mutex; use std::fmt; use std::path::PathBuf; use std::sync::Arc; -use async_lock::Mutex; pub use crate::providers::internal::auth::{DeviceCodeHandler, DeviceCodePrompt}; diff --git a/crates/rig-core/src/streaming/mod.rs b/crates/rig-core/src/streaming/mod.rs index 646ed59daa..0172a3d098 100644 --- a/crates/rig-core/src/streaming/mod.rs +++ b/crates/rig-core/src/streaming/mod.rs @@ -14,12 +14,12 @@ use crate::message::{ }; use crate::wasm_compat::WasmCompatSend; use futures::stream::{AbortHandle, Abortable}; +use futures::task::AtomicWaker; use futures::{Stream, StreamExt}; pub use identity::{MintKind, StreamPartId, SyntheticIds, WireId}; use parts::PartsAccumulator; use serde::{Deserialize, Serialize}; use std::pin::Pin; -use futures::task::AtomicWaker; use std::sync::Arc; use std::sync::atomic::{AtomicBool, Ordering}; use std::task::{Context, Poll}; From 16aa2a6e7f1acbba59ce69e9e3b0cf2f6dbb8c7a Mon Sep 17 00:00:00 2001 From: stephen Date: Fri, 21 Aug 2026 13:24:21 -0700 Subject: [PATCH 3/3] refactor: use futures::lock::Mutex for the auth caches instead of async-lock The futures crate is already a dependency and its async mutex is a drop-in for this slow, low-contention path (token refresh), so the PR now adds zero new direct dependencies. --- Cargo.lock | 1 - Cargo.toml | 1 - crates/rig-core/Cargo.toml | 1 - crates/rig-core/src/providers/chatgpt/auth/mod.rs | 2 +- crates/rig-core/src/providers/copilot/auth/mod.rs | 2 +- 5 files changed, 2 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 83e5775924..39daa3dec9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9925,7 +9925,6 @@ dependencies = [ "anyhow", "as-any", "assert_fs", - "async-lock", "async-stream", "base64 0.22.1", "bytes", diff --git a/Cargo.toml b/Cargo.toml index c29484d8ec..2218be38ed 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -110,7 +110,6 @@ anyhow = "1" arrow-array = "58" as-any = "0.3" assert_fs = "1" -async-lock = "3" async-stream = "0.3" axum = "0.8" aws-config = { version = "1", default-features = false } diff --git a/crates/rig-core/Cargo.toml b/crates/rig-core/Cargo.toml index ba584fa968..823a232a20 100644 --- a/crates/rig-core/Cargo.toml +++ b/crates/rig-core/Cargo.toml @@ -25,7 +25,6 @@ doctest = true [dependencies] as-any = { workspace = true } -async-lock = { workspace = true } async-stream = { workspace = true } base64 = { workspace = true } bytes = { workspace = true } diff --git a/crates/rig-core/src/providers/chatgpt/auth/mod.rs b/crates/rig-core/src/providers/chatgpt/auth/mod.rs index 34aac1c413..c400bff173 100644 --- a/crates/rig-core/src/providers/chatgpt/auth/mod.rs +++ b/crates/rig-core/src/providers/chatgpt/auth/mod.rs @@ -1,6 +1,6 @@ //! Shared ChatGPT authentication types and target-specific dispatch. -use async_lock::Mutex; +use futures::lock::Mutex; use std::fmt; use std::path::PathBuf; use std::sync::Arc; diff --git a/crates/rig-core/src/providers/copilot/auth/mod.rs b/crates/rig-core/src/providers/copilot/auth/mod.rs index 3b142df717..58936291f4 100644 --- a/crates/rig-core/src/providers/copilot/auth/mod.rs +++ b/crates/rig-core/src/providers/copilot/auth/mod.rs @@ -1,4 +1,4 @@ -use async_lock::Mutex; +use futures::lock::Mutex; use std::fmt; use std::path::PathBuf; use std::sync::Arc;