fix(aiguard): scope framework collision avoidance to request and response phases

.cursor/rules/ai-guard.mdc

@@ -58,7 +58,7 @@ ddtrace/aiguard/
 ├── _initialization.py      # load_ai_guard(): lazy, gated on DD_AI_GUARD_ENABLED
 ├── _listener.py            # ai_guard_listen(): registers all core.on(...) hooks,
 │                           #   gated per integration by DD_AI_GUARD_<LLM>_ENABLED
-├── _context.py             # contextvar depth counter to avoid double-scanning
+├── _context.py             # phase-scoped (REQUEST/RESPONSE) claims to avoid double-scanning
 ├── _common.py              # _get(), wrap_abort_error(), evaluate_auto() helpers
 ├── _redaction.py           # applies the service's redaction_replacements
 ├── _trace_utils.py         # _aiguard_manual_keep(): force-keep the AI Guard trace
@@ -135,7 +135,7 @@ Events currently consumed:
 | OpenAI Responses | `openai.responses.create.before` / `.after` |
 | OpenAI (stream wrappers) | `openai.patch` / `.unpatch` |
 | Anthropic | `anthropic.messages.create.before` / `.after`, `anthropic.patch` / `.unpatch` |
-| LangChain | `langchain.{chatmodel,llm}.{generate,agenerate,stream}.before` / `.finally`, `langchain.{chatmodel,llm}.stream.started`, `langchain.patch` / `.unpatch` |
+| LangChain | `langchain.{chatmodel,llm}.{generate,agenerate}.before`, `langchain.{chatmodel,llm}.{generate,agenerate}.after`, `langchain.patch` / `.unpatch` |
 | HTTP (client IP) | `set_http_meta_for_asm` |
 
 ### 3. Provider converters → canonical `Message` shape
@@ -172,10 +172,79 @@ messages are the redacted ones — never slice on them.
 ### 5. Double-scan avoidance (`_context.py`)
 
 A framework integration (LangChain, Strands) and a provider integration
