Fix(llm): Do not call block_on in case of durable exec inside golem

Makefile.toml

@@ -98,12 +98,26 @@ command = "cargo"
 args = ["clippy", "--all-targets", "--", "--no-deps", "-Dwarnings"]
 
 [tasks.check-clippy-no-default-features]
-description = "Runs clippy with --no-default-features to catch breakage in the durability-off code path. Scoped to the stt crates for now; other crates still have stale passthrough paths that need the same cleanup before they can be included here."
+description = "Runs clippy with --no-default-features to catch breakage in the durability-off code path. Scoped to the llm and stt crates for now; other crates still have stale passthrough paths that need the same cleanup before they can be included here."
 install_crate = "clippy"
 command = "cargo"
 args = [
     "clippy",
     "-p",
+    "golem-ai-llm",
+    "-p",
+    "golem-ai-llm-openai",
+    "-p",
+    "golem-ai-llm-anthropic",
+    "-p",
+    "golem-ai-llm-grok",
+    "-p",
+    "golem-ai-llm-ollama",
+    "-p",
+    "golem-ai-llm-openrouter",
+    "-p",
+    "golem-ai-llm-bedrock",
+    "-p",
     "golem-ai-stt",
     "-p",
     "golem-ai-stt-whisper",

llm/anthropic/src/lib.rs

@@ -278,20 +278,23 @@ impl Anthropic {
 impl LlmProvider for Anthropic {
     type ChatStream = LlmChatStream<AnthropicChatStream>;
 
-    fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
+    async fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
         let anthropic_api_key = get_config_key(Self::ENV_VAR_NAME)?;
         let client = MessagesApi::new(anthropic_api_key);
         let request = events_to_request(events, config)?;
         Self::request(client, request)
     }
 
-    fn stream(events: Vec<Event>, config: Config) -> ChatStream {
-        ChatStream::new(Self::unwrapped_stream(events, config))
+    async fn stream(events: Vec<Event>, config: Config) -> ChatStream {
+        ChatStream::new(Self::unwrapped_stream(events, config).await)
     }
 }
 
 impl ExtendedLlmProvider for Anthropic {
-    fn unwrapped_stream(events: Vec<Event>, config: Config) -> LlmChatStream<AnthropicChatStream> {
+    async fn unwrapped_stream(
+        events: Vec<Event>,
+        config: Config,
+    ) -> LlmChatStream<AnthropicChatStream> {
         with_config_key(
             Self::ENV_VAR_NAME,
             AnthropicChatStream::failed,

llm/bedrock/Cargo.toml

@@ -17,6 +17,9 @@ aws-sdk-bedrockruntime = { version = "1.56.0", default-features = false }
 aws-smithy-types = { version = "1.3.1" }
 aws-smithy-runtime-api = "1.8.3"
 
+async-trait = "0.1.89"
+futures = { workspace = true }
+
 wasi = { workspace = true }
 
 # To infer mime types of downloaded images before passing to bedrock

llm/bedrock/src/async_utils.rs

@@ -1,20 +1,4 @@
 use std::future::Future;
-use wstd::runtime;
-
-pub fn get_async_runtime() -> AsyncRuntime {
-    AsyncRuntime
-}
-
-pub struct AsyncRuntime;
-
-impl AsyncRuntime {
-    pub fn block_on<F>(self, f: F) -> F::Output
-    where
-        F: Future,
-    {
-        runtime::block_on(f)
-    }
-}
 
 #[derive(Clone)]
 pub struct UnsafeFuture<Fut> {

llm/bedrock/src/lib.rs

@@ -1,4 +1,3 @@
-use async_utils::get_async_runtime;
 use client::Bedrock as BedrockClient;
 use golem_ai_llm::durability::{DurableLLM, ExtendedLlmProvider};
 use golem_ai_llm::model::{
@@ -20,32 +19,22 @@ pub struct Bedrock;
 impl LlmProvider for Bedrock {
     type ChatStream = BedrockChatStream;
 
-    fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
-        let runtime = get_async_runtime();
-
-        runtime.block_on(async {
-            let client = get_bedrock_client().await?;
-            client.converse(events, config).await
-        })
+    async fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
+        let client = get_bedrock_client().await?;
+        client.converse(events, config).await
     }
 
-    fn stream(events: Vec<Event>, config: Config) -> ChatStream {
-        ChatStream::new(Self::unwrapped_stream(events, config))
+    async fn stream(events: Vec<Event>, config: Config) -> ChatStream {
+        ChatStream::new(Self::unwrapped_stream(events, config).await)
     }
 }
 
 impl ExtendedLlmProvider for Bedrock {
-    fn unwrapped_stream(messages: Vec<Event>, config: Config) -> Self::ChatStream {
-        let runtime = get_async_runtime();
-
-        runtime.block_on(async {
-            let bedrock = get_bedrock_client().await;
-
-            match bedrock {
-                Ok(client) => client.converse_stream(messages, config).await,
-                Err(err) => BedrockChatStream::failed(err),
-            }
-        })
+    async fn unwrapped_stream(messages: Vec<Event>, config: Config) -> Self::ChatStream {
+        match get_bedrock_client().await {
+            Ok(client) => client.converse_stream(messages, config).await,
+            Err(err) => BedrockChatStream::failed(err),
+        }
     }
 
     fn retry_prompt(

llm/bedrock/src/stream.rs

@@ -1,44 +1,39 @@
-use crate::{
-    async_utils,
-    conversions::{converse_stream_output_to_stream_event, custom_error, merge_metadata},
-};
+use crate::conversions::{converse_stream_output_to_stream_event, custom_error, merge_metadata};
+use async_trait::async_trait;
 use aws_sdk_bedrockruntime::{
     self as bedrock, primitives::event_stream::EventReceiver,
     types::error::ConverseStreamOutputError,
 };
+use futures::lock::Mutex;
 use golem_ai_llm::{model as llm, ChatStreamInterface};
-use std::cell::{RefCell, RefMut};
+use std::cell::RefCell;
 
 type BedrockEventSource =
     EventReceiver<bedrock::types::ConverseStreamOutput, ConverseStreamOutputError>;
 
 pub struct BedrockChatStream {
-    stream: RefCell<Option<BedrockEventSource>>,
+    stream: Mutex<Option<BedrockEventSource>>,
     failure: Option<llm::Error>,
     finished: RefCell<bool>,
 }
 
 impl BedrockChatStream {
     pub fn new(stream: BedrockEventSource) -> BedrockChatStream {
         BedrockChatStream {
-            stream: RefCell::new(Some(stream)),
+            stream: Mutex::new(Some(stream)),
             failure: None,
             finished: RefCell::new(false),
         }
     }
 
     pub fn failed(error: llm::Error) -> BedrockChatStream {
         BedrockChatStream {
-            stream: RefCell::new(None),
+            stream: Mutex::new(None),
             failure: Some(error),
             finished: RefCell::new(true),
         }
     }
 
-    fn stream_mut(&self) -> RefMut<'_, Option<BedrockEventSource>> {
-        self.stream.borrow_mut()
-    }
-
     fn failure(&self) -> &Option<llm::Error> {
         &self.failure
     }
@@ -50,65 +45,68 @@ impl BedrockChatStream {
     fn set_finished(&self) {
         *self.finished.borrow_mut() = true;
     }
-    fn get_single_event(&self) -> Option<Result<llm::StreamEvent, llm::Error>> {
-        if let Some(stream) = self.stream_mut().as_mut() {
-            let runtime = async_utils::get_async_runtime();
+    async fn get_single_event(&self) -> Option<Result<llm::StreamEvent, llm::Error>> {
+        let mut stream_guard = self.stream.lock().await;
+        if let Some(stream) = stream_guard.as_mut() {
+            let token = stream.recv().await;
+            drop(stream_guard);
 
-            runtime.block_on(async move {
-                let token = stream.recv().await;
-                log::trace!("Bedrock stream event: {token:?}");
-
-                match token {
-                    Ok(Some(output)) => {
-                        log::trace!("Processing bedrock stream event: {output:?}");
-                        converse_stream_output_to_stream_event(output).map(Ok)
-                    }
-                    Ok(None) => {
-                        log::trace!("running set_finished on stream due to None event received");
-                        self.set_finished();
-                        None
-                    }
-                    Err(error) => {
-                        log::trace!("running set_finished on stream due to error: {error:?}");
-                        self.set_finished();
-                        Some(Err(custom_error(
-                            llm::ErrorCode::InternalError,
-                            format!("An error occurred while reading event stream: {error}"),
-                        )))
-                    }
+            log::trace!("Bedrock stream event: {token:?}");
+            match token {
+                Ok(Some(output)) => {
+                    log::trace!("Processing bedrock stream event: {output:?}");
+                    converse_stream_output_to_stream_event(output).map(Ok)
+                }
+                Ok(None) => {
+                    log::trace!("running set_finished on stream due to None event received");
+                    self.set_finished();
+                    None
+                }
+                Err(error) => {
+                    log::trace!("running set_finished on stream due to error: {error:?}");
+                    self.set_finished();
+                    Some(Err(custom_error(
+                        llm::ErrorCode::InternalError,
+                        format!("An error occurred while reading event stream: {error}"),
+                    )))
                 }
-            })
-        } else if let Some(error) = self.failure() {
-            self.set_finished();
-            Some(Err(error.clone()))
+            }
         } else {
-            None
+            drop(stream_guard);
+
+            if let Some(error) = self.failure() {
+                self.set_finished();
+                Some(Err(error.clone()))
+            } else {
+                None
+            }
         }
     }
 }
 
+#[async_trait(?Send)]
 impl ChatStreamInterface for BedrockChatStream {
-    fn poll_next(&self) -> Option<Vec<Result<llm::StreamEvent, llm::Error>>> {
+    async fn poll_next(&self) -> Option<Vec<Result<llm::StreamEvent, llm::Error>>> {
         if self.is_finished() {
             return Some(vec![]);
         }
-        self.get_single_event().map(|event| {
-            if let Ok(llm::StreamEvent::Finish(metadata)) = &event {
-                if let Some(Ok(llm::StreamEvent::Finish(final_metadata))) = self.get_single_event()
-                {
-                    return vec![Ok(llm::StreamEvent::Finish(merge_metadata(
-                        metadata.clone(),
-                        final_metadata,
-                    )))];
-                }
+        let event = self.get_single_event().await?;
+        if let Ok(llm::StreamEvent::Finish(metadata)) = &event {
+            if let Some(Ok(llm::StreamEvent::Finish(final_metadata))) =
+                self.get_single_event().await
+            {
+                return Some(vec![Ok(llm::StreamEvent::Finish(merge_metadata(
+                    metadata.clone(),
+                    final_metadata,
+                )))]);
             }
-            vec![event]
-        })
+        }
+        Some(vec![event])
     }
 
-    fn get_next(&self) -> Vec<Result<llm::StreamEvent, llm::Error>> {
+    async fn get_next(&self) -> Vec<Result<llm::StreamEvent, llm::Error>> {
         loop {
-            if let Some(events) = self.poll_next() {
+            if let Some(events) = self.poll_next().await {
                 return events;
             }
         }

llm/grok/src/lib.rs

@@ -156,20 +156,23 @@ impl Grok {
 impl LlmProvider for Grok {
     type ChatStream = LlmChatStream<GrokChatStream>;
 
-    fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
+    async fn send(events: Vec<Event>, config: Config) -> Result<Response, Error> {
         let xai_api_key = get_config_key(Self::ENV_VAR_NAME)?;
         let client = CompletionsApi::new(xai_api_key);
         let request = events_to_request(events, config)?;
         Self::request(client, request)
     }
 
-    fn stream(messages: Vec<Event>, config: Config) -> ChatStream {
-        ChatStream::new(Self::unwrapped_stream(messages, config))
+    async fn stream(messages: Vec<Event>, config: Config) -> ChatStream {
+        ChatStream::new(Self::unwrapped_stream(messages, config).await)
     }
 }
 
 impl ExtendedLlmProvider for Grok {
-    fn unwrapped_stream(messages: Vec<Event>, config: Config) -> LlmChatStream<GrokChatStream> {
+    async fn unwrapped_stream(
+        messages: Vec<Event>,
+        config: Config,
+    ) -> LlmChatStream<GrokChatStream> {
         with_config_key(Self::ENV_VAR_NAME, GrokChatStream::failed, |xai_api_key| {
             let client = CompletionsApi::new(xai_api_key);
 

llm/llm/Cargo.toml

@@ -11,6 +11,7 @@ description = "Library for working with LLM APIs on Golem Cloud"
 path = "src/lib.rs"
 
 [dependencies]
+async-trait = "0.1.89"
 golem-rust = { workspace = true }
 
 golem-wasi-http = { workspace = true }

llm/llm/src/chat_stream.rs

@@ -1,6 +1,7 @@
 use crate::event_source::{Event, EventSource, MessageEvent};
 use crate::model::{Error, ErrorCode, StreamEvent};
 use crate::ChatStreamInterface;
+use async_trait::async_trait;
 use golem_rust::golem_wasm::Pollable;
 use std::cell::{Ref, RefMut};
 use std::task::Poll;
@@ -32,8 +33,9 @@ impl<T: LlmChatStreamState> LlmChatStream<T> {
     }
 }
 
+#[async_trait(?Send)]
 impl<T: LlmChatStreamState> ChatStreamInterface for LlmChatStream<T> {
-    fn poll_next(&self) -> Option<Vec<Result<StreamEvent, Error>>> {
+    async fn poll_next(&self) -> Option<Vec<Result<StreamEvent, Error>>> {
         if self.implementation.is_finished() {
             return Some(vec![]);
         }
@@ -96,11 +98,11 @@ impl<T: LlmChatStreamState> ChatStreamInterface for LlmChatStream<T> {
         }
     }
 
-    fn get_next(&self) -> Vec<Result<StreamEvent, Error>> {
+    async fn get_next(&self) -> Vec<Result<StreamEvent, Error>> {
         let pollable = self.subscribe();
         loop {
             pollable.block();
-            if let Some(events) = self.poll_next() {
+            if let Some(events) = self.poll_next().await {
                 return events;
             }
         }