From 619d907104ff1bbf77aa9342af5300158d0b81a8 Mon Sep 17 00:00:00 2001
From: Sakshar Thakkar <sakshar.thakkar@uipath.com>
Date: Fri, 19 Dec 2025 17:17:40 -0800
Subject: [PATCH] feat: add execution_callbacks context var for per-execution
 callback injection
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Add execution_callbacks ContextVar to allow wrappers to inject callbacks at execution time
- Modify _get_graph_config to copy self.callbacks (prevent mutation) and extend with context callbacks
- Export execution_callbacks from runtime module

This enables clean callback injection without monkey-patching, solving the duplicate callback issue when runtime instances are reused across multiple executions (chat/debug mode).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 src/uipath_langchain/runtime/__init__.py |  3 ++-
 src/uipath_langchain/runtime/runtime.py  | 15 ++++++++++++++-
 2 files changed, 16 insertions(+), 2 deletions(-)

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