diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py b/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py index def67c9af..7d81b37d0 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/__init__.py @@ -51,6 +51,10 @@ - No forked plugin required - uses standard OpenAIAgentsPlugin """ +from agentex.lib.core.temporal.plugins.openai_agents.run import ( + OpenAIAgentsTurnResult, + run_turn, +) from agentex.lib.core.temporal.plugins.openai_agents.hooks.hooks import ( TemporalStreamingHooks, ) @@ -81,4 +85,6 @@ "streaming_parent_span_id", "TemporalStreamingHooks", "stream_lifecycle_content", -] \ No newline at end of file + "run_turn", + "OpenAIAgentsTurnResult", +] diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/hooks/hooks.py b/src/agentex/lib/core/temporal/plugins/openai_agents/hooks/hooks.py index 758b0db27..30d358cc9 100644 --- a/src/agentex/lib/core/temporal/plugins/openai_agents/hooks/hooks.py +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/hooks/hooks.py @@ -1,9 +1,30 @@ """Temporal streaming hooks for OpenAI Agents SDK lifecycle events. This module provides a convenience class for streaming agent lifecycle events -to the AgentEx UI via Temporal activities. +to the AgentEx UI via Temporal activities, and (optionally) tracing tool calls +to SGP with both inputs and outputs. + +Two responsibilities, independently switchable: + +1. UI message emission, split into tool requests / tool responses / handoffs + (each default True). Leave all on for the non-streaming model provider, which + does not emit them itself. When pairing with ``TemporalStreamingModelProvider`` + set ``emit_tool_requests=False`` — that model already streams the tool REQUEST + from the model output, so emitting it here double-posts. But keep + ``emit_tool_responses=True``: the streaming model does NOT emit a function + tool's response, so ``on_tool_end`` is its only source (disabling it makes the + tool-call "done" events vanish). ``run_turn`` wires this correctly for you. + +2. SGP tracing (enabled when ``trace_id`` is provided): opens a span named after + the tool on tool start with the tool ARGUMENTS as its input and closes it on + tool end with the result as its output, parented to ``parent_span_id``. Token + usage metrics are always emitted via ``LLMMetricsHooks`` regardless of these + flags. """ +from __future__ import annotations + +import json import logging from typing import Any, override from datetime import timedelta @@ -19,6 +40,23 @@ logger = logging.getLogger(__name__) +# Best-effort tracing budget — a tracing outage must never break a tool call. +_TRACE_TIMEOUT = timedelta(seconds=5) +# Cap tool-result span output so a large payload can't bloat the trace. +_MAX_SPAN_OUTPUT_CHARS = 2000 + + +def _get_adk() -> Any: + """Lazily import the adk facade for workflow-safe tracing. + + Kept lazy (not a module-level import) so this core hooks module does not pull + the full adk surface — and its optional deps — at import time. Only invoked + when a tool span is actually created (i.e. when ``trace_id`` is set). + """ + from agentex.lib import adk + + return adk + class TemporalStreamingHooks(LLMMetricsHooks): """Convenience hooks class for streaming OpenAI Agent lifecycle events to the AgentEx UI. @@ -27,10 +65,14 @@ class TemporalStreamingHooks(LLMMetricsHooks): AgentEx UI via Temporal activities. It subclasses the OpenAI Agents SDK's RunHooks to intercept lifecycle events and forward them for real-time UI updates. - Lifecycle events streamed: - - Tool requests (on_tool_start): Streams when a tool is about to be invoked - - Tool responses (on_tool_end): Streams the tool's execution result - - Agent handoffs (on_handoff): Streams when control transfers between agents + Lifecycle events streamed (each gated by its own flag, all default True): + - Tool requests (on_tool_start, ``emit_tool_requests``): when a tool is invoked + - Tool responses (on_tool_end, ``emit_tool_responses``): the tool's result + - Agent handoffs (on_handoff, ``emit_handoffs``): when control transfers + + Tracing (when ``trace_id`` is provided): + - One SGP span per tool call, named after the tool, with the tool + arguments as the span input and the tool result as the span output. Usage: Basic usage - streams all lifecycle events:: @@ -40,6 +82,18 @@ class TemporalStreamingHooks(LLMMetricsHooks): hooks = TemporalStreamingHooks(task_id="abc123") result = await Runner.run(agent, input, hooks=hooks) + Paired with the streaming model provider (it already streams the tool + REQUEST, so suppress that here — but keep responses, which the model does + not emit). Prefer ``run_turn`` which wires this for you:: + + hooks = TemporalStreamingHooks( + task_id="abc123", + emit_tool_requests=False, + emit_tool_responses=True, + trace_id=trace_id, + parent_span_id=parent_span_id, + ) + Advanced - subclass for custom behavior:: class MyCustomHooks(TemporalStreamingHooks): @@ -62,22 +116,82 @@ async def on_agent_start(self, context, agent): Attributes: task_id: The AgentEx task ID for routing streamed events timeout: Timeout for streaming activity calls (default: 10 seconds) + emit_tool_requests: Whether to stream the ToolRequestContent on tool start + emit_tool_responses: Whether to stream the ToolResponseContent on tool end + emit_handoffs: Whether to stream the handoff text message + trace_id: When set, tool calls are traced to SGP (input + output) + parent_span_id: Parent span for the per-tool spans """ def __init__( self, task_id: str, timeout: timedelta = timedelta(seconds=10), + *, + emit_tool_requests: bool = True, + emit_tool_responses: bool = True, + emit_handoffs: bool = True, + trace_id: str | None = None, + parent_span_id: str | None = None, ): """Initialize the streaming hooks. + Request and response emission are independently switchable because the + ``TemporalStreamingModelProvider`` emits a function tool's REQUEST from + the model output but NOT its response — the function result only ever + surfaces here via ``on_tool_end``. So when pairing with that provider, + set ``emit_tool_requests=False`` (the model already posted the request) + but keep ``emit_tool_responses=True`` (otherwise the tool-call "done" + events disappear). ``run_turn`` wires this correctly for you. + Args: task_id: AgentEx task ID for routing streamed events to the correct UI session timeout: Timeout for streaming activity invocations (default: 10 seconds) + emit_tool_requests: When True (default) stream a ToolRequestContent on + tool start. Set False when a streaming model provider already + emits the request, to avoid double-posting it. + emit_tool_responses: When True (default) stream a ToolResponseContent + on tool end. Keep True with the streaming model provider — it does + NOT emit function-tool responses, so this is their only source. + emit_handoffs: When True (default) stream a handoff text message. + trace_id: When provided, open an SGP span per tool call (named after + the tool) with the arguments as input and the result as output. When None, + no tool spans are created (token-usage metrics still emit). + parent_span_id: Parent span id the per-tool spans attach to. """ super().__init__() self.task_id = task_id self.timeout = timeout + self.emit_tool_requests = emit_tool_requests + self.emit_tool_responses = emit_tool_responses + self.emit_handoffs = emit_handoffs + self.trace_id = trace_id + self.parent_span_id = parent_span_id + # tool_call_id -> open SGP span, so on_tool_end closes the right one. + self._tool_spans: dict[str, Any] = {} + + @staticmethod + def _tool_call_id(context: RunContextWrapper, tool: Tool) -> str: + tool_context = context if isinstance(context, ToolContext) else None + return getattr(tool_context, "tool_call_id", None) or f"call_{id(tool)}" + + @staticmethod + def _parse_tool_arguments(context: RunContextWrapper) -> dict[str, Any]: + """Parse the JSON ``tool_arguments`` off a ToolContext into a dict. + + Returns an empty dict for a non-ToolContext or unparseable arguments — + a tool call must never fail because its args could not be displayed. + """ + tool_context = context if isinstance(context, ToolContext) else None + raw = getattr(tool_context, "tool_arguments", None) + if not raw: + return {} + try: + parsed = json.loads(raw) + except (json.JSONDecodeError, TypeError): + logger.warning(f"Failed to parse tool arguments: {raw!r}") + return {} + return parsed if isinstance(parsed, dict) else {"value": parsed} @override async def on_agent_start(self, context: RunContextWrapper, agent: Agent) -> None: # noqa: ARG002 @@ -102,100 +216,101 @@ async def on_agent_end(self, context: RunContextWrapper, agent: Agent, output: A agent: The agent that completed output: The agent's output """ - logger.debug(f"[TemporalStreamingHooks] Agent '{agent.name}' completed execution with output type: {type(output).__name__}") + logger.debug( + f"[TemporalStreamingHooks] Agent '{agent.name}' completed execution with output type: {type(output).__name__}" + ) @override async def on_tool_start(self, context: RunContextWrapper, agent: Agent, tool: Tool) -> None: # noqa: ARG002 - """Stream tool request when a tool starts execution. + """Stream the tool request (optional) and open a traced span (optional). - Extracts the tool_call_id and tool_arguments from the context and streams a - ToolRequestContent message to the UI showing that the tool is about to execute. + Streams a ToolRequestContent message when ``emit_tool_requests`` is True, + and opens an SGP span named after the tool (input = arguments) when + ``trace_id`` is set. Both read the same parsed arguments. Args: - context: The run context wrapper (will be a ToolContext with tool_call_id and tool_arguments) + context: The run context wrapper (a ToolContext with tool_call_id and tool_arguments) agent: The agent executing the tool tool: The tool being executed """ - import json - - tool_context = context if isinstance(context, ToolContext) else None - tool_call_id = tool_context.tool_call_id if tool_context else f"call_{id(tool)}" - - # Extract tool arguments from context - tool_arguments = {} - if tool_context and hasattr(tool_context, 'tool_arguments'): - try: - # tool_arguments is a JSON string, parse it - tool_arguments = json.loads(tool_context.tool_arguments) - except (json.JSONDecodeError, TypeError): - # If parsing fails, log and use empty dict - logger.warning(f"Failed to parse tool arguments: {tool_context.tool_arguments}") - tool_arguments = {} - - await workflow.execute_activity( - stream_lifecycle_content, - args=[ - self.task_id, - ToolRequestContent( - author="agent", - tool_call_id=tool_call_id, - name=tool.name, - arguments=tool_arguments, - ).model_dump(), - ], - start_to_close_timeout=self.timeout, - ) + tool_call_id = self._tool_call_id(context, tool) + tool_arguments = self._parse_tool_arguments(context) + + if self.emit_tool_requests: + await workflow.execute_activity( + stream_lifecycle_content, + args=[ + self.task_id, + ToolRequestContent( + author="agent", + tool_call_id=tool_call_id, + name=tool.name, + arguments=tool_arguments, + ).model_dump(), + ], + start_to_close_timeout=self.timeout, + ) + + await self._maybe_start_tool_span(tool_call_id, tool.name, tool_arguments) @override async def on_tool_end( - self, context: RunContextWrapper, agent: Agent, tool: Tool, result: str # noqa: ARG002 + self, + context: RunContextWrapper, + agent: Agent, # noqa: ARG002 + tool: Tool, + result: str, ) -> None: - """Stream tool response when a tool completes execution. + """Stream the tool response (optional) and close the traced span (optional). - Extracts the tool_call_id and streams a ToolResponseContent message to the UI - showing the tool's execution result. + Streams a ToolResponseContent message when ``emit_tool_responses`` is True, + and closes the matching tool span (output = result) when one was + opened in on_tool_start. Args: - context: The run context wrapper (will be a ToolContext with tool_call_id) + context: The run context wrapper (a ToolContext with tool_call_id) agent: The agent that executed the tool tool: The tool that was executed result: The tool's execution result """ - tool_context = context if isinstance(context, ToolContext) else None - tool_call_id = ( - getattr(tool_context, "tool_call_id", f"call_{id(tool)}") - if tool_context - else f"call_{id(tool)}" - ) - - await workflow.execute_activity( - stream_lifecycle_content, - args=[ - self.task_id, - ToolResponseContent( - author="agent", - tool_call_id=tool_call_id, - name=tool.name, - content=result, - ).model_dump(), - ], - start_to_close_timeout=self.timeout, - ) + tool_call_id = self._tool_call_id(context, tool) + + if self.emit_tool_responses: + await workflow.execute_activity( + stream_lifecycle_content, + args=[ + self.task_id, + ToolResponseContent( + author="agent", + tool_call_id=tool_call_id, + name=tool.name, + content=result, + ).model_dump(), + ], + start_to_close_timeout=self.timeout, + ) + + await self._maybe_end_tool_span(tool_call_id, result) @override async def on_handoff( - self, context: RunContextWrapper, from_agent: Agent, to_agent: Agent # noqa: ARG002 + self, + context: RunContextWrapper, + from_agent: Agent, + to_agent: Agent, # noqa: ARG002 ) -> None: """Stream handoff message when control transfers between agents. Sends a text message to the UI indicating that one agent is handing off - to another agent. + to another agent. No-op when ``emit_handoffs`` is False. Args: context: The run context wrapper from_agent: The agent transferring control to_agent: The agent receiving control """ + if not self.emit_handoffs: + return await workflow.execute_activity( stream_lifecycle_content, args=[ @@ -208,3 +323,73 @@ async def on_handoff( ], start_to_close_timeout=self.timeout, ) + + async def _maybe_start_tool_span(self, tool_call_id: str, tool_name: str, arguments: dict[str, Any]) -> None: + """Open a span named after the tool with the arguments as input. + + The span name is the bare ``tool_name`` (no prefix) to match the shared + unified-harness span reducer (``core/harness/span_derivation.py``), so + OpenAI Temporal traces look the same as every other harness. + + Best-effort: tracing must never break a tool call, so any failure is + logged and swallowed. No-op when ``trace_id`` is not set. + """ + if not self.trace_id: + return + try: + span = await _get_adk().tracing.start_span( + trace_id=self.trace_id, + parent_id=self.parent_span_id, + name=tool_name, + input={"arguments": arguments}, + start_to_close_timeout=_TRACE_TIMEOUT, + ) + if span is not None: + self._tool_spans[tool_call_id] = span + except Exception as e: # noqa: BLE001 - tracing is best-effort + logger.warning(f"[tracing] tool start_span failed (non-fatal): {e}") + + async def _maybe_end_tool_span(self, tool_call_id: str, result: Any) -> None: + """Close the span opened for ``tool_call_id`` with the result as output.""" + span = self._tool_spans.pop(tool_call_id, None) + if span is None or not self.trace_id: + return + try: + span.output = {"result": str(result)[:_MAX_SPAN_OUTPUT_CHARS]} + await _get_adk().tracing.end_span( + trace_id=self.trace_id, + span=span, + start_to_close_timeout=_TRACE_TIMEOUT, + ) + except Exception as e: # noqa: BLE001 - tracing is best-effort + logger.warning(f"[tracing] tool end_span failed (non-fatal): {e}") + + async def close_open_tool_spans(self) -> None: + """Close any tool spans still open because ``on_tool_end`` never fired. + + ``on_tool_start`` opens a span that ``on_tool_end`` is expected to close. + If the runner terminates mid-tool (max-turns exceeded, cancellation, an + unexpected SDK exception), the matching ``on_tool_end`` never runs and the + span would otherwise stay open forever — orphaned in the tracing backend. + Call this from a ``finally`` around ``Runner.run`` to drain the leftovers. + + Best-effort, like the rest of tracing: each span is closed with an + ``incomplete`` marker and any failure is logged and swallowed. + """ + if not self._tool_spans: + return + orphaned = list(self._tool_spans.items()) + self._tool_spans.clear() + for tool_call_id, span in orphaned: + logger.warning( + f"[tracing] tool span for {tool_call_id} left open (on_tool_end never fired); closing as incomplete" + ) + try: + span.output = {"result": None, "status": "incomplete"} + await _get_adk().tracing.end_span( + trace_id=self.trace_id, + span=span, + start_to_close_timeout=_TRACE_TIMEOUT, + ) + except Exception as e: # noqa: BLE001 - tracing is best-effort + logger.warning(f"[tracing] orphan tool end_span failed (non-fatal): {e}") diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/run.py b/src/agentex/lib/core/temporal/plugins/openai_agents/run.py new file mode 100644 index 000000000..0fb21bfe4 --- /dev/null +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/run.py @@ -0,0 +1,161 @@ +"""``run_turn`` — the unified entry point for the OpenAI Agents Temporal harness. + +This is the ``Runner.run`` analogue of the CLI harness's +``UnifiedEmitter.auto_send_turn``: it owns the repeatable per-turn concerns so +agents don't hand-roll them. + +What it does: + +1. Runs the agent via ``Runner.run`` with hooks that emit each tool call exactly + ONCE. The ``TemporalStreamingModelProvider`` already streams the tool-call + message from the model output, so the hooks are wired with + ``emit_messages=False`` to avoid the double-post; they still trace tool calls + (input + output) and emit token-usage metrics. +2. Normalizes token usage off the run result into a harness-independent + ``TurnUsage`` so callers can attach it to the turn span / task metadata, + matching what the CLI harness reports. + +What it deliberately does NOT do: sandboxing. Sandbox provisioning is a +composable concern carried on ``RunConfig`` (the SDK's ``SandboxRunConfig``) and +is passed straight through. Agent-specific lifecycle UI (e.g. surfacing sandbox +provisioning as a tool card) belongs in a caller-supplied ``hooks`` subclass, +not here. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any +from dataclasses import dataclass + +from agents import Runner + +from agentex.lib.utils.logging import make_logger +from agentex.lib.core.harness.types import TurnUsage +from agentex.lib.core.temporal.plugins.openai_agents.hooks.hooks import TemporalStreamingHooks + +if TYPE_CHECKING: + from agents import RunHooks, RunConfig + from agents.result import RunResult + +logger = make_logger(__name__) + +# Mirror the OpenAI Agents SDK default; callers can override per turn. +_DEFAULT_MAX_TURNS = 10 + + +@dataclass +class OpenAIAgentsTurnResult: + """The raw SDK run result plus normalized agentex usage. + + The raw ``result`` is kept so callers retain ``final_output``, + ``to_input_list()`` and any provider extras (e.g. sandbox resume state); + ``usage`` is the harness-independent token/cost summary for the turn span. + """ + + result: "RunResult" + usage: TurnUsage + + @property + def final_output(self) -> Any: + return self.result.final_output + + +def _extract_turn_usage(result: "RunResult", *, model: str | None = None) -> TurnUsage: + """Map the SDK's aggregated ``context_wrapper.usage`` onto ``TurnUsage``. + + Tolerant of a missing/partial Usage shape (non-OpenAI providers routed via + litellm may omit the nested token details) — absent fields stay None. + """ + usage = getattr(getattr(result, "context_wrapper", None), "usage", None) + if usage is None: + return TurnUsage(model=model) + + input_details = getattr(usage, "input_tokens_details", None) + output_details = getattr(usage, "output_tokens_details", None) + return TurnUsage( + model=model, + input_tokens=getattr(usage, "input_tokens", None), + output_tokens=getattr(usage, "output_tokens", None), + total_tokens=getattr(usage, "total_tokens", None), + cached_input_tokens=getattr(input_details, "cached_tokens", None), + reasoning_tokens=getattr(output_details, "reasoning_tokens", None), + num_llm_calls=getattr(usage, "requests", None), + ) + + +async def run_turn( + starting_agent: Any, + input: Any, + *, + task_id: str, + trace_id: str | None = None, + parent_span_id: str | None = None, + run_config: "RunConfig | None" = None, + hooks: "RunHooks | None" = None, + model: str | None = None, + max_turns: int = _DEFAULT_MAX_TURNS, +) -> OpenAIAgentsTurnResult: + """Run one agent turn and return the result plus normalized usage. + + Args: + starting_agent: The agent to run. + input: The input list / string passed to ``Runner.run``. + task_id: AgentEx task id for streaming. + trace_id: When set, tool calls are traced to SGP (input + output). Only + applied when ``hooks`` is omitted (it flows into the default + ``TemporalStreamingHooks``). Ignored when you pass your own ``hooks`` + — see ``hooks`` below. + parent_span_id: Parent span for the per-tool spans (typically the turn + span). Same caveat as ``trace_id``: only applied to the default hooks. + run_config: Forwarded to ``Runner.run`` verbatim (carries the model + provider and any ``SandboxRunConfig``). Left untouched here. + hooks: Optional hooks override. When omitted, a default + ``TemporalStreamingHooks(emit_tool_requests=False, ...)`` is used so + the streaming model is the sole tool-REQUEST emitter while the hooks + still emit tool RESPONSES (the model does not), and ``trace_id`` / + ``parent_span_id`` are forwarded into it. When you pass your own + subclass (also with ``emit_tool_requests=False``) to add agent-specific + lifecycle behavior such as a sandbox-ready card, ``trace_id`` and + ``parent_span_id`` are NOT applied for you — pass them to your + subclass's constructor yourself if you want tool spans traced. + model: Model name recorded on the returned usage; derived from the agent + when not supplied. + max_turns: Forwarded to ``Runner.run``. + + Returns: + OpenAIAgentsTurnResult with the raw run result and normalized usage. + """ + if hooks is None: + hooks = TemporalStreamingHooks( + task_id=task_id, + # The streaming model already posts the tool REQUEST, so suppress it + # here (no double-post) — but keep responses, which the model does not + # emit for function tools (on_tool_end is their only source). + emit_tool_requests=False, + emit_tool_responses=True, + trace_id=trace_id, + parent_span_id=parent_span_id, + ) + + run_kwargs: dict[str, Any] = {"hooks": hooks, "max_turns": max_turns} + if run_config is not None: + run_kwargs["run_config"] = run_config + + try: + result = await Runner.run(starting_agent, input, **run_kwargs) + finally: + # If the runner terminated mid-tool (max-turns, cancellation, SDK error), + # on_tool_end never fired for the in-flight call, leaving its span open. + # Drain any leftovers so they don't orphan in the tracing backend. + if isinstance(hooks, TemporalStreamingHooks): + await hooks.close_open_tool_spans() + + resolved_model = model + if resolved_model is None: + agent_model = getattr(starting_agent, "model", None) + resolved_model = str(agent_model) if agent_model else None + + return OpenAIAgentsTurnResult( + result=result, + usage=_extract_turn_usage(result, model=resolved_model), + ) diff --git a/src/agentex/lib/core/temporal/plugins/openai_agents/tests/test_run_turn_and_hooks.py b/src/agentex/lib/core/temporal/plugins/openai_agents/tests/test_run_turn_and_hooks.py new file mode 100644 index 000000000..244182ac5 --- /dev/null +++ b/src/agentex/lib/core/temporal/plugins/openai_agents/tests/test_run_turn_and_hooks.py @@ -0,0 +1,247 @@ +"""Tests for the unified OpenAI-Agents turn surface. + +Covers: +- ``TemporalStreamingHooks`` message-emission gating (``emit_messages``), so the + streaming model can be the sole tool-message emitter (no double-post). +- ``TemporalStreamingHooks`` input-bearing tool spans (input = arguments, + output = result) when a ``trace_id`` is provided. +- ``run_turn`` usage extraction and default-hooks wiring. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from agents.tool_context import ToolContext + +from agentex.lib.core.temporal.plugins.openai_agents import run as run_mod +from agentex.lib.core.temporal.plugins.openai_agents.hooks import hooks as hooks_mod + +TemporalStreamingHooks = hooks_mod.TemporalStreamingHooks + + +def _tool_context(args: str = '{"query": "hi"}') -> ToolContext: + return ToolContext( + context=None, + tool_name="search", + tool_call_id="call_abc", + tool_arguments=args, + ) + + +def _tool() -> MagicMock: + tool = MagicMock() + tool.name = "search" + return tool + + +# --------------------------------------------------------------------------- # +# Argument parsing +# --------------------------------------------------------------------------- # + + +def test_parse_tool_arguments_valid_dict(): + assert TemporalStreamingHooks._parse_tool_arguments(_tool_context('{"a": 1}')) == {"a": 1} + + +def test_parse_tool_arguments_garbage_is_empty(): + assert TemporalStreamingHooks._parse_tool_arguments(_tool_context("not json")) == {} + + +def test_parse_tool_arguments_non_tool_context_is_empty(): + assert TemporalStreamingHooks._parse_tool_arguments(SimpleNamespace()) == {} + + +# --------------------------------------------------------------------------- # +# Message emission gating (the double-post fix + the response-survival guard) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_defaults_stream_tool_request(monkeypatch): + exec_activity = AsyncMock() + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", exec_activity) + + hooks = TemporalStreamingHooks(task_id="t1") # all emit flags default True + await hooks.on_tool_start(_tool_context(), MagicMock(), _tool()) + + exec_activity.assert_awaited_once() + # args=[task_id, ToolRequestContent.model_dump()] + _, kwargs = exec_activity.call_args + payload = kwargs["args"][1] + assert payload["name"] == "search" + assert payload["arguments"] == {"query": "hi"} + + +@pytest.mark.asyncio +async def test_requests_off_skips_request_but_keeps_response(monkeypatch): + """The streaming-model pairing: suppress the duplicate REQUEST, but the + RESPONSE must still emit (the model never emits function-tool responses).""" + exec_activity = AsyncMock() + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", exec_activity) + + hooks = TemporalStreamingHooks(task_id="t1", emit_tool_requests=False, emit_tool_responses=True) + await hooks.on_tool_start(_tool_context(), MagicMock(), _tool()) + exec_activity.assert_not_awaited() # request suppressed + + await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "the result") + exec_activity.assert_awaited_once() # response still emitted + _, kwargs = exec_activity.call_args + payload = kwargs["args"][1] + assert payload["name"] == "search" + assert payload["content"] == "the result" + + +@pytest.mark.asyncio +async def test_responses_off_skips_response(monkeypatch): + exec_activity = AsyncMock() + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", exec_activity) + + hooks = TemporalStreamingHooks(task_id="t1", emit_tool_responses=False) + await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "result") + + exec_activity.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_emit_handoffs_false_skips_handoff(monkeypatch): + exec_activity = AsyncMock() + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", exec_activity) + + hooks = TemporalStreamingHooks(task_id="t1", emit_handoffs=False) + await hooks.on_handoff(MagicMock(), MagicMock(name="from"), MagicMock(name="to")) + + exec_activity.assert_not_awaited() + + +# --------------------------------------------------------------------------- # +# Input-bearing tool spans (the "traces have outputs but no inputs" fix) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_tool_span_carries_input_and_output(monkeypatch): + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", AsyncMock()) + span = SimpleNamespace(output=None) + start_span = AsyncMock(return_value=span) + end_span = AsyncMock() + fake_adk = SimpleNamespace(tracing=SimpleNamespace(start_span=start_span, end_span=end_span)) + monkeypatch.setattr(hooks_mod, "_get_adk", lambda: fake_adk) + + hooks = TemporalStreamingHooks( + task_id="t1", emit_tool_requests=False, trace_id="trace-1", parent_span_id="parent-1" + ) + await hooks.on_tool_start(_tool_context(), MagicMock(), _tool()) + + start_span.assert_awaited_once() + _, kwargs = start_span.call_args + assert kwargs["name"] == "search" + assert kwargs["parent_id"] == "parent-1" + assert kwargs["input"] == {"arguments": {"query": "hi"}} + + await hooks.on_tool_end(_tool_context(), MagicMock(), _tool(), "the answer") + end_span.assert_awaited_once() + assert span.output == {"result": "the answer"} + + +@pytest.mark.asyncio +async def test_no_trace_id_means_no_span(monkeypatch): + monkeypatch.setattr(hooks_mod.workflow, "execute_activity", AsyncMock()) + start_span = AsyncMock() + fake_adk = SimpleNamespace(tracing=SimpleNamespace(start_span=start_span)) + monkeypatch.setattr(hooks_mod, "_get_adk", lambda: fake_adk) + + hooks = TemporalStreamingHooks(task_id="t1", emit_tool_requests=False, trace_id=None) + await hooks.on_tool_start(_tool_context(), MagicMock(), _tool()) + + start_span.assert_not_awaited() + + +# --------------------------------------------------------------------------- # +# Usage extraction +# --------------------------------------------------------------------------- # + + +def _result_with_usage() -> SimpleNamespace: + usage = SimpleNamespace( + requests=3, + input_tokens=100, + output_tokens=40, + total_tokens=140, + input_tokens_details=SimpleNamespace(cached_tokens=20), + output_tokens_details=SimpleNamespace(reasoning_tokens=10), + ) + return SimpleNamespace(context_wrapper=SimpleNamespace(usage=usage), final_output="done") + + +def test_extract_turn_usage_maps_fields(): + usage = run_mod._extract_turn_usage(_result_with_usage(), model="openai/gpt-5.5") + assert usage.model == "openai/gpt-5.5" + assert usage.input_tokens == 100 + assert usage.output_tokens == 40 + assert usage.total_tokens == 140 + assert usage.cached_input_tokens == 20 + assert usage.reasoning_tokens == 10 + assert usage.num_llm_calls == 3 + + +def test_extract_turn_usage_missing_usage_is_tolerant(): + usage = run_mod._extract_turn_usage(SimpleNamespace(), model="m") + assert usage.model == "m" + assert usage.input_tokens is None + assert usage.num_llm_calls is None + + +# --------------------------------------------------------------------------- # +# run_turn +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_run_turn_returns_usage_and_passes_through_result(monkeypatch): + fake_result = _result_with_usage() + runner_run = AsyncMock(return_value=fake_result) + monkeypatch.setattr(run_mod.Runner, "run", runner_run) + + agent = SimpleNamespace(model="openai/gpt-5.5") + out = await run_mod.run_turn( + agent, + [{"role": "user", "content": "hi"}], + task_id="t1", + trace_id="trace-1", + parent_span_id="parent-1", + ) + + assert isinstance(out, run_mod.OpenAIAgentsTurnResult) + assert out.final_output == "done" + assert out.usage.total_tokens == 140 + assert out.usage.model == "openai/gpt-5.5" + + # Default hooks must be wired so the streaming model is the sole tool-REQUEST + # emitter, while the hooks still emit tool RESPONSES (the model does not). + runner_run.assert_awaited_once() + _, kwargs = runner_run.call_args + hooks = kwargs["hooks"] + assert hooks.emit_tool_requests is False + assert hooks.emit_tool_responses is True + assert hooks.trace_id == "trace-1" + assert hooks.parent_span_id == "parent-1" + + +@pytest.mark.asyncio +async def test_run_turn_respects_supplied_hooks(monkeypatch): + runner_run = AsyncMock(return_value=_result_with_usage()) + monkeypatch.setattr(run_mod.Runner, "run", runner_run) + + custom_hooks = TemporalStreamingHooks(task_id="t1", emit_tool_requests=False) + await run_mod.run_turn( + SimpleNamespace(model="m"), + "hi", + task_id="t1", + hooks=custom_hooks, + ) + + _, kwargs = runner_run.call_args + assert kwargs["hooks"] is custom_hooks