fix(openai-agents): address governance review — trace_id, counters, graph walk

Review findings (Viswa) for PR #358:

- Drop the per-hooks uuid trace_id: minting one at install time was
  identical across every call and diverged across handoff nodes. Trace
  correlation is owned by the layer below (OTel span / HTTP resolve),
  matching the LangChain adapter. Requires uipath-core >= 0.5.20, which
  removed trace_id from EvaluatorProtocol — bumped in the lock.
- Count llm/tool calls only after governance passes: a DENY raised
  before the bump, inflating the counter on blocked calls.
- Walk the handoff graph by isinstance(node, Agent) instead of
  duck-typing hasattr(node, "hooks") — the SDK type-checks the slot, so
  non-Agent objects could slip through.
- Log (not silently drop) when the graph walk hits the node cap.
- Report the live executing agent's name rather than the install-time
  entrypoint name, so attribution is correct after a handoff.
- Treat an item as a function call only when explicitly typed
  function_call or it carries arguments (no bare-name misclassification).
- Cap _stringify output so an oversized tool result / arg blob can't
  hand a multi-megabyte string to the evaluator.
- Trailing newline (W292).

Tests: FakeAgent now subclasses agents.Agent (graph walk isinstance-checks
Agent); added live-agent-name and no-inflation-on-block coverage; made the
caplog assertions propagation-independent.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: aditik0303

packages/uipath-openai-agents/src/uipath_openai_agents/governance/hooks.py

@@ -50,7 +50,6 @@
 import json
 import logging
 from typing import Any, Dict, List
-from uuid import uuid4
 
 from agents import Agent, AgentHooks
 from uipath.core.adapters import EvaluatorProtocol
@@ -65,6 +64,10 @@
 # full prompt — see :func:`_latest_input_text`.
 _BEFORE_MODEL_TEXT_CAP = 64000
 
