fix(agent-framework): address governance review — streaming, try/finally

Review findings (Viswa) for PR #361:

- AFTER_MODEL / AFTER_TOOL now run inside a try/finally around call_next,
  so they still fire (audit + rules observe the turn) when the model or
  tool call raises. GovernanceBlockException from before_* still aborts
  before call_next; the underlying error still propagates after the finally.
- Streaming AFTER_MODEL: on a streaming ChatContext, context.result is a
  ResponseStream after call_next (reading it yielded empty text). We now
  register a stream_result_hook so AFTER_MODEL runs on the finalized
  ChatResponse the framework assembles once the stream is consumed.
- Drop the per-callbacks uuid trace_id (identical for every call); trace
  correlation is owned by the layer below, matching LangChain. Requires
  uipath-core >= 0.5.20 (removed trace_id from EvaluatorProtocol) — bumped.
- Count llm/tool calls only after governance passes: a DENY raised before
  the bump, inflating the counter on blocked calls.
- Cap _stringify output so an oversized tool result / message can't hand a
  multi-megabyte string to the evaluator.
- Trailing newline (W292).

Note: agents are not cached by the factory (each runtime gets a fresh
instance to avoid concurrent-Workflow reuse), so the cached-agent
re-attach concern does not apply here — the idempotency guard is only a
double-install safety net. The PR description's entry-point claim is stale
(factory-evaluator uses no governance entry-point) and will be corrected.

Tests: added finally-on-error (model + tool), streaming result-hook, and
no-inflation-on-block coverage.

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

Author: aditik0303

packages/uipath-agent-framework/src/uipath_agent_framework/governance/middleware.py

@@ -45,7 +45,6 @@
 import logging
 from collections.abc import Mapping
 from typing import Any, Awaitable, Callable, Dict, List
