agent-loop-guard-runtime 0.6.0a2__py3-none-any.whl

app/core/security.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from fastapi import HTTPException, Request
+
+
+def bearer_token(request: Request) -> str:
+    auth = request.headers.get("authorization", "")
+    if auth.lower().startswith("bearer "):
+        return auth.split(" ", 1)[1].strip()
+    api_key = request.headers.get("x-api-key") or request.headers.get("api-key")
+    if api_key:
+        return api_key.strip()
+    raise HTTPException(
+        status_code=401,
+        detail={"error": {"type": "agent_loop_guard_auth", "message": "Missing gateway key."}},
+    )
+
+
+def external_session_id(request: Request, protocol: str) -> str | None:
+    explicit = request.headers.get("x-alg-session-id")
+    if explicit:
+        return explicit.strip()
+    if protocol == "anthropic":
+        claude_id = request.headers.get("x-claude-code-session-id")
+        if claude_id:
+            return claude_id.strip()
+    return None
+
+
+def filtered_upstream_headers(headers: dict[str, str]) -> dict[str, str]:
+    hop_by_hop = {
+        "host",
+        "content-length",
+        "connection",
+        "keep-alive",
+        "proxy-authenticate",
+        "proxy-authorization",
+        "te",
+        "trailer",
+        "transfer-encoding",
+        "upgrade",
+    }
+    result: dict[str, str] = {}
+    for key, value in headers.items():
+        lower = key.lower()
+        if lower in hop_by_hop:
+            continue
+        if lower.startswith("x-alg-"):
+            continue
+        if lower in {"authorization", "x-api-key", "api-key"}:
+            continue
+        result[key] = value
+    return result
+

app/core/token_meter.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def _strings(value: Any) -> list[str]:
+    if isinstance(value, str):
+        return [value]
+    if isinstance(value, dict):
+        result: list[str] = []
+        for item in value.values():
+            result.extend(_strings(item))
+        return result
+    if isinstance(value, list):
+        result = []
+        for item in value:
+            result.extend(_strings(item))
+        return result
+    return []
+
+
+def estimate_tokens(value: Any) -> int:
+    chars = sum(len(item) for item in _strings(value))
+    return max(1, (chars + 3) // 4)
+
+
+def usage_from_response(protocol: str, response_json: dict[str, Any] | None) -> dict[str, int | bool] | None:
+    if not response_json:
+        return None
+    usage = response_json.get("usage")
+    if not isinstance(usage, dict):
+        return None
+    if protocol == "anthropic":
+        input_tokens = int(usage.get("input_tokens") or 0)
+        output_tokens = int(usage.get("output_tokens") or 0)
+        return {
+            "input_tokens": input_tokens,
+            "output_tokens": output_tokens,
+            "total_tokens": input_tokens + output_tokens,
+            "estimated": False,
+        }
+    input_tokens = int(usage.get("input_tokens") or usage.get("prompt_tokens") or 0)
+    output_tokens = int(usage.get("output_tokens") or usage.get("completion_tokens") or 0)
+    total = int(usage.get("total_tokens") or input_tokens + output_tokens)
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total,
+        "estimated": False,
+    }
+
+
+def usage_or_estimate(
+    protocol: str, request_body: Any, response_json: dict[str, Any] | None = None
+) -> dict[str, int | bool]:
+    usage = usage_from_response(protocol, response_json)
+    if usage is not None and int(usage["total_tokens"]) > 0:
+        return usage
+    input_tokens = estimate_tokens(request_body)
+    output_tokens = estimate_tokens(response_json or "")
+    return {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": input_tokens + output_tokens,
+        "estimated": True,
+    }
+

app/db/models.py