-(OpenAI, Anthropic) can both fire for the same call. To avoid scanning twice,
-the framework wraps its dispatch + LLM call with the `aiguard_context()`
-context manager (a contextvar depth counter), and provider listeners
-short-circuit when `is_aiguard_context_active()` returns `True`.
+(OpenAI, Anthropic) can both fire for the same call. The framework **claims the
+phases it evaluates itself** (`Phase.REQUEST`, `Phase.RESPONSE`), and each
+provider check skips only the phase it guards:
+
+| Integration | Claims |
+|-------------|--------|
+| LangChain `generate` / `agenerate` | `REQUEST` + `RESPONSE` for the whole call |
+| LangChain model `_stream` / `_astream` | `REQUEST` + `RESPONSE` while the model stream is read (request and response evaluated by the buffer, see below); with stream analysis off, only for the first read |
+| Strands | `REQUEST` + `RESPONSE` |
+
+| Provider check | Phase |
+|----------------|-------|
+| Request evaluation (OpenAI Chat / Responses, Anthropic) | `is_aiguard_context_active(Phase.REQUEST)` |
+| Non-streaming response evaluation | `is_aiguard_context_active(Phase.RESPONSE)` |
+| Buffered stream response evaluation (`_streaming.py`) | `is_aiguard_context_active(Phase.RESPONSE)` |
+
+Never claim a phase the framework does not actually evaluate: claiming
+`RESPONSE` on a path that cannot scan the response silences the provider's
+check and leaves the response unevaluated (APPSEC-70286).
+
+AI Guard takes the LangChain claims in its own wrappers around `BaseChatModel` /
+`BaseLLM` `generate` / `agenerate`, installed by `_langchain_patch` at
+`langchain.patch`. The LangChain contrib only dispatches product-agnostic
+`.before` / `.after` events and carries no AI Guard state.
+
+Streamed responses are buffered around each model class's own `_stream` /
+`_astream` (`_buffered_stream`), not in `stream()` and not by the provider's
+buffer. Chat models that implement the langchain-core 1.4 protocol-event hooks
+(`_stream_chat_model_events` / `_astream_chat_model_events`) get the same buffer
+on those hooks, because `stream_events` / `astream_events` v3 read them instead
+of `_stream`. Subclasses define those methods themselves, so the wrappers are
+installed per class: at `langchain.patch` for every existing subclass, and from
+an `__init_subclass__` hook on `BaseChatModel` / `BaseLLM` for classes defined
+later (`_buffer_model_classes`). Installing on first use of `generate` or
+`stream` is not enough: paths such as `stream_events(version="v3")` read
+`_stream` directly. A buffer skips only a `super()._stream` call on the model it
+is already reading, on every read of that model, not only the first; another model streamed during that read (a router model's
+inner model) gets its own buffer. The layer matters:
+
+- **Below LangChain's callbacks.** `stream()`, `astream_events()`, LangGraph's
+  `stream_mode="messages"` and generate's internal streaming all hand each
+  token to callbacks after reading `_stream`. A buffer on `stream()` would
+  deliver tokens to them before the verdict, and would finish the model run
+  before `astream_events()` could emit any stream event. Providers that report
+  tokens to the run manager from inside `_stream` get a deferring run manager
+  (`_DeferredRunManager`) whose tokens are replayed after the verdict.
+- **Above the provider.** LangChain wraps every provider read inside `_astream`
+  in its own read timeout, so a provider buffer that drains the whole stream on
+  the first read turns a long generation into a timeout.
+- **Request first.** The buffer evaluates the request before it claims and
+  reads the model, unless that model's own `generate` is running
+  (`_GENERATING_MODEL`), whose `.before` listener already did. A direct
+  `_stream` read has no other request check, and the claim makes the provider
+  skip its own. AI Guard does not listen to `.stream.before`.
+- **Once per call.** When generate streams internally, the buffer records the
+  payload it evaluated and the `.generate.after` listener skips that exact
+  payload. The record is scoped to the generate call, so a cached response is
+  still evaluated.
+
+The buffer also raises the `BaseException`-based `AIGuardAbortError`, which
+`with_fallbacks` and retry policies cannot swallow the way they swallow the
+provider's `Exception`-based abort errors.
+
+Claims are shared objects, not per-context counters, so a release from a
+different asyncio task is seen by the claiming task (APPSEC-70282). Keep the
+handle returned by `set_aiguard_context_active(...)` and release it with
+`reset_aiguard_context_active(handle)`. Prefer `aiguard_context(...)` in an AI
+Guard-owned wrapper, which claims and releases in one frame (LangChain). Only
+when a framework's own hooks split claim and release, carry the handle on
+per-call state the framework provides (Strands `invocation_state`); never add
+AI Guard state to a contrib integration to carry it. Never release by searching
+the current context for "the latest claim": that misses claims made in another
+task and takes an inner call's claim when calls nest.
 
 ## Configuration
 
@@ -239,10 +308,12 @@ must be imported from `ddtrace.aiguard.integrations.strands`.
   silently dropping data.
 - **`allow_raise=True` is required** at dispatch sites that must be able to
   block, so `AIGuardAbortError` can propagate out of the patched call.
-- **Respect the context counter.** When adding a framework integration that
-  may nest over a provider integration, wrap with `aiguard_context()` and have
-  the provider listener honor `is_aiguard_context_active()`.
-- **Preserve invariant comments** for lazy import timing, stream counter
+- **Respect phase-scoped claims.** When adding a framework integration that
+  may nest over a provider integration, claim only the phases it evaluates
+  (`aiguard_context(Phase.REQUEST, ...)` or `set_aiguard_context_active(...)`
+  with the handle kept for release), and have every provider check name its
+  phase (`is_aiguard_context_active(Phase.REQUEST)` / `(Phase.RESPONSE)`).
+- **Preserve invariant comments** for lazy import timing, stream claim
   lifecycle, and legacy message translation when editing nearby code. Do not add
   new `AIDEV-*` labels (deprecated — see AGENTS.md).
 - **Tests** live under `tests/aiguard/`. Use the `run-tests` skill;

