From 3e5ecded042791fd434d9597256d3d0bae67f488 Mon Sep 17 00:00:00 2001 From: BingCHuanJ Date: Sat, 24 Oct 2026 18:14:48 +0800 Subject: [PATCH 1/2] fix(mobile-web): keep model selection scoped to remote sessions Co-authored-by: OpenBitFun <318544290+bitfun-ai@users.noreply.github.com> --- .github/workflows/ci.yml | 3 +- src/apps/cli/src/account.rs | 210 +++--- src/apps/cli/src/dispatch/mod.rs | 2 + .../desktop/src/api/remote_connect_api.rs | 136 ++-- .../core/src/service/remote_connect/mod.rs | 4 + .../src/remote_connect.rs | 326 ++++++++- .../src/remote_connect/model_write_fence.rs | 184 +++++ .../src/remote_connect/rpc_execution.rs | 189 +++++ .../tests/remote_connect_contracts.rs | 2 + src/mobile-web/AGENTS.md | 1 + src/mobile-web/package.json | 1 + .../src/components/ChatModelControls.tsx | 66 +- src/mobile-web/src/i18n/messages.ts | 12 + src/mobile-web/src/pages/ChatPage.tsx | 147 ++-- .../src/services/RelayHttpClient.ts | 17 +- .../src/services/RemoteSessionManager.ts | 214 +++++- .../tests/fixtures/session-model.tsx | 246 +++++++ .../tests/session-model-browser.test.mjs | 655 ++++++++++++++++++ 18 files changed, 2067 insertions(+), 348 deletions(-) create mode 100644 src/crates/services/services-integrations/src/remote_connect/model_write_fence.rs create mode 100644 src/crates/services/services-integrations/src/remote_connect/rpc_execution.rs create mode 100644 src/mobile-web/tests/fixtures/session-model.tsx create mode 100644 src/mobile-web/tests/session-model-browser.test.mjs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 64f2f54dac..ed76d6ad43 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -700,12 +700,13 @@ jobs: timeout-minutes: 5 run: pnpm --dir src/mobile-web run test:account-browser - - name: Test mobile web composer and rollback layouts + - name: Test mobile web chat controls and layouts if: needs.build-impact.outputs.frontend_required != 'false' timeout-minutes: 3 run: | pnpm --dir src/mobile-web run test:prompt-suggestion-browser pnpm --dir src/mobile-web run test:rollback-browser + pnpm --dir src/mobile-web run test:session-model-browser - name: Upload mobile input layout screenshots if: always() && needs.build-impact.outputs.frontend_required != 'false' diff --git a/src/apps/cli/src/account.rs b/src/apps/cli/src/account.rs index cd8a875e4d..ea12020adf 100644 --- a/src/apps/cli/src/account.rs +++ b/src/apps/cli/src/account.rs @@ -332,22 +332,25 @@ impl CliAccountRoutingHost { relay_url, cancellation: Some(cancelled.clone()), }; - tokio::select! { - biased; - _ = cancelled.changed() => { tracing::debug!("Retired pending RPC after account routing transition"); } - _ = async { - ACTIVE_PEER_FANOUT_OWNER - .scope(fanout_owner, async { - self.handle_current_relay_event( - event, - relay_client, - account_generation, - expected_token, - ) - .await; - }) - .await; - } => {} + let device_rpc = matches!(&event, RelayEvent::DeviceMessageReceived { .. }); + let mut event_cancellation = cancelled.clone(); + let handling = ACTIVE_PEER_FANOUT_OWNER.scope( + fanout_owner, + self.handle_current_relay_event( + event, + relay_client, + account_generation, + expected_token, + &mut event_cancellation, + ), + ); + if device_rpc { + // RPC execution has its own cancellation boundary. Do not drop an + // accepted model write or lose its task-local owner and routing lease. + handling.await; + } else { + let _ = + remote_connect::rpc_execution::run_rpc_stage(&mut cancelled, false, handling).await; } } @@ -357,6 +360,7 @@ impl CliAccountRoutingHost { relay_client: &Arc, account_generation: u64, expected_token: &str, + cancellation: &mut tokio::sync::watch::Receiver, ) { let runtime = match self.runtime() { Ok(runtime) => runtime, @@ -438,54 +442,63 @@ impl CliAccountRoutingHost { encrypted_data, nonce, } => { - let Ok((session, relay_url)) = runtime - .read_account_context_for_generation(account_generation) - .await - else { - return; - }; - if session.token != expected_token - || !self - .routing_loop_is_current(account_generation, relay_client) - .await - { - return; - } - let plaintext = match session - .decrypt_from_peer(&relay_url, &source_device_id, &encrypted_data, &nonce) - .await - { - Ok(plaintext) => plaintext, - Err(error) => { - tracing::warn!("Failed to decrypt device message: {error}"); - return; - } - }; use remote_connect::remote_server::{RemoteCommand, RemoteResponse}; - let command: RemoteCommand = match serde_json::from_str(&plaintext) { - Ok(command) => command, - Err(error) => { - tracing::warn!("Could not parse device command: {error}"); - return; + use remote_connect::rpc_execution::run_rpc_stage; + + let prepare = async { + let Ok((session, relay_url)) = runtime + .read_account_context_for_generation(account_generation) + .await + else { + return None; + }; + if session.token != expected_token + || !self + .routing_loop_is_current(account_generation, relay_client) + .await + { + return None; } + let plaintext = match session + .decrypt_from_peer(&relay_url, &source_device_id, &encrypted_data, &nonce) + .await + { + Ok(plaintext) => plaintext, + Err(error) => { + tracing::warn!("Failed to decrypt device message: {error}"); + return None; + } + }; + let command: RemoteCommand = match serde_json::from_str(&plaintext) { + Ok(command) => command, + Err(error) => { + tracing::warn!("Could not parse device command: {error}"); + return None; + } + }; + tracing::info!( + "Device command received: source={source_device_id} corr={correlation_id}" + ); + let peer_key = match session + .peer_message_key(&relay_url, &source_device_id) + .await + { + Ok(key) => key, + Err(error) => { + tracing::warn!("Failed to resolve peer message key: {error}"); + return None; + } + }; + Some((session, command, peer_key)) }; - tracing::info!( - "Device command received: source={source_device_id} corr={correlation_id}" - ); - let peer_key = match session - .peer_message_key(&relay_url, &source_device_id) - .await - { - Ok(key) => key, - Err(error) => { - tracing::warn!("Failed to resolve peer message key: {error}"); - return; - } + let Some((session, command, peer_key)) = + run_rpc_stage(cancellation, false, prepare).await.flatten() + else { + return; }; - // Host streams are answered from this host's memory for the - // requesting device; every other command goes to the dispatcher. - let hub = self.host_stream_hub().await; - let response = + let model_write = matches!(&command, RemoteCommand::SetSessionModel { .. }); + let execute = async { + let hub = self.host_stream_hub().await; match openbitfun_core::service::remote_connect::handle_host_stream_command( hub.as_ref(), &source_device_id, @@ -499,50 +512,55 @@ impl CliAccountRoutingHost { crate::peer_host::handle_host_invoke(command, args.clone()).await } RemoteCommand::DeviceEvent { event, payload } => { - // A host is also a controller of its peers: stream - // hints addressed to this device feed local subscribers. session.deliver_device_event(&source_device_id, event, payload); crate::peer_host::handle_device_event_command() } other => RemoteServer::new(peer_key).dispatch(other).await, }, - }; - if !self - .routing_loop_is_current(account_generation, relay_client) - .await - { - return; - } - let response_json = serde_json::to_string(&response).unwrap_or_else(|error| { - serde_json::to_string(&RemoteResponse::Error { - message: format!("failed to serialize RPC response: {error}"), - }) - .unwrap_or_else(|_| { - r#"{"resp":"error","message":"serialize failed"}"#.to_string() - }) - }); - let Ok((encrypted_response, response_nonce)) = - encryption::encrypt_to_base64(&peer_key, &response_json) - else { - tracing::warn!("Failed to encrypt RPC response"); + } + }; + let Some(response) = run_rpc_stage(cancellation, model_write, execute).await else { return; }; - let reply_target = if source_device_id == "rpc" { - "rpc" - } else { - source_device_id.as_str() + let reply = async { + if !self + .routing_loop_is_current(account_generation, relay_client) + .await + { + return; + } + let response_json = serde_json::to_string(&response).unwrap_or_else(|error| { + serde_json::to_string(&RemoteResponse::Error { + message: format!("failed to serialize RPC response: {error}"), + }) + .unwrap_or_else(|_| { + r#"{"resp":"error","message":"serialize failed"}"#.to_string() + }) + }); + let Ok((encrypted_response, response_nonce)) = + encryption::encrypt_to_base64(&peer_key, &response_json) + else { + tracing::warn!("Failed to encrypt RPC response"); + return; + }; + let reply_target = if source_device_id == "rpc" { + "rpc" + } else { + source_device_id.as_str() + }; + if let Err(error) = relay_client + .send_device_message( + reply_target, + &correlation_id, + &encrypted_response, + &response_nonce, + ) + .await + { + tracing::warn!("Failed to send RPC response: {error}"); + } }; - if let Err(error) = relay_client - .send_device_message( - reply_target, - &correlation_id, - &encrypted_response, - &response_nonce, - ) - .await - { - tracing::warn!("Failed to send RPC response: {error}"); - } + let _ = run_rpc_stage(cancellation, false, reply).await; } RelayEvent::Disconnected => { tracing::info!("Device routing disconnected"); diff --git a/src/apps/cli/src/dispatch/mod.rs b/src/apps/cli/src/dispatch/mod.rs index 00e3bf4753..c61d8368a3 100644 --- a/src/apps/cli/src/dispatch/mod.rs +++ b/src/apps/cli/src/dispatch/mod.rs @@ -724,6 +724,8 @@ async fn inspect_model_readiness() -> Result { default_models: Default::default(), models_dev_reasoning_catalog: None, reasoning_preset_selection_supported: true, + session_model_write_fence_supported: false, + session_model_write_token: None, session_model_id: None, session_reasoning_preset: None, }); diff --git a/src/apps/desktop/src/api/remote_connect_api.rs b/src/apps/desktop/src/api/remote_connect_api.rs index ab96ce7e7e..c97a3af3a6 100644 --- a/src/apps/desktop/src/api/remote_connect_api.rs +++ b/src/apps/desktop/src/api/remote_connect_api.rs @@ -2855,79 +2855,81 @@ async fn connect_route_devices(route: RelayRoute) -> Result {}, - _ = async move { - // A transition cancels this future before - // taking the write lease. Captured account - // ownership cannot cross into its replacement. + use openbitfun_core::service::remote_connect::rpc_execution::run_rpc_stage; + + // Keep the routing lease until an accepted + // model write finishes, even after retirement. let _routing_effect = routing_effect; - // Host streams are answered from this host's - // memory for the requesting device; every - // other command goes to the local dispatcher. - let hub = route_host_stream_hub(route).await; - let execution = match openbitfun_core::service::remote_connect::handle_host_stream_command( - hub.as_ref(), - &source_device_id, - &cmd, - ) - .await - { - Some(response) => serde_json::to_value(response) - .map_err(anyhow::Error::from), - None => execute_local_remote_command(&cmd, route, &source_device_id).await, + let model_write = + matches!(&cmd, RemoteCommand::SetSessionModel { .. }); + let execute = async { + let hub = route_host_stream_hub(route).await; + match openbitfun_core::service::remote_connect::handle_host_stream_command( + hub.as_ref(), + &source_device_id, + &cmd, + ) + .await + { + Some(response) => serde_json::to_value(response) + .map_err(anyhow::Error::from), + None => execute_local_remote_command(&cmd, route, &source_device_id).await, + } }; - // Returning drops this reply only. The loop - // re-checks ownership at the top of every - // iteration, so a stale connection is still - // retired there — just not from in here. - if !device_routing_owner_is_current(&rpc_owner).await { + let Some(execution) = + run_rpc_stage(&mut cancelled, model_write, execute) + .await + else { return; - } - match execution { - Ok(resp_value) => { - let control_ping = - is_successful_control_ping(&cmd, &resp_value); - let sent = send_rpc_envelope( - &rpc_owner, - &rpc_session, - &source_device_id, - &correlation_id, - resp_value, - ) - .await; - if let Some(generation) = - ping_generation.filter(|_| sent && control_ping) - { - record_control_ping_if_owner( - &rpc_owner, - generation, - ping_received_at, - match &cmd { - RemoteCommand::Ping { client } => { - client.as_ref() - } - _ => None, - }, + }; + let reply = async { + if !device_routing_owner_is_current(&rpc_owner).await { + return; + } + match execution { + Ok(resp_value) => { + let control_ping = is_successful_control_ping( + &cmd, + &resp_value, ); + let sent = send_rpc_envelope( + &rpc_owner, + &rpc_session, + &source_device_id, + &correlation_id, + resp_value, + ) + .await; + if let Some(generation) = ping_generation + .filter(|_| sent && control_ping) + { + record_control_ping_if_owner( + &rpc_owner, + generation, + ping_received_at, + match &cmd { + RemoteCommand::Ping { client } => { + client.as_ref() + } + _ => None, + }, + ); + } + } + Err(e) => { + log::warn!("RPC: execute command failed: {e}"); + send_rpc_error( + &rpc_owner, + &rpc_session, + &source_device_id, + &correlation_id, + format!("RPC execute failed: {e}"), + ) + .await; } } - Err(e) => { - log::warn!("RPC: execute command failed: {e}"); - send_rpc_error( - &rpc_owner, - &rpc_session, - &source_device_id, - &correlation_id, - format!("RPC execute failed: {e}"), - ) - .await; - } - } - - } => {} - } + }; + let _ = run_rpc_stage(&mut cancelled, false, reply).await; }); } Err(e) => { diff --git a/src/crates/assembly/core/src/service/remote_connect/mod.rs b/src/crates/assembly/core/src/service/remote_connect/mod.rs index e8166e3ef9..4262e4c603 100644 --- a/src/crates/assembly/core/src/service/remote_connect/mod.rs +++ b/src/crates/assembly/core/src/service/remote_connect/mod.rs @@ -51,6 +51,10 @@ pub mod host_stream_subscriber { pub use openbitfun_services_integrations::remote_connect::host_stream_subscriber::*; } +pub mod rpc_execution { + pub use openbitfun_services_integrations::remote_connect::rpc_execution::*; +} + pub mod session_records { pub use openbitfun_services_integrations::remote_connect::session_records::*; } diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index dece1716f0..e7b0aad2c2 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -17,6 +17,7 @@ pub mod file_projection; pub mod host_stream; pub mod host_stream_subscriber; mod lan; +mod model_write_fence; mod page_upload; pub mod pairing; pub mod qr_generator; @@ -25,6 +26,7 @@ mod realtime_payload; pub mod relay_client; pub mod relay_failure; mod relay_http; +pub mod rpc_execution; pub mod session_records; pub mod session_store; @@ -1648,9 +1650,43 @@ where Err(message) => RemoteResponse::Error { message }, } } - RemoteCommand::GetModelCatalog { session_id } => { + RemoteCommand::GetModelCatalog { + session_id, + prepare_model_write, + model_write_fence, + } => { + // Only a targeted barrier retires a prepared write. Ordinary readers + // wait for active commits without revoking other controllers' intents. + if session_id.is_none() && (*prepare_model_write || model_write_fence.is_some()) { + return RemoteResponse::Error { + message: "A model write barrier requires a session".to_string(), + }; + } + let gate = match session_id + .as_deref() + .map(model_write_fence::session) + .transpose() + { + Ok(gate) => gate, + Err(message) => return RemoteResponse::Error { message }, + }; + let mut fence = match gate.as_ref() { + Some(gate) => Some(gate.lock().await), + None => None, + }; + if let (Some(fence), Some(token)) = (fence.as_mut(), model_write_fence.as_deref()) { + fence.retire(token); + } match host.load_model_catalog(session_id.as_deref()).await { - Ok(catalog) => RemoteResponse::ModelCatalog { catalog }, + Ok(mut catalog) => { + catalog.session_model_write_fence_supported = true; + catalog.session_model_write_token = if *prepare_model_write { + fence.as_mut().map(|fence| fence.issue()) + } else { + None + }; + RemoteResponse::ModelCatalog { catalog } + } Err(message) => RemoteResponse::Error { message }, } } @@ -1658,17 +1694,30 @@ where session_id, model_id, reasoning_preset, - } => match host - .update_session_model_selection( - session_id, - model_id, - reasoning_preset.as_ref().map(|preset| preset.as_deref()), - ) - .await - { - Ok(selection) => remote_session_model_updated_response(session_id.clone(), selection), - Err(message) => RemoteResponse::Error { message }, - }, + model_write_token, + } => { + let gate = match model_write_fence::session(session_id) { + Ok(gate) => gate, + Err(message) => return RemoteResponse::Error { message }, + }; + let mut fence = gate.lock().await; + if let Err(message) = fence.claim(model_write_token.as_deref()) { + return RemoteResponse::Error { message }; + } + match host + .update_session_model_selection( + session_id, + model_id, + reasoning_preset.as_ref().map(|preset| preset.as_deref()), + ) + .await + { + Ok(selection) => { + remote_session_model_updated_response(session_id.clone(), selection) + } + Err(message) => RemoteResponse::Error { message }, + } + } RemoteCommand::UpdateSessionTitle { session_id, title } => { if let Err(message) = host.ensure_session_loaded(session_id).await { return RemoteResponse::Error { message }; @@ -1983,6 +2032,11 @@ pub struct RemoteModelCatalog { pub default_models: RemoteDefaultModelsConfig, #[serde(default)] pub reasoning_preset_selection_supported: bool, + /// Offers per-intent preparation and targeted barriers on catalog reads. + #[serde(default)] + pub session_model_write_fence_supported: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub session_model_write_token: Option, #[serde(skip_serializing_if = "Option::is_none")] pub session_model_id: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -2094,6 +2148,8 @@ pub fn build_remote_model_catalog(facts: RemoteModelCatalogFacts) -> RemoteModel models_dev_reasoning_catalog: facts.models_dev_reasoning_catalog, default_models: facts.default_models, reasoning_preset_selection_supported: true, + session_model_write_fence_supported: false, + session_model_write_token: None, session_model_id: facts.session_model_id, session_reasoning_preset: facts.session_reasoning_preset, } @@ -2706,12 +2762,18 @@ pub enum RemoteCommand { }, GetModelCatalog { session_id: Option, + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + prepare_model_write: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + model_write_fence: Option, }, SetSessionModel { session_id: String, model_id: String, #[serde(default, deserialize_with = "deserialize_present_nullable")] reasoning_preset: Option>, + #[serde(default, skip_serializing_if = "Option::is_none")] + model_write_token: Option, }, UpdateSessionTitle { session_id: String, @@ -4739,6 +4801,7 @@ mod tests { created_requests: Mutex>, list_identities: Mutex>, model_updates: Mutex>)>>, + model_update_barrier: Option>, removed_trackers: Mutex>, history_error: Option, rollback_requests: Mutex)>>, @@ -4867,8 +4930,16 @@ mod tests { async fn load_model_catalog( &self, - _session_id: Option<&str>, + session_id: Option<&str>, ) -> Result { + let selection = self + .model_updates + .lock() + .unwrap() + .iter() + .rev() + .find(|(id, _, _)| Some(id.as_str()) == session_id) + .cloned(); Ok(RemoteModelCatalog { version: 1, models: Vec::new(), @@ -4876,8 +4947,10 @@ mod tests { models_dev_reasoning_catalog: None, default_models: RemoteDefaultModelsConfig::default(), reasoning_preset_selection_supported: true, - session_model_id: None, - session_reasoning_preset: None, + session_model_write_fence_supported: false, + session_model_write_token: None, + session_model_id: selection.as_ref().map(|(_, model, _)| model.clone()), + session_reasoning_preset: selection.and_then(|(_, _, preset)| preset.flatten()), }) } @@ -4887,6 +4960,10 @@ mod tests { model_id: &str, reasoning_preset: Option>, ) -> Result { + if let Some(barrier) = &self.model_update_barrier { + barrier.0.notify_one(); + barrier.1.notified().await; + } self.model_updates.lock().unwrap().push(( session_id.to_string(), model_id.to_string(), @@ -5198,6 +5275,7 @@ mod tests { session_id: "session-a".to_string(), model_id: "model-1".to_string(), reasoning_preset: Some(Some("high".to_string())), + model_write_token: None, }, ) .await; @@ -5220,6 +5298,222 @@ mod tests { ); } + async fn fenced_catalog(host: &FakeSessionHost, session_id: &str) -> RemoteModelCatalog { + match handle_remote_session_command( + host, + &RemoteCommand::GetModelCatalog { + session_id: Some(session_id.to_string()), + prepare_model_write: true, + model_write_fence: None, + }, + ) + .await + { + RemoteResponse::ModelCatalog { catalog } => catalog, + response => panic!("Unexpected catalog response: {response:?}"), + } + } + + #[tokio::test] + async fn remote_session_model_fence_rejects_delayed_and_replayed_wire_writes() { + let host = FakeSessionHost::default(); + let session_id = uuid::Uuid::new_v4().to_string(); + let old = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let current = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let barrier: RemoteCommand = serde_json::from_value(serde_json::json!({ + "cmd": "get_model_catalog", "session_id": session_id, "model_write_fence": old, + })) + .unwrap(); + assert!(matches!( + handle_remote_session_command(&host, &barrier).await, + RemoteResponse::ModelCatalog { .. } + )); + let wire_write = |token: &str| { + serde_json::from_value::(serde_json::json!({ + "cmd": "set_session_model", "session_id": session_id, + "model_id": "model-c", "reasoning_preset": null, "model_write_token": token, + })) + .unwrap() + }; + assert!(matches!( + handle_remote_session_command(&host, &wire_write(&old)).await, + RemoteResponse::Error { .. } + )); + assert!( + host.model_updates.lock().unwrap().is_empty(), + "a retired request must never reach the model writer" + ); + assert!(matches!( + handle_remote_session_command(&host, &wire_write(¤t)).await, + RemoteResponse::SessionModelUpdated { .. } + )); + assert!(matches!( + handle_remote_session_command(&host, &wire_write(¤t)).await, + RemoteResponse::Error { .. } + )); + assert_eq!(host.model_updates.lock().unwrap().len(), 1); + assert_eq!( + fenced_catalog(&host, &session_id) + .await + .session_model_id + .as_deref(), + Some("model-c") + ); + } + + #[tokio::test] + async fn remote_session_model_fence_read_waits_for_the_real_host_write() { + let barrier = Arc::new((tokio::sync::Notify::new(), tokio::sync::Notify::new())); + let host = Arc::new(FakeSessionHost { + model_update_barrier: Some(barrier.clone()), + ..Default::default() + }); + let session_id = uuid::Uuid::new_v4().to_string(); + let token = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let command = RemoteCommand::SetSessionModel { + session_id: session_id.clone(), + model_id: "model-c".to_string(), + reasoning_preset: Some(None), + model_write_token: Some(token), + }; + let writing_host = host.clone(); + let write = tokio::spawn(async move { + handle_remote_session_command(writing_host.as_ref(), &command).await + }); + barrier.0.notified().await; + let read = fenced_catalog(&host, &session_id); + tokio::pin!(read); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!( + std::future::Future::poll(read.as_mut(), &mut context).is_pending(), + "a transport timeout must not let a read overtake an executing write" + ); + barrier.1.notify_one(); + assert!(matches!( + write.await.unwrap(), + RemoteResponse::SessionModelUpdated { .. } + )); + assert_eq!(read.await.session_model_id.as_deref(), Some("model-c")); + } + + #[tokio::test] + async fn remote_session_model_fence_plain_read_preserves_other_preparations() { + let host = FakeSessionHost::default(); + let session_id = uuid::Uuid::new_v4().to_string(); + let first = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let second = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let read: RemoteCommand = serde_json::from_value(serde_json::json!({ + "cmd": "get_model_catalog", "session_id": session_id, + })) + .unwrap(); + match handle_remote_session_command(&host, &read).await { + RemoteResponse::ModelCatalog { catalog } => { + assert!(catalog.session_model_write_fence_supported); + assert!(catalog.session_model_write_token.is_none()); + } + response => panic!("Unexpected catalog response: {response:?}"), + } + for token in [first, second] { + let command = RemoteCommand::SetSessionModel { + session_id: session_id.clone(), + model_id: "model-c".to_string(), + reasoning_preset: None, + model_write_token: Some(token), + }; + assert!(matches!( + handle_remote_session_command(&host, &command).await, + RemoteResponse::SessionModelUpdated { .. } + )); + } + assert_eq!(host.model_updates.lock().unwrap().len(), 2); + } + + #[tokio::test] + async fn remote_session_model_fence_requires_session_scope() { + let host = FakeSessionHost::default(); + for (prepare_model_write, model_write_fence) in + [(true, None), (false, Some("token".to_string()))] + { + assert!(matches!( + handle_remote_session_command( + &host, + &RemoteCommand::GetModelCatalog { + session_id: None, + prepare_model_write, + model_write_fence, + } + ) + .await, + RemoteResponse::Error { .. } + )); + } + } + + #[tokio::test] + async fn remote_session_model_fence_cancelled_execution_drains_before_read_and_replacement() { + let barrier = Arc::new((tokio::sync::Notify::new(), tokio::sync::Notify::new())); + let host = Arc::new(FakeSessionHost { + model_update_barrier: Some(barrier.clone()), + ..Default::default() + }); + let session_id = uuid::Uuid::new_v4().to_string(); + let token = fenced_catalog(&host, &session_id) + .await + .session_model_write_token + .unwrap(); + let command = RemoteCommand::SetSessionModel { + session_id: session_id.clone(), + model_id: "model-c".to_string(), + reasoning_preset: Some(None), + model_write_token: Some(token), + }; + let lifecycle = Arc::new(tokio::sync::RwLock::new(())); + let lease = lifecycle.clone().read_owned().await; + let (cancel, mut cancelled) = tokio::sync::watch::channel(0); + let writing_host = host.clone(); + let write = tokio::spawn(async move { + let _lease = lease; + rpc_execution::run_rpc_stage( + &mut cancelled, + true, + handle_remote_session_command(writing_host.as_ref(), &command), + ) + .await + }); + barrier.0.notified().await; + cancel.send_replace(1); + let read = fenced_catalog(&host, &session_id); + tokio::pin!(read); + let replacement = lifecycle.write(); + tokio::pin!(replacement); + let mut context = std::task::Context::from_waker(std::task::Waker::noop()); + assert!(std::future::Future::poll(read.as_mut(), &mut context).is_pending()); + assert!(std::future::Future::poll(replacement.as_mut(), &mut context).is_pending()); + barrier.1.notify_one(); + assert!( + write.await.unwrap().is_none(), + "a cancelled owner cannot reply" + ); + assert_eq!(read.await.session_model_id.as_deref(), Some("model-c")); + let _replacement = replacement.await; + assert_eq!(host.model_updates.lock().unwrap().len(), 1); + } + #[tokio::test] async fn remote_session_handler_removes_tracker_after_delete_success() { let host = FakeSessionHost::default(); diff --git a/src/crates/services/services-integrations/src/remote_connect/model_write_fence.rs b/src/crates/services/services-integrations/src/remote_connect/model_write_fence.rs new file mode 100644 index 0000000000..a944d52b94 --- /dev/null +++ b/src/crates/services/services-integrations/src/remote_connect/model_write_fence.rs @@ -0,0 +1,184 @@ +//! Targeted barriers fence delayed model writes without revoking other readers. + +use std::collections::{HashMap, VecDeque}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::Instant; +use tokio::sync::Mutex as AsyncMutex; + +const MAX_SESSION_FENCES: usize = 256; +const MAX_PREPARED_WRITES: usize = 64; + +#[derive(Default)] +pub(super) struct ModelWriteFence { + tokens: VecDeque, +} + +impl ModelWriteFence { + pub(super) fn issue(&mut self) -> String { + let token = uuid::Uuid::new_v4().to_string(); + if self.tokens.len() >= MAX_PREPARED_WRITES { + self.tokens.pop_front(); + } + self.tokens.push_back(token.clone()); + token + } + + pub(super) fn retire(&mut self, token: &str) -> bool { + if let Some(index) = self.tokens.iter().position(|candidate| candidate == token) { + self.tokens.remove(index); + return true; + } + false + } + + pub(super) fn claim(&mut self, token: Option<&str>) -> Result<(), String> { + match token { + Some(token) if !self.retire(token) => Err( + "Session model selection expired; reload the session model before retrying" + .to_string(), + ), + Some(_) => Ok(()), + None => { + // Legacy writes invalidate outstanding preparations, but an + // ordinary catalog read never revokes another caller's intent. + self.tokens.clear(); + Ok(()) + } + } + } +} + +struct Entry { + gate: Arc>, + last_used: Instant, +} + +struct Registry { + entries: Mutex>, + capacity: usize, +} + +impl Registry { + fn new(capacity: usize) -> Self { + Self { + entries: Mutex::new(HashMap::new()), + capacity, + } + } + + fn session(&self, session_id: &str) -> Result>, String> { + if session_id.is_empty() || session_id.len() > 512 { + return Err("Invalid session identity for model selection".to_string()); + } + let mut entries = self + .entries + .lock() + .map_err(|_| "Session model fence unavailable".to_string())?; + if let Some(entry) = entries.get_mut(session_id) { + entry.last_used = Instant::now(); + return Ok(entry.gate.clone()); + } + if entries.len() >= self.capacity { + // A held or queued guard owns another Arc. Never evict it: a second + // gate for that session would let a read overtake an active write. + let idle = entries + .iter() + .filter(|(_, entry)| Arc::strong_count(&entry.gate) == 1) + .min_by_key(|(_, entry)| entry.last_used) + .map(|(id, _)| id.clone()); + if let Some(id) = idle { + entries.remove(&id); + } else { + return Err("Session model settings are busy; retry later".to_string()); + } + } + let gate = Arc::new(AsyncMutex::new(ModelWriteFence::default())); + entries.insert( + session_id.to_string(), + Entry { + gate: gate.clone(), + last_used: Instant::now(), + }, + ); + Ok(gate) + } +} + +pub(super) fn session(session_id: &str) -> Result>, String> { + static REGISTRY: OnceLock = OnceLock::new(); + REGISTRY + .get_or_init(|| Registry::new(MAX_SESSION_FENCES)) + .session(session_id) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn a_targeted_barrier_retires_only_its_write_and_tokens_are_single_use() { + let mut fence = ModelWriteFence::default(); + let old = fence.issue(); + let current = fence.issue(); + assert!(fence.retire(&old)); + assert!(fence.claim(Some(&old)).is_err()); + assert!(fence.claim(Some(¤t)).is_ok()); + assert!(fence.claim(Some(¤t)).is_err()); + } + + #[test] + fn prepared_writes_are_independent_and_bounded() { + let mut fence = ModelWriteFence::default(); + let first = fence.issue(); + let second = fence.issue(); + assert!(fence.claim(Some(&first)).is_ok()); + assert!(fence.claim(Some(&second)).is_ok()); + let retired = fence.issue(); + for _ in 0..MAX_PREPARED_WRITES { + fence.issue(); + } + assert_eq!(fence.tokens.len(), MAX_PREPARED_WRITES); + assert!(fence.claim(Some(&retired)).is_err()); + } + + #[test] + fn legacy_writes_retire_outstanding_tokens() { + let mut fence = ModelWriteFence::default(); + let token = fence.issue(); + fence.claim(None).unwrap(); + assert!(fence.claim(Some(&token)).is_err()); + } + + #[tokio::test] + async fn a_catalog_read_waits_for_the_actual_write_not_its_transport_owner() { + let registry = Registry::new(1); + let gate = registry.session("A").unwrap(); + let mut write = gate.lock().await; + let token = write.issue(); + write.claim(Some(&token)).unwrap(); + let reading_gate = registry.session("A").unwrap(); + assert!(reading_gate.try_lock().is_err()); + assert!( + registry.session("B").is_err(), + "an active gate cannot be evicted" + ); + drop(write); + let fresh = reading_gate.lock().await.issue(); + assert!(gate.lock().await.claim(Some(&token)).is_err()); + assert!(gate.lock().await.claim(Some(&fresh)).is_ok()); + } + + #[tokio::test] + async fn eviction_cannot_revalidate_a_delayed_write() { + let registry = Registry::new(1); + let token = registry.session("A").unwrap().lock().await.issue(); + drop(registry.session("B").unwrap()); + assert!(registry + .session("A") + .unwrap() + .lock() + .await + .claim(Some(&token)) + .is_err()); + } +} diff --git a/src/crates/services/services-integrations/src/remote_connect/rpc_execution.rs b/src/crates/services/services-integrations/src/remote_connect/rpc_execution.rs new file mode 100644 index 0000000000..69c93848b0 --- /dev/null +++ b/src/crates/services/services-integrations/src/remote_connect/rpc_execution.rs @@ -0,0 +1,189 @@ +//! Cancellation boundaries for account-owned remote RPC execution. + +use std::future::Future; +use std::sync::atomic::{AtomicBool, Ordering}; +use tokio::sync::watch; + +/// Runs one RPC stage while its caller retains the current routing lease. +/// +/// Only model mutation execution may finish after cancellation. Preparation, +/// ownership revalidation and replies must remain separate cancellable stages. +/// A mutation accepted by its first poll is drained in this task, preserving +/// task-local ownership and keeping the lease until its commit or failure ends. +pub async fn run_rpc_stage( + cancellation: &mut watch::Receiver, + finish_started: bool, + operation: F, +) -> Option +where + F: Future, +{ + if cancellation.has_changed().unwrap_or(true) { + return None; + } + let admitted = AtomicBool::new(false); + let checkpoint = cancellation.clone(); + let execution = async { + if checkpoint.has_changed().unwrap_or(true) { + return None; + } + admitted.store(true, Ordering::Relaxed); + Some(operation.await) + }; + tokio::pin!(execution); + tokio::select! { + biased; + _ = cancellation.changed() => { + if finish_started && admitted.load(Ordering::Relaxed) { + let _ = execution.await; + } + None + } + result = &mut execution => result, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + use tokio::sync::{Notify, RwLock}; + + tokio::task_local! { + static OWNER: u64; + } + + #[tokio::test] + async fn cancellation_before_admission_never_executes() { + let (cancel, mut cancelled) = watch::channel(0); + cancel.send_replace(1); + let executed = AtomicBool::new(false); + assert_eq!( + run_rpc_stage(&mut cancelled, true, async { + executed.store(true, Ordering::Relaxed); + }) + .await, + None + ); + assert!(!executed.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn a_closed_owner_never_executes() { + let (cancel, mut cancelled) = watch::channel(0); + drop(cancel); + let executed = AtomicBool::new(false); + assert_eq!( + run_rpc_stage(&mut cancelled, true, async { + executed.store(true, Ordering::Relaxed); + }) + .await, + None + ); + assert!(!executed.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn completed_execution_returns_its_result() { + let (_cancel, mut cancelled) = watch::channel(0); + assert_eq!( + run_rpc_stage(&mut cancelled, true, async { 42 }).await, + Some(42) + ); + } + + async fn cancelled_write_retains_lease_and_owner(fails: bool) { + let (cancel, mut cancelled) = watch::channel(0); + let lifecycle = Arc::new(RwLock::new(())); + let lease = lifecycle.clone().read_owned().await; + let admitted = Arc::new(Notify::new()); + let finish = Arc::new(Notify::new()); + let committed = Arc::new(AtomicBool::new(false)); + let task_admitted = admitted.clone(); + let task_finish = finish.clone(); + let task_committed = committed.clone(); + let task = tokio::spawn(OWNER.scope(7, async move { + let _lease = lease; + let result = run_rpc_stage(&mut cancelled, true, async { + assert_eq!(OWNER.get(), 7); + task_admitted.notify_one(); + task_finish.notified().await; + assert_eq!(OWNER.get(), 7); + task_committed.store(true, Ordering::Relaxed); + if fails { + Err("commit failed") + } else { + Ok(()) + } + }) + .await; + assert!( + result.is_none(), + "a retired RPC cannot proceed to its reply" + ); + })); + admitted.notified().await; + cancel.send_replace(1); + // A replacement owner must wait, even after cancellation was observed. + let replacement = lifecycle.write(); + tokio::pin!(replacement); + assert!( + std::future::poll_fn(|cx| std::task::Poll::Ready( + replacement.as_mut().poll(cx).is_pending() + )) + .await + ); + assert!(!committed.load(Ordering::Relaxed)); + finish.notify_one(); + task.await.unwrap(); + let _replacement = replacement.await; + assert!(committed.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn cancelled_model_write_finishes_before_replacement_without_replying() { + cancelled_write_retains_lease_and_owner(false).await; + } + + #[tokio::test] + async fn cancelled_failed_write_finishes_before_replacement_without_replying() { + cancelled_write_retains_lease_and_owner(true).await; + } + + #[tokio::test] + async fn retirement_after_execution_suppresses_ownership_checks_and_reply() { + let (cancel, mut cancelled) = watch::channel(0); + let result = run_rpc_stage(&mut cancelled, true, async { + cancel.send_replace(1); + 42 + }) + .await; + assert_eq!(result, Some(42)); + let replied = AtomicBool::new(false); + assert_eq!( + run_rpc_stage(&mut cancelled, false, async { + replied.store(true, Ordering::Relaxed); + }) + .await, + None + ); + assert!(!replied.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn other_stages_remain_cancellable() { + let (cancel, mut cancelled) = watch::channel(0); + let admitted = Arc::new(Notify::new()); + let task_admitted = admitted.clone(); + let task = tokio::spawn(async move { + run_rpc_stage(&mut cancelled, false, async { + task_admitted.notify_one(); + std::future::pending::<()>().await; + }) + .await + }); + admitted.notified().await; + cancel.send_replace(1); + assert_eq!(task.await.unwrap(), None); + } +} diff --git a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs index 64778cdc32..34e7c2039f 100644 --- a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs +++ b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs @@ -2965,6 +2965,8 @@ fn sample_remote_model_catalog(version: u64) -> RemoteModelCatalog { ..RemoteDefaultModelsConfig::default() }, reasoning_preset_selection_supported: true, + session_model_write_fence_supported: false, + session_model_write_token: None, session_model_id: Some("model-1".to_string()), session_reasoning_preset: None, } diff --git a/src/mobile-web/AGENTS.md b/src/mobile-web/AGENTS.md index eebab366f5..62fbd35dd8 100644 --- a/src/mobile-web/AGENTS.md +++ b/src/mobile-web/AGENTS.md @@ -59,6 +59,7 @@ pnpm --dir src/mobile-web run test:host-queue # idempotent outbox and real brows pnpm --dir src/mobile-web run test:host-stream # HostStream reader: hints, reconnects, host restart, unsupported hosts; catalog subscription; next-input suggestions pnpm --dir src/mobile-web run test:prompt-suggestion-browser # Real composer and stream projection: compact/wide, original text, offline/sending gates, no automatic submission pnpm --dir src/mobile-web run test:rollback-browser # Real chat edit/rollback sheets: full-width fields, long text, live resize, reduced viewport, cancel without host mutations +pnpm --dir src/mobile-web run test:session-model-browser # Per-session host model/reasoning selection, legacy hosts, read-only navigation and stale responses pnpm --dir src/mobile-web run test:account-login # account login without an online desktop pnpm --dir src/mobile-web run test:account-browser # real Chrome tabs, persistence, migration, races; simulated Relay pnpm --dir src/mobile-web run test:images # image preparation and upload limits diff --git a/src/mobile-web/package.json b/src/mobile-web/package.json index c3a9ed81a8..1ae5680fbb 100644 --- a/src/mobile-web/package.json +++ b/src/mobile-web/package.json @@ -23,6 +23,7 @@ "test:host-stream": "node --test tests/host-stream.test.mjs tests/host-catalog-subscription.test.mjs tests/prompt-suggestion.test.mjs", "test:prompt-suggestion-browser": "node --test tests/prompt-suggestion-browser.test.mjs", "test:rollback-browser": "node --test tests/rollback-browser.test.mjs", + "test:session-model-browser": "node --test tests/session-model-browser.test.mjs", "test:session-stream-browser": "node --test tests/session-stream-browser.test.mjs", "test:runtime-files": "node --test tests/runtime-file-upload.test.mjs tests/runtime-file-download.test.mjs", "test:interaction-mailbox": "node --test tests/interaction-mailbox-browser.test.mjs", diff --git a/src/mobile-web/src/components/ChatModelControls.tsx b/src/mobile-web/src/components/ChatModelControls.tsx index a1d5a1d22f..4a3b0da43d 100644 --- a/src/mobile-web/src/components/ChatModelControls.tsx +++ b/src/mobile-web/src/components/ChatModelControls.tsx @@ -5,8 +5,6 @@ import { useI18n } from '../i18n'; import type { RemoteModelCatalog, RemoteModelConfig } from '../services/RemoteSessionManager'; import ComposerSelectionSurface from './ComposerSelectionSurface'; -const MOBILE_LAST_SELECTED_MODEL_ID_KEY = 'openbitfun.mobile.last_selected_model_id'; - const SparklesIcon: React.FC<{ className?: string; size?: number }> = ({ className, size = 10 }) => (