@@ -0,0 +1,296 @@
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, Integer, String, Text
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
+
+
+def utcnow() -> datetime:
+    return datetime.now(UTC)
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Project(Base):
+    __tablename__ = "projects"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    name: Mapped[str] = mapped_column(String(160), nullable=False)
+    mode: Mapped[str] = mapped_column(String(32), nullable=False, default="shadow")
+    provider: Mapped[str] = mapped_column(String(32), nullable=False, default="mock")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    agents: Mapped[list[Agent]] = relationship(back_populates="project")
+    policies: Mapped[list[Policy]] = relationship(back_populates="project")
+
+
+class Agent(Base):
+    __tablename__ = "agents"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    project_id: Mapped[str] = mapped_column(ForeignKey("projects.id"), nullable=False)
+    name: Mapped[str] = mapped_column(String(160), nullable=False)
+    protocol: Mapped[str] = mapped_column(String(32), nullable=False, default="openai")
+    key_hash: Mapped[str] = mapped_column(String(96), nullable=False, index=True)
+    enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
+    paused: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    project: Mapped[Project] = relationship(back_populates="agents")
+
+
+class Policy(Base):
+    __tablename__ = "policies"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    project_id: Mapped[str] = mapped_column(ForeignKey("projects.id"), nullable=False, index=True)
+    rule_type: Mapped[str] = mapped_column(String(64), nullable=False)
+    threshold: Mapped[int] = mapped_column(Integer, nullable=False)
+    window_size: Mapped[int] = mapped_column(Integer, nullable=False, default=8)
+    action: Mapped[str] = mapped_column(String(32), nullable=False, default="block")
+    enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
+
+    project: Mapped[Project] = relationship(back_populates="policies")
+
+
+class GuardSession(Base):
+    __tablename__ = "sessions"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    project_id: Mapped[str] = mapped_column(ForeignKey("projects.id"), nullable=False, index=True)
+    agent_id: Mapped[str] = mapped_column(ForeignKey("agents.id"), nullable=False, index=True)
+    external_session_id: Mapped[str | None] = mapped_column(String(160), nullable=True, index=True)
+    status: Mapped[str] = mapped_column(String(32), nullable=False, default="active")
+    name: Mapped[str | None] = mapped_column(String(200), nullable=True)
+    note: Mapped[str | None] = mapped_column(Text, nullable=True)
+    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    request_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    input_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    output_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    total_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    flagged_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    blocked_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+
+    requests: Mapped[list[RequestRecord]] = relationship(back_populates="session")
+    events: Mapped[list[Event]] = relationship(back_populates="session")
+
+
+class RequestRecord(Base):
+    __tablename__ = "requests"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    session_id: Mapped[str] = mapped_column(ForeignKey("sessions.id"), nullable=False, index=True)
+    protocol: Mapped[str] = mapped_column(String(32), nullable=False)
+    endpoint: Mapped[str] = mapped_column(String(160), nullable=False)
+    model: Mapped[str | None] = mapped_column(String(200), nullable=True)
+    fingerprint: Mapped[str] = mapped_column(String(96), nullable=False, index=True)
+    error_fingerprint: Mapped[str | None] = mapped_column(String(96), nullable=True, index=True)
+    status: Mapped[int] = mapped_column(Integer, nullable=False)
+    latency_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    input_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    output_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    total_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    estimated_tokens: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
+    decision: Mapped[str] = mapped_column(String(32), nullable=False, default="allow")
+    reason_code: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    request_preview: Mapped[str | None] = mapped_column(Text, nullable=True)
+    response_preview: Mapped[str | None] = mapped_column(Text, nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    session: Mapped[GuardSession] = relationship(back_populates="requests")
+    tool_calls: Mapped[list[ToolCall]] = relationship(back_populates="request")
+
+
+class Event(Base):
+    __tablename__ = "events"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    session_id: Mapped[str] = mapped_column(ForeignKey("sessions.id"), nullable=False, index=True)
+    request_id: Mapped[str | None] = mapped_column(ForeignKey("requests.id"), nullable=True)
+    type: Mapped[str] = mapped_column(String(64), nullable=False)
+    rule_id: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    severity: Mapped[str] = mapped_column(String(32), nullable=False, default="info")
+    reason_code: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    mode: Mapped[str] = mapped_column(String(32), nullable=False, default="shadow")
+    payload_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    session: Mapped[GuardSession] = relationship(back_populates="events")
+
+
+class ToolCall(Base):
+    __tablename__ = "tool_calls"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    request_id: Mapped[str] = mapped_column(ForeignKey("requests.id"), nullable=False, index=True)
+    tool_name: Mapped[str] = mapped_column(String(160), nullable=False)
+    args_hash: Mapped[str] = mapped_column(String(96), nullable=False, index=True)
+    result_status: Mapped[str | None] = mapped_column(String(64), nullable=True)
+
+    request: Mapped[RequestRecord] = relationship(back_populates="tool_calls")
+
+
+class TraceRun(Base):
+    __tablename__ = "trace_runs"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    source_session_id: Mapped[str | None] = mapped_column(
+        ForeignKey("sessions.id"), nullable=True, index=True
+    )
+    project_id: Mapped[str] = mapped_column(String(64), nullable=False, index=True, default="default")
+    task_id: Mapped[str | None] = mapped_column(String(160), nullable=True, index=True)
+    task_fingerprint: Mapped[str | None] = mapped_column(String(96), nullable=True, index=True)
+    agent_name: Mapped[str | None] = mapped_column(String(160), nullable=True)
+    model: Mapped[str | None] = mapped_column(String(200), nullable=True)
+    status: Mapped[str] = mapped_column(String(32), nullable=False, default="running", index=True)
+    failure_tag: Mapped[str | None] = mapped_column(String(80), nullable=True, index=True)
+    input_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    output_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    total_tokens: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    total_cost_micros: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    duration_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    span_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    event_count: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    pinned: Mapped[bool] = mapped_column(Boolean, default=False, nullable=False)
+    attributes_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    spans: Mapped[list[TraceSpan]] = relationship(
+        back_populates="run", cascade="all, delete-orphan"
+    )
+    events: Mapped[list[TraceEvent]] = relationship(
+        back_populates="run", cascade="all, delete-orphan"
+    )
+    artifacts: Mapped[list[TraceArtifact]] = relationship(
+        back_populates="run", cascade="all, delete-orphan"
+    )
+
+
+class TraceSpan(Base):
+    __tablename__ = "trace_spans"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    trace_id: Mapped[str] = mapped_column(ForeignKey("trace_runs.id"), nullable=False, index=True)
+    parent_span_id: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
+    name: Mapped[str] = mapped_column(String(120), nullable=False, index=True)
+    status: Mapped[str] = mapped_column(String(32), nullable=False, default="ok", index=True)
+    start_ns: Mapped[int] = mapped_column(BigInteger, nullable=False)
+    end_ns: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
+    duration_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    attributes_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+
+    run: Mapped[TraceRun] = relationship(back_populates="spans")
+    events: Mapped[list[TraceEvent]] = relationship(
+        back_populates="span", cascade="all, delete-orphan"
+    )
+
+
+class TraceEvent(Base):
+    __tablename__ = "trace_events"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    trace_id: Mapped[str] = mapped_column(ForeignKey("trace_runs.id"), nullable=False, index=True)
+    span_id: Mapped[str | None] = mapped_column(ForeignKey("trace_spans.id"), nullable=True, index=True)
+    name: Mapped[str] = mapped_column(String(120), nullable=False, index=True)
+    severity: Mapped[str] = mapped_column(String(32), nullable=False, default="info")
+    timestamp_ns: Mapped[int] = mapped_column(BigInteger, nullable=False)
+    attributes_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+
+    run: Mapped[TraceRun] = relationship(back_populates="events")
+    span: Mapped[TraceSpan | None] = relationship(back_populates="events")
+
+
+class TraceArtifact(Base):
+    __tablename__ = "trace_artifacts"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    trace_id: Mapped[str] = mapped_column(ForeignKey("trace_runs.id"), nullable=False, index=True)
+    path: Mapped[str] = mapped_column(String(500), nullable=False)
+    mime_type: Mapped[str | None] = mapped_column(String(120), nullable=True)
+    size: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    sha256: Mapped[str | None] = mapped_column(String(96), nullable=True, index=True)
+    attributes_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+    run: Mapped[TraceRun] = relationship(back_populates="artifacts")
+
+
+class MCPServer(Base):
+    __tablename__ = "mcp_servers"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    name: Mapped[str] = mapped_column(String(160), nullable=False)
+    transport: Mapped[str] = mapped_column(String(32), nullable=False, default="stdio")
+    target: Mapped[str] = mapped_column(String(500), nullable=False)
+    fingerprint: Mapped[str] = mapped_column(String(96), nullable=False, index=True)
+    enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
+    first_seen_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+
+class MCPSession(Base):
+    __tablename__ = "mcp_sessions"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    server_id: Mapped[str] = mapped_column(ForeignKey("mcp_servers.id"), nullable=False, index=True)
+    client_name: Mapped[str | None] = mapped_column(String(160), nullable=True)
+    protocol_version: Mapped[str | None] = mapped_column(String(32), nullable=True)
+    mode: Mapped[str] = mapped_column(String(32), nullable=False, default="enforce")
+    started_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+    ended_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+
+
+class MCPToolSnapshot(Base):
+    __tablename__ = "mcp_tool_snapshots"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    session_id: Mapped[str] = mapped_column(ForeignKey("mcp_sessions.id"), nullable=False, index=True)
+    server_id: Mapped[str] = mapped_column(ForeignKey("mcp_servers.id"), nullable=False, index=True)
+    tool_name: Mapped[str] = mapped_column(String(160), nullable=False, index=True)
+    schema_hash: Mapped[str] = mapped_column(String(96), nullable=False)
+    schema_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+    risk_tags_json: Mapped[str] = mapped_column(Text, nullable=False, default="[]")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+
+class MCPDecisionEvent(Base):
+    __tablename__ = "mcp_decision_events"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    session_id: Mapped[str | None] = mapped_column(
+        ForeignKey("mcp_sessions.id"), nullable=True, index=True
+    )
+    server_id: Mapped[str] = mapped_column(ForeignKey("mcp_servers.id"), nullable=False, index=True)
+    trace_id: Mapped[str | None] = mapped_column(String(64), nullable=True, index=True)
+    request_id: Mapped[str | None] = mapped_column(String(160), nullable=True, index=True)
+    tool_name: Mapped[str] = mapped_column(String(160), nullable=False, index=True)
+    argument_hash: Mapped[str] = mapped_column(String(96), nullable=False, index=True)
+    policy_version: Mapped[str] = mapped_column(String(64), nullable=False)
+    action: Mapped[str] = mapped_column(String(32), nullable=False, index=True)
+    reason: Mapped[str] = mapped_column(String(300), nullable=False)
+    rule_id: Mapped[str | None] = mapped_column(String(160), nullable=True)
+    mode: Mapped[str] = mapped_column(String(32), nullable=False, default="enforce")
+    latency_ms: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
+    result_status: Mapped[str | None] = mapped_column(String(64), nullable=True)
+    attributes_json: Mapped[str] = mapped_column(Text, nullable=False, default="{}")
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)
+
+
+class MCPApproval(Base):
+    __tablename__ = "mcp_approvals"
+
+    id: Mapped[str] = mapped_column(String(64), primary_key=True)
+    decision_event_id: Mapped[str] = mapped_column(
+        ForeignKey("mcp_decision_events.id"), nullable=False, index=True
+    )
+    status: Mapped[str] = mapped_column(String(32), nullable=False, default="pending", index=True)
+    scope: Mapped[str] = mapped_column(String(32), nullable=False, default="once")
+    argument_hash: Mapped[str] = mapped_column(String(96), nullable=False)
+    expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False)
+    decided_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
+    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=utcnow)