ajeetraina777/datadog-ai-guard-kit:latest

schemaVersion: "2"
kind: mixin
name: datadog-ai-guard
displayName: Datadog AI Guard
description: Installs the Datadog AI Guard SDKs (Python ddtrace + Node dd-trace) and wires agentless credentials/env so AI apps and agents built inside the sandbox can screen LLM prompts, tool calls, and outputs for prompt injection, jailbreaks, tool misuse, and sensitive-data exfiltration via client.evaluate(...). Datadog API/APP keys are proxy-injected and never enter the container. Sets DD_* environment variables (last-wins if a later --kit overrides them).
licenses:
    - Apache-2.0
args:
    apm:
        default: "false"
        description: Send AI Guard evaluations as APM traces so they appear in the Datadog UI ("AI Guard / Submit your first trace"). "true" runs a Datadog Agent container in the sandbox (trace intake on 127.0.0.1:8126) that forwards through the credential-injecting proxy. Default "false" keeps the kit fully agentless — evaluate() still returns live verdicts for inline enforcement, they just won't show in the trace UI.
        pattern: ^(true|false)$
    env:
        default: sandbox
        description: Value for DD_ENV (deployment environment tag).
        pattern: ^[A-Za-z0-9._:/-]+$
    service:
        default: sbx-ai-guard
        description: Value for DD_SERVICE (service name tag).
        pattern: ^[A-Za-z0-9._:/-]+$
    site:
        default: datadoghq.com
        description: 'Datadog site (DD_SITE): datadoghq.com | datadoghq.eu | us3.datadoghq.com | us5.datadoghq.com | ap1.datadoghq.com'
        pattern: ^[a-z0-9.-]+\.[a-z]{2,}$
