fix(autogen): attribute bridged tools to live turn steps

integrations/autogen/src/agentflow_autogen/adapter.py

@@ -16,14 +16,14 @@
 from app.models.run import RunStatus
 from autogen_agentchat.base import TaskResult
 from autogen_agentchat.messages import (
-    BaseChatMessage,
     ModelClientStreamingChunkEvent,
     TextMessage,
     ToolCallExecutionEvent,
     ToolCallRequestEvent,
 )
 
 from agentflow_autogen.tools import (
+    BridgedToolExecution,
     build_function_tools,
     inject_tools,
 )
@@ -68,8 +68,9 @@ async def run(self, ctx: AdapterContext) -> AdapterResult:
                         bridged_tools, bridged_names = build_function_tools(
                             surface,
                             ctx,
-                            step_index=handler.current_step_index,
+                            step_index=lambda: handler.current_step_index,
                             on_tool_error=bridged_failure.append,
+                            on_tool_execution=handler.record_bridged_tool,
                         )
                     else:
                         bridged_tools = []
@@ -132,6 +133,8 @@ def __init__(
         self._final_reply = ""
         self._replies: list[str] = []
         self._active_tool_calls: dict[str, tuple[str, str, float]] = {}
+        self._bridged_tool_executions: list[BridgedToolExecution] = []
+        self._pending_bridged_calls: dict[str, tuple[str, dict[str, Any]]] = {}
         self._pending_user_prompt: str | None = None
 
     @property
@@ -149,15 +152,30 @@ async def begin_step(self, node: str) -> None:
         self._started_at = time.monotonic()
         self._tokens_in = 0
         self._tokens_out = 0
+        self._final_reply = ""
         await self.ctx.emit_step_started(index=self._current_step_index, node=self._current_node)
 
+    async def ensure_step(self, node: str) -> None:
+        """Start or rotate the active step when an AutoGen turn changes source."""
+        if not self.step_started or self.step_completed:
+            await self.begin_step(node)
+            return
+        if self.per_turn_steps and node != self._current_node:
+            await self.complete_step(output={"reply": self._final_reply})
+            await self.begin_step(node)
+
     async def emit_user_prompt(self, prompt: str) -> None:
         self._pending_user_prompt = prompt
         step_index = self._current_step_index if self.step_started else None
         await self.ctx.emit_message(role="user", content=prompt, step_index=step_index)
 
+    def record_bridged_tool(self, execution: BridgedToolExecution) -> None:
+        """Hold tool events until AutoGen emits an event with the owning source."""
+        self._bridged_tool_executions.append(execution)
+
     async def handle(self, event: Any, *, bridged_names: frozenset[str]) -> None:
         if isinstance(event, ModelClientStreamingChunkEvent):
+            await self.ensure_step(str(event.source or self.default_node))
             if self.stream_tokens and isinstance(event.content, str):
                 await self.ctx.emit_token_delta(
                     step_index=self._current_step_index,
@@ -191,13 +209,7 @@ async def _handle_text_message(self, message: TextMessage) -> None:
             )
             return
 
-        if self.per_turn_steps:
-            if self.step_started and not self.step_completed:
-                await self.complete_step(output={"reply": self._final_reply})
-            await self.begin_step(str(message.source or self.default_node))
-
-        if not self.step_started:
-            await self.begin_step(str(message.source or self.default_node))
+        await self.ensure_step(str(message.source or self.default_node))
 
         self._accumulate_usage(message)
         content = str(message.content or "")
@@ -217,11 +229,19 @@ async def _handle_tool_request(
         *,
         bridged_names: frozenset[str],
     ) -> None:
+        await self.ensure_step(str(event.source or self.default_node))
         self._accumulate_usage(event)
         for call in event.content:
+            arguments = _parse_tool_arguments(call.arguments)
             if call.name in bridged_names:
+                emitted = await self._emit_bridged_tool(
+                    name=call.name,
+                    arguments=arguments,
+                    call_id=call.id,
+                )
+                if not emitted:
+                    self._pending_bridged_calls[call.id] = (call.name, arguments)
                 continue
-            arguments = _parse_tool_arguments(call.arguments)
             call_id = await self.ctx.emit_tool_call_started(
                 step_index=self._current_step_index,
                 name=call.name,
@@ -235,8 +255,16 @@ async def _handle_tool_execution(
         *,
         bridged_names: frozenset[str],
     ) -> None:
+        await self.ensure_step(str(event.source or self.default_node))
         for result in event.content:
             if result.name in bridged_names:
+                pending = self._pending_bridged_calls.pop(result.call_id, None)
+                name, arguments = pending or (result.name, {})
+                await self._emit_bridged_tool(
+                    name=name,
+                    arguments=arguments,
+                    call_id=result.call_id,
+                )
                 continue
             active = self._active_tool_calls.pop(result.call_id, None)
             if active is None:
@@ -258,6 +286,50 @@ async def _handle_tool_execution(
                 latency_ms=int((time.monotonic() - started) * 1000),
             )
 
+    async def _emit_bridged_tool(
+        self,
+        *,
+        name: str,
+        arguments: dict[str, Any],
+        call_id: str,
+    ) -> bool:
+        execution_index = next(
+            (
+                index
+                for index, execution in enumerate(self._bridged_tool_executions)
+                if execution.name == name and execution.arguments == arguments
+            ),
+            None,
+        )
+        if execution_index is None:
+            execution_index = next(
+                (
+                    index
+                    for index, execution in enumerate(self._bridged_tool_executions)
+                    if execution.name == name
+                ),
+                None,
+            )
+        if execution_index is None:
+            return False
+
+        execution = self._bridged_tool_executions.pop(execution_index)
+        emitted_call_id = await self.ctx.emit_tool_call_started(
+            step_index=self._current_step_index,
+            name=execution.name,
+            arguments=execution.arguments,
+            call_id=call_id,
+        )
+        await self.ctx.emit_tool_call_completed(
+            step_index=self._current_step_index,
+            name=execution.name,
+            call_id=emitted_call_id,
+            result=execution.result,
+            error=execution.error,
+            latency_ms=execution.latency_ms,
+        )
+        return True
+
     async def _handle_task_result(self, result: TaskResult) -> None:
         for message in result.messages:
             if isinstance(message, TextMessage) and message.source != USER_SOURCE:

integrations/autogen/src/agentflow_autogen/tools.py

@@ -3,20 +3,35 @@
 from __future__ import annotations
 
 import json
-from typing import Annotated, Any, Callable
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Annotated, Any
 
 from app.adapters.adapter_tools import AdapterToolSurface
 from app.adapters.base import AdapterContext
 from app.adapters.tool_registry import ToolDefinition
 from autogen_core.tools import FunctionTool
 
 
+@dataclass(frozen=True, slots=True)
+class BridgedToolExecution:
+    """A tool result waiting for AutoGen to reveal its owning turn."""
+
+    name: str
+    arguments: dict[str, Any]
+    result: dict[str, Any] | None
+    error: str | None
+    latency_ms: int
+
+
 def build_function_tools(
     surface: AdapterToolSurface,
     ctx: AdapterContext,
     *,
-    step_index: int,
+    step_index: int | Callable[[], int],
     on_tool_error: Callable[[str], None] | None = None,
+    on_tool_execution: Callable[[BridgedToolExecution], None] | None = None,
 ) -> tuple[list[FunctionTool], frozenset[str]]:
     """Create AutoGen tools that delegate execution to ``AdapterToolSurface``."""
     tools: list[FunctionTool] = []
@@ -29,6 +44,7 @@ def build_function_tools(
             step_index=step_index,
             definition=definition,
             on_tool_error=on_tool_error,
+            on_tool_execution=on_tool_execution,
         )
         tools.append(
             FunctionTool(
@@ -46,50 +62,80 @@ def _make_handler(
     surface: AdapterToolSurface,
     ctx: AdapterContext,
     *,
-    step_index: int,
+    step_index: int | Callable[[], int],
     definition: ToolDefinition,
     on_tool_error: Callable[[str], None] | None = None,
+    on_tool_execution: Callable[[BridgedToolExecution], None] | None = None,
 ):
     properties = (definition.parameters or {}).get("properties") or {}
 
     async def _execute(arguments: dict[str, Any]) -> str:
         try:
-            result = await surface.execute(
-                ctx,
-                step_index=step_index,
-                name=definition.name,
-                arguments=arguments,
-            )
+            if on_tool_execution is None:
+                current_step_index = step_index() if callable(step_index) else step_index
+                result = await surface.execute(
+                    ctx,
+                    step_index=current_step_index,
+                    name=definition.name,
+                    arguments=arguments,
+                )
+            else:
+                started = time.monotonic()
+                try:
+                    result = await surface.lookup(definition.name).handler(arguments)
+                except Exception as exc:
+                    error = str(exc) or type(exc).__name__
+                    on_tool_execution(
+                        BridgedToolExecution(
+                            name=definition.name,
+                            arguments=arguments,
+                            result=None,
+                            error=error,
+                            latency_ms=int((time.monotonic() - started) * 1000),
+                        )
+                    )
+                    raise
+                if not isinstance(result, dict):
+                    result = {"result": result}
+                on_tool_execution(
+                    BridgedToolExecution(
+                        name=definition.name,
+                        arguments=arguments,
+                        result=result,
+                        error=None,
+                        latency_ms=int((time.monotonic() - started) * 1000),
+                    )
+                )
             return _stringify_result(result)
         except Exception as exc:
             if on_tool_error is not None:
                 on_tool_error(str(exc) or type(exc).__name__)
             return str(exc) or type(exc).__name__
 
     if not properties:
-        async def handler() -> str:
+        async def no_arguments_handler() -> str:
             return await _execute({})
 
-        handler.__name__ = definition.name
-        return handler
+        no_arguments_handler.__name__ = definition.name
+        return no_arguments_handler
 
     if set(properties) == {"text"}:
-        async def handler(text: Annotated[str, "Tool input text."] = "") -> str:
+        async def text_handler(text: Annotated[str, "Tool input text."] = "") -> str:
             return await _execute({"text": text})
 
-        handler.__name__ = definition.name
-        return handler
+        text_handler.__name__ = definition.name
+        return text_handler
 
-    async def handler(input_json: Annotated[str, "JSON-encoded tool arguments."] = "{}") -> str:
+    async def json_handler(input_json: Annotated[str, "JSON-encoded tool arguments."] = "{}") -> str:
         try:
             parsed = json.loads(input_json or "{}")
         except json.JSONDecodeError:
             parsed = {"raw": input_json}
         arguments = parsed if isinstance(parsed, dict) else {"value": parsed}
         return await _execute(arguments)
 
-    handler.__name__ = definition.name
-    return handler
+    json_handler.__name__ = definition.name
+    return json_handler
 
 
 def _stringify_result(result: dict[str, Any]) -> str: