openscript-sdk 0.1.0__py3-none-any.whl

openscript_sdk/server/routers/guard.py

@@ -0,0 +1,231 @@
+from __future__ import annotations
+
+import time
+from typing import Any, Literal
+
+from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.encoders import jsonable_encoder
+from pydantic import BaseModel, Field
+
+from openscript_sdk.contracts.types import ActionBlockedError, ActionContext
+from openscript_sdk.server.auth import require_api_key
+from openscript_sdk.server.guard import SECTIONS, Checkpoint
+
+router = APIRouter(prefix="/v1", tags=["guard"])
+
+
+class GuardRequest(BaseModel):
+    """One phase of one action, asked by a client that runs the action itself."""
+
+    phase: Literal["input", "output"]
+    #: "invoke" for an agent's input and output, "stream" for the same when
+    #: the answer is streamed (SecureAgent.stream's action), "tool_call" for a
+    #: tool's {"name", "args"} and its result. Tools go to a different list.
+    action: Literal["invoke", "stream", "tool_call"]
+    session_id: str = "default"
+    agent_id: str = "default"
+    input: dict[str, Any]
+    #: Required for phase "output": what the agent or tool returned.
+    output: dict[str, Any] | None = None
+    #: The caller's role, for tool rules restricted by role.
+    role: str | None = None
+    #: An approved approval to redeem in the input phase.
+    approval_id: str = ""
+
+
+class GuardResponse(BaseModel):
+    decision: Literal["allow", "deny", "require_approval"]
+    reason: str = ""
+    blocked_by: str = ""
+    #: Set when decision is "require_approval": decide it through
+    #: /v1/approvals, then ask again with it.
+    approval_id: str = ""
+    #: The input as the policies left it. Give the agent or tool this, not
+    #: what was sent: a redacting policy replaces what it redacts.
+    input: dict[str, Any]
+    #: The output as the policies left it (phase "output" only).
+    output: dict[str, Any] | None = None
+    risk_score: float = 0.0
+    risk_categories: dict[str, float] = Field(default_factory=dict)
+    #: What each policy wrote about this phase, keyed by its category.
+    findings: dict[str, Any] = Field(default_factory=dict)
+    #: Time spent in the pipeline, measured inside the server, so a client can
+    #: tell the policies' cost from the network's.
+    pipeline_ms: float = 0.0
+    #: Action "stream", phase "input" only: the guards that scan a guarded
+    #: stream as it goes, and the tail a client must hold back for them.
+    stream_guards: list[str] | None = None
+    stream_window: int | None = None
+
+
+class StreamScanRequest(BaseModel):
+    """The text a client is holding back in a guarded stream, to be scanned."""
+
+    session_id: str = "default"
+    agent_id: str = "default"
+    #: The input as the input phase left it. Only for the record of a refusal.
+    input: dict[str, Any] = Field(default_factory=dict)
+    #: The held tail plus whatever has arrived since.
+    text: str
+    #: True at the end of the stream: matches at the edge are now definite.
+    final: bool = False
+
+
+class StreamScanResponse(BaseModel):
+    decision: Literal["allow", "deny"]
+    #: The text as the guards left it (redacted); "" when denied.
+    text: str = ""
+    reason: str = ""
+    blocked_by: str = ""
+    labels: list[str] = Field(default_factory=list)
+    risk: float = 0.0
+    pipeline_ms: float = 0.0
+
+
+def _checkpoint(request: Request, action: str) -> Checkpoint:
+    checkpoints: dict[str, Checkpoint] | None = getattr(request.app.state, "guard", None)
+    if not checkpoints:
+        raise HTTPException(
+            status_code=503,
+            detail="no policies configured; start the server with OPENSCRIPT_GUARD_CONFIG",
+        )
+    checkpoint = checkpoints.get(action)
+    if checkpoint is None:
+        raise HTTPException(
+            status_code=503,
+            detail=f"the guard config has no '{SECTIONS[action]}' section for action '{action}'",
+        )
+    return checkpoint
+
+
+def _ms(start: float) -> float:
+    return round((time.perf_counter() - start) * 1000, 3)
+
+
+def _response(
+    context: ActionContext,
+    phase: str,
+    decision: str,
+    start: float,
+    blocked_by: str = "",
+    approval_id: str = "",
+) -> GuardResponse:
+    findings = {
+        key: value
+        for key, value in context.metadata.items()
+        if isinstance(value, dict) and ("risk" in value or "category" in value)
+    }
+    return GuardResponse(
+        decision=decision,  # type: ignore[arg-type]
+        reason=context.decision_reason,
+        blocked_by=blocked_by,
+        approval_id=approval_id,
+        input=jsonable_encoder(context.input_data),
+        # Withheld when refused: what a policy blocked is not handed back to be used.
+        output=(
+            jsonable_encoder(context.output_data)
+            if phase == "output" and decision == "allow"
+            else None
+        ),
+        risk_score=round(context.risk_score, 4),
+        risk_categories=context.risk_categories,
+        findings=jsonable_encoder(findings),
+        pipeline_ms=_ms(start),
+    )
+
+
+@router.post("/guard", dependencies=[Depends(require_api_key)])
+async def guard(req: GuardRequest, request: Request) -> GuardResponse:
+    """Run one phase of the configured policies over an action the caller runs.
+
+    A refusal is an answer, not an error: it comes back with status 200 and
+    decision "deny" or "require_approval". A policy that fails closed is an
+    error (500), and so is a server with nothing configured (503). A client
+    must treat every error as a refusal.
+    """
+    checkpoint = _checkpoint(request, req.action)
+    if req.phase == "output" and req.output is None:
+        raise HTTPException(status_code=422, detail="phase 'output' needs 'output'")
+
+    metadata: dict[str, Any] = {}
+    if req.action == "tool_call":
+        name, args = req.input.get("name"), req.input.get("args")
+        if not isinstance(name, str) or not name or not isinstance(args, dict):
+            raise HTTPException(
+                status_code=422,
+                detail="a tool_call's input is {'name': <tool>, 'args': {...}}",
+            )
+        # Every argument came from the caller: there are no defaults to tell apart.
+        metadata["defaulted_args"] = []
+    if req.role is not None:
+        metadata["role"] = req.role
+
+    context = ActionContext(
+        action=req.action,
+        agent_id=req.agent_id,
+        session_id=req.session_id,
+        input_data=req.input,
+        output_data=req.output or {},
+        metadata=metadata,
+    )
+    start = time.perf_counter()
+    try:
+        if req.phase == "input":
+            context = await checkpoint.check_input(context, req.approval_id)
+        else:
+            context = await checkpoint.check_output(context)
+    except ActionBlockedError as blocked:
+        context = blocked.context or context
+        decision = "require_approval" if blocked.approval_id else "deny"
+        return _response(
+            context, req.phase, decision, start, blocked.interceptor, blocked.approval_id
+        )
+    except RuntimeError as exc:
+        # A FAIL_CLOSED policy that raised. The action is refused either way.
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+    response = _response(context, req.phase, "allow", start)
+    if req.action == "stream" and req.phase == "input":
+        guards = checkpoint.stream_guards()
+        response.stream_guards = [guard.name for guard in guards]
+        # What SecureAgent holds back: no match a guard can make is longer.
+        response.stream_window = max((guard.max_match_len for guard in guards), default=0)
+    return response
+
+
+@router.post("/guard/stream", dependencies=[Depends(require_api_key)])
+async def guard_stream(req: StreamScanRequest, request: Request) -> StreamScanResponse:
+    """Scan the held-back text of a guarded stream, as SecureAgent does per chunk.
+
+    The client releases everything but the last stream_window characters of
+    what comes back, and sends the rest again with the next chunk. A deny
+    comes back as an answer with no text: nothing of the match is released.
+    """
+    checkpoint = _checkpoint(request, "stream")
+    context = ActionContext(
+        action="stream",
+        agent_id=req.agent_id,
+        session_id=req.session_id,
+        input_data=req.input,
+    )
+    start = time.perf_counter()
+    try:
+        text = await checkpoint.scan_stream(context, req.text, req.final)
+    except ActionBlockedError as blocked:
+        found = context.metadata.get("stream_guard", {})
+        return StreamScanResponse(
+            decision="deny",
+            reason=blocked.reason,
+            blocked_by=blocked.interceptor,
+            labels=found.get("labels", []),
+            risk=found.get("risk", 0.0),
+            pipeline_ms=_ms(start),
+        )
+    found = context.metadata.get("stream_guard", {})
+    return StreamScanResponse(
+        decision="allow",
+        text=text,
+        labels=found.get("labels", []),
+        risk=found.get("risk", 0.0),
+        pipeline_ms=_ms(start),
+    )

