Skip to content

feat(aiguard): send evaluations through the Agent when the service has no keys - #21569

Draft
eliottness wants to merge 1 commit into
DataDog:mainfrom
eliottness:eliott.bouhana/aiguard-agent-data-path
Draft

eliottness wants to merge 1 commit into
DataDog:mainfrom
eliottness:eliott.bouhana/aiguard-agent-data-path

Conversation

@eliottness

Copy link
Copy Markdown
Contributor

Description

A service instrumented by Single Step Instrumentation never receives DD_API_KEY and DD_APP_KEY, and the AI Guard SDK needs both to call the evaluate API, so AI Guard simply stayed off for those services. This adds the SDK half of the Agent data path RFC: when the service has no credentials and the local Agent advertises /ai_guard/proxy/, the client sends evaluations to the Agent, which attaches its own API key and application key.

The Agent side is DataDog/datadog-agent#59716.

Transport selection follows the RFC and adds no new setting:

  1. Both DD_API_KEY and DD_APP_KEY set: direct, exactly as before, and DD_AI_GUARD_ENDPOINT still applies.
  2. Otherwise, if the Agent's /info lists /ai_guard/proxy/: the Agent transport, sending the same body and DD-AI-GUARD-* headers with no credentials.
  3. Otherwise: no transport, and evaluations fail like any other failed AI Guard call, saying why.

Details worth a reviewer's attention:

  • Setup never fails on a missing route. The route is resolved per evaluation, not when the client is built, because ai_guard_listen() constructs the client before registering any listener and load_ai_guard() clears its one-shot flag even on failure. Raising here would mean a keyless service whose Agent is not ready yet never installs AI Guard and never recovers.
  • Discovery is throttled and shared. While the route is absent the Agent is re-checked at most once per AGENT_CHECK_INTERVAL, behind a fork-safe lock, with the first check on a background thread so startup never blocks on /info. A generation counter makes a probe that finishes late unable to resurrect a route that was invalidated meanwhile.
  • A 404 without Datadog-Agent-Version means the route vanished, so the client forgets it and goes back to discovery; a versioned 404 stays an ordinary service error.
  • 401/403 in Agent mode are reported as the Agent's credentials, not the service's, since that is who owns them.
  • Unix sockets work. The connection base is kept whole and the request path passed separately, because urlparse puts a socket path in path and leaves netloc empty, so folding the route into the URL would produce unix://.
  • Telemetry adds the two metrics the RFC defines: an ai_guard.transport gauge tagged transport:agent|direct, and an ai_guard.agent.requests count per evaluation sent to the Agent, tagged status_code, error_type and agent_error (the Agent's own Datadog-Agent-Proxy-Error code, carried out of the transport through a ContextVar because Response keeps no headers and one client is shared across threads).

A single credential keeps working as before: each key is sent on its own, so a caller constructing AIGuardClient with only one is unaffected.

Testing

tests/aiguard/api/: 475 passed. The new cases cover the selection matrix (both keys, one key, none, endpoint override), the TCP and Unix connection bases with the outgoing path, body and absence of credential headers, the throttled re-check, rediscovery on an unversioned 404 versus an ordinary versioned 404, the 401/403 reporting, and both metrics with exact call counts rather than presence checks. Each new assertion was mutation-checked: reverting the behavior it names makes it fail.

Validated end to end against the Datadog prod org, with a locally built trace-agent from DataDog/datadog-agent#59716 and an app holding no credential: a benign prompt returned ALLOW, a prompt-injection attempt raised AIGuardAbortError with DENY, and both appear in the backend as ai_guard spans with blocked=false and blocked=true.

Risks

Low for existing users: a service with both keys takes the same path as before and never calls /info. The behavior change is that new_ai_guard_client() no longer raises ValueError when credentials are missing; it returns a client whose evaluations fail with AIGuardClientError instead, which is what lets AI Guard install and recover on an Agent that becomes ready later.

Not included, and deliberately: ai_guard.transport is emitted per evaluation rather than on every telemetry flush, because the telemetry writer is a thin shim over the native worker and exposes no flush hook. A service that initializes and flushes without ever evaluating therefore reports no transport gauge.

Additional Notes

Draft: the Agent-side PR is not merged yet, so the Agent transport only has a route to discover once it ships.

…s no keys

A service instrumented by Single Step Instrumentation never gets DD_API_KEY and DD_APP_KEY, so AI Guard stayed off. When both keys are missing and the Agent advertises /ai_guard/proxy/ in its /info endpoint list, the client now sends evaluations to that route and the Agent attaches its own credentials. Services that set both keys keep calling the API directly.

Adds the two SDK metrics the RFC defines: an ai_guard.transport gauge tagged with the transport in use, and an ai_guard.agent.requests count per evaluation sent to the Agent, tagged with the status code, the error type and the Agent's own error code.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant