diff --git a/crates/python/src/py_api/mod.rs b/crates/python/src/py_api/mod.rs index 9d37efa45..b2f055dd4 100644 --- a/crates/python/src/py_api/mod.rs +++ b/crates/python/src/py_api/mod.rs @@ -41,6 +41,7 @@ use nemo_relay::api::tool::ToolAttributes; use nemo_relay::codec::response::AnnotatedLlmResponse; use nemo_relay::codec::traits::{LlmCodec, LlmResponseCodec}; use nemo_relay::error::{FlowError, Result as FlowResult}; +use pyo3::IntoPyObjectExt; use pyo3::prelude::*; use tokio_stream::StreamExt; use uuid::Uuid; @@ -69,6 +70,87 @@ fn python_event_loop_running(py: Python<'_>) -> PyResult { } } +#[pyclass] +struct SafeFutureCanceller { + sender: Option>, +} + +#[pymethods] +impl SafeFutureCanceller { + fn __call__(&mut self, future: &Bound<'_, PyAny>) -> PyResult<()> { + if future.call_method0("cancelled")?.is_truthy()? + && let Some(sender) = self.sender.take() + { + let _ = sender.send(()); + } + Ok(()) + } +} + +#[pyclass] +struct SafeFutureCompleter { + future: Py, + result: Option>>, +} + +#[pymethods] +impl SafeFutureCompleter { + fn __call__(&mut self, py: Python<'_>) -> PyResult<()> { + let future = self.future.bind(py); + if future.call_method0("cancelled")?.is_truthy()? { + return Ok(()); + } + match self.result.take().expect("future completion result") { + Ok(value) => future.call_method1("set_result", (value.bind(py),))?, + Err(error) => future.call_method1("set_exception", (error.into_bound_py_any(py)?,))?, + }; + Ok(()) + } +} + +fn safe_future_into_py<'py, F>(py: Python<'py>, future: F) -> PyResult> +where + F: Future>> + Send + 'static, +{ + let locals = pyo3_async_runtimes::tokio::get_current_locals(py)?; + let event_loop = locals.event_loop(py).unbind(); + let python_future = event_loop.bind(py).call_method0("create_future")?.unbind(); + let (cancel_sender, mut cancel_receiver) = tokio::sync::oneshot::channel(); + python_future.call_method1( + py, + "add_done_callback", + (SafeFutureCanceller { + sender: Some(cancel_sender), + },), + )?; + let completion_future = python_future.clone_ref(py); + pyo3_async_runtimes::tokio::get_runtime().spawn(async move { + let result = tokio::select! { + result = pyo3_async_runtimes::tokio::scope(locals, future) => Some(result), + _ = &mut cancel_receiver => None, + }; + let Some(result) = result else { return }; + Python::attach(|py| { + let loop_is_closed = event_loop + .bind(py) + .call_method0("is_closed") + .and_then(|closed| closed.is_truthy()) + .unwrap_or(true); + if loop_is_closed { + return; + } + let _ = event_loop.bind(py).call_method1( + "call_soon_threadsafe", + (SafeFutureCompleter { + future: completion_future, + result: Some(result), + },), + ); + }); + }); + Ok(python_future.into_bound(py)) +} + fn with_python_publication_context(f: impl FnOnce() -> T) -> T { with_publication_context(py_callable::capture_python_publication_context(), f) } @@ -716,7 +798,7 @@ fn tool_call_execute<'py>( let scope_stack = current_scope_stack_handle(); let publication_context = py_callable::capture_python_publication_context(); let publication_buffer = capture_nested_publication_buffer(); - pyo3_async_runtimes::tokio::future_into_py(py, async move { + safe_future_into_py(py, async move { with_task_nested_publication_buffer( publication_buffer, with_task_publication_context(