openscript_sdk/server/routers/sessions.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends, Request
+from sqlalchemy import case, func, select
+
+from openscript_sdk.contracts.server_types import EventType
+from openscript_sdk.events.models import EventRow
+from openscript_sdk.events.store import EventStore
+from openscript_sdk.server.auth import require_api_key
+
+router = APIRouter(prefix="/sessions", tags=["sessions"])
+
+
+def _get_event_store(request: Request) -> EventStore:
+    return request.app.state.event_store  # type: ignore[no-any-return]
+
+
+@router.get("", dependencies=[Depends(require_api_key)])
+async def list_sessions(store: EventStore = Depends(_get_event_store)) -> dict:
+    """List all sessions with aggregate metadata."""
+    async with store._session_factory() as session:
+        stmt = (
+            select(
+                EventRow.session_id,
+                func.count(EventRow.id).label("event_count"),
+                func.max(EventRow.timestamp).label("last_activity"),
+                func.sum(
+                    case(
+                        (EventRow.event_type == EventType.THREAT_DETECTED.value, 1),
+                        else_=0,
+                    )
+                ).label("threat_count"),
+            )
+            .group_by(EventRow.session_id)
+            .order_by(func.max(EventRow.timestamp).desc())
+        )
+        result = await session.execute(stmt)
+        rows = result.all()
+
+    sessions = [
+        {
+            "session_id": row.session_id,
+            "event_count": row.event_count,
+            "threat_count": int(row.threat_count or 0),
+            "last_activity": row.last_activity.isoformat() if row.last_activity else None,
+            "flagged": int(row.threat_count or 0) > 0,
+        }
+        for row in rows
+    ]
+    return {"count": len(sessions), "sessions": sessions}
+
+
+@router.get("/{session_id}/stats", dependencies=[Depends(require_api_key)])
+async def session_stats(
+    session_id: str,
+    store: EventStore = Depends(_get_event_store),
+) -> dict:
+    """Return per-event-type counts and threat info for a session."""
+    events = await store.get_events(session_id=session_id, limit=10_000)
+
+    type_counts: dict[str, int] = {}
+    threat_scores: list[float] = []
+
+    for event in events:
+        type_counts[event.event_type.value] = type_counts.get(event.event_type.value, 0) + 1
+        if event.event_type == EventType.THREAT_DETECTED:
+            score = event.payload.get("score")
+            if isinstance(score, (int, float)):
+                threat_scores.append(float(score))
+
+    return {
+        "session_id": session_id,
+        "event_count": len(events),
+        "event_type_counts": type_counts,
+        "threat_count": len(threat_scores),
+        "max_threat_score": round(max(threat_scores), 4) if threat_scores else 0.0,
+        "flagged": len(threat_scores) > 0,
+        "first_event": events[0].timestamp.isoformat() if events else None,
+        "last_event": events[-1].timestamp.isoformat() if events else None,
+    }

