diff --git a/src/channels/wasm/host.rs b/src/channels/wasm/host.rs index 89b2a313..03a6170f 100644 --- a/src/channels/wasm/host.rs +++ b/src/channels/wasm/host.rs @@ -300,6 +300,51 @@ impl ChannelHostState { } } +/// In-memory workspace store for WASM channels. +/// +/// Persists workspace writes across callback invocations within a single +/// channel lifetime. This allows WASM channels to maintain state (e.g., +/// Telegram polling offsets) between poll ticks without requiring a +/// full database-backed workspace. +/// +/// Uses `std::sync::RwLock` (not tokio) because WASM execution runs +/// inside `spawn_blocking`. +pub struct ChannelWorkspaceStore { + data: std::sync::RwLock>, +} + +impl ChannelWorkspaceStore { + /// Create a new empty workspace store. + pub fn new() -> Self { + Self { + data: std::sync::RwLock::new(std::collections::HashMap::new()), + } + } + + /// Commit pending writes from a callback execution into the store. + pub fn commit_writes(&self, writes: &[PendingWorkspaceWrite]) { + if writes.is_empty() { + return; + } + if let Ok(mut data) = self.data.write() { + for write in writes { + tracing::debug!( + path = %write.path, + content_len = write.content.len(), + "Committing workspace write to channel store" + ); + data.insert(write.path.clone(), write.content.clone()); + } + } + } +} + +impl crate::tools::wasm::WorkspaceReader for ChannelWorkspaceStore { + fn read(&self, path: &str) -> Option { + self.data.read().ok()?.get(path).cloned() + } +} + /// Rate limiter for channel message emission. /// /// Tracks emission rates across multiple executions. @@ -497,4 +542,56 @@ mod tests { assert_eq!(state.channel_name(), "telegram"); } + + #[test] + fn test_channel_workspace_store_commit_and_read() { + use crate::channels::wasm::host::{ChannelWorkspaceStore, PendingWorkspaceWrite}; + use crate::tools::wasm::WorkspaceReader; + + let store = ChannelWorkspaceStore::new(); + + // Initially empty + assert!(store.read("channels/telegram/offset").is_none()); + + // Commit some writes + let writes = vec![ + PendingWorkspaceWrite { + path: "channels/telegram/offset".to_string(), + content: "103".to_string(), + }, + PendingWorkspaceWrite { + path: "channels/telegram/state.json".to_string(), + content: r#"{"ok":true}"#.to_string(), + }, + ]; + store.commit_writes(&writes); + + // Should be readable + assert_eq!( + store.read("channels/telegram/offset"), + Some("103".to_string()) + ); + assert_eq!( + store.read("channels/telegram/state.json"), + Some(r#"{"ok":true}"#.to_string()) + ); + + // Overwrite a value + let writes2 = vec![PendingWorkspaceWrite { + path: "channels/telegram/offset".to_string(), + content: "200".to_string(), + }]; + store.commit_writes(&writes2); + assert_eq!( + store.read("channels/telegram/offset"), + Some("200".to_string()) + ); + + // Empty writes are a no-op + store.commit_writes(&[]); + assert_eq!( + store.read("channels/telegram/offset"), + Some("200".to_string()) + ); + } } diff --git a/src/channels/wasm/wrapper.rs b/src/channels/wasm/wrapper.rs index 212334a6..e28599a1 100644 --- a/src/channels/wasm/wrapper.rs +++ b/src/channels/wasm/wrapper.rs @@ -42,7 +42,9 @@ use wasmtime_wasi::{ResourceTable, WasiCtx, WasiCtxBuilder, WasiView}; use crate::channels::wasm::capabilities::ChannelCapabilities; use crate::channels::wasm::error::WasmChannelError; -use crate::channels::wasm::host::{ChannelEmitRateLimiter, ChannelHostState, EmittedMessage}; +use crate::channels::wasm::host::{ + ChannelEmitRateLimiter, ChannelHostState, ChannelWorkspaceStore, EmittedMessage, +}; use crate::channels::wasm::router::RegisteredEndpoint; use crate::channels::wasm::runtime::{PreparedChannelModule, WasmChannelRuntime}; use crate::channels::wasm::schema::ChannelConfig; @@ -547,6 +549,10 @@ pub struct WasmChannel { /// Pairing store for DM pairing (guest access control). pairing_store: Arc, + + /// In-memory workspace store persisting writes across callback invocations. + /// Ensures WASM channels can maintain state (e.g., polling offsets) between ticks. + workspace_store: Arc, } impl WasmChannel { @@ -577,6 +583,7 @@ impl WasmChannel { credentials: Arc::new(RwLock::new(HashMap::new())), typing_task: RwLock::new(None), pairing_store, + workspace_store: Arc::new(ChannelWorkspaceStore::new()), } } @@ -634,6 +641,26 @@ impl WasmChannel { self.endpoints.read().await.clone() } + /// Inject the workspace store as the reader into a capabilities clone. + /// + /// Ensures `workspace_read` capability is present with the store as its reader, + /// so WASM callbacks can read previously written workspace state. + fn inject_workspace_reader( + capabilities: &ChannelCapabilities, + store: &Arc, + ) -> ChannelCapabilities { + let mut caps = capabilities.clone(); + let ws_cap = caps + .tool_capabilities + .workspace_read + .get_or_insert_with(|| crate::tools::wasm::WorkspaceCapability { + allowed_prefixes: Vec::new(), + reader: None, + }); + ws_cap.reader = Some(Arc::clone(store) as Arc); + caps + } + /// Add channel host functions to the linker using generated bindings. /// /// Uses the wasmtime::component::bindgen! generated `add_to_linker` function @@ -765,12 +792,13 @@ impl WasmChannel { let runtime = Arc::clone(&self.runtime); let prepared = Arc::clone(&self.prepared); - let capabilities = self.capabilities.clone(); + let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store); let config_json = self.config_json.read().await.clone(); let timeout = self.runtime.config().callback_timeout; let channel_name = self.name.clone(); let credentials = self.get_credentials().await; let pairing_store = self.pairing_store.clone(); + let workspace_store = self.workspace_store.clone(); // Execute in blocking task with timeout let result = tokio::time::timeout(timeout, async move { @@ -801,8 +829,13 @@ impl WasmChannel { } }; - let host_state = + let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); + + // Commit pending workspace writes to the persistent store + let pending_writes = host_state.take_pending_writes(); + workspace_store.commit_writes(&pending_writes); + Ok((config, host_state)) }) .await @@ -897,10 +930,11 @@ impl WasmChannel { let runtime = Arc::clone(&self.runtime); let prepared = Arc::clone(&self.prepared); - let capabilities = self.capabilities.clone(); + let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store); let timeout = self.runtime.config().callback_timeout; let credentials = self.get_credentials().await; let pairing_store = self.pairing_store.clone(); + let workspace_store = self.workspace_store.clone(); // Prepare request data let method = method.to_string(); @@ -940,8 +974,13 @@ impl WasmChannel { .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?; let response = convert_http_response(wit_response); - let host_state = + let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); + + // Commit pending workspace writes to the persistent store + let pending_writes = host_state.take_pending_writes(); + workspace_store.commit_writes(&pending_writes); + Ok((response, host_state)) }) .await @@ -989,11 +1028,12 @@ impl WasmChannel { let runtime = Arc::clone(&self.runtime); let prepared = Arc::clone(&self.prepared); - let capabilities = self.capabilities.clone(); + let capabilities = Self::inject_workspace_reader(&self.capabilities, &self.workspace_store); let timeout = self.runtime.config().callback_timeout; let channel_name = self.name.clone(); let credentials = self.get_credentials().await; let pairing_store = self.pairing_store.clone(); + let workspace_store = self.workspace_store.clone(); // Execute in blocking task with timeout let result = tokio::time::timeout(timeout, async move { @@ -1013,8 +1053,13 @@ impl WasmChannel { .call_on_poll(&mut store) .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?; - let host_state = + let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); + + // Commit pending workspace writes to the persistent store + let pending_writes = host_state.take_pending_writes(); + workspace_store.commit_writes(&pending_writes); + Ok(((), host_state)) }) .await @@ -1501,6 +1546,7 @@ impl WasmChannel { let credentials = self.credentials.clone(); let pairing_store = self.pairing_store.clone(); let callback_timeout = self.runtime.config().callback_timeout; + let workspace_store = self.workspace_store.clone(); tokio::spawn(async move { let mut interval_timer = tokio::time::interval(interval); @@ -1523,6 +1569,7 @@ impl WasmChannel { &credentials, pairing_store.clone(), callback_timeout, + &workspace_store, ).await; match result { @@ -1565,7 +1612,10 @@ impl WasmChannel { /// Execute a single poll callback with a fresh WASM instance. /// - /// Returns any emitted messages from the callback. + /// Returns any emitted messages from the callback. Pending workspace writes + /// are committed to the shared `ChannelWorkspaceStore` so state persists + /// across poll ticks (e.g., Telegram polling offset). + #[allow(clippy::too_many_arguments)] async fn execute_poll( channel_name: &str, runtime: &Arc, @@ -1574,6 +1624,7 @@ impl WasmChannel { credentials: &RwLock>, pairing_store: Arc, timeout: Duration, + workspace_store: &Arc, ) -> Result, WasmChannelError> { // Skip if no WASM bytes (testing mode) if prepared.component_bytes.is_empty() { @@ -1586,9 +1637,10 @@ impl WasmChannel { let runtime = Arc::clone(runtime); let prepared = Arc::clone(prepared); - let capabilities = capabilities.clone(); + let capabilities = Self::inject_workspace_reader(capabilities, workspace_store); let credentials_snapshot = credentials.read().await.clone(); let channel_name_owned = channel_name.to_string(); + let workspace_store = Arc::clone(workspace_store); // Execute in blocking task with timeout let result = tokio::time::timeout(timeout, async move { @@ -1608,8 +1660,13 @@ impl WasmChannel { .call_on_poll(&mut store) .map_err(|e| Self::map_wasm_error(e, &prepared.name, prepared.limits.fuel))?; - let host_state = + let mut host_state = Self::extract_host_state(&mut store, &prepared.name, &capabilities); + + // Commit pending workspace writes to the persistent store + let pending_writes = host_state.take_pending_writes(); + workspace_store.commit_writes(&pending_writes); + Ok(host_state) }) .await @@ -2230,6 +2287,8 @@ mod tests { let credentials = Arc::new(tokio::sync::RwLock::new(std::collections::HashMap::new())); let timeout = std::time::Duration::from_secs(5); + let workspace_store = Arc::new(crate::channels::wasm::host::ChannelWorkspaceStore::new()); + let result = WasmChannel::execute_poll( "poll-test", &runtime, @@ -2238,6 +2297,7 @@ mod tests { &credentials, Arc::new(PairingStore::new()), timeout, + &workspace_store, ) .await;