feat: add nopoll feature flag to stop remove pollable when using async runtime

iambenkay · iambenkay · commit 4ee8617a1ca2 · 2025-06-13T21:54:13.000+01:00
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,6 +17,8 @@ lto = true
 opt-level = 's'
 
 [workspace.dependencies]
+golem-llm = { path = "llm", version = "0.0.0", default-features = false }
+
 golem-rust = "1.6.0"
 log = "0.4.27"
 reqwest = { git = "https://github.com/golemcloud/reqwest", branch = "update-may-2025", features = [
diff --git a/llm-anthropic/Cargo.toml b/llm-anthropic/Cargo.toml
@@ -16,11 +16,11 @@ default = ["durability"]
 durability = ["golem-rust/durability", "golem-llm/durability"]
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 
 golem-rust = { workspace = true }
 log = { workspace = true }
-reqwest =  { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 wit-bindgen-rt = { workspace = true }
@@ -40,4 +40,4 @@ path = "wit"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit/deps/golem-llm" }
-"wasi:io" = { path = "wit/deps/wasi:io"}
+"wasi:io" = { path = "wit/deps/wasi:io" }
diff --git a/llm-anthropic/wit/anthropic.wit b/llm-anthropic/wit/anthropic.wit
diff --git a/llm-bedrock/Cargo.toml b/llm-bedrock/Cargo.toml
@@ -8,7 +8,7 @@ repository = "https://github.com/golemcloud/golem-llm"
 description = "WebAssembly component for working with Amazon Bedrock APIs, with special support for Golem Cloud"
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 
 # AWS SDK crates
 aws-config = { version = "1.5.19", default-features = false, features = [
@@ -38,7 +38,11 @@ path = "src/lib.rs"
 
 [features]
 default = ["durability"]
-durability = ["golem-rust/durability", "golem-llm/durability"]
+durability = [
+    "golem-rust/durability",
+    "golem-llm/durability",
+    "golem-llm/nopoll",
+]
 
 [package.metadata.component]
 package = "golem:llm-bedrock"
diff --git a/llm-bedrock/src/lib.rs b/llm-bedrock/src/lib.rs
@@ -62,10 +62,6 @@ impl ExtendedGuest for BedrockComponent {
 
         bedrock.unwrap().converse_stream(messages, config)
     }
-
-    fn subscribe(stream: &Self::ChatStream) -> golem_rust::wasm_rpc::Pollable {
-        stream.subscribe()
-    }
 }
 
 type DurableBedrockComponent = DurableLLM<BedrockComponent>;
diff --git a/llm-bedrock/src/stream.rs b/llm-bedrock/src/stream.rs
@@ -3,7 +3,6 @@ use aws_sdk_bedrockruntime::{
     types::error::ConverseStreamOutputError,
 };
 use golem_llm::golem::llm::llm;
-use golem_rust::wasm_rpc::wasi::io::poll::Pollable;
 use std::cell::{RefCell, RefMut};
 
 use crate::{
@@ -52,10 +51,6 @@ impl BedrockChatStream {
     fn set_finished(&self) {
         *self.finished.borrow_mut() = true;
     }
-
-    pub fn subscribe(&self) -> Pollable {
-        golem_rust::bindings::wasi::clocks::monotonic_clock::subscribe_duration(0)
-    }
 }
 
 impl llm::GuestChatStream for BedrockChatStream {
@@ -97,10 +92,8 @@ impl llm::GuestChatStream for BedrockChatStream {
     }
 
     fn blocking_get_next(&self) -> Vec<llm::StreamEvent> {
-        let pollable = self.subscribe();
         let mut result = Vec::new();
         loop {
-            pollable.block();
             match self.get_next() {
                 Some(events) => {
                     result.extend(events);
diff --git a/llm-grok/Cargo.toml b/llm-grok/Cargo.toml
@@ -16,11 +16,11 @@ default = ["durability"]
 durability = ["golem-rust/durability", "golem-llm/durability"]
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 
 golem-rust = { workspace = true }
 log = { workspace = true }
-reqwest =  { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 wit-bindgen-rt = { workspace = true }
@@ -40,4 +40,4 @@ path = "wit"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit/deps/golem-llm" }
-"wasi:io" = { path = "wit/deps/wasi:io"}
+"wasi:io" = { path = "wit/deps/wasi:io" }
diff --git a/llm-ollama/Cargo.toml b/llm-ollama/Cargo.toml
@@ -17,7 +17,7 @@ default = ["durability"]
 durability = ["golem-rust/durability", "golem-llm/durability"]
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 base64 = "0.21"
 mime_guess = "2.0"
 url = "2.4"
diff --git a/llm-openai/Cargo.toml b/llm-openai/Cargo.toml
@@ -16,11 +16,11 @@ default = ["durability"]
 durability = ["golem-rust/durability", "golem-llm/durability"]
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 
 golem-rust = { workspace = true }
 log = { workspace = true }
-reqwest =  { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 wit-bindgen-rt = { workspace = true }
@@ -40,4 +40,4 @@ path = "wit"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit/deps/golem-llm" }
-"wasi:io" = { path = "wit/deps/wasi:io"}
+"wasi:io" = { path = "wit/deps/wasi:io" }
diff --git a/llm-openrouter/Cargo.toml b/llm-openrouter/Cargo.toml
@@ -16,11 +16,11 @@ default = ["durability"]
 durability = ["golem-rust/durability", "golem-llm/durability"]
 
 [dependencies]
-golem-llm = { path = "../llm", version = "0.0.0", default-features = false }
+golem-llm = { workspace = true }
 
 golem-rust = { workspace = true }
 log = { workspace = true }
-reqwest =  { workspace = true }
+reqwest = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 wit-bindgen-rt = { workspace = true }
@@ -40,4 +40,4 @@ path = "wit"
 
 [package.metadata.component.target.dependencies]
 "golem:llm" = { path = "wit/deps/golem-llm" }
-"wasi:io" = { path = "wit/deps/wasi:io"}
+"wasi:io" = { path = "wit/deps/wasi:io" }
diff --git a/llm/Cargo.toml b/llm/Cargo.toml
@@ -24,3 +24,4 @@ wit-bindgen = { version = "0.40.0" }
 [features]
 default = ["durability"]
 durability = ["golem-rust/durability"]
+nopoll = []
diff --git a/llm/src/durability.rs b/llm/src/durability.rs
@@ -1,4 +1,5 @@
 use crate::golem::llm::llm::{Config, ContentPart, Guest, Message, Role, StreamDelta};
+#[cfg(not(feature = "nopoll"))]
 use golem_rust::wasm_rpc::Pollable;
 use std::marker::PhantomData;
 
@@ -57,7 +58,7 @@ pub trait ExtendedGuest: Guest + 'static {
         });
         extended_messages
     }
-
+    #[cfg(not(feature = "nopoll"))]
     fn subscribe(stream: &Self::ChatStream) -> Pollable;
 }
 
@@ -105,9 +106,9 @@ mod durable_impl {
         ChatEvent, ChatStream, Config, Guest, GuestChatStream, Message, StreamDelta, StreamEvent,
         ToolCall, ToolResult,
     };
-    use golem_rust::bindings::golem::durability::durability::{
-        DurableFunctionType, LazyInitializedPollable,
-    };
+    use golem_rust::bindings::golem::durability::durability::DurableFunctionType;
+    #[cfg(not(feature = "nopoll"))]
+    use golem_rust::bindings::golem::durability::durability::LazyInitializedPollable;
     use golem_rust::durability::Durability;
     use golem_rust::wasm_rpc::Pollable;
     use golem_rust::{with_persistence_level, FromValueAndType, IntoValue, PersistenceLevel};
@@ -197,11 +198,13 @@ mod durable_impl {
     enum DurableChatStreamState<Impl: ExtendedGuest> {
         Live {
             stream: Impl::ChatStream,
+            #[cfg(not(feature = "nopoll"))]
             pollables: Vec<LazyInitializedPollable>,
         },
         Replay {
             original_messages: Vec<Message>,
             config: Config,
+            #[cfg(not(feature = "nopoll"))]
             pollables: Vec<LazyInitializedPollable>,
             partial_result: Vec<StreamDelta>,
             finished: bool,
@@ -218,6 +221,7 @@ mod durable_impl {
             Self {
                 state: RefCell::new(Some(DurableChatStreamState::Live {
                     stream,
+                    #[cfg(not(feature = "nopoll"))]
                     pollables: Vec::new(),
                 })),
                 subscription: RefCell::new(None),
@@ -229,14 +233,15 @@ mod durable_impl {
                 state: RefCell::new(Some(DurableChatStreamState::Replay {
                     original_messages,
                     config,
+                    #[cfg(not(feature = "nopoll"))]
                     pollables: Vec::new(),
                     partial_result: Vec::new(),
                     finished: false,
                 })),
                 subscription: RefCell::new(None),
             }
         }
-
+        #[cfg(not(feature = "nopoll"))]
         fn subscribe(&self) -> Pollable {
             let mut state = self.state.borrow_mut();
             match &mut *state {
@@ -257,17 +262,25 @@ mod durable_impl {
     impl<Impl: ExtendedGuest> Drop for DurableChatStream<Impl> {
         fn drop(&mut self) {
             let _ = self.subscription.take();
+
             match self.state.take() {
                 Some(DurableChatStreamState::Live {
+                    #[cfg(not(feature = "nopoll"))]
                     mut pollables,
                     stream,
                 }) => {
                     with_persistence_level(PersistenceLevel::PersistNothing, move || {
+                        #[cfg(not(feature = "nopoll"))]
                         pollables.clear();
                         drop(stream);
                     });
                 }
-                Some(DurableChatStreamState::Replay { mut pollables, .. }) => {
+                Some(DurableChatStreamState::Replay {
+                    #[cfg(not(feature = "nopoll"))]
+                    mut pollables,
+                    ..
+                }) => {
+                    #[cfg(not(feature = "nopoll"))]
                     pollables.clear();
                 }
                 None => {}
@@ -295,6 +308,7 @@ mod durable_impl {
                     Some(DurableChatStreamState::Replay {
                         original_messages,
                         config,
+                        #[cfg(not(feature = "nopoll"))]
                         pollables,
                         partial_result,
                         finished,
@@ -311,7 +325,7 @@ mod durable_impl {
                                         extended_messages,
                                         config.clone(),
                                     );
-
+                                    #[cfg(not(feature = "nopoll"))]
                                     for lazy_initialized_pollable in pollables {
                                         lazy_initialized_pollable.set(Impl::subscribe(&stream));
                                     }
@@ -330,14 +344,19 @@ mod durable_impl {
                 };
 
                 if let Some(stream) = new_live_stream {
+                    #[cfg(not(feature = "nopoll"))]
                     let pollables = match state.take() {
                         Some(DurableChatStreamState::Live { pollables, .. }) => pollables,
                         Some(DurableChatStreamState::Replay { pollables, .. }) => pollables,
                         None => {
                             unreachable!()
                         }
                     };
-                    *state = Some(DurableChatStreamState::Live { stream, pollables });
+                    *state = Some(DurableChatStreamState::Live {
+                        stream,
+                        #[cfg(not(feature = "nopoll"))]
+                        pollables,
+                    });
                 }
 
                 result
@@ -378,13 +397,17 @@ mod durable_impl {
         }
 
         fn blocking_get_next(&self) -> Vec<StreamEvent> {
+            #[cfg(not(feature = "nopoll"))]
             let mut subscription = self.subscription.borrow_mut();
+            #[cfg(not(feature = "nopoll"))]
             if subscription.is_none() {
                 *subscription = Some(self.subscribe());
             }
+            #[cfg(not(feature = "nopoll"))]
             let subscription = subscription.as_mut().unwrap();
             let mut result = Vec::new();
             loop {
+                #[cfg(not(feature = "nopoll"))]
                 subscription.block();
                 match self.get_next() {
                     Some(events) => {

Original file line number	Diff line number	Diff line change
`@@ -62,10 +62,6 @@ impl ExtendedGuest for BedrockComponent {`
`62`	`62`
`63`	`63`	`bedrock.unwrap().converse_stream(messages, config)`
`64`	`64`	`}`
`65`		`-`
`66`		`- fn subscribe(stream: &Self::ChatStream) -> golem_rust::wasm_rpc::Pollable {`
`67`		`- stream.subscribe()`
`68`		`- }`
`69`	`65`	`}`
`70`	`66`
`71`	`67`	`type DurableBedrockComponent = DurableLLM<BedrockComponent>;`