-from uuid import uuid4
 
 from agent_framework._middleware import (
     ChatContext,
@@ -158,22 +157,26 @@ def __init__(
         self._evaluator = evaluator
         self._agent_name = agent_name
         self._session_id = session_id
-        self._trace_id = str(uuid4())
+        # ``trace_id`` is intentionally NOT held here. A single uuid minted at
+        # install time would be identical for every call. Trace correlation is
+        # owned by the layer below (OTel span / HTTP resolve at call time),
+        # matching the LangChain adapter.
         self._session_state: Dict[str, Any] = {"tool_calls": 0, "llm_calls": 0}
 
     # ----- Model --------------------------------------------------------
 
     def before_model(self, messages: Any) -> None:
         """Evaluate BEFORE_MODEL on the latest message only (see ADK rationale)."""
         try:
-            self._session_state["llm_calls"] = (
-                self._session_state.get("llm_calls", 0) + 1
-            )
             self._evaluator.evaluate_before_model(
                 model_input=self._latest_message_text(messages),
                 agent_name=self._agent_name,
                 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
@@ -187,7 +190,6 @@ def after_model(self, result: Any) -> None:
                 model_output=self._response_text(result),
                 agent_name=self._agent_name,
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
             )
         except GovernanceBlockException:
             raise
@@ -199,17 +201,18 @@ def after_model(self, result: Any) -> None:
     def before_tool(self, function: Any, arguments: Any) -> None:
         """Evaluate TOOL_CALL with the tool name + arguments."""
         try:
-            self._session_state["tool_calls"] = (
-                self._session_state.get("tool_calls", 0) + 1
-            )
             self._evaluator.evaluate_tool_call(
                 tool_name=getattr(function, "name", None) or "unknown",
                 tool_args=_coerce_args(arguments),
                 agent_name=self._agent_name,
                 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
@@ -223,7 +226,6 @@ def after_tool(self, function: Any, result: Any) -> None:
                 tool_result="" if result is None else _stringify(result),
                 agent_name=self._agent_name,
                 runtime_id=self._session_id,
-                trace_id=self._trace_id,
             )
         except GovernanceBlockException:
             raise
@@ -281,8 +283,39 @@ async def process(
         self, context: ChatContext, call_next: Callable[[], Awaitable[None]]
     ) -> None:
         self._cb.before_model(getattr(context, "messages", None))
-        await call_next()
-        self._cb.after_model(getattr(context, "result", None))
+
+        if getattr(context, "stream", False):
+            # Streaming: after ``call_next`` ``context.result`` is a
+            # ResponseStream, not finalized text — reading it for AFTER_MODEL
+            # yields nothing. Register a result hook so AFTER_MODEL runs on the
+            # finalized ``ChatResponse`` the framework assembles once the stream
+            # is consumed.
+            hooks = getattr(context, "stream_result_hooks", None)
+            if isinstance(hooks, list):
+                hooks.append(self._govern_streamed_result)
+            else:  # pragma: no cover - defensive: framework always provides it
+                logger.debug(
+                    "ChatContext has no stream_result_hooks; AFTER_MODEL will "
+                    "not run for this streamed response"
+                )
+            await call_next()
+            return
+
+        try:
+            await call_next()
+        finally:
+            # AFTER_MODEL must run even if the model call raised, so audit and
+            # rules still observe whatever result is present.
+            self._cb.after_model(getattr(context, "result", None))
+
+    def _govern_streamed_result(self, response: Any) -> Any:
+        """``stream_result_hook``: govern the finalized streamed ``ChatResponse``.
+
+        Returns the response unchanged (governance observes, it does not
+        rewrite). A DENY raised here still propagates to abort the run.
+        """
+        self._cb.after_model(response)
+        return response
 
 
 class GovernanceFunctionMiddleware(FunctionMiddleware):
@@ -298,8 +331,11 @@ async def process(
     ) -> None:
         function = getattr(context, "function", None)
         self._cb.before_tool(function, getattr(context, "arguments", None))
-        await call_next()
-        self._cb.after_tool(function, getattr(context, "result", None))
+        try:
+            await call_next()
+        finally:
+            # AFTER_TOOL must run even if the tool call raised.
+            self._cb.after_tool(function, getattr(context, "result", None))
 
 
 # Tuple used for isinstance idempotency / detach checks.
@@ -328,11 +364,17 @@ def _coerce_args(arguments: Any) -> Dict[str, Any]:
     return {}
 
 
-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.
+
+    Bounded by ``cap`` so an oversized tool result or message payload can't
+    hand a multi-megabyte string to the evaluator. Callers that slice the
+    result again (the ``_message_text`` / ``_response_text`` fallbacks) are
+    unaffected.
+    """
     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-agent-framework/tests/governance/test_middleware.py

@@ -108,7 +108,9 @@ async def _noop_next() -> None:
 
 def test_install_governance_appends_both_middleware():
     agent = FakeAgent()
-    returned = install_governance(agent, FakeEvaluator(), agent_name="x", session_id="s")
+    returned = install_governance(
+        agent, FakeEvaluator(), agent_name="x", session_id="s"
+    )
     assert returned is agent
     kinds = [type(m) for m in agent.middleware]
     assert GovernanceChatMiddleware in kinds
@@ -158,20 +160,28 @@ def _factory_without_init():
         UiPathAgentFrameworkRuntimeFactory,
     )
 
-    return UiPathAgentFrameworkRuntimeFactory.__new__(UiPathAgentFrameworkRuntimeFactory)
+    return UiPathAgentFrameworkRuntimeFactory.__new__(
+        UiPathAgentFrameworkRuntimeFactory
+    )
 
 
 def _stub_factory_runtime(monkeypatch, factory_mod):
     """Stub storage + runtime constructions so only the governance branch runs."""
     monkeypatch.setattr(factory_mod, "ScopedCheckpointStorage", lambda *a, **k: None)
-    monkeypatch.setattr(factory_mod, "UiPathAgentFrameworkRuntime", lambda **kw: SimpleNamespace(**kw))
-    monkeypatch.setattr(factory_mod, "UiPathResumableRuntime", lambda **kw: SimpleNamespace(**kw))
+    monkeypatch.setattr(
+        factory_mod, "UiPathAgentFrameworkRuntime", lambda **kw: SimpleNamespace(**kw)
+    )
+    monkeypatch.setattr(
+        factory_mod, "UiPathResumableRuntime", lambda **kw: SimpleNamespace(**kw)
+    )
     monkeypatch.setattr(factory_mod, "UiPathResumeTriggerHandler", lambda *a, **k: None)
 
     async def _storage(self):
         return SimpleNamespace(checkpoint_storage=object())
 
-    monkeypatch.setattr(factory_mod.UiPathAgentFrameworkRuntimeFactory, "_get_storage", _storage)
+    monkeypatch.setattr(
+        factory_mod.UiPathAgentFrameworkRuntimeFactory, "_get_storage", _storage
+    )
 
 
 async def test_factory_installs_governance_when_evaluator_supplied(monkeypatch):
@@ -231,6 +241,45 @@ async def test_chat_middleware_caps_text():
     assert len(ev.calls[0][1]["model_input"]) <= _BEFORE_MODEL_TEXT_CAP
 
 
+async def test_after_model_runs_even_when_model_call_raises():
+    """AFTER_MODEL must fire from the finally so audit/rules still observe the
+    turn, and the underlying error must still propagate."""
+    ev = FakeEvaluator()
+    mw = GovernanceChatMiddleware(_make_callbacks(ev))
+
+    async def boom_next() -> None:
+        raise RuntimeError("model exploded")
+
+    context = SimpleNamespace(messages=[_msg("hi")], result=SimpleNamespace(text=""))
+    with pytest.raises(RuntimeError, match="model exploded"):
+        await mw.process(context, boom_next)
+    assert [h for h, _ in ev.calls] == ["before_model", "after_model"]
+
+
+async def test_streaming_governs_finalized_response_via_result_hook():
+    """Streaming: context.result is a ResponseStream after call_next, so
+    AFTER_MODEL runs from a stream_result_hook on the finalized ChatResponse."""
+    ev = FakeEvaluator()
+    mw = GovernanceChatMiddleware(_make_callbacks(ev))
+    context = SimpleNamespace(
+        messages=[_msg("the question")],
+        stream=True,
+        stream_result_hooks=[],
+        result=None,
+    )
+    await mw.process(context, _noop_next)
+
+    # BEFORE_MODEL fired; AFTER_MODEL deferred to the registered hook.
+    assert [h for h, _ in ev.calls] == ["before_model"]
+    assert len(context.stream_result_hooks) == 1
+
+    finalized = SimpleNamespace(text="the streamed answer")
+    returned = context.stream_result_hooks[0](finalized)
+    assert returned is finalized  # hook returns the response unchanged
+    assert [h for h, _ in ev.calls] == ["before_model", "after_model"]
+    assert ev.calls[-1][1]["model_output"] == "the streamed answer"
+
+
 # --------------------------------------------------------------------------
 # FunctionMiddleware → TOOL_CALL / AFTER_TOOL
 # --------------------------------------------------------------------------
@@ -270,6 +319,29 @@ async def test_function_middleware_coerces_pydantic_args():
     assert ev.calls[1][1]["tool_result"] == ""  # None result → ""
 
 
+async def test_after_tool_runs_even_when_tool_call_raises():
+    ev = FakeEvaluator()
+    mw = GovernanceFunctionMiddleware(_make_callbacks(ev))
+
+    async def boom_next() -> None:
+        raise RuntimeError("tool exploded")
+
+    context = SimpleNamespace(function=FakeTool("t"), arguments={}, result=None)
+    with pytest.raises(RuntimeError, match="tool exploded"):
+        await mw.process(context, boom_next)
+    assert [h for h, _ in ev.calls] == ["tool_call", "after_tool"]
+
+
+def test_blocked_before_tool_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_callbacks(ev)
+    with pytest.raises(GovernanceBlockException):
+        cb.before_tool(FakeTool("t"), {})
+    assert ev.calls[-1][1]["session_state"]["tool_calls"] == 0
+    assert cb._session_state["tool_calls"] == 0
+
+
 # --------------------------------------------------------------------------
 # enforcement semantics
 # --------------------------------------------------------------------------
@@ -312,4 +384,4 @@ def evaluate_before_model(self, **_: Any) -> None:
     mw = GovernanceChatMiddleware(cb)
     with caplog.at_level(logging.WARNING):
         await mw.process(SimpleNamespace(messages=[_msg("x")], result=None), _noop_next)
-    assert any("governance check failed" in r.message for r in caplog.records)
+    assert any("governance check failed" in r.message for r in caplog.records)