(burr) branch main updated: fix(core): pass __context/__tracer to async and streaming class-based actions (#937)

[eb...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

elijahbenizzy pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/burr.git


The following commit(s) were added to refs/heads/main by this push:
     new c17ddbf0 fix(core): pass __context/__tracer to async and streaming class-based actions (#937)
c17ddbf0 is described below

commit c17ddbf05d86035fdef0b76deedfdb63870d00f5
Author: breken <su...@brekfuz.com>
AuthorDate: Sat Oct 17 23:34:01 2026 +0800

    fix(core): pass __context/__tracer to async and streaming class-based actions (#937)
    
    * Pass __context/__tracer to async and streaming class-based actions
    
    A class-based action that declares `__context` (or `__tracer`) in its
    run signature has that parameter name-mangled by Python, e.g. to
    `_MyAction__context`. The sync run path remaps the injected input to
    the mangled name, but the async run path and both multi-step streaming
    paths passed `__context=` as-is, so the action failed with
    `TypeError: run() got an unexpected keyword argument '__context'`.
    
    Apply the same remap in `_arun_function` and in the sync and async
    multi-step streaming runners.
    
    Generated-by: Claude Code (Claude Opus 5.5)
    Co-Authored-By: Claude Opus 5.5 (1M context) <no...@anthropic.com>
    
    * Note the missing __context/__tracer remap in single-step runners
    
    The single-step runners call run_and_update / stream_run_and_update with
    the injected inputs as-is, so a SingleStepAction that names __context in
    its signature still gets a TypeError. SingleStepAction is internal and
    in-tree subclasses take **run_kwargs, so leave a TODO at each site rather
    than change behavior here.
    
    Generated-by: Claude Code (Claude Opus 5.5)
    Co-Authored-By: Claude Opus 5.5 <no...@anthropic.com>
    
    ---------
    
    Co-authored-by: breken-ai <31...@users.noreply.github.com>
    Co-authored-by: Claude Opus 5.5 (1M context) <no...@anthropic.com>
---
 burr/core/application.py       | 21 ++++++---
 tests/core/test_application.py | 99 ++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 115 insertions(+), 5 deletions(-)

diff --git a/burr/core/application.py b/burr/core/application.py
index 5d54636b..91ef8b1b 100644
--- a/burr/core/application.py
+++ b/burr/core/application.py
@@ -158,6 +158,14 @@ def _remap_dunder_parameters(
     return inputs
 
 
+def _remap_injected_inputs(method: Callable, inputs: Dict[str, Any]) -> Dict[str, Any]:
+    """Remaps the framework-injected ``__context``/``__tracer`` inputs to the name-mangled
+    parameter names a class-based action's method may use. See ``_remap_dunder_parameters``."""
+    if "__context" in inputs or "__tracer" in inputs:
+        return _remap_dunder_parameters(method, inputs, ["__context", "__tracer"])
+    return inputs
+
+
 def _run_function(function: Function, state: State, inputs: Dict[str, Any], name: str) -> dict:
     """Runs a function, returning the result of running the function.
     Note this restricts the keys in the state to only those that the
@@ -176,9 +184,7 @@ def _run_function(function: Function, state: State, inputs: Dict[str, Any], name
         )
     state_to_use = state.subset(*function.reads)
     function.validate_inputs(inputs)
-    if "__context" in inputs or "__tracer" in inputs:
-        # potentially need to remap the __context & __tracer variables
-        inputs = _remap_dunder_parameters(function.run, inputs, ["__context", "__tracer"])
+    inputs = _remap_injected_inputs(function.run, inputs)
     result = function.run(state_to_use, **inputs)
     _validate_result(result, name)
     return result
@@ -191,6 +197,7 @@ async def _arun_function(
     Async version of the above."""
     state_to_use = state.subset(*function.reads)
     function.validate_inputs(inputs)
+    inputs = _remap_injected_inputs(function.run, inputs)
     result = await function.run(state_to_use, **inputs)
     _validate_result(result, name)
     return result
@@ -311,6 +318,7 @@ def _run_single_step_action(
     """
     # TODO -- guard all reads/writes with a subset of the state
     action.validate_inputs(inputs)
+    # TODO: __context/__tracer are not remapped here (see _remap_injected_inputs).
     result, new_state = _adjust_single_step_output(
         action.run_and_update(state, **inputs), action.name, action.schema
     )
@@ -334,6 +342,7 @@ def _run_single_step_streaming_action(
     action.validate_inputs(inputs)
     stream_initialize_time = system.now()
     first_stream_start_time = None
+    # TODO: __context/__tracer are not remapped here (see _remap_injected_inputs).
     generator = action.stream_run_and_update(state, **inputs)
     result = None
     state_update = None
@@ -402,6 +411,7 @@ async def _arun_single_step_streaming_action(
     action.validate_inputs(inputs)
     stream_initialize_time = system.now()
     first_stream_start_time = None
+    # TODO: __context/__tracer are not remapped here (see _remap_injected_inputs).
     generator = action.stream_run_and_update(state, **inputs)
     result = None
     state_update = None
@@ -477,7 +487,7 @@ def _run_multi_step_streaming_action(
     """
     action.validate_inputs(inputs)
     stream_initialize_time = system.now()
-    generator = action.stream_run(state, **inputs)
+    generator = action.stream_run(state, **_remap_injected_inputs(action.stream_run, inputs))
     result = None
     first_stream_start_time = None
     count = 0
@@ -535,7 +545,7 @@ async def _arun_multi_step_streaming_action(
     """Runs a multi-step streaming action in async. See the synchronous version for more details."""
     action.validate_inputs(inputs)
     stream_initialize_time = system.now()
-    generator = action.stream_run(state, **inputs)
+    generator = action.stream_run(state, **_remap_injected_inputs(action.stream_run, inputs))
     result = None
     first_stream_start_time = None
     count = 0
@@ -587,6 +597,7 @@ async def _arun_single_step_action(
     """Runs a single step action in async. See the synchronous version for more details."""
     state_to_use = state
     action.validate_inputs(inputs)
+    # TODO: __context/__tracer are not remapped here (see _remap_injected_inputs).
     result, new_state = _adjust_single_step_output(
         await action.run_and_update(state_to_use, **inputs), action.name, action.schema
     )
diff --git a/tests/core/test_application.py b/tests/core/test_application.py
index 5a252f77..efb7d205 100644
--- a/tests/core/test_application.py
+++ b/tests/core/test_application.py
@@ -4256,6 +4256,105 @@ def test_remap_context_variable_without_mangled_context():
     assert _remap_dunder_parameters(_action.run, inputs, ["__context", "__tracer"]) == expected
 
 
+class AsyncActionWithContext(Action):
+    """Class-based async action whose run() takes the injected ``__context``.
+    Python name-mangles the parameter to ``_AsyncActionWithContext__context``."""
+
+    @property
+    def reads(self) -> list[str]:
+        return []
+
+    @property
+    def writes(self) -> list[str]:
+        return ["app_id"]
+
+    @property
+    def inputs(self) -> list[str]:
+        return ["__context"]
+
+    async def run(self, state: State, __context: ApplicationContext) -> dict:
+        return {"app_id": __context.app_id}
+
+    def update(self, result: dict, state: State) -> State:
+        return state.update(**result)
+
+
+class StreamingActionWithContext(StreamingAction):
+    @property
+    def reads(self) -> list[str]:
+        return []
+
+    @property
+    def writes(self) -> list[str]:
+        return ["app_id"]
+
+    @property
+    def inputs(self) -> list[str]:
+        return ["__context"]
+
+    def stream_run(
+        self, state: State, __context: ApplicationContext
+    ) -> Generator[dict, None, None]:
+        yield {"app_id": __context.app_id}
+
+    def update(self, result: dict, state: State) -> State:
+        return state.update(**result)
+
+
+class AsyncStreamingActionWithContext(AsyncStreamingAction):
+    @property
+    def reads(self) -> list[str]:
+        return []
+
+    @property
+    def writes(self) -> list[str]:
+        return ["app_id"]
+
+    @property
+    def inputs(self) -> list[str]:
+        return ["__context"]
+
+    async def stream_run(self, state: State, __context: ApplicationContext) -> AsyncGenerator:
+        yield {"app_id": __context.app_id}
+
+    def update(self, result: dict, state: State) -> State:
+        return state.update(**result)
+
+
+def _build_context_app(action_: Action) -> Application:
+    return (
+        ApplicationBuilder()
+        .with_actions(ctx_action=action_)
+        .with_transitions()
+        .with_entrypoint("ctx_action")
+        .with_identifiers(app_id="context-app-id")
+        .build()
+    )
+
+
+async def test_astep_class_based_action_receives_context():
+    app = _build_context_app(AsyncActionWithContext())
+    _, result, state = await app.astep()
+    assert result == {"app_id": "context-app-id"}
+    assert state["app_id"] == "context-app-id"
+
+
+def test_stream_result_class_based_streaming_action_receives_context():
+    app = _build_context_app(StreamingActionWithContext())
+    _, container = app.stream_result(halt_after=["ctx_action"])
+    result, state = container.get()
+    assert result == {"app_id": "context-app-id"}
+    assert state["app_id"] == "context-app-id"
+
+
+async def test_astream_result_class_based_async_streaming_action_receives_context():
+    app = _build_context_app(AsyncStreamingActionWithContext())
+    _, container = await app.astream_result(halt_after=["ctx_action"])
+    result, state = await container.get()
+    assert result == {"app_id": "context-app-id"}
+    assert state["app_id"] == "context-app-id"
+
+
 async def test_async_application_builder_initialize_raises_on_broken_persistor():
     """Persisters should return None when there is no state to be loaded and the default used."""
     await asyncio.sleep(0.00001)