agentInstructions:
    content: |
        # Datadog AI Guard

        This sandbox has the Datadog AI Guard SDKs pre-installed so any AI app or
        agent you build here can screen LLM interactions in real time for **prompt
        injection, jailbreaks, tool misuse, and sensitive-data exfiltration**.

        Call `evaluate(...)` before you act on user input, before you run a tool the
        model requested, and (optionally) on model output. Pass the full conversation
        so far; AI Guard returns an action — typically ALLOW / DENY / ABORT — and with
        `block=True/`{ block: true }` it raises when the interaction should be blocked.

        ## Already configured

        - **Credentials**: `DD_API_KEY` / `DD_APP_KEY` are a proxy-managed placeholder,
          never the real value. Do **not** try to read or print them — the sbx proxy
          swaps in the real values on outbound calls to the AI Guard endpoint
          (`app.$DD_SITE` on base sites, the bare `$DD_SITE` on us3/us5/ap1). They never
          exist in the container.
        - **Proxy routing**: outbound HTTPS must go through the sbx proxy for the swap
          to happen. Python is handled automatically (a `.pth` auto-imports the
          `_sbx_proxy_tunnel` shim). **Node**: `import '~/.datadog/sbx_proxy_tunnel.mjs'`
          **before** dd-trace makes any call, otherwise it bypasses the proxy and 401s.
        - **Environment**: `DD_AI_GUARD_ENABLED=true`, `DD_SITE`, `DD_ENV`,
          `DD_SERVICE` are exported. APM tracing and instrumentation telemetry are
          disabled by default (agentless mode, no local Datadog Agent).
        - **Trace UI (optional)**: `evaluate()` always returns the verdict for inline
          enforcement. To also see evaluations in the Datadog AI Guard UI, launch the
          kit with `--kit-arg apm=true` — it runs an in-sandbox Datadog Agent that
          forwards the evaluation spans through the proxy. Start/restart it from a
          shell with `sh ~/.datadog/start-agent.sh`. Not needed for enforcement.
        - **Network**: egress is allowed to the AI Guard endpoint (`app.$DD_SITE` /
          `$DD_SITE`) plus the pip/npm registries.

        ## Python (ddtrace >= 3.19.0)

        ```python
        from ddtrace.aiguard import new_ai_guard_client, Message, Options

        client = new_ai_guard_client()
        result = client.evaluate(
            messages=[
                Message(role="system", content="You are an AI Assistant"),
                Message(role="user", content=user_input),
            ],
            options=Options(block=True),  # raises if the interaction should be blocked
        )
        ```

        ## Node.js (dd-trace >= 5.69.0)

        dd-trace is installed globally. In your project either `npm install dd-trace`
        (network is already allowed) or resolve the global copy. Import the proxy
        helper first so dd-trace's HTTPS is routed through the injecting proxy:

        ```javascript
        import '~/.datadog/sbx_proxy_tunnel.mjs';  // MUST be first; else 401
        import tracer from 'dd-trace';

        const result = await tracer.aiguard.evaluate(
          [
            { role: 'system', content: 'You are an AI Assistant' },
            { role: 'user', content: userInput },
          ],
          { block: true },
        );
        ```

        Ready-to-run examples are in `~/.datadog/`. An operational runbook is in
        `~/runbooks/datadog-ai-guard.md`. Full docs:
        https://docs.datadoghq.com/security/ai_guard/
permissions:
    network:
        allow:
            - app.${{ kit.args.site }}
            - ${{ kit.args.site }}
            - '*.${{ kit.args.site }}'
            - pypi.org
            - files.pythonhosted.org
            - registry.npmjs.org
            - trace.agent.${{ kit.args.site }}
            - registry-1.docker.io
            - auth.docker.io
            - production.cloudflare.docker.com
credentials:
    - service: datadogapi
      description: Datadog API key (sent as the DD-API-KEY header).
      required: true
      apiKey:
        name: DD_API_KEY
        inject:
            - domain: app.${{ kit.args.site }}
              header: DD-API-KEY
              format: '%s'
            - domain: ${{ kit.args.site }}
              header: DD-API-KEY
              format: '%s'
            - domain: trace.agent.${{ kit.args.site }}
              header: DD-Api-Key
              format: '%s'
            - domain: api.${{ kit.args.site }}
              header: DD-API-KEY
              format: '%s'
    - service: datadogapp
      description: Datadog application key (sent as the DD-APPLICATION-KEY header).
      required: true
      apiKey:
        name: DD_APP_KEY
        inject:
            - domain: app.${{ kit.args.site }}
              header: DD-APPLICATION-KEY
              format: '%s'
            - domain: ${{ kit.args.site }}
              header: DD-APPLICATION-KEY
              format: '%s'
environment:
    variables:
        DD_AI_GUARD_ENABLED: "true"
        DD_APM_TRACING_ENABLED: ${{ kit.args.apm }}
        DD_ENV: ${{ kit.args.env }}
        DD_INSTRUMENTATION_TELEMETRY_ENABLED: "false"
        DD_SERVICE: ${{ kit.args.service }}
        DD_SITE: ${{ kit.args.site }}
setup:
    install:
        - command: 'command -v python3 >/dev/null 2>&1 && (python3 -c ''import ddtrace'' 2>/dev/null || python3 -m pip install --break-system-packages ''ddtrace>=3.19.0'' || python3 -m pip install ''ddtrace>=3.19.0'') || echo ''datadog-ai-guard: python3 not found, skipping ddtrace'''
          description: Install the Datadog Python SDK (ddtrace) for AI Guard
        - command: |
            command -v python3 >/dev/null 2>&1 && python3 -c '
            import glob, os, sysconfig, site, shutil
            cands = glob.glob("/home/*/.datadog/_sbx_proxy_tunnel.py") + ["/root/.datadog/_sbx_proxy_tunnel.py"]
            src = next((c for c in cands if os.path.isfile(c)), None)
            if not src:
                print("datadog-ai-guard: proxy shim source not found; SDK calls may 401"); raise SystemExit(0)
            dirs = set()
            try:
                dirs.update(p for p in site.getsitepackages() if p)
            except Exception:
                pass
            for k in ("purelib", "platlib"):
                p = sysconfig.get_paths().get(k)
                if p:
                    dirs.add(p)
            done = []
            for d in dirs:
                try:
                    os.makedirs(d, exist_ok=True)
                    shutil.copyfile(src, os.path.join(d, "_sbx_proxy_tunnel.py"))
                    with open(os.path.join(d, "aaa_sbx_proxy_tunnel.pth"), "w") as f:
                        f.write("import _sbx_proxy_tunnel\n")
                    done.append(d)
                except Exception:
                    pass
            print("datadog-ai-guard: proxy tunnel shim installed ->", done or "(no writable site dir)")
            ' || echo "datadog-ai-guard: python3 not found, skipping proxy shim"
          description: Route Python stdlib HTTPS through the sbx proxy so the AI Guard SDK works
        - command: 'command -v npm >/dev/null 2>&1 && (npm ls -g dd-trace >/dev/null 2>&1 || npm install -g ''dd-trace@^5.69.0'') || echo ''datadog-ai-guard: npm not found, skipping dd-trace'''
          description: Install the Datadog Node SDK (dd-trace) for AI Guard
        - command: |
            if [ "${{ kit.args.apm }}" != "true" ]; then
              echo "datadog-ai-guard: apm=false (agentless) — skipping Datadog Agent; evaluate() still returns live verdicts."
            elif ! command -v docker >/dev/null 2>&1; then
              echo "datadog-ai-guard: apm=true but docker is unavailable — cannot run the trace Agent."
            else
              s="$(ls /home/*/.datadog/start-agent.sh /root/.datadog/start-agent.sh 2>/dev/null | head -1)"
              [ -n "$s" ] && chmod +x "$s" 2>/dev/null || true
              # Warm the image so the first runtime start is fast; non-fatal if it fails.
              docker pull "${DD_AGENT_IMAGE:-gcr.io/datadoghq/agent:7}" >/dev/null 2>&1 \
                || echo "datadog-ai-guard: agent image pre-pull deferred (will pull on first start)."
              # Best-effort start now; if the runtime credential env is not yet present
              # at setup time, start it from a shell later: sh ~/.datadog/start-agent.sh
              if [ -n "$s" ] && [ -n "${DD_API_KEY:-}" ]; then
                DD_APM_TRACING_ENABLED=true \
                DD_SITE="${{ kit.args.site }}" \
                HTTPS_PROXY="${HTTPS_PROXY:-http://gateway.docker.internal:3128}" \
                sh "$s" || echo "datadog-ai-guard: agent start deferred — run 'sh ~/.datadog/start-agent.sh' from a shell."
              else
                echo "datadog-ai-guard: start the trace Agent from a shell -> sh ~/.datadog/start-agent.sh"
              fi
            fi
          description: Start an in-sandbox Datadog Agent for the AI Guard trace UI (apm=true only)