Conversation
Codeowners resolved asResolved from the full PR diff against |
Circular import analysis
|
Dependency direction analysis
|
❌ ErrorsYour PR has failed checks. Please review the issues below and take necessary action before merging. 🚦 2 Pipeline jobs failed
ℹ️ InfoNo other issues found (see more)🧪 All tests passed Useful? React with 👍 / 👎 This comment will be updated automatically if new data arrives.🔗 Commit SHA: d118016 | Docs | View more details | Give us feedback! |
8a41031 to
7d263a4
Compare
…onse phases AI Guard used one all-or-nothing depth counter to stop a framework and a provider integration from evaluating the same call twice. While it was raised, provider listeners skipped both their request and their response evaluation. LangChain streaming raised that counter on iteration entry, because it evaluates the request itself. But it has no stream after-event, so it could not evaluate the response -- and the counter had already switched off the provider's buffered-stream evaluation, which was the only thing left that could. A streamed LangChain response was therefore scanned by nobody, even with DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true. Track the two phases independently so a framework suppresses only what it actually covers: LangChain generate/agenerate REQUEST + RESPONSE (evaluates both) LangChain streaming REQUEST only (provider covers the response) Strands REQUEST + RESPONSE (before/after model call) Provider before-listeners now guard on REQUEST, after-listeners and the buffered stream on RESPONSE. set_aiguard_context_active still returns a pairing handle, and no arguments still claims every phase, so existing callers and the thread / task isolation and nesting tests keep exercising the new implementation unchanged. Also make reset_aiguard_context_active tolerate a token created in a different Context. Two phases means two tokens and two chances for ContextVar.reset to raise ValueError into a framework's cleanup path; it now falls back to a plain decrement, which is correct whether the context inherited the claim or never saw it. That fixes the Strands before/after-invocation crash in APPSEC-70282 for every caller. APPSEC-70286 APPSEC-70282 Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The phase-scoping tests only observed the context flags, so nothing proved a streamed LangChain response is actually evaluated. Run ChatOpenAI.stream with the OpenAI integration patched underneath and assert the buffered stream evaluates the response, blocks it before any chunk is delivered, and stays off when the flag is disabled. The OpenAI integration injects stream_options.include_usage, so the request bodies differ from the existing recordings; the two new cassettes reuse a recorded stream response. APPSEC-70286 Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
7d263a4 to
bf3cb54
Compare
The Phase import moved _report_converter_error's add_error_log call from line 59 to 60, so its line-based exemption in the constant-log checker no longer matched and the error-log-check CI job failed. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5282bf50cd
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0b1d72be7a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Carry the .stream.started claim handle on the stream handler so .stream.finally releases it even when another asyncio task finalizes the stream. Update the AI Guard guide for phase-scoped claims. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…coped-collision-avoidance # Conflicts: # .cursor/rules/ai-guard.mdc
…ease generate claims by handle Evaluate streamed LangChain responses in a buffer installed on BaseChatModel/BaseLLM stream/astream instead of the provider's buffer, so LangChain's per-chunk read timeout still applies, blocks raise the BaseException-based AIGuardAbortError that with_fallbacks cannot swallow, streamed tool calls are recorded as evaluated for the agent hooks, and any provider is covered. LangChain streams claim both phases again. Generate now carries its claim handle from .before to .finally in a per-call state dict; the tokenless reset_aiguard_context_active_current is removed. The Anthropic after event skips streamed raw responses, the stream handler releases its claim even if on_span_finish raises, and is_aiguard_context_active is fast again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…coped-collision-avoidance
…e contrib product-agnostic AI Guard now takes and releases its LangChain claims in one frame of its own wrappers around BaseChatModel/BaseLLM generate, agenerate, stream and astream: generate holds the claim for the whole call, streams only while each chunk is pulled. No claim handle has to travel through the contrib, so the LangChain contrib drops the AI Guard state dicts, the aiguard_*_event options and the finally/started events that existed only to release claims. It keeps its product-agnostic before/after events. Claiming per chunk also leaves the caller's loop body unclaimed, so direct SDK calls made inside a LangChain stream loop are evaluated (APPSEC-70480). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… the callbacks The stream buffer sat on stream() / astream(), above the point where LangChain hands tokens to callbacks, so callback handlers, astream_events() and LangGraph messages streaming saw the response before the verdict, and astream_events() on a bare model emitted no stream events with stream analysis on. Buffer around each model class's own _stream / _astream instead, installed per class on first use. Tokens a model reports to the run manager from inside _stream are deferred until the verdict. This also covers streaming inside invoke() (APPSEC-70474); .generate.after skips a payload the buffer already evaluated in the same call. One claim per stream instead of per chunk, and the chunks are merged in one add_ai_message_chunks call. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rs at patch time The nested-buffer marker was a plain flag, so any other model streamed while a buffer read its model (a router model's inner model) skipped its own buffer and its callbacks got tokens before the verdict. The marker now holds the model being read, and only a super()._stream call on that model skips the buffer. Buffers were installed on first use of generate or stream, so a path reading _stream directly (stream_events v3) on a fresh class was not buffered. Install them at langchain.patch for existing subclasses and from an __init_subclass__ hook for later ones; unpatch removes both. The stream / astream wrappers only installed buffers and are gone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 04cf299cd2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not aiguard_config._ai_guard_analyze_stream_responses_enabled: | ||
| iterator = None | ||
| try: | ||
| with aiguard_context(Phase.REQUEST, Phase.RESPONSE): |
There was a problem hiding this comment.
When _stream is consumed directly, as in the event-stream path this per-class hook is intended to cover, no langchain.*.stream.before listener has evaluated the prompt. Claiming Phase.REQUEST here therefore makes the nested OpenAI/Anthropic before-hook skip its evaluation; with stream analysis disabled (the default), this branch never performs any replacement evaluation at all, and with it enabled the request is only included in the post-response evaluation after it has already reached the model. Evaluate the request before taking the request claim, or leave that phase unclaimed unless an outer LangChain before-hook has run.
Useful? React with 👍 / 👎.
| async for chunk in func(*args, **kwargs): | ||
| events.append((True, chunk)) |
There was a problem hiding this comment.
When a task is cancelled or times out while this enabled buffer is draining the model response, control exits the async for without closing the iterator returned by _astream. Because buffering yields nothing until the drain completes, cancellation is the caller's main way to stop a slow stream, and this can leave the provider connection and its model-run resources open; the flag-disabled branch already avoids this by retaining the iterator and awaiting aclose() in finally. Apply the same cleanup around the buffered drain.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c9af054706
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…ffer install lock Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…asses can be collected Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2 participants
Description
Jira: APPSEC-70286, APPSEC-70474, APPSEC-70480 and APPSEC-70282
AI Guard used a single all-or-nothing depth counter to stop a framework integration and a provider integration from evaluating the same call twice. While it was raised, provider listeners skipped both their request and their response evaluation. LangChain streaming raised it on iteration entry but never evaluated the streamed response, so with
DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=truea streamed LangChain response was scanned by nobody.Phase-scoped claims
The counter is replaced by independent REQUEST and RESPONSE claims, so a framework only suppresses the provider checks it performs itself:
generate/agenerate.before, response in.after(#20457)_stream/_astream.stream.before/.generate.before, response in the model-stream buffer (below)Provider
.beforelisteners check REQUEST; provider.afterlisteners and the buffered streams check RESPONSE.AI Guard owns its LangChain logic; the contrib stays product-agnostic
The LangChain contrib (
ddtrace/contrib/internal/langchain/) carries no AI Guard state or naming. It only dispatches the product-agnosticlangchain.{chatmodel,llm}.{generate,agenerate,stream}.beforeand.{generate,agenerate}.afterevents, from inside the LLM span, so a listener's block is still recorded on that span and the LLMObs output is kept. Againstmainthe contrib diff is removals only, plus renaming the stream optionaiguard_before_eventtobefore_event: the AI Guard state dicts, theaiguard_*_eventoptions, the stream handler'sstart_streamoverride and the.finally/.stream.started/.stream.finallyevents are gone. They existed only so AI Guard could release a claim in a different listener from the one that took it.Instead, at
langchain.patchAI Guard installs its own wrappers (outermost, around the contrib's) onBaseChatModel/BaseLLM:generate/agenerate:with aiguard_context(REQUEST, RESPONSE)around the call._stream/_astreamof each model class: the stream buffer below, which also takes the stream's claim. Installed atlangchain.patchfor every existing subclass and from an__init_subclass__hook for classes defined later, so a path that reads_streamdirectly (stream_events(version="v3")) is covered on first use.Every claim is taken and released in one frame, so no handle travels through the contrib and nothing can leave a claim behind: not a nested inner stream, not an early break, not a close from another task.
The claim covers only the read of the model's stream, so the caller's loop body and LangChain callbacks are never inside it. This fixes APPSEC-70480: direct OpenAI / Anthropic calls made while iterating a LangChain stream are evaluated. Previously the claim spanned the whole iteration and those calls skipped AI Guard.
Streamed LangChain responses are buffered around each model's
_stream/_astreamWith
DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true, AI Guard wraps the_stream/_astreameach model class defines. The whole model stream is read, the chunks are merged into the message LangChain would build, request + response are evaluated, and only then are the chunks replayed. Subclasses define these methods themselves, so the wrappers are installed per class (see above). The layer matters:stream(),astream_events(), LangGraphstream_mode="messages"and generate's internal streaming all hand each token to callbacks after reading_stream. No callback sees a token before the verdict, andastream_events()still emits its stream events because the model run is still open when the first chunk is replayed. Providers that report tokens to the run manager from inside_stream(LLMs,ChatOpenAI(streaming=True)) get a deferring run manager whose tokens are replayed after the verdict.invoke()(APPSEC-70474). When LangChain streams inside generate (a streaming callback handler,streaming=True,stream=True), the buffer evaluates the response before any token is reported. It records the payload it evaluated, and.generate.afterskips that exact payload, so the response is evaluated once. The record is scoped to the generate call, so a cached response is still evaluated._astreamin its own 120 swait_for. The buffer drains_astreamfrom outside, so every provider read keeps its own timeout.BaseException-basedAIGuardAbortError. The provider'sOpenAIAIGuardAbortErroris anException, whichwith_fallbacksand retry policies caught and routed to an unevaluated fallback.response_metadata; chunk addition carries it into the aggregate, so the legacyAgentExecutor(which streams by default) and the langgraph hooks skip them.add_ai_message_chunkscall instead of pairwise.It is passthrough when the flag is off. Conversion errors fail open; only a block raises. A subclass
_streamthat callssuper()._streamis buffered once: the buffer skips only a nested read of the model it is already reading. Another model streamed during that read (a router model's inner model) gets its own buffer.Claims are shared objects, released by handle (APPSEC-70282)
An asyncio task works on a copy of its parent's Context, so a counter lowered from another task left the claiming task covered for good, and later OpenAI/Anthropic calls in it skipped AI Guard. Claims are now shared objects stored in a
ContextVar: a release from any task is seen by all, never raises (the oldContextVar.resetraisedValueErrorinto the framework's cleanup path), and never drops an unrelated claim.Every claim is released by its handle, never by searching the current context for "the latest claim", which misses claims made in another task and takes an inner call's claim when calls nest (the tokenless helper is removed). LangChain claims and releases in one frame via
aiguard_context. Strands, whose own hooks split claim and release, keeps the handle oninvocation_state.Smaller fixes
.afterevent skips streamed requests: awith_raw_responsestream reaches it with its body unread, and converting it raised and reported a converter error on every call.is_aiguard_context_activehas an early return and a plain loop (it runs on every provider call)..cursor/rules/ai-guard.mdc) documents which phase each framework claims and each provider check reads, never claiming a phase you do not evaluate, the AI Guard-owned LangChain wrappers, why the stream buffer sits at_stream, and release by handle (never add AI Guard state to a contrib).Known limitations
langchain-core 1.4 adds native protocol events (
_stream_chat_model_events), which its v2 event path reads instead of_stream. No model in langchain-openai implements them yet; a model that does is not buffered on that path.Out of scope
Pre-existing issues raised in review are tracked separately in Jira and are not fixed here: REQUEST claimed when LangChain skips its request check, the Strands invocation-wide RESPONSE claim over Strands' own models, the Strands
structured_output_asyncclaim leak, the async raw-responseparse()coroutine, async evaluation blocking the event loop, Strands Graph parallel nodes sharing one handle, the double request evaluation whenstream()falls back toinvoke(), and making the phase a required argument.Testing
test_context.py): phase scoping; cross-task release (no raise, clears the claiming task, keeps unrelated claims); a generator closed from another task releases its claim.test_langchain.py): streamed chat (sync, async) and LLM responses evaluated before delivery, with the evaluated content checked against the delivered chunks; a response block delivers zero chunks;with_fallbacksdoes not swallow a block; a model with no provider integration is covered; a streamingAgentExecutor(stream_runnable=True) evaluates its tool call once; streamed tool calls marked evaluated for the agent hooks; claims observed at the provider's HTTP request for streamed and non-streamed calls; the caller's loop body is unclaimed with and without the buffer (APPSEC-70480); a_generatereturning with an inner stream still open leaves no claim; anastream()closed from another task leaves no claim; conversion errors fail open.New in this revision: callback tokens and
on_llm_endarrive only after the verdict, and none arrive on a block;astream_events()on a bare model emits every stream event, after the verdict; a chain'sastream_events()emits no model output on a block;invoke(stream=True)is evaluated once, before its tokens (APPSEC-70474); tokens a model reports from inside_stream/_astreamwait for the verdict and are dropped on a block (chat sync, async and LLM); asuper()._streamsubclass is buffered once; one claim per stream, with and without the buffer; a stream under an outer framework claim is still evaluated; the one-call chunk merge matches chunk addition; an inner model streamed by a router model is buffered on its own, so its callbacks see nothing on a block;_streamread directly is buffered for classes defined before and after patching; unpatch removes the hook and the wrappers.test_streaming.py): unchanged coverage of evaluate-then-replay and blocking.Cassettes: 2 files in
tests/cassettes/openai/for the OpenAI-patched streaming requests (stream_options.include_usagechanges the body); response bodies reuse a recorded stream plus a usage chunk.Local results on the final code:
aiguard::ai_guard_langchainlangchain 1.x (py3.13)aiguard::ai_guard_langchainlangchain 0.2.17 (py3.12)aiguard::ai_guard_langchainlangchain 0.1.20 (py3.11)aiguard::ai_guard_openaiaiguard::ai_guard_anthropicaiguard::ai_guard_strandsaiguard::ai_guard_apiaiguard::ai_guard_litellm_guardrailllmobs::langchainSkips are the langchain-version-gated agent tests, plus the
astream_events(version="v2")tests on langchain-core 0.1.tests/llmobs/test_base_stream_handler.py: 94 passed locally (the.stream.finallypairing test was removed with the event); its async tests need pytest-asyncio, which the local llmobs venv lacks (CI covers them).Lint clean:
fmt,typing,spelling; ast-grep'sno-rst-double-backtickswarnings in the touched packages dropped from 58 to 49, none on added lines.Risks
Behaviour change for streamed LangChain responses when
DD_AI_GUARD_ANALYZE_STREAM_RESPONSES_ENABLED=true. They are now buffered and evaluated where previously they streamed through unevaluated: higher time-to-first-token, for the caller and for LangChain callbacks (astream_events(), LangGraph messages streaming, models streaming insideinvoke()), and one additional evaluation per streamed LangChain call. The flag is off by default, so nobody is affected without opting in. Called out in the release note.The LangChain contrib change is removals only: it no longer dispatches the AI Guard-only
.finally/.stream.started/.stream.finallyevents (AI Guard was their only listener) and no longer carries AI Guard options._context.pyis private; its no-argument defaults are kept for existing callers and will be removed in a follow-up.AI Guard wraps LangChain model methods directly, as it already does for agent
planand langgraphToolNode. The_stream/_astreamwrappers and the__init_subclass__hook are installed atlangchain.patchand removed on unpatch. CoreExecutionContexts were considered instead, but a context entered at stream start and exited at finalize would change the caller's current core context for the whole loop, and would stay stuck in the starting task when an async stream is finalized in another task.Additional Notes
The OpenAI buffered-stream wrappers are installed at
openai.patchtime only if the flag is already on (existing behaviour); the OpenAI-patched LangChain fixture enables it before patching. The LangChain model-stream buffer checks the flag per call.🤖 Generated with Claude Code