fix: persist WASM channel workspace writes across callbacks (#264)

* fix: persist WASM channel workspace writes across callbacks

WASM channel callbacks (polling, webhooks, on_start) call
workspace_write() to persist state, but the host code never committed
these writes — take_pending_writes() was never called. Additionally,
no WorkspaceReader was injected into channel capabilities, so
workspace_read() always returned None.

This caused Telegram's polling offset to reset to 0 on every tick,
making getUpdates re-deliver already-processed messages and producing
2-4 duplicate LLM responses per user message.

Add ChannelWorkspaceStore (Arc-wrapped HashMap with std::sync::RwLock)
that persists across callback invocations within a channel's lifetime.
Inject it as the WorkspaceReader and commit pending writes after every
callback execution (on_start, on_poll, on_http_request, execute_poll).

Co-Authored-By: Claude Opus 4.6 <[email protected]>

* style: fix formatting

Co-Authored-By: Claude Opus 4.6 <[email protected]>

---------

Co-authored-by: Claude Opus 4.6 <[email protected]>
This commit is contained in:
firat.sertgoz
2026-02-20 15:52:21 +00:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 3829d81269
commit 7df356c109
2 changed files with 167 additions and 10 deletions
+97
View File
@@ -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<std::collections::HashMap<String, String>>,
}
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<String> {
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())
);
}
}
+70 -10
View File
@@ -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<PairingStore>,
/// In-memory workspace store persisting writes across callback invocations.
/// Ensures WASM channels can maintain state (e.g., polling offsets) between ticks.
workspace_store: Arc<ChannelWorkspaceStore>,
}
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<ChannelWorkspaceStore>,
) -> 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<dyn crate::tools::wasm::WorkspaceReader>);
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<WasmChannelRuntime>,
@@ -1574,6 +1624,7 @@ impl WasmChannel {
credentials: &RwLock<HashMap<String, String>>,
pairing_store: Arc<PairingStore>,
timeout: Duration,
workspace_store: &Arc<ChannelWorkspaceStore>,
) -> Result<Vec<EmittedMessage>, 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;