diff --git a/src/uipath_langchain/runtime/__init__.py b/src/uipath_langchain/runtime/__init__.py index cd8f015f..1e4d7308 100644 --- a/src/uipath_langchain/runtime/__init__.py +++ b/src/uipath_langchain/runtime/__init__.py @@ -5,7 +5,7 @@ ) from uipath_langchain.runtime.factory import UiPathLangGraphRuntimeFactory -from uipath_langchain.runtime.runtime import UiPathLangGraphRuntime +from uipath_langchain.runtime.runtime import UiPathLangGraphRuntime, execution_callbacks from uipath_langchain.runtime.schema import ( get_entrypoints_schema, get_graph_schema, @@ -33,4 +33,5 @@ def create_factory( "get_graph_schema", "UiPathLangGraphRuntimeFactory", "UiPathLangGraphRuntime", + "execution_callbacks", ] diff --git a/src/uipath_langchain/runtime/runtime.py b/src/uipath_langchain/runtime/runtime.py index 10c41114..c44e3c72 100644 --- a/src/uipath_langchain/runtime/runtime.py +++ b/src/uipath_langchain/runtime/runtime.py @@ -1,5 +1,6 @@ import logging import os +from contextvars import ContextVar from typing import Any, AsyncGenerator from uuid import uuid4 @@ -31,6 +32,12 @@ logger = logging.getLogger(__name__) +# Context var for per-execution callback injection (set by wrappers like TelemetryRuntimeWrapper) +# This allows injecting callbacks at execution time without mutating self.callbacks +execution_callbacks: ContextVar[list[BaseCallbackHandler] | None] = ContextVar( + "uipath_execution_callbacks", default=None +) + class UiPathLangGraphRuntime: """ @@ -197,9 +204,15 @@ async def get_schema(self) -> UiPathRuntimeSchema: def _get_graph_config(self) -> RunnableConfig: """Build graph execution configuration.""" + # Copy callbacks to prevent mutation, extend with context callbacks + callbacks = list(self.callbacks) + ctx_callbacks = execution_callbacks.get() + if ctx_callbacks: + callbacks.extend(ctx_callbacks) + graph_config: RunnableConfig = { "configurable": {"thread_id": self.runtime_id}, - "callbacks": self.callbacks, + "callbacks": callbacks, } # Add optional config from environment