openscript_sdk/server/routers/sse.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+import json
+
+from fastapi import APIRouter, Depends, Query, Request
+from fastapi.responses import StreamingResponse
+
+from openscript_sdk.events.store import EventStore
+from openscript_sdk.server.auth import require_api_key
+
+router = APIRouter(tags=["streaming"])
+
+_POLL_INTERVAL = 0.5  # seconds between EventStore polls
+_KEEPALIVE_INTERVAL = 15  # seconds between SSE keepalive comments
+
+
+def _get_event_store(request: Request) -> EventStore:
+    return request.app.state.event_store  # type: ignore[no-any-return]
+
+
+@router.get(
+    "/events/{session_id}/stream",
+    dependencies=[Depends(require_api_key)],
+    response_class=StreamingResponse,
+)
+async def stream_events(
+    session_id: str,
+    after_sequence: int = Query(default=0, ge=0),
+    store: EventStore = Depends(_get_event_store),
+) -> StreamingResponse:
+    """Server-Sent Events stream for a session.
+
+    Clients pass ?after_sequence=N to resume from a known position.
+    The server polls EventStore every 0.5 s and pushes new events.
+    A keepalive comment is emitted every 15 s to prevent proxy timeouts.
+    """
+
+    async def _generate():
+        cursor = after_sequence
+        keepalive_ticks = 0
+
+        while True:
+            events = await store.get_events(
+                session_id=session_id,
+                after_sequence=cursor,
+                limit=100,
+            )
+
+            for event in events:
+                data = json.dumps(event.model_dump(mode="json"))
+                yield f"data: {data}\n\n"
+                cursor = max(cursor, event.sequence_num)
+
+            keepalive_ticks += 1
+            if keepalive_ticks * _POLL_INTERVAL >= _KEEPALIVE_INTERVAL:
+                yield ": keepalive\n\n"
+                keepalive_ticks = 0
+
+            await asyncio.sleep(_POLL_INTERVAL)
+
+    return StreamingResponse(
+        _generate(),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no",
+        },
+    )