ddtrace/aiguard/_context.py

@@ -1,77 +1,128 @@
-"""Active-flag tracking for AI Guard collision avoidance.
-
-When a framework integration (e.g. LangChain, Strands) is already evaluating
-messages through AI Guard, provider-level integrations (e.g. OpenAI) must
-skip their own evaluation to avoid double-scanning. The framework calls
-``set_aiguard_context_active()`` around its dispatch + LLM call block and
-the provider listener calls ``is_aiguard_context_active()`` to decide
-whether to short-circuit
+"""Phase-scoped collision avoidance for AI Guard.
+
+A framework integration (LangChain, Strands) and a provider integration (OpenAI,
+Anthropic) can both fire for the same call. The framework marks the phases it
+evaluates itself, and a provider listener skips only the phase it is asked about
+-- so a framework that covers the request but not the response leaves the
+provider's response protection in place.
+
+The phases are tracked independently so a framework never suppresses a check it
+does not perform itself. A single all-or-nothing flag once let LangChain
+streaming switch off the provider's buffered-stream evaluation without doing its
+own, and left streamed responses unevaluated (APPSEC-70286).
+
+Who claims what:
+
+- LangChain generate / agenerate: REQUEST and RESPONSE for the whole call (it
+  evaluates both).
+- LangChain model streams: REQUEST and RESPONSE while the model's own _stream /
+  _astream is read. The response is buffered and evaluated there, below
+  LangChain's callbacks and above the provider, so the provider's buffered
+  stream stays passthrough and cannot trip LangChain's per-chunk timeout.
+- Strands: REQUEST and RESPONSE (before- and after-model-call hooks).
+
+Claims are shared objects rather than per-context counters: an asyncio task
+works on a copy of its parent's Context, so a counter lowered from another task
+would leave the claiming task covered for good. Every claim is released by the
+handle set_aiguard_context_active returned, never by searching the current
+context, so a release can neither miss its claim nor take someone else's. Prefer
+aiguard_context, which claims and releases in one frame; carry the handle only
+when a framework's hooks split claim and release (Strands).
 """
 
 from collections.abc import Iterator
 import contextlib
 import contextvars
+from enum import Enum
 from typing import Optional
 
 
-_AI_GUARD_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar("ai_guard_active_depth", default=0)
+class Phase(Enum):
+    """Half of a model call a framework may take responsibility for."""
 
+    REQUEST = "request"
+    RESPONSE = "response"
 
-def is_aiguard_context_active() -> bool:
-    """Return ``True`` if a framework-level AI Guard evaluation is in progress."""
-    return _AI_GUARD_DEPTH.get() > 0
 
+ALL_PHASES: tuple[Phase, ...] = (Phase.REQUEST, Phase.RESPONSE)
 
-def set_aiguard_context_active() -> contextvars.Token[int]:
-    """Mark the current execution context as already under AI Guard evaluation.
 
-    Returns an opaque :class:`contextvars.Token` to pair with
-    :func:`reset_aiguard_context_active`. Nested set / reset pairs increment
-    and decrement the same depth counter, so reads return ``True`` until every
-    set is matched by a reset.
+class _Claim:
+    """One phase held by one framework call.
+
+    Every Context copied while the claim is held shares this object, so a
+    release from any of them -- including a different asyncio task -- is seen by
+    all, rather than only lowering the releasing task's own copy (APPSEC-70282).
     """
-    return _AI_GUARD_DEPTH.set(_AI_GUARD_DEPTH.get() + 1)
 
+    __slots__ = ("phase", "released")
 
