Repository navigation
feat(aiguard): send evaluations through the Agent when the service has no keys - #21569
Draft
eliottness wants to merge 1 commit into
Draft
eliottness wants to merge 1 commit into
eliottness wants to merge 1 commit into
Conversation
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
A service instrumented by Single Step Instrumentation never receives
DD_API_KEYandDD_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:
DD_API_KEYandDD_APP_KEYset: direct, exactly as before, andDD_AI_GUARD_ENDPOINTstill applies./infolists/ai_guard/proxy/: the Agent transport, sending the same body andDD-AI-GUARD-*headers with no credentials.Details worth a reviewer's attention:
ai_guard_listen()constructs the client before registering any listener andload_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.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.Datadog-Agent-Versionmeans the route vanished, so the client forgets it and goes back to discovery; a versioned 404 stays an ordinary service error.urlparseputs a socket path inpathand leavesnetlocempty, so folding the route into the URL would produceunix://.ai_guard.transportgauge taggedtransport:agent|direct, and anai_guard.agent.requestscount per evaluation sent to the Agent, taggedstatus_code,error_typeandagent_error(the Agent's ownDatadog-Agent-Proxy-Errorcode, carried out of the transport through aContextVarbecauseResponsekeeps 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
AIGuardClientwith 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 raisedAIGuardAbortErrorwithDENY, and both appear in the backend asai_guardspans withblocked=falseandblocked=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 thatnew_ai_guard_client()no longer raisesValueErrorwhen credentials are missing; it returns a client whose evaluations fail withAIGuardClientErrorinstead, which is what lets AI Guard install and recover on an Agent that becomes ready later.Not included, and deliberately:
ai_guard.transportis 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.