openscript_sdk/server/routers/threat.py

@@ -0,0 +1,119 @@
+from __future__ import annotations
+
+from fastapi import APIRouter, Depends
+
+from openscript_sdk.contracts.server_types import (
+    PolicyEvalRequest,
+    PolicyEvalResponse,
+    ThreatScoreRequest,
+    ThreatScoreResponse,
+    ToolValidateRequest,
+    ToolValidateResponse,
+)
+from openscript_sdk.interceptors.threat import score_text
+from openscript_sdk.policies.tool_firewall import ToolRule, ToolRules, validate_tool_call
+from openscript_sdk.server.auth import require_api_key
+
+router = APIRouter(prefix="/v1", tags=["security"])
+
+# Policy rules: maps capability flag → list of actions that require it.
+# An action is denied when the agent lacks a required capability.
+_CAPABILITY_RULES: dict[str, list[str]] = {
+    "external_write": ["send_email", "write_file", "post_webhook", "create_record"],
+    "sensitive_access": ["read_database", "read_secrets", "access_pii"],
+    "untrusted_input": [],  # always allowed — flag is informational
+}
+
+# BuildToolRules from _CAPABILITY_RULES so /v1/tools/validate enforces
+# the same capability-gated actions the rest of the sever knows about
+_TOOL_RULES = ToolRules(
+    default_deny=False,
+    rules={
+        action: ToolRule(
+            allowed_roles=["*"],
+            # flag actions that require sensitive_access as needing approval
+            requires_approval=(capability == "sensitive_access"),
+        )
+        for capability, actions in _CAPABILITY_RULES.items()
+        for action in actions
+    },
+)
+
+
+@router.post("/threat/score", dependencies=[Depends(require_api_key)])
+async def threat_score(req: ThreatScoreRequest) -> ThreatScoreResponse:
+    """Score arbitrary text for prompt-injection signals.
+
+    Stateless — no database access. Safe to call on every LLM turn.
+    """
+    text = req.text
+    if req.context.get("include_turn_number") and req.turn_number > 0:
+        # Slightly lower threshold tolerance for later turns (heuristic)
+        pass
+
+    score, signals = score_text(text)
+    flagged = score >= 0.7
+    reason = ""
+    if flagged:
+        top = max(signals, key=lambda k: signals[k]) if signals else "unknown"
+        reason = f"score {score:.2f} — top signal: {top}"
+
+    return ThreatScoreResponse(
+        session_id=req.session_id,
+        score=round(score, 4),
+        signals={k: round(v, 4) for k, v in signals.items()},
+        flagged=flagged,
+        reason=reason,
+    )
+
+
+@router.post("/policy/evaluate", dependencies=[Depends(require_api_key)])
+async def policy_evaluate(req: PolicyEvalRequest) -> PolicyEvalResponse:
+    """Evaluate an action against capability-based policy rules.
+
+    Returns ALLOW or DENY with the list of policies evaluated.
+    """
+    evaluated: list[str] = []
+    denied_reasons: list[str] = []
+
+    for capability, restricted_actions in _CAPABILITY_RULES.items():
+        if req.action in restricted_actions:
+            policy_name = f"requires_{capability}"
+            evaluated.append(policy_name)
+            agent_has_cap = req.capabilities.get(capability, False)
+            if not agent_has_cap:
+                denied_reasons.append(f"action '{req.action}' requires capability '{capability}'")
+
+    if denied_reasons:
+        return PolicyEvalResponse(
+            session_id=req.session_id,
+            decision="deny",
+            reason="; ".join(denied_reasons),
+            policies_evaluated=evaluated,
+        )
+
+    return PolicyEvalResponse(
+        session_id=req.session_id,
+        decision="allow",
+        reason="all policies passed",
+        policies_evaluated=evaluated or ["default_allow"],
+    )
+
+
+@router.post("/tools/validate", dependencies=[Depends(require_api_key)])
+async def validate_tool(req: ToolValidateRequest) -> ToolValidateResponse:
+    """Validate a tool call against the server-level firewall rules.
+
+    Uses the same rules engine as ToolFirewallPolicy so SDK and server
+    behavior are always consistent.
+    """
+    result = validate_tool_call(
+        {"name": req.tool_name, "args": req.args},
+        role=req.role,
+        rules=_TOOL_RULES,
+    )
+    return ToolValidateResponse(
+        allowed=result["allowed"],
+        reason=result["reason"],
+        requires_approval=result.get("requires_approval", False),
+    )