-def reset_aiguard_context_active(token: Optional[contextvars.Token[int]]) -> None:
-    """Restore the depth counter to its value before the matching ``set``.
+    def __init__(self, phase: Phase) -> None:
+        self.phase = phase
+        self.released = False
+
+
+_CLAIMS: contextvars.ContextVar[tuple[_Claim, ...]] = contextvars.ContextVar("ai_guard_claims", default=())
+
+# Opaque pairing handle returned by set / consumed by reset.
+PhaseTokens = tuple[_Claim, ...]
 
-    A ``None`` token is a defensive no-op (e.g. cleanup paths that may run
-    without a prior ``set``).
-    """
-    if token is None:
-        return
-    _AI_GUARD_DEPTH.reset(token)
 
+def _prune() -> None:
+    """Drop released claims from the current context's stack."""
+    claims = _CLAIMS.get()
+    live = tuple(claim for claim in claims if not claim.released)
+    if len(live) != len(claims):
+        _CLAIMS.set(live)
 
-def reset_aiguard_context_active_current() -> None:
-    """Tokenless companion to :func:`reset_aiguard_context_active`.
 
-    Decrements the depth counter for the current context. Used when the
-    original token is not accessible — e.g. a framework's ``.after``
-    listener releasing the counter that the matching ``.before`` listener
-    bumped, since the dispatch infrastructure does not thread the token
-    through to the after-event.
+def is_aiguard_context_active(phase: Optional[Phase] = None) -> bool:
+    """Return whether a framework already covers phase.
 
-    Safe to call when the counter is already zero (no-op): the ``.after``
-    event may fire without a matching ``.before`` if dispatch is
-    reconfigured at runtime.
+    Omitting phase asks whether any phase is covered. Callers that guard a
+    specific evaluation should always name their phase; the phase-less form
+    exists for callers that only need to know an evaluation is in flight.
     """
-    depth = _AI_GUARD_DEPTH.get()
-    if depth > 0:
-        _AI_GUARD_DEPTH.set(depth - 1)
+    claims = _CLAIMS.get()
+    if not claims:
+        return False
+    # Runs on every provider call: a plain loop is several times faster than any() over a generator.
+    for claim in claims:
+        if not claim.released and (phase is None or claim.phase is phase):
+            return True
+    return False
+
+
+def set_aiguard_context_active(*phases: Phase) -> PhaseTokens:
+    """Claim phases for the current execution context.
+
+    No arguments claims every phase. Returns a handle to pair with
+    reset_aiguard_context_active; nested claims stack, so reads stay true until
+    every claim is released.
+    """
+    tokens = tuple(_Claim(phase) for phase in (phases or ALL_PHASES))
+    _CLAIMS.set(tuple(claim for claim in _CLAIMS.get() if not claim.released) + tokens)
+    return tokens
 
 
-@contextlib.contextmanager
-def aiguard_context() -> Iterator[None]:
-    """Mark the current task as under AI Guard evaluation for the block's duration.
+def reset_aiguard_context_active(tokens: Optional[PhaseTokens]) -> None:
+    """Release the claims in tokens, from any context. A falsy handle is a no-op.
 
-    Framework integrations (LangChain, Strands) wrap their dispatch + LLM
-    call block with this so nested provider-level integrations (e.g. OpenAI)
-    skip their own evaluation.
+    Only these claims are released, so a release that lands in another task
+    can neither leave the claiming task covered nor drop an unrelated claim.
     """
-    token = set_aiguard_context_active()
+    if not tokens:
+        return
+    for claim in tokens:
+        claim.released = True
+    _prune()
+
+
+@contextlib.contextmanager
+def aiguard_context(*phases: Phase) -> Iterator[None]:
+    """Claim phases for the duration of the block. No arguments claims every phase."""
+    tokens = set_aiguard_context_active(*phases)
     try:
         yield
     finally:
-        reset_aiguard_context_active(token)
+        reset_aiguard_context_active(tokens)