diff --git a/.github/workflows/system-tests.yml b/.github/workflows/system-tests.yml index f3afbcb1cd6..e0a2a864d12 100644 --- a/.github/workflows/system-tests.yml +++ b/.github/workflows/system-tests.yml @@ -95,7 +95,7 @@ jobs: persist-credentials: false repository: 'DataDog/system-tests' # Automatically managed, use scripts/update-system-tests-version to update - ref: '0526d7eb6ad321c14f7d4c7574cf8970089757c6' + ref: '24c2e96e7703a8128d5a29cae2e95776d86ef790' - name: Download wheels to binaries directory uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1 @@ -148,7 +148,7 @@ jobs: persist-credentials: false repository: 'DataDog/system-tests' # Automatically managed, use scripts/update-system-tests-version to update - ref: '0526d7eb6ad321c14f7d4c7574cf8970089757c6' + ref: '24c2e96e7703a8128d5a29cae2e95776d86ef790' - name: Build runner uses: ./.github/actions/install_runner @@ -458,4 +458,4 @@ jobs: needs.integration-frameworks-system-tests.result == 'cancelled'|| needs.tracer-release.result == 'failure' || needs.tracer-release.result == 'cancelled' - run: exit 1 + run: exit 1 \ No newline at end of file diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index 3958e232627..6884140d662 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -16,7 +16,7 @@ variables: DD_VPA_TEMPLATE: "vpa-template-cpu-p70-10percent-2x-oom-min-cap" # CI_DEBUG_SERVICES: "true" # Automatically managed, use scripts/update-system-tests-version to update - SYSTEM_TESTS_REF: "0526d7eb6ad321c14f7d4c7574cf8970089757c6" + SYSTEM_TESTS_REF: "24c2e96e7703a8128d5a29cae2e95776d86ef790" # Profiling native build image (built from dd/images/dd-trace-py/profiling_native) PROFILING_NATIVE_IMAGE: "registry.ddbuild.io/dd-trace-py:v103334885-be1888c-profiling_native" diff --git a/.sg/rules/core-raising-dispatch.yml b/.sg/rules/core-raising-dispatch.yml deleted file mode 100644 index e9e86575b93..00000000000 --- a/.sg/rules/core-raising-dispatch.yml +++ /dev/null @@ -1,55 +0,0 @@ -id: core-raising-dispatch -message: Avoid using `raising_dispatch()` - prefer `core.dispatch()` with explicit try/except -severity: error -language: python -ignores: - # Implementation and re-export - - "ddtrace/internal/core/__init__.py" - - "ddtrace/internal/core/event_hub.py" - # Allow-listed existing callsites - - "ddtrace/contrib/internal/langchain/patch.py" - - "ddtrace/contrib/internal/openai/patch.py" -rule: - any: - # Match function calls to raising_dispatch - - pattern: core.raising_dispatch($$$ARGS) - - pattern: raising_dispatch($$$ARGS) - - pattern: event_hub.raising_dispatch($$$ARGS) - # Match imports of raising_dispatch - - pattern: from ddtrace.internal.core import raising_dispatch - - pattern: from ddtrace.internal.core import raising_dispatch as $ALIAS - - pattern: from ddtrace.internal.core import $$$IMPORTS, raising_dispatch - - pattern: from ddtrace.internal.core import $$$IMPORTS, raising_dispatch as $ALIAS - - pattern: from ddtrace.internal.core import raising_dispatch, $$$IMPORTS - - pattern: from ddtrace.internal.core import raising_dispatch as $ALIAS, $$$IMPORTS - - pattern: from ddtrace.internal.core.event_hub import raising_dispatch - - pattern: from ddtrace.internal.core.event_hub import raising_dispatch as $ALIAS - - pattern: from ddtrace.internal.core.event_hub import $$$IMPORTS, raising_dispatch - - pattern: from ddtrace.internal.core.event_hub import $$$IMPORTS, raising_dispatch as $ALIAS - - pattern: from ddtrace.internal.core.event_hub import raising_dispatch, $$$IMPORTS - - pattern: from ddtrace.internal.core.event_hub import raising_dispatch as $ALIAS, $$$IMPORTS -constraints: - ALIAS: - regex: ".*" -note: | - `raising_dispatch()` is built on top of `dispatch_with_results()`, which is - itself discouraged (see the `core-dispatch-with-results` rule). New callsites - should not be introduced; the long-term plan is to retire this helper in - favor of `core.dispatch(..., allow_raise=True)` once that pattern lands. - - Prefer letting the listener raise through `core.dispatch()` and handle the - exception at the callsite: - - try: - core.dispatch("my.event", (args,)) - except MyException: - # handle interruption - pass - - Note: try/except will block all other listeners from running after the - exception is raised. - - This rule does not apply to: - - The implementation in ddtrace/internal/core/event_hub.py - - The export in ddtrace/internal/core/__init__.py - - The allow-listed contrib callsites in langchain/patch.py and openai/patch.py diff --git a/.sg/tests/__snapshots__/core-raising-dispatch-snapshot.yml b/.sg/tests/__snapshots__/core-raising-dispatch-snapshot.yml deleted file mode 100644 index 622f9ffd078..00000000000 --- a/.sg/tests/__snapshots__/core-raising-dispatch-snapshot.yml +++ /dev/null @@ -1,108 +0,0 @@ -id: core-raising-dispatch -snapshots: - ? | - # Using raising_dispatch in a flow - core.raising_dispatch("my.event", (x,)) - print("done") - : labels: - - source: core.raising_dispatch("my.event", (x,)) - style: primary - start: 35 - end: 74 - core.raising_dispatch("my.event", (arg1, arg2)): - labels: - - source: core.raising_dispatch("my.event", (arg1, arg2)) - style: primary - start: 0 - end: 47 - ? | - core.raising_dispatch("my.event", (arg1, arg2)) - : labels: - - source: core.raising_dispatch("my.event", (arg1, arg2)) - style: primary - start: 0 - end: 47 - event_hub.raising_dispatch("my.event", ()): - labels: - - source: event_hub.raising_dispatch("my.event", ()) - style: primary - start: 0 - end: 42 - ? | - from ddtrace.internal import core - core.raising_dispatch("my.event", (arg1, arg2)) - : labels: - - source: core.raising_dispatch("my.event", (arg1, arg2)) - style: primary - start: 34 - end: 81 - from ddtrace.internal.core import dispatch, raising_dispatch: - labels: - - source: from ddtrace.internal.core import dispatch, raising_dispatch - style: primary - start: 0 - end: 60 - from ddtrace.internal.core import dispatch, raising_dispatch, on: - labels: - - source: from ddtrace.internal.core import dispatch, raising_dispatch, on - style: primary - start: 0 - end: 64 - from ddtrace.internal.core import raising_dispatch: - labels: - - source: from ddtrace.internal.core import raising_dispatch - style: primary - start: 0 - end: 50 - ? | - from ddtrace.internal.core import raising_dispatch - raising_dispatch("my.event", (arg1,)) - : labels: - - source: from ddtrace.internal.core import raising_dispatch - style: primary - start: 0 - end: 50 - from ddtrace.internal.core import raising_dispatch as rd: - labels: - - source: from ddtrace.internal.core import raising_dispatch as rd - style: primary - start: 0 - end: 56 - from ddtrace.internal.core import raising_dispatch, on: - labels: - - source: from ddtrace.internal.core import raising_dispatch, on - style: primary - start: 0 - end: 54 - from ddtrace.internal.core.event_hub import dispatch, raising_dispatch: - labels: - - source: from ddtrace.internal.core.event_hub import dispatch, raising_dispatch - style: primary - start: 0 - end: 70 - from ddtrace.internal.core.event_hub import raising_dispatch: - labels: - - source: from ddtrace.internal.core.event_hub import raising_dispatch - style: primary - start: 0 - end: 60 - ? | - from ddtrace.internal.core.event_hub import raising_dispatch - raising_dispatch("my.event", ()) - : labels: - - source: from ddtrace.internal.core.event_hub import raising_dispatch - style: primary - start: 0 - end: 60 - from ddtrace.internal.core.event_hub import raising_dispatch as rd: - labels: - - source: from ddtrace.internal.core.event_hub import raising_dispatch as rd - style: primary - start: 0 - end: 66 - raising_dispatch("my.event", (arg1,)): - labels: - - source: raising_dispatch("my.event", (arg1,)) - style: primary - start: 0 - end: 37 diff --git a/.sg/tests/core-raising-dispatch-test.yml b/.sg/tests/core-raising-dispatch-test.yml deleted file mode 100644 index 96781a33069..00000000000 --- a/.sg/tests/core-raising-dispatch-test.yml +++ /dev/null @@ -1,48 +0,0 @@ -id: core-raising-dispatch -valid: - # These should NOT trigger the rule (valid code) - - core.dispatch("my.event", (arg1, arg2)) - - | - from ddtrace.internal.core import dispatch - dispatch("my.event", (arg1,)) - - | - # Using try/except for flow control is OK - try: - core.dispatch("my.event", (args,)) - except MyException: - pass - - | - # Other functions with similar names are OK - my_module.raising_dispatch(args) - raising_dispatch_other(args) - -invalid: - # These should trigger the rule (warnings) - # Direct function calls - - core.raising_dispatch("my.event", (arg1, arg2)) - - raising_dispatch("my.event", (arg1,)) - - event_hub.raising_dispatch("my.event", ()) - - | - core.raising_dispatch("my.event", (arg1, arg2)) - - | - # Using raising_dispatch in a flow - core.raising_dispatch("my.event", (x,)) - print("done") - # Import statements - - from ddtrace.internal.core import raising_dispatch - - from ddtrace.internal.core import raising_dispatch as rd - - from ddtrace.internal.core import dispatch, raising_dispatch - - from ddtrace.internal.core import raising_dispatch, on - - from ddtrace.internal.core import dispatch, raising_dispatch, on - - from ddtrace.internal.core.event_hub import raising_dispatch - - from ddtrace.internal.core.event_hub import raising_dispatch as rd - - from ddtrace.internal.core.event_hub import dispatch, raising_dispatch - - | - from ddtrace.internal.core import raising_dispatch - raising_dispatch("my.event", (arg1,)) - - | - from ddtrace.internal.core.event_hub import raising_dispatch - raising_dispatch("my.event", ()) - - | - from ddtrace.internal import core - core.raising_dispatch("my.event", (arg1, arg2)) diff --git a/ddtrace/appsec/_ai_guard/_langchain.py b/ddtrace/appsec/_ai_guard/_langchain.py index 225004357d9..a068ddac9d5 100644 --- a/ddtrace/appsec/_ai_guard/_langchain.py +++ b/ddtrace/appsec/_ai_guard/_langchain.py @@ -224,8 +224,8 @@ def _langchain_generate_finally(*args, **kwargs): Releases the AI Guard active counter that the matching ``.before`` listener bumped. Dispatched from the contrib's ``finally`` block so it - fires on every exit path — success, block (``_raising_dispatch`` raises - out of ``.before``), or exception inside the underlying LLM call. The + fires on every exit path — success, block (``core.dispatch(..., allow_raise=True)`` + raises out of ``.before``), or exception inside the underlying LLM call. The counter reset is a no-op when the counter is already zero, so listener invocations that don't pair with a ``.before`` set are safe. """ @@ -247,12 +247,12 @@ def _langchain_llm_stream_before(client: AIGuardClient, instance, args, kwargs): def _evaluate_langchain_messages(client: AIGuardClient, messages): - """Evaluate the prompt and surface a block as a returned ``AIGuardAbortError``. + """Evaluate the prompt and re-raise ``AIGuardAbortError`` on a block. - Returns the abort error so the contrib's dispatcher can re-raise it; the - ``AIGuardClient`` already gates on ``ai_guard_config._ai_guard_block``, - so a returned error always represents a blocking decision. - Allow / skip paths return ``None``. + Re-raises so the contrib's ``core.dispatch(..., allow_raise=True)`` + propagates the abort. The ``AIGuardClient`` already gates on + ``ai_guard_config._ai_guard_block``, so a raised error always represents + a blocking decision. Allow / skip paths return ``None``. """ from langchain_core.messages import HumanMessage @@ -260,8 +260,8 @@ def _evaluate_langchain_messages(client: AIGuardClient, messages): if len(messages) > 0 and isinstance(messages[-1], HumanMessage): try: client.evaluate(_convert_messages(messages), Options(block=ai_guard_config._ai_guard_block)) - except AIGuardAbortError as e: - return e + except AIGuardAbortError: + raise except Exception: logger.debug("Failed to evaluate chat model prompt", exc_info=True) return None diff --git a/ddtrace/appsec/_ai_guard/_openai.py b/ddtrace/appsec/_ai_guard/_openai.py index b6019dc824c..2bfee941deb 100644 --- a/ddtrace/appsec/_ai_guard/_openai.py +++ b/ddtrace/appsec/_ai_guard/_openai.py @@ -66,6 +66,19 @@ class OpenAIAIGuardAbortError(openai.UnprocessableEntityError, AIGuardAbortError Catchable as either ``openai.APIError`` / ``openai.UnprocessableEntityError`` (idiomatic OpenAI error handling, no retry on 422) or ``AIGuardAbortError`` (Datadog-specific, exposes ``action`` / ``reason``). + + AIDEV-NOTE: catchability asymmetry vs plain ``AIGuardAbortError``. + ``AIGuardAbortError`` derives from ``DDBlockException(BaseException)`` + so a generic ``except Exception:`` does NOT catch it (by design — a + block decision must not be silently swallowed). However, + ``OpenAIAIGuardAbortError`` *also* inherits from + ``openai.UnprocessableEntityError``, which is ``Exception``-derived, + so via MRO this subclass IS catchable by ``except Exception:``. + That asymmetry is intentional: the OpenAI contrib must keep + ``except openai.APIError:`` blocks working unchanged for users + migrating from non-AI-Guard error handling. Code that wants + uniform block detection across providers should branch on + ``isinstance(e, AIGuardAbortError)``. """ def __init__(self, action, reason, tags=None, sds=None, tag_probs=None): @@ -285,7 +298,7 @@ def _openai_chat_completion_before(client, kwargs): try: client.evaluate(ai_guard_messages, Options(block=ai_guard_config._ai_guard_block)) except AIGuardAbortError as e: - return _wrap_abort_error(e) + raise _wrap_abort_error(e) except Exception: logger.debug("Failed to evaluate OpenAI chat completion request", exc_info=True) return None @@ -298,7 +311,7 @@ def _openai_chat_completion_after(client, kwargs, resp): returns. Skips streaming responses (handled separately) and when a framework evaluation is already active. - On block: returns an ``OpenAIAIGuardAbortError`` (or plain + On block: raises an ``OpenAIAIGuardAbortError`` (or plain ``AIGuardAbortError`` when the OpenAI SDK is not importable). Allow / skip paths return ``None``. """ @@ -316,7 +329,7 @@ def _openai_chat_completion_after(client, kwargs, resp): try: client.evaluate(all_messages, Options(block=ai_guard_config._ai_guard_block)) except AIGuardAbortError as e: - return _wrap_abort_error(e) + raise _wrap_abort_error(e) except Exception: logger.debug("Failed to evaluate OpenAI chat completion response", exc_info=True) return None diff --git a/ddtrace/appsec/ai_guard/_api_client.py b/ddtrace/appsec/ai_guard/_api_client.py index a817defc820..8e0fa9a8339 100644 --- a/ddtrace/appsec/ai_guard/_api_client.py +++ b/ddtrace/appsec/ai_guard/_api_client.py @@ -14,6 +14,7 @@ from ddtrace.ext import http from ddtrace.internal import core from ddtrace.internal import telemetry +from ddtrace.internal._exceptions import DDBlockException import ddtrace.internal.logger as ddlogger from ddtrace.internal.settings.asm import ai_guard_config from ddtrace.internal.telemetry import TELEMETRY_NAMESPACE @@ -94,8 +95,13 @@ def __init__(self, message: Optional[str], status: int = 0, errors: Optional[lis super().__init__(message) -class AIGuardAbortError(Exception): - """Exception to abort current execution due to security policy.""" +class AIGuardAbortError(DDBlockException): + """Exception to abort current execution due to security policy. + + Inherits from ``DDBlockException`` (which is ``BaseException``-derived) so + that a generic ``except Exception:`` in user code does not accidentally + swallow an AI Guard block decision. + """ def __init__( self, diff --git a/ddtrace/contrib/internal/langchain/patch.py b/ddtrace/contrib/internal/langchain/patch.py index a084a100580..dea0c821738 100644 --- a/ddtrace/contrib/internal/langchain/patch.py +++ b/ddtrace/contrib/internal/langchain/patch.py @@ -9,6 +9,7 @@ from ddtrace.contrib.internal.trace_utils import unwrap from ddtrace.contrib.internal.trace_utils import wrap from ddtrace.internal import core +from ddtrace.internal._exceptions import DDBlockException from ddtrace.internal.compat import is_wrapted from ddtrace.internal.logger import get_logger from ddtrace.internal.utils import ArgumentError @@ -82,10 +83,10 @@ def traced_llm_generate(func, instance, args, kwargs): integration.llmobs_set_prompt_tag(instance, span) try: - core.raising_dispatch("langchain.llm.generate.before", (prompts,)) + core.dispatch("langchain.llm.generate.before", (prompts,), allow_raise=True) completions = func(*args, **kwargs) core.dispatch("langchain.llm.generate.after", (prompts, completions)) - except Exception: + except (DDBlockException, Exception): span.set_exc_info(*sys.exc_info()) raise finally: @@ -115,10 +116,10 @@ async def traced_llm_agenerate(func, instance, args, kwargs): completions = None try: - core.raising_dispatch("langchain.llm.agenerate.before", (prompts,)) + core.dispatch("langchain.llm.agenerate.before", (prompts,), allow_raise=True) completions = await func(*args, **kwargs) core.dispatch("langchain.llm.agenerate.after", (prompts, completions)) - except Exception: + except (DDBlockException, Exception): span.set_exc_info(*sys.exc_info()) raise finally: @@ -147,10 +148,10 @@ def traced_chat_model_generate(func, instance, args, kwargs): chat_completions = None try: - core.raising_dispatch("langchain.chatmodel.generate.before", (chat_messages,)) + core.dispatch("langchain.chatmodel.generate.before", (chat_messages,), allow_raise=True) chat_completions = func(*args, **kwargs) core.dispatch("langchain.chatmodel.generate.after", (chat_messages, chat_completions)) - except Exception: + except (DDBlockException, Exception): span.set_exc_info(*sys.exc_info()) raise finally: @@ -179,10 +180,10 @@ async def traced_chat_model_agenerate(func, instance, args, kwargs): chat_completions = None try: - core.raising_dispatch("langchain.chatmodel.agenerate.before", (chat_messages,)) + core.dispatch("langchain.chatmodel.agenerate.before", (chat_messages,), allow_raise=True) chat_completions = await func(*args, **kwargs) core.dispatch("langchain.chatmodel.agenerate.after", (chat_messages, chat_completions)) - except Exception: + except (DDBlockException, Exception): span.set_exc_info(*sys.exc_info()) raise finally: @@ -313,7 +314,7 @@ def traced_chat_stream(func, instance, args, kwargs): llm_provider = instance._llm_type model = _extract_model_name(instance) - core.raising_dispatch("langchain.chatmodel.stream.before", (instance, args, kwargs)) + core.dispatch("langchain.chatmodel.stream.before", (instance, args, kwargs), allow_raise=True) def _on_span_started(span: Span): integration.record_instance(instance, span) @@ -347,13 +348,14 @@ def traced_llm_stream(func, instance, args, kwargs): llm_provider = instance._llm_type model = _extract_model_name(instance) - core.raising_dispatch( + core.dispatch( "langchain.llm.stream.before", ( instance, args, kwargs, ), + allow_raise=True, ) def _on_span_start(span: Span): diff --git a/ddtrace/contrib/internal/openai/patch.py b/ddtrace/contrib/internal/openai/patch.py index 6146d03e060..978b8fb8532 100644 --- a/ddtrace/contrib/internal/openai/patch.py +++ b/ddtrace/contrib/internal/openai/patch.py @@ -10,6 +10,7 @@ from ddtrace.contrib.trace_utils import unwrap from ddtrace.contrib.trace_utils import wrap from ddtrace.internal import core +from ddtrace.internal._exceptions import DDBlockException from ddtrace.internal.logger import get_logger from ddtrace.internal.utils.formats import deep_getattr from ddtrace.internal.utils.version import parse_version @@ -245,7 +246,7 @@ def patched_endpoint(func, instance, args, kwargs): integration = openai._datadog_integration is_chat = patch_hook in _CHAT_COMPLETION_HOOKS if is_chat: - core.raising_dispatch("openai.chat.completions.create.before", (kwargs,)) + core.dispatch("openai.chat.completions.create.before", (kwargs,), allow_raise=True) g = _traced_endpoint(patch_hook, integration, instance, args, kwargs) g.send(None) @@ -265,7 +266,7 @@ def patched_endpoint(func, instance, args, kwargs): override_return = e.value if is_chat and not kwargs.get("stream") and resp is not None and err is None: - core.raising_dispatch("openai.chat.completions.create.after", (kwargs, resp)) + core.dispatch("openai.chat.completions.create.after", (kwargs, resp), allow_raise=True) if override_return is not None: return override_return @@ -412,8 +413,8 @@ async def async_wrapper(): # Guard even when the caller never awaits the returned coroutine. if is_chat: try: - core.raising_dispatch("openai.chat.completions.create.before", (kwargs,)) - except BaseException: + core.dispatch("openai.chat.completions.create.before", (kwargs,), allow_raise=True) + except DDBlockException: # AI Guard blocked the request — discard the unstarted SDK # coroutine so Python doesn't emit a "coroutine was never # awaited" warning for it. @@ -440,7 +441,7 @@ async def async_wrapper(): override_return = e.value if is_chat and not kwargs.get("stream") and resp is not None and err is None: - core.raising_dispatch("openai.chat.completions.create.after", (kwargs, resp)) + core.dispatch("openai.chat.completions.create.after", (kwargs, resp), allow_raise=True) if override_return is not None: if resp is not send_resp and override_return is not None: diff --git a/ddtrace/internal/_exceptions.py b/ddtrace/internal/_exceptions.py index c197b7ce30a..e2836e036be 100644 --- a/ddtrace/internal/_exceptions.py +++ b/ddtrace/internal/_exceptions.py @@ -2,10 +2,19 @@ from typing import TypeVar -class BlockingException(BaseException): +class DDBlockException(BaseException): + """ + Base class for any in-tree decision to abort the current operation + (web request blocking, AI Guard policy abort, future product blocks). + Inherits from BaseException so a generic ``except Exception:`` handler in + user code does not accidentally swallow a blocking decision. + """ + + +class BlockingException(DDBlockException): """ Exception raised when a request is blocked by ASM - It derives from BaseException to avoid being caught by the general Exception handler + It derives from BaseException (via DDBlockException) to avoid being caught by the general Exception handler """ diff --git a/ddtrace/internal/core/__init__.py b/ddtrace/internal/core/__init__.py index fc53688883d..88dac0220b1 100644 --- a/ddtrace/internal/core/__init__.py +++ b/ddtrace/internal/core/__init__.py @@ -133,7 +133,6 @@ def done_callback(f): from .event_hub import dispatch_with_results # noqa:F401 from .event_hub import has_listeners # noqa:F401 from .event_hub import on # noqa:F401 -from .event_hub import raising_dispatch # noqa:F401 from .event_hub import reset as reset_listeners # noqa:F401 from .events import EventType diff --git a/ddtrace/internal/core/event_hub.py b/ddtrace/internal/core/event_hub.py index 20746f376d2..f7bd66a55c0 100644 --- a/ddtrace/internal/core/event_hub.py +++ b/ddtrace/internal/core/event_hub.py @@ -73,9 +73,14 @@ def reset(event_id: Optional[str] = None, callback: Optional[Callable[..., Any]] del _listeners[event_id] -def dispatch_event(event) -> None: +def dispatch_event(event, allow_raise: bool = False) -> None: """Call all hooks for the provided event. + When ``allow_raise=True``, listener ``Exception``s propagate to the caller + (the first listener to raise wins; subsequent listeners are skipped). + ``BaseException``-derived exceptions (including ``DDBlockException``) + always propagate regardless of ``allow_raise``. + PERF: Avoid calling `dispatch` to reduce function calls/overhead of this function. """ global _listeners @@ -88,12 +93,18 @@ def dispatch_event(event) -> None: try: local_hook(event) except Exception: - if config._raise: + if allow_raise or config._raise: raise -def dispatch(event_id: str, args: tuple[Any, ...] = ()) -> None: - """Call all hooks for the provided event_id with the provided args""" +def dispatch(event_id: str, args: tuple[Any, ...] = (), allow_raise: bool = False) -> None: + """Call all hooks for the provided event_id with the provided args. + + When ``allow_raise=True``, listener ``Exception``s propagate to the caller + (the first listener to raise wins; subsequent listeners are skipped). + ``BaseException``-derived exceptions (including ``DDBlockException``) + always propagate regardless of ``allow_raise``. + """ global _listeners if event_id not in _listeners: @@ -103,7 +114,7 @@ def dispatch(event_id: str, args: tuple[Any, ...] = ()) -> None: try: local_hook(*args) except Exception: - if config._raise: + if allow_raise or config._raise: raise @@ -126,12 +137,3 @@ def dispatch_with_results(event_id: str, args: tuple[Any, ...] = ()) -> EventRes results[name] = EventResult(ResultType.RESULT_EXCEPTION, None, e) return results - - -def raising_dispatch(event_id: str, args: tuple[Any, ...] = ()) -> None: - """Deprecated: use ``dispatch`` with try/except instead.""" - results = dispatch_with_results(event_id, args) - for event in results.values(): - # we explicitly set the exception as a value to prevent caught exceptions from leaking - if isinstance(event.value, Exception): - raise event.value diff --git a/releasenotes/notes/aiguard-abort-baseexception-deprecate-raising-dispatch-3a5c7f61ee4d2e1b.yaml b/releasenotes/notes/aiguard-abort-baseexception-deprecate-raising-dispatch-3a5c7f61ee4d2e1b.yaml new file mode 100644 index 00000000000..b6b2c99fe76 --- /dev/null +++ b/releasenotes/notes/aiguard-abort-baseexception-deprecate-raising-dispatch-3a5c7f61ee4d2e1b.yaml @@ -0,0 +1,8 @@ +--- +upgrade: + - | + AI Guard: ``ddtrace.appsec.ai_guard.AIGuardAbortError`` now derives from + ``ddtrace.internal._exceptions.DDBlockException`` (a ``BaseException`` + subclass) instead of ``Exception``. This brings AI Guard block decisions in + line with how ASM blocks are surfaced and prevents a generic + ``except Exception:`` in user code from silently swallowing a block. diff --git a/tests/appsec/ai_guard/langchain/test_langchain.py b/tests/appsec/ai_guard/langchain/test_langchain.py index ffc26a03044..6cdf53c85d3 100644 --- a/tests/appsec/ai_guard/langchain/test_langchain.py +++ b/tests/appsec/ai_guard/langchain/test_langchain.py @@ -471,3 +471,61 @@ async def test_streamed_llm_async_block(mock_execute_request, langchain_openai, pass mock_execute_request.assert_called_once() + + +# --------------------------------------------------------------------------- +# Span observability on AI Guard block (non-streaming) +# +# The four non-stream LangChain wrappers (``traced_llm_generate`` / +# ``traced_chat_model_generate`` and async variants) catch ``DDBlockException`` +# explicitly and call ``span.set_exc_info`` before re-raising. Without that +# explicit arm, ``AIGuardAbortError`` (a ``BaseException`` subclass via +# ``DDBlockException``) would slip past ``except Exception:`` and the LLM span +# would finish with no error info — leaving a hole between the AI Guard span +# (block decision) and the LLM span (no link back to the abort). +# --------------------------------------------------------------------------- + + +def _find_llm_span_with_error(test_spans): + """Return the non-AI-Guard span (the LLM span) and its error tags.""" + from ddtrace.appsec._constants import AI_GUARD + + spans = test_spans.spans + llm_span = next((s for s in spans if s.name != AI_GUARD.RESOURCE_TYPE), None) + assert llm_span is not None, f"No LLM span found among: {[s.name for s in spans]}" + return llm_span + + +@patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request") +def test_chat_block_tags_llm_span_with_set_exc_info( + mock_execute_request, langchain_openai, openai_url, tracer, test_spans +): + """On AI Guard block, the LLM span must carry ``error == 1`` and an + ``error.type`` containing ``AIGuardAbortError``. + """ + mock_execute_request.return_value = mock_evaluate_response("DENY") + + chat = langchain_openai.ChatOpenAI(temperature=0, max_tokens=256, n=1, base_url=openai_url) + with pytest.raises(AIGuardAbortError): + chat.invoke(input=[HumanMessage(content="When do you use 'whom' instead of 'who'?")]) + + llm_span = _find_llm_span_with_error(test_spans) + assert llm_span.error == 1 + assert "AIGuardAbortError" in (llm_span.get_tag("error.type") or "") + + +@pytest.mark.asyncio +@patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request") +async def test_chat_async_block_tags_llm_span_with_set_exc_info( + mock_execute_request, langchain_openai, openai_url, tracer, test_spans +): + """Async variant of ``test_chat_block_tags_llm_span_with_set_exc_info``.""" + mock_execute_request.return_value = mock_evaluate_response("DENY") + + chat = langchain_openai.ChatOpenAI(temperature=0, max_tokens=256, n=1, base_url=openai_url) + with pytest.raises(AIGuardAbortError): + await chat.ainvoke(input=[HumanMessage(content="When do you use 'whom' instead of 'who'?")]) + + llm_span = _find_llm_span_with_error(test_spans) + assert llm_span.error == 1 + assert "AIGuardAbortError" in (llm_span.get_tag("error.type") or "") diff --git a/tests/appsec/ai_guard/openai/test_openai.py b/tests/appsec/ai_guard/openai/test_openai.py index 60112dfd3ba..d2cc40415be 100644 --- a/tests/appsec/ai_guard/openai/test_openai.py +++ b/tests/appsec/ai_guard/openai/test_openai.py @@ -9,8 +9,10 @@ """ import asyncio +import gc import threading from unittest.mock import patch +import warnings import pytest @@ -632,6 +634,31 @@ async def test_chat_async_unawaited_coro_does_not_evaluate(mock_execute_request, mock_execute_request.assert_not_called() +@pytest.mark.asyncio +@patch("ddtrace.appsec.ai_guard._api_client.AIGuardClient._execute_request") +async def test_chat_async_block_closes_unstarted_sdk_coroutine(mock_execute_request, async_openai_client): + """When AI Guard blocks at await time, the unstarted SDK coroutine MUST be + closed so Python doesn't emit ``RuntimeWarning: coroutine ... was never + awaited`` for it. Force GC inside the ``catch_warnings`` block so any + unclosed-coroutine warning surfaces while ``simplefilter("error")`` is + still active — the warning fires from ``coroutine.__del__``, not from + the raise itself, so a naive context exit would let it leak past the + assertion. + """ + mock_execute_request.return_value = mock_evaluate_response("DENY") + + with warnings.catch_warnings(record=True) as recorded: + warnings.simplefilter("always") + with pytest.raises(AIGuardAbortError): + await async_openai_client.chat.completions.create(messages=_user_messages(), **CHAT_PARAMS) + gc.collect() + + never_awaited = [ + w for w in recorded if issubclass(w.category, RuntimeWarning) and "was never awaited" in str(w.message) + ] + assert not never_awaited, f"Unstarted SDK coroutine leaked: {[str(w.message) for w in never_awaited]}" + + # --------------------------------------------------------------------------- # After-model evaluation: ALLOW before, DENY after # --------------------------------------------------------------------------- diff --git a/tests/internal/test_context_events_api.py b/tests/internal/test_context_events_api.py index d5d1572fa63..b18ac2e5d6b 100644 --- a/tests/internal/test_context_events_api.py +++ b/tests/internal/test_context_events_api.py @@ -239,6 +239,95 @@ def on_type_error(*_): with core.context_with_data("my.cool.context"): pass + @with_config_raise_value(raise_value=False) + def test_core_dispatch_allow_raise_propagates_exception(self): + def on_runtime_error(*_): + raise RuntimeError("OH NO!") + + core.on("my.cool.event", on_runtime_error) + + # Default allow_raise=False: Exception is swallowed + assert core.dispatch("my.cool.event", (1,)) is None + + # allow_raise=True: Exception propagates + with pytest.raises(RuntimeError): + core.dispatch("my.cool.event", (1,), allow_raise=True) + + @with_config_raise_value(raise_value=False) + def test_core_dispatch_allow_raise_short_circuits_listeners(self): + calls = [] + + def first(*_): + calls.append("first") + raise RuntimeError("first failed") + + def second(*_): + calls.append("second") + + core.on("my.cool.event", first, "first") + core.on("my.cool.event", second, "second") + + with pytest.raises(RuntimeError): + core.dispatch("my.cool.event", (), allow_raise=True) + assert calls == ["first"] + + @with_config_raise_value(raise_value=False) + def test_core_dispatch_ddblockexception_always_propagates(self): + from ddtrace.internal._exceptions import DDBlockException + + class FakeBlock(DDBlockException): + pass + + def on_block(*_): + raise FakeBlock() + + core.on("my.cool.event", on_block) + + # allow_raise=False (default): BaseException-derived block still propagates + # (it's not caught by the internal `except Exception:`) + with pytest.raises(DDBlockException): + core.dispatch("my.cool.event", ()) + + # allow_raise=True: also propagates + with pytest.raises(DDBlockException): + core.dispatch("my.cool.event", (), allow_raise=True) + + @with_config_raise_value(raise_value=False) + def test_core_dispatch_event_allow_raise_propagates(self): + from ddtrace.internal.core.event_hub import dispatch_event + + class FakeEvent: + event_name = "my.cool.event" + + def listener(_): + raise RuntimeError("evt boom") + + core.on("my.cool.event", listener) + + # Default allow_raise=False: swallowed + assert dispatch_event(FakeEvent()) is None + + # allow_raise=True: propagates + with pytest.raises(RuntimeError): + dispatch_event(FakeEvent(), allow_raise=True) + + def test_ddblockexception_inheritance(self): + from ddtrace.internal._exceptions import BlockingException + from ddtrace.internal._exceptions import DDBlockException + + # DDBlockException is BaseException-derived only, not Exception + assert issubclass(DDBlockException, BaseException) + assert not issubclass(DDBlockException, Exception) + + # BlockingException inherits the new base class + assert issubclass(BlockingException, DDBlockException) + + # AIGuardAbortError inherits from DDBlockException + from ddtrace.appsec.ai_guard._api_client import AIGuardAbortError + + assert issubclass(AIGuardAbortError, DDBlockException) + assert not issubclass(AIGuardAbortError, Exception) + def test_core_dispatch_context_ended(self): context_id = "my.cool.context" event_name = "context.ended.%s" % context_id