+# Hard cap on how many nodes the handoff-graph walk visits, guarding against
+# cyclic or pathologically deep agent graphs. Hitting it is logged, not silent.
+_MAX_GRAPH_NODES = 1000
+
 
 def install_governance(
     agent: Agent,
@@ -108,30 +111,47 @@ def install_governance(
 
 
 def _iter_agents(root: Any) -> List[Any]:
-    """Return every agent node reachable through the ``handoffs`` graph.
-
-    A node qualifies if it exposes the ``hooks`` slot. Handoff targets may be
-    ``Agent`` instances or ``Handoff`` objects that carry the target on
-    ``.agent``; both are followed so a multi-agent app is governed end to end.
-    Cycles and pathological depth are bounded by an id-visited set and a hard
-    cap.
+    """Return every ``Agent`` reachable through the ``handoffs`` graph.
+
+    A node qualifies only if it is a real :class:`agents.Agent`. The SDK
+    type-checks the ``hooks`` slot, so duck-typing on ``hasattr(node, "hooks")``
+    could let non-Agent objects through — we isinstance-check instead. Handoff
+    targets may be ``Agent`` instances or ``Handoff`` objects that carry the
+    target on ``.agent``; both are followed so a multi-agent app is governed end
+    to end. Cycles and pathological depth are bounded by an id-visited set and a
+    hard cap (``_MAX_GRAPH_NODES``), which logs rather than silently truncating.
+
+    Not walked: agents reachable only as tools (``agent.as_tool()``) or embedded
+    in input/output guardrail functions — the SDK closes over those behind
+    opaque callables, so they are governed by their own runtime rather than this
+    graph walk.
     """
     found: List[Any] = []
     seen: set[int] = set()
     stack: List[Any] = [root]
-    while stack and len(seen) < 1000:
+    capped = False
+    while stack:
+        if len(seen) >= _MAX_GRAPH_NODES:
+            capped = True
+            break
         node = stack.pop()
         if node is None or id(node) in seen:
             continue
         seen.add(id(node))
-        if hasattr(node, "hooks"):
+        if isinstance(node, Agent):
             found.append(node)
         handoffs = getattr(node, "handoffs", None)
         if isinstance(handoffs, (list, tuple)):
             for h in handoffs:
                 # A Handoff wraps its target agent on ``.agent``; a bare Agent
                 # is itself the target.
                 stack.append(getattr(h, "agent", h))
+    if capped:
+        logger.warning(
+            "install_governance stopped walking the agent graph at the %d-node "
+            "cap; agents beyond it will not be governed",
+            _MAX_GRAPH_NODES,
+        )
     return found
 
 
@@ -158,9 +178,24 @@ def __init__(
         self._agent_name = agent_name
         self._session_id = session_id
         self._inner = inner
-        self._trace_id = str(uuid4())
+        # ``trace_id`` is intentionally NOT held here. A single uuid minted at
+        # install time would be identical for every model/tool call and would
+        # diverge across handoff nodes (each carries its own hooks). Trace
+        # correlation is owned by the layer below: OTel-backed sinks read the
+        # live span on the caller's thread, HTTP consumers resolve the canonical
+        # id at call time. This matches the LangChain adapter.
         self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0}
 
+    def _resolve_agent_name(self, agent: Any) -> str:
+        """Prefer the live executing agent's name over the install-time name.
+
+        After a handoff the running node may differ from the graph entrypoint
+        the factory named us with; reporting the actual agent gives governance
+        accurate attribution. Falls back to the install-time name.
+        """
+        name = getattr(agent, "name", None)
+        return name if isinstance(name, str) and name else self._agent_name
+
     # ----- Model hooks -----------------------------------------------------
 
     async def on_llm_start(
@@ -180,31 +215,33 @@ async def on_llm_start(
         the prompt for context.
         """
         try:
-            self._session_state["llm_calls"] = (
-                self._session_state.get("llm_calls", 0) + 1
-            )
             model_input = _latest_input_text(input_items)
             self._evaluator.evaluate_before_model(
                 model_input=model_input,
-                agent_name=self._agent_name,
+                agent_name=self._resolve_agent_name(agent),
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
+            )
+            # Count only calls that passed governance — a DENY raises above, so
+            # a blocked call must not inflate the counter.
+            self._session_state["llm_calls"] = (
+                self._session_state.get("llm_calls", 0) + 1
             )
         except GovernanceBlockException:
             raise
         except Exception as e:  # noqa: BLE001 - governance must not break the run
             logger.warning("on_llm_start governance check failed (continuing): %s", e)
-        await _delegate(self._inner, "on_llm_start", context, agent, system_prompt, input_items)
+        await _delegate(
+            self._inner, "on_llm_start", context, agent, system_prompt, input_items
+        )
 
     async def on_llm_end(self, context: Any, agent: Any, response: Any) -> None:
         """Evaluate AFTER_MODEL rules immediately after the LLM response."""
         try:
             model_output = _model_response_text(response)
             self._evaluator.evaluate_after_model(
                 model_output=model_output,
-                agent_name=self._agent_name,
+                agent_name=self._resolve_agent_name(agent),
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
             )
         except GovernanceBlockException:
             raise
@@ -223,18 +260,19 @@ async def on_tool_start(self, context: Any, agent: Any, tool: Any) -> None:
         at the model layer where the call's arguments are visible in the output.
         """
         try:
-            self._session_state["tool_calls"] = (
-                self._session_state.get("tool_calls", 0) + 1
-            )
             tool_name = getattr(tool, "name", None) or "unknown"
             self._evaluator.evaluate_tool_call(
                 tool_name=tool_name,
                 tool_args={},
-                agent_name=self._agent_name,
+                agent_name=self._resolve_agent_name(agent),
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
                 session_state=self._session_state,
             )
+            # Count only calls that passed governance; the evaluator saw the
+            # count of prior tool calls, and a DENY raises before this bump.
+            self._session_state["tool_calls"] = (
+                self._session_state.get("tool_calls", 0) + 1
+            )
         except GovernanceBlockException:
             raise
         except Exception as e:  # noqa: BLE001
@@ -256,9 +294,8 @@ async def on_tool_end(
             self._evaluator.evaluate_after_tool(
                 tool_name=tool_name,
                 tool_result=tool_result,
-                agent_name=self._agent_name,
+                agent_name=self._resolve_agent_name(agent),
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
             )
         except GovernanceBlockException:
             raise
@@ -341,10 +378,13 @@ def _item_text(item: Any) -> str:
 
     pieces: List[str] = []
 
-    # A function/tool call carries its intent in name + arguments.
+    # A function/tool call carries its intent in name + arguments. Treat an
+    # item as a call only when it is explicitly typed ``function_call`` or it
+    # actually carries arguments — a bare ``name`` on some other item type (a
+    # named message part) is not a tool call.
     name = _get(item, "name")
     arguments = _get(item, "arguments")
-    if name and (_get(item, "type") in (None, "function_call") or arguments is not None):
+    if name and (_get(item, "type") == "function_call" or arguments is not None):
         if isinstance(name, str):
             pieces.append(name)
         if arguments is not None:
@@ -419,11 +459,15 @@ def _get(obj: Any, attr: str) -> Any:
     return getattr(obj, attr, None)
 
 
-def _stringify(value: Any) -> str:
-    """Render a dict / object payload as compact, scannable text."""
+def _stringify(value: Any, cap: int = _BEFORE_MODEL_TEXT_CAP) -> str:
+    """Render a dict / object payload as compact, scannable text, capped.
+
+    The result is bounded by ``cap`` so an oversized tool result or argument
+    blob can't hand a multi-megabyte string to the evaluator.
+    """
     if isinstance(value, str):
-        return value
+        return value[:cap]
     try:
-        return json.dumps(value, default=str, ensure_ascii=False)
+        return json.dumps(value, default=str, ensure_ascii=False)[:cap]
     except (TypeError, ValueError):
-        return str(value)
+        return str(value)[:cap]

packages/uipath-openai-agents/tests/governance/test_hooks.py

@@ -14,10 +14,12 @@
 from __future__ import annotations
 
 import logging
+from contextlib import contextmanager
 from types import SimpleNamespace
-from typing import Any, List
+from typing import Any, Iterator, List
 
 import pytest
+from agents import Agent
 from uipath.core.governance.exceptions import GovernanceBlockException
 
 from uipath_openai_agents.governance.hooks import (
@@ -62,14 +64,14 @@ def evaluate_after_tool(self, **kwargs: Any) -> None:
         self._record("after_tool", **kwargs)
 
 
-class FakeAgent:
-    """Minimal stand-in for ``agents.Agent`` (duck-typed by the adapter)."""
+class FakeAgent(Agent):  # type: ignore[type-arg]
+    """A real ``agents.Agent`` — the graph walk isinstance-checks ``Agent``, so
+    a bare duck-typed stand-in would be (correctly) skipped by
+    ``install_governance``. Subclassing keeps the construction lightweight while
+    remaining a genuine ``Agent`` instance."""
 
     def __init__(self, name: str = "agent", handoffs: List[Any] | None = None):
-        self.name = name
-        self.hooks: Any = None
-        self.tools: List[Any] = []
-        self.handoffs = handoffs or []
+        super().__init__(name=name, handoffs=handoffs or [])
 
 
 class FakeTool:
@@ -132,9 +134,7 @@ def test_install_governance_installs_on_all_agents_in_handoff_graph():
     leaf_b = FakeAgent("b")
     root = FakeAgent("root", handoffs=[leaf_a, leaf_b])
 
-    returned = install_governance(
-        root, FakeEvaluator(), agent_name="x", session_id="s"
-    )
+    returned = install_governance(root, FakeEvaluator(), agent_name="x", session_id="s")
 
     assert returned is root  # original returned, not a proxy
     for node in (root, leaf_a, leaf_b):
@@ -167,8 +167,30 @@ def test_install_governance_chains_existing_hooks():
     assert agent.hooks._inner is user_hooks
 
 
+_HOOKS_LOGGER = "uipath_openai_agents.governance.hooks"
+
+
+@contextmanager
+def _capture_hooks_logs(caplog: Any) -> Iterator[None]:
+    """Attach caplog's handler straight to the hooks logger.
+
+    Some sibling suites configure an ancestor ``uipath*`` logger with
+    ``propagate=False``, which silently breaks caplog's default root-handler
+    capture. Attaching directly to the target logger is propagation-independent.
+    """
+    logger = logging.getLogger(_HOOKS_LOGGER)
+    logger.addHandler(caplog.handler)
+    prev = logger.level
+    logger.setLevel(logging.WARNING)
+    try:
+        yield
+    finally:
+        logger.removeHandler(caplog.handler)
+        logger.setLevel(prev)
+
+
 def test_install_governance_warns_when_no_agent(caplog):
-    with caplog.at_level(logging.WARNING):
+    with _capture_hooks_logs(caplog):
         install_governance(object(), FakeEvaluator(), agent_name="x", session_id="s")  # type: ignore[arg-type]
     assert any("no Agent" in r.message for r in caplog.records)
 
@@ -282,6 +304,26 @@ async def test_on_tool_end_none_result():
     assert ev.calls[-1][1]["tool_result"] == ""
 
 
+async def test_reports_live_agent_name_not_install_time_name():
+    """After a handoff the executing agent differs from the graph entrypoint
+    the factory named us with; governance should attribute the live agent."""
+    ev = FakeEvaluator()
+    cb = _make_hooks(ev)  # install-time name is "agent-1"
+    await cb.on_llm_start(None, FakeAgent("billing_specialist"), None, [_msg("hi")])
+    assert ev.calls[-1][1]["agent_name"] == "billing_specialist"
+
+
+async def test_blocked_call_does_not_increment_counter():
+    """A DENY raises before the counter bump, so the count is not inflated."""
+    ev = FakeEvaluator(block_on="tool_call")
+    cb = _make_hooks(ev)
+    with pytest.raises(GovernanceBlockException):
+        await cb.on_tool_start(None, FakeAgent(), FakeTool("t"))
+    # evaluator saw the pre-call count (0) and the block prevented the bump
+    assert ev.calls[-1][1]["session_state"]["tool_calls"] == 0
+    assert cb._session_state["tool_calls"] == 0
+
+
 # --------------------------------------------------------------------------
 # chaining to user hooks
 # --------------------------------------------------------------------------
@@ -305,10 +347,19 @@ async def test_governance_delegates_to_inner_hooks():
 @pytest.mark.parametrize(
     "hook,invoke",
     [
-        ("before_model", lambda cb: cb.on_llm_start(None, FakeAgent(), None, [_msg("hi")])),
-        ("after_model", lambda cb: cb.on_llm_end(None, FakeAgent(), SimpleNamespace(output=[]))),
+        (
+            "before_model",
+            lambda cb: cb.on_llm_start(None, FakeAgent(), None, [_msg("hi")]),
+        ),
+        (
+            "after_model",
+            lambda cb: cb.on_llm_end(None, FakeAgent(), SimpleNamespace(output=[])),
+        ),
         ("tool_call", lambda cb: cb.on_tool_start(None, FakeAgent(), FakeTool("t"))),
-        ("after_tool", lambda cb: cb.on_tool_end(None, FakeAgent(), FakeTool("t"), {"r": 1})),
+        (
+            "after_tool",
+            lambda cb: cb.on_tool_end(None, FakeAgent(), FakeTool("t"), {"r": 1}),
+        ),
     ],
 )
 async def test_block_exception_propagates(hook, invoke):
@@ -327,7 +378,7 @@ def evaluate_before_model(self, **_: Any) -> None:
         agent_name="a",
         session_id="s",
     )
-    with caplog.at_level(logging.WARNING):
+    with _capture_hooks_logs(caplog):
         # must NOT raise — a governance bug can't break the agent run
         await cb.on_llm_start(None, FakeAgent(), None, [_msg("x")])
     assert any("governance check failed" in r.message for r in caplog.records)
@@ -357,7 +408,9 @@ async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch):
     from uipath_openai_agents.runtime import factory as factory_mod
 
     # Stub the runtime so we don't introspect a real Agent.
-    monkeypatch.setattr(factory_mod, "UiPathOpenAIAgentRuntime", lambda **kw: SimpleNamespace(**kw))
+    monkeypatch.setattr(
+        factory_mod, "UiPathOpenAIAgentRuntime", lambda **kw: SimpleNamespace(**kw)
+    )
     agent = FakeAgent()
     await _factory_without_init()._create_runtime_instance(
         agent=agent, runtime_id="r", entrypoint="e", evaluator=FakeEvaluator()
@@ -368,9 +421,11 @@ async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch):
 async def test_factory_skips_governance_without_evaluator(monkeypatch):
     from uipath_openai_agents.runtime import factory as factory_mod
 
-    monkeypatch.setattr(factory_mod, "UiPathOpenAIAgentRuntime", lambda **kw: SimpleNamespace(**kw))
+    monkeypatch.setattr(
+        factory_mod, "UiPathOpenAIAgentRuntime", lambda **kw: SimpleNamespace(**kw)
+    )
     agent = FakeAgent()
     await _factory_without_init()._create_runtime_instance(
         agent=agent, runtime_id="r", entrypoint="e"
     )
-    assert agent.hooks is None
+    assert agent.hooks is None