znyx-runtime 1.0.0__py3-none-any.whl

znyx_runtime/api/routes.py

@@ -0,0 +1,284 @@
+"""
+Runtime API routes - evaluation endpoints only.
+No DB dependency, no admin routes.
+"""
+import logging
+import os
+from typing import Optional
+
+from fastapi import APIRouter, Depends, HTTPException, Header
+
+from znyx_core.core.models import (
+    EvaluationRequest, EvaluationResponse, HealthResponse, ToolEvaluationRequest,
+    RetrievalEvaluationRequest, AgentPlanEvaluationRequest, AgentStepEvaluationRequest,
+    MemoryWriteEvaluationRequest,
+)
+
+logger = logging.getLogger(__name__)
+router = APIRouter()
+
+
+def _is_prod() -> bool:
+    from znyx_core.utils.env import is_production
+    return is_production()
+
+
+def _runtime_auth_required() -> bool:
+    """Default: on. Only explicitly disabled when RUNTIME_REQUIRE_AUTH=false,
+    and only honoured outside production — production always requires auth."""
+    raw = os.getenv("RUNTIME_REQUIRE_AUTH", "true").strip().lower()
+    if _is_prod():
+        return True
+    return raw not in ("0", "false", "no", "off")
+
+
+async def _optional_runtime_auth(
+    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
+    authorization: Optional[str] = Header(default=None),
+) -> None:
+    """Enforce runtime auth. Default-on; only disablable in non-production."""
+    if not _runtime_auth_required():
+        return
+
+    expected = os.getenv("RUNTIME_API_KEY", "").strip()
+    if not expected:
+        raise HTTPException(
+            status_code=500,
+            detail="Server misconfiguration: RUNTIME_REQUIRE_AUTH is enabled but RUNTIME_API_KEY is not set.",
+        )
+
+    # Accept either X-API-Key or Authorization: Bearer
+    token = x_api_key
+    if not token and authorization and authorization.startswith("Bearer "):
+        token = authorization[7:]
+
+    if not token:
+        raise HTTPException(
+            status_code=401,
+            detail="Authentication required. Set X-API-Key or Authorization: Bearer header.",
+        )
+
+    # Constant-time compare to prevent timing attacks.
+    import hmac as _hmac
+
+    if not _hmac.compare_digest(token, expected):
+        raise HTTPException(status_code=401, detail="Invalid authentication token")
+
+
+def _get_evaluator():
+    from znyx_runtime.main import evaluator
+    return evaluator
+
+
+def _get_bundle_manager():
+    from znyx_runtime.main import bundle_manager
+    return bundle_manager
+
+
+def _get_heartbeat():
+    from znyx_runtime.main import heartbeat
+    return heartbeat
+
+
+def _get_runtime_judge():
+    from znyx_runtime.main import runtime_judge
+    return runtime_judge
+
+
+def _runtime_judge_ctx(policy, scoped):
+    """Build the per-request judge context for a runtime evaluation, or None
+    when judges aren't used / judge audit is off. Stamps the request's tenant as org_scope
+    so the CP drain can attribute the spooled rows."""
+    rj = _get_runtime_judge()
+    if rj is None:
+        return None
+    return rj.context_for(policy, env=getattr(scoped, "env", None) or "prod",
+                          org_scope=getattr(scoped, "tenant_id", None))
+
+
+# Public aliases consumed by stream_routes and any other runtime-scoped code.
+# These intentionally shadow the private underscore names so callers can do:
+#   from znyx_runtime.api.routes import verify_runtime_auth, get_evaluator
+verify_runtime_auth = _optional_runtime_auth
+get_evaluator = _get_evaluator
+get_heartbeat = _get_heartbeat
+
+
+@router.post("/v1/evaluate/input", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_input(request: EvaluationRequest) -> EvaluationResponse:
+    """Evaluate input text before sending to LLM."""
+    try:
+        bm = _get_bundle_manager()
+        if bm.effective_env:
+            request = request.model_copy(update={"env": bm.effective_env})
+        policy = bm.get_policy(
+            tenant_id=request.tenant_id,
+            app_id=request.app_id,
+            agent_id=request.agent_id,
+            env=request.env,
+        )
+        ev = _get_evaluator()
+        result = await ev.evaluate(request, context="input", policy=policy,
+                                   judge_ctx=_runtime_judge_ctx(policy, request))
+        hb = _get_heartbeat()
+        if hb:
+            hb.increment_eval_count()
+        return result
+    except RuntimeError as e:
+        logger.error(f"Policy unavailable: {e}")
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.post("/v1/evaluate/output", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_output(request: EvaluationRequest) -> EvaluationResponse:
+    """Evaluate output text from LLM before returning to user."""
+    try:
+        bm = _get_bundle_manager()
+        if bm.effective_env:
+            request = request.model_copy(update={"env": bm.effective_env})
+        policy = bm.get_policy(
+            tenant_id=request.tenant_id,
+            app_id=request.app_id,
+            agent_id=request.agent_id,
+            env=request.env,
+        )
+        ev = _get_evaluator()
+        result = await ev.evaluate(request, context="output", policy=policy,
+                                   judge_ctx=_runtime_judge_ctx(policy, request))
+        hb = _get_heartbeat()
+        if hb:
+            hb.increment_eval_count()
+        return result
+    except RuntimeError as e:
+        logger.error(f"Policy unavailable: {e}")
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.post("/v1/evaluate/tool", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_tool(request: ToolEvaluationRequest) -> EvaluationResponse:
+    """Evaluate tool invocation against governance policies."""
+    try:
+        bm = _get_bundle_manager()
+        if bm.effective_env:
+            request = request.model_copy(update={"env": bm.effective_env})
+        policy = bm.get_policy(
+            tenant_id=request.tenant_id,
+            app_id=request.app_id,
+            agent_id=request.agent_id,
+            env=request.env,
+        )
+        ev = _get_evaluator()
+        result = await ev.evaluate_tool(request, policy=policy,
+                                        judge_ctx=_runtime_judge_ctx(policy, request))
+        hb = _get_heartbeat()
+        if hb:
+            hb.increment_eval_count()
+        return result
+    except RuntimeError as e:
+        logger.error(f"Policy unavailable: {e}")
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Tool evaluation failed")
+        raise HTTPException(status_code=500, detail="Tool evaluation failed")
+
+
+# ── New-stage evaluate endpoints (generalized stage dispatch) ──────
+
+async def _evaluate_stage_runtime(scoped, stage: str) -> EvaluationResponse:
+    """Shared runtime handler for the per-stage endpoints: resolve the bundle
+    policy for the request scope and dispatch through the generalized stage pipeline."""
+    bm = _get_bundle_manager()
+    if bm.effective_env:
+        scoped = scoped.model_copy(update={"env": bm.effective_env})
+    policy = bm.get_policy(
+        tenant_id=scoped.tenant_id,
+        app_id=scoped.app_id,
+        agent_id=scoped.agent_id,
+        env=scoped.env,
+    )
+    ev = _get_evaluator()
+    result = await ev.evaluate_stage(scoped, stage, policy=policy,
+                                     judge_ctx=_runtime_judge_ctx(policy, scoped))
+    hb = _get_heartbeat()
+    if hb:
+        hb.increment_eval_count()
+    return result
+
+
+@router.post("/v1/evaluate/retrieval", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_retrieval(request: RetrievalEvaluationRequest) -> EvaluationResponse:
+    """Evaluate retrieved RAG chunks for indirect prompt injection before they enter context (LLM01)."""
+    try:
+        return await _evaluate_stage_runtime(request, "retrieval")
+    except RuntimeError:
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Retrieval evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.post("/v1/evaluate/agent-plan", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_agent_plan(request: AgentPlanEvaluationRequest) -> EvaluationResponse:
+    """Evaluate a proposed multi-step agent plan for excessive agency (LLM06)."""
+    try:
+        return await _evaluate_stage_runtime(request, "agent_plan")
+    except RuntimeError:
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Agent-plan evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.post("/v1/evaluate/agent-step", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_agent_step(request: AgentStepEvaluationRequest) -> EvaluationResponse:
+    """Evaluate a single agent-loop iteration for budget/depth caps (LLM10)."""
+    try:
+        return await _evaluate_stage_runtime(request, "agent_loop")
+    except RuntimeError:
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Agent-step evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.post("/v1/evaluate/memory-write", response_model=EvaluationResponse, dependencies=[Depends(_optional_runtime_auth)])
+async def evaluate_memory_write(request: MemoryWriteEvaluationRequest) -> EvaluationResponse:
+    """Evaluate text being written to agent memory for persistent injection (LLM01)."""
+    try:
+        return await _evaluate_stage_runtime(request, "memory_write")
+    except RuntimeError:
+        raise HTTPException(status_code=503, detail="Policy unavailable (fail-closed)")
+    except Exception:
+        logger.exception("Memory-write evaluation failed")
+        raise HTTPException(status_code=500, detail="Evaluation failed")
+
+
+@router.get("/healthz", response_model=HealthResponse)
+async def health_check() -> HealthResponse:
+    """Health check - always returns 200."""
+    return HealthResponse(status="ok")
+
+
+@router.get("/readyz", response_model=HealthResponse)
+async def readiness_check() -> HealthResponse:
+    """Readiness check - returns 200 only if a policy is loaded."""
+    bm = _get_bundle_manager()
+    if bm.is_ready:
+        return HealthResponse(status="ready")
+    raise HTTPException(status_code=503, detail="Not ready - no policy loaded")
+
+
+@router.get("/v1/bundle/status")
+async def bundle_status():
+    """Return metadata about the currently loaded policy bundle."""
+    bm = _get_bundle_manager()
+    return {
+        "ready": bm.is_ready,
+        **bm.bundle_info,
+    }

znyx_runtime/api/stream_routes.py

@@ -0,0 +1,101 @@
+"""Streaming guardrails SSE endpoint - evaluate LLM output in real-time."""
+import json
+import logging
+from typing import Any, Dict, List, Optional
+from fastapi import APIRouter, Depends
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel, Field
+
+from znyx_runtime.api.routes import verify_runtime_auth, get_evaluator, get_heartbeat
+
+logger = logging.getLogger(__name__)
+router = APIRouter(tags=["Runtime API"])
+
+
+class StreamEvaluateRequest(BaseModel):
+    """Request body for streaming evaluation."""
+    request_id: str = "stream-0"
+    tenant_id: str = "default"
+    app_id: str = "default"
+    context: str = Field(default="output", description="input or output")
+    chunks: List[str] = Field(..., description="Text chunks to evaluate in order")
+    policy: Optional[Dict[str, Any]] = Field(default=None, description="Inline policy (optional)")
+    window_size: int = Field(default=200, ge=50, le=2000)
+    overlap: int = Field(default=40, ge=0, le=500)
+
+
+@router.post(
+    "/v1/evaluate/stream",
+    operation_id="runtime.evaluateStream",
+    summary="Evaluate streaming LLM output as it arrives",
+)
+async def evaluate_stream(
+    body: StreamEvaluateRequest,
+    _auth=Depends(verify_runtime_auth),
+    evaluator=Depends(get_evaluator),
+):
+    """Evaluate text chunks via Server-Sent Events.
+
+    Accepts a list of text chunks and streams back SSE events as each window
+    is evaluated. Events:
+      - ``chunk``: forwarded text chunk
+      - ``guardrail``: window evaluation result (non-blocking)
+      - ``block``: window triggered a BLOCK decision
+      - ``done``: final summary with aggregate metrics
+    """
+    from znyx_core.engine.streaming import StreamingEvaluator
+    from znyx_core.core.models import EvaluationRequest
+
+    # Resolve policy
+    policy = body.policy
+    if policy is None and evaluator.policy_resolver:
+        policy = evaluator.policy_resolver.resolve(
+            tenant_id=body.tenant_id,
+            app_id=body.app_id,
+            agent_id="default",
+            env="prod",
+        )
+    if policy is None:
+        policy = {}
+
+    request = EvaluationRequest(
+        request_id=body.request_id,
+        tenant_id=body.tenant_id,
+        app_id=body.app_id,
+        text="",  # text comes in chunks
+    )
+
+    stream_eval = StreamingEvaluator(
+        policy=policy,
+        context=body.context,
+        window_size=body.window_size,
+        overlap=body.overlap,
+        request=request,
+    )
+
+    async def event_generator():
+        for chunk in body.chunks:
+            events = stream_eval.push(chunk)
+            for event in events:
+                yield f"event: {event['event']}\ndata: {json.dumps(event['data'])}\n\n"
+                if event["event"] == "block":
+                    # Still send done after block
+                    break
+            if stream_eval.is_blocked:
+                break
+
+        summary = stream_eval.flush()
+        yield f"event: {summary['event']}\ndata: {json.dumps(summary['data'])}\n\n"
+        hb = get_heartbeat()
+        if hb:
+            hb.increment_eval_count()
+
+    return StreamingResponse(
+        event_generator(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )

znyx_runtime/audit_sink.py

@@ -0,0 +1,195 @@
+"""Egress audit sink.
+
+The zero-DB runtime cannot write control-plane tables, but the egress contract requires that every
+boundary-crossing call (remote_llm / remote_api / a hosted inference sidecar) is
+audited — and that a *failed* audit write never silently permits an un-audited
+egress. This module defines the runtime-side abstraction the egress gate calls
+*before* any such call:
+
+    event = EgressAuditEvent(...metadata only...)
+    await sink.record(event)        # raises (fail-closed) if it cannot durably record
+    # ...only now make the egress call...
+
+Backends (selectable, see ``make_audit_sink``):
+  * ``SpoolAuditSink``  — durable append-only JSON-lines spool on local disk
+        (default ``~/.znyx/egress-audit.spool``). Works without a database, so it
+        is valid in OSS/local mode. The control plane drains the spool into
+        ``egress_events`` rows and hash-chain-signs them.
+  * ``NoopAuditSink``   — explicit opt-out (records nothing).
+
+Failure semantics: in ``fail_mode="closed"`` (the default) a write failure raises
+``AuditWriteError`` so the caller denies the egress. In ``fail_mode="open"`` the
+failure is logged and swallowed (the egress proceeds un-audited) — only for
+operators who explicitly accept that risk.
+
+Privacy: events carry METADATA ONLY (destination host/region, byte count, redaction
+flag, detector key, model/rubric ids) — never raw prompt/response content, mirroring
+the telemetry posture. Audit is NOT gated on GDPR telemetry consent (different lawful
+basis: security/compliance logging).
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import os
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import List, Optional
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_SPOOL = Path.home() / ".znyx" / "egress-audit.spool"
+
+
+class AuditWriteError(RuntimeError):
+    """Raised by a fail-closed sink when an audit event cannot be durably recorded."""
+
+
+@dataclass
+class EgressAuditEvent:
+    """A single boundary-crossing audit record. Metadata only — no raw content."""
+    occurred_at: str                                  # ISO-8601 timestamp (caller-supplied)
+    mode: str                                         # remote_llm | remote_api | local_ml(hosted) | ...
+    detector_key: str                                 # e.g. "jailbreak"
+    destination_host: str                             # host the call left the boundary to
+    org_scope: Optional[str] = None                   # tenant/org identifier (no PII)
+    trace_id: Optional[str] = None
+    destination_region: Optional[str] = None
+    bytes_out: int = 0                                # size of the payload sent (count, not content)
+    redacted: bool = False                            # was redact_before_egress applied
+    allowlisted: bool = False                         # destination matched the per-detector allowlist
+    model_version: Optional[str] = None               # model_id@revision the call targeted
+    event_id: Optional[str] = None                    # stable UUID → egress_events PK (idempotent drain)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class AuditSink(ABC):
+    """Records egress audit events. ``record`` MUST raise in fail-closed mode if it
+    cannot durably persist the event, so the caller can deny the egress."""
+
+    def __init__(self, fail_mode: str = "closed"):
+        self.fail_mode = fail_mode if fail_mode in ("closed", "open") else "closed"
+
+    @abstractmethod
+    async def record(self, event: EgressAuditEvent) -> None:
+        ...
+
+    @abstractmethod
+    def record_sync(self, event: EgressAuditEvent) -> None:
+        """Synchronous durable record, for the (synchronous) escalation gate which
+        cannot await. MUST raise in fail-closed mode if it cannot durably persist, so
+        the gate denies the egress."""
+        ...
+
+    def _handle_failure(self, exc: Exception) -> None:
+        if self.fail_mode == "closed":
+            raise AuditWriteError(f"egress audit write failed (fail-closed): {exc}") from exc
+        logger.error("egress audit write failed (fail-open — egress NOT audited): %s", exc)
+
+
+class NoopAuditSink(AuditSink):
+    """Records nothing. Explicit opt-out only."""
+
+    async def record(self, event: EgressAuditEvent) -> None:
+        return None
+
+    def record_sync(self, event: EgressAuditEvent) -> None:
+        return None
+
+
+class SpoolAuditSink(AuditSink):
+    """Durable append-only JSON-lines spool on local disk.
+
+    Each ``record`` appends one JSON line and fsyncs before returning, so a recorded
+    event survives a crash. The control plane drains and signs the spool.
+    """
+
+    def __init__(self, spool_path: Optional[str] = None, fail_mode: str = "closed"):
+        super().__init__(fail_mode=fail_mode)
+        self.path = Path(spool_path) if spool_path else _DEFAULT_SPOOL
+        self._lock = asyncio.Lock()
+
+    def _append_sync(self, line: str) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        # Append + flush + fsync so the record is durable before we return.
+        with open(self.path, "a", encoding="utf-8") as fh:
+            fh.write(line + "\n")
+            fh.flush()
+            os.fsync(fh.fileno())
+
+    async def record(self, event: EgressAuditEvent) -> None:
+        line = json.dumps(event.to_dict(), separators=(",", ":"), sort_keys=True)
+        try:
+            async with self._lock:
+                await asyncio.to_thread(self._append_sync, line)
+        except Exception as exc:  # noqa: BLE001 — any IO failure must hit the fail-mode gate
+            self._handle_failure(exc)
+
+    def record_sync(self, event: EgressAuditEvent) -> None:
+        line = json.dumps(event.to_dict(), separators=(",", ":"), sort_keys=True)
+        try:
+            self._append_sync(line)
+        except Exception as exc:  # noqa: BLE001 — any IO failure must hit the fail-mode gate
+            self._handle_failure(exc)
+
+    def read_all(self) -> List[dict]:
+        """Read every spooled event (used by the control-plane drainer / tests).
+
+        Resilient to a single corrupt/partial line (e.g. an interrupted write): a
+        line that won't parse is logged and skipped so one bad record can't abort the
+        whole drain."""
+        if not self.path.exists():
+            return []
+        out: List[dict] = []
+        for raw in self.path.read_text(encoding="utf-8").splitlines():
+            raw = raw.strip()
+            if not raw:
+                continue
+            try:
+                out.append(json.loads(raw))
+            except json.JSONDecodeError as exc:
+                logger.warning("skipping unparseable egress spool line: %s", exc)
+        return out
+
+
+def make_audit_sink(
+    mode: str = "spool",
+    fail_mode: str = "closed",
+    spool_path: Optional[str] = None,
+) -> AuditSink:
+    """Construct an audit sink from runtime configuration."""
+    if mode == "noop":
+        return NoopAuditSink(fail_mode=fail_mode)
+    if mode == "spool":
+        return SpoolAuditSink(spool_path=spool_path, fail_mode=fail_mode)
+    logger.warning("Unknown audit_sink_mode %r — defaulting to durable spool", mode)
+    return SpoolAuditSink(spool_path=spool_path, fail_mode=fail_mode)
+
+
+def make_audit_egress_sink(sink: "AuditSink"):
+    """Adapt an ``AuditSink`` to the (synchronous) egress-gate callback the
+    escalation engine calls before any boundary-crossing call.
+
+    The escalation engine passes an ``znyx_core.engine.egress.EgressEvent``; we map
+    it to an ``EgressAuditEvent`` and ``record_sync`` it. A fail-closed sink raises
+    on a failed write, which the gate turns into a denial (no un-audited egress)."""
+    def _sink(ev) -> None:
+        sink.record_sync(EgressAuditEvent(
+            occurred_at=ev.occurred_at,
+            mode=ev.mode,
+            detector_key=ev.detector_key,
+            destination_host=ev.destination_host or "",
+            org_scope=ev.org_scope,
+            trace_id=ev.trace_id,
+            destination_region=ev.destination_region,
+            bytes_out=ev.bytes_out,
+            redacted=ev.redacted,
+            allowlisted=ev.allowlisted,
+            model_version=ev.model_version,
+            event_id=ev.event_id,
+        ))
+    return _sink