secondshift 0.1.0b12__py3-none-any.whl

secondshift/gateway/proxy.py

@@ -0,0 +1,369 @@
+"""Model gateway: the only way a run reaches a model provider.
+
+A run authenticates with a short-lived token bound to that run. The gateway
+checks the path, model and token budget, injects the workspace's API key and
+streams the provider's answer back, recording token usage. The run never sees
+the provider key.
+"""
+
+import asyncio
+import json
+import logging
+import weakref
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+from asgiref.sync import sync_to_async
+from django.conf import settings
+from django.http import HttpRequest, HttpResponse, StreamingHttpResponse
+
+from secondshift.config.constants import TOKEN_BUDGET_MESSAGE, ModelFamily
+from secondshift.connections.model_providers import api_root, provider_family
+from secondshift.connections.models import Connection
+from secondshift.connections.outbound import (
+    OutboundUrlError,
+    PinnedTransport,
+    validate_outbound_url,
+)
+from secondshift.gateway.metering import StreamMeter, record_usage, usage_from_json
+from secondshift.gateway.models import RunToken
+from secondshift.gateway.tokens import aauthenticate, request_token
+
+logger = logging.getLogger(__name__)
+
+ANTHROPIC_VERSION = "2023-06-01"
+
+
+@dataclass(frozen=True)
+class FamilyRoutes:
+    # Paths a run may call, relative to the provider's API root. Only the
+    # billed ones are metered.
+    allowed: frozenset[str]
+    metered: frozenset[str]
+    forwarded_headers: tuple[str, ...] = ()
+
+
+OPENAI_PATHS = frozenset({"chat/completions", "responses"})
+ROUTES = {
+    ModelFamily.ANTHROPIC: FamilyRoutes(
+        allowed=frozenset({"messages", "messages/count_tokens"}),
+        metered=frozenset({"messages"}),
+        forwarded_headers=("anthropic-version", "anthropic-beta"),
+    ),
+    ModelFamily.OPENAI: FamilyRoutes(allowed=OPENAI_PATHS, metered=OPENAI_PATHS),
+}
+FORWARDED_RESPONSE_HEADERS = (
+    "request-id",
+    "x-request-id",
+    "retry-after",
+)
+
+
+# One pooled client per event loop: connections to providers are reused across
+# requests, and a client is never shared between loops.
+clients: "weakref.WeakKeyDictionary[asyncio.AbstractEventLoop, httpx.AsyncClient]" = (
+    weakref.WeakKeyDictionary()
+)
+
+
+def upstream_client() -> httpx.AsyncClient:
+    loop = asyncio.get_running_loop()
+    client = clients.get(loop)
+    if client is None:
+        client = httpx.AsyncClient(
+            transport=PinnedTransport(),
+            follow_redirects=False,
+            timeout=httpx.Timeout(
+                settings.OUTBOUND_TIMEOUT_SECONDS,
+                read=settings.GATEWAY_READ_TIMEOUT_SECONDS,
+            ),
+        )
+        clients[loop] = client
+    return client
+
+
+def error_response(
+    family: ModelFamily, status: int, kind: str, message: str
+) -> HttpResponse:
+    """An error in the provider's own shape, so agent SDKs surface it cleanly."""
+    if family == ModelFamily.ANTHROPIC:
+        body: dict[str, Any] = {
+            "type": "error",
+            "error": {"type": kind, "message": message},
+        }
+    else:
+        body = {"error": {"message": message, "type": kind, "code": kind}}
+    return HttpResponse(
+        json.dumps(body), status=status, content_type="application/json"
+    )
+
+
+def upstream_headers(
+    request: HttpRequest, family: ModelFamily, connection: Connection
+) -> dict[str, str]:
+    headers = {
+        "Content-Type": "application/json",
+        "Accept": request.headers.get("accept", "application/json"),
+        "User-Agent": "SecondShift-Gateway/1",
+    }
+    for name in ROUTES[family].forwarded_headers:
+        if value := request.headers.get(name):
+            headers[name] = value
+    if family == ModelFamily.ANTHROPIC:
+        headers.setdefault("anthropic-version", ANTHROPIC_VERSION)
+        headers["x-api-key"] = connection.access_token
+    elif connection.access_token:
+        headers["Authorization"] = f"Bearer {connection.access_token}"
+    return headers
+
+
+def copy_headers(upstream: httpx.Response, response: HttpResponse) -> None:
+    for name in FORWARDED_RESPONSE_HEADERS:
+        if value := upstream.headers.get(name):
+            response[name] = value
+
+
+class GatewayRefusal(Exception):
+    """A request the gateway answers itself, without calling the provider."""
+
+    def __init__(
+        self, family: ModelFamily, status: int, kind: str, message: str
+    ) -> None:
+        super().__init__(message)
+        self.response = error_response(family, status, kind, message)
+
+
+async def forward(request: HttpRequest, family_name: str, path: str) -> HttpResponse:
+    if family_name not in ROUTES:
+        return error_response(
+            ModelFamily.OPENAI, 404, "not_found_error", "Unknown API family."
+        )
+    family = ModelFamily(family_name)
+    try:
+        run_token = await authorize(request, family, path)
+        raw, payload = await read_payload(request, family)
+        model, metered = check_request(run_token, family, path, payload)
+        raw = with_stream_usage(family, path, payload, raw)
+        upstream = await send_upstream(request, run_token, family, path, raw)
+        response = await relay(upstream, run_token, family, model, metered)
+    except GatewayRefusal as refusal:
+        return refusal.response
+    copy_headers(upstream, response)
+    return response
+
+
+async def authorize(request: HttpRequest, family: ModelFamily, path: str) -> RunToken:
+    """The run token, if it may call this provider's endpoint."""
+    run_token = await aauthenticate(request_token(request))
+    if run_token is None:
+        raise GatewayRefusal(
+            family, 401, "authentication_error", "Invalid or expired run token."
+        )
+    if provider_family(run_token.connection) != family:
+        raise GatewayRefusal(
+            family,
+            404,
+            "not_found_error",
+            "This run's model provider does not serve this API.",
+        )
+    if request.method != "POST" or path not in ROUTES[family].allowed:
+        raise GatewayRefusal(
+            family, 404, "not_found_error", "This endpoint is not available."
+        )
+    return run_token
+
+
+async def read_payload(
+    request: HttpRequest, family: ModelFamily
+) -> tuple[bytes, dict[str, Any]]:
+    """The raw body and its JSON, refused when too large or without a model."""
+    limit = settings.GATEWAY_MAX_BODY_BYTES
+    try:
+        declared = int(request.headers.get("content-length") or 0)
+    except ValueError:
+        declared = 0
+    # The ASGI handler has already spooled the body. Read it off the event loop.
+    raw = (
+        b""
+        if declared > limit
+        else await sync_to_async(request.read, thread_sensitive=False)(limit + 1)
+    )
+    if declared > limit or len(raw) > limit:
+        raise GatewayRefusal(
+            family, 413, "request_too_large", "The request body is too large."
+        )
+    try:
+        payload = json.loads(raw)
+    except ValueError as error:
+        raise GatewayRefusal(
+            family, 400, "invalid_request_error", "The body must be JSON."
+        ) from error
+    if not isinstance(payload, dict) or not isinstance(payload.get("model"), str):
+        raise GatewayRefusal(
+            family, 400, "invalid_request_error", "The request must name a model."
+        )
+    return raw, payload
+
+
+def check_request(
+    run_token: RunToken, family: ModelFamily, path: str, payload: dict[str, Any]
+) -> tuple[str, bool]:
+    """The model, and whether the call counts against the token budget."""
+    model = payload["model"]
+    allowed = run_token.connection.allowed_models
+    if allowed and model not in allowed:
+        raise GatewayRefusal(
+            family,
+            403,
+            "permission_error",
+            f"Model {model} is not allowed for this workspace.",
+        )
+    metered = path in ROUTES[family].metered
+    if metered and run_token.used_tokens >= run_token.budget_tokens:
+        refusal = GatewayRefusal(family, 429, "rate_limit_error", TOKEN_BUDGET_MESSAGE)
+        # Provider SDKs retry a 429 unless told not to. Waiting never helps here.
+        refusal.response["x-should-retry"] = "false"
+        raise refusal
+    return model, metered
+
+
+def with_stream_usage(
+    family: ModelFamily, path: str, payload: dict[str, Any], raw: bytes
+) -> bytes:
+    """Without this, streamed OpenAI chat completions never report usage."""
+    if not (
+        family == ModelFamily.OPENAI
+        and path == "chat/completions"
+        and payload.get("stream")
+    ):
+        return raw
+    options = payload.get("stream_options")
+    payload["stream_options"] = {
+        **(options if isinstance(options, dict) else {}),
+        "include_usage": True,
+    }
+    return json.dumps(payload).encode()
+
+
+async def send_upstream(
+    request: HttpRequest,
+    run_token: RunToken,
+    family: ModelFamily,
+    path: str,
+    raw: bytes,
+) -> httpx.Response:
+    connection = run_token.connection
+    root = api_root(connection)
+    try:
+        # Scheme and userinfo rules. The transport re-checks addresses on connect.
+        validate_outbound_url(root, "provider URL", resolve=False)
+    except OutboundUrlError as error:
+        raise GatewayRefusal(family, 502, "api_error", str(error)) from error
+    url = f"{root}/{path}"
+    client = upstream_client()
+    try:
+        return await client.send(
+            client.build_request(
+                "POST",
+                url,
+                content=raw,
+                headers=upstream_headers(request, family, connection),
+            ),
+            stream=True,
+        )
+    except httpx.TimeoutException as error:
+        raise GatewayRefusal(
+            family, 504, "api_error", "The provider timed out."
+        ) from error
+    except httpx.HTTPError as error:
+        logger.warning(
+            "gateway request for run %s failed: %s",
+            run_token.run_id,
+            type(error).__name__,
+        )
+        raise GatewayRefusal(
+            family, 502, "api_error", "The provider could not be reached."
+        ) from error
+
+
+async def relay(
+    upstream: httpx.Response,
+    run_token: RunToken,
+    family: ModelFamily,
+    model: str,
+    metered: bool,
+) -> HttpResponse:
+    """Pass the provider's answer on, metering its token use."""
+    content_type = upstream.headers.get("content-type", "application/json")
+    if content_type.startswith("text/event-stream"):
+        response = StreamingHttpResponse(
+            relay_stream(upstream, run_token, family, model, metered),
+            status=upstream.status_code,
+            content_type=content_type,
+        )
+        response["Cache-Control"] = "no-cache"
+        response["X-Accel-Buffering"] = "no"
+        return response
+    try:
+        content = await upstream.aread()
+    except httpx.HTTPError as error:
+        raise GatewayRefusal(
+            family, 502, "api_error", "The provider response was interrupted."
+        ) from error
+    finally:
+        await upstream.aclose()
+    if metered:
+        try:
+            body = json.loads(content)
+        except ValueError:
+            body = None
+        await record_usage(
+            run_token,
+            family,
+            model,
+            upstream.status_code,
+            usage_from_json(family, body),
+        )
+    return HttpResponse(content, status=upstream.status_code, content_type=content_type)
+
+
+async def relay_stream(
+    upstream: httpx.Response,
+    run_token: RunToken,
+    family: ModelFamily,
+    model: str,
+    metered: bool,
+) -> AsyncIterator[bytes]:
+    meter = StreamMeter(family)
+    try:
+        # aiter_bytes yields each chunk as it arrives, decompressed.
+        async for chunk in upstream.aiter_bytes():
+            meter.feed(chunk)
+            yield chunk
+    except httpx.HTTPError as error:
+        logger.warning(
+            "gateway stream for run %s ended early: %s",
+            run_token.run_id,
+            type(error).__name__,
+        )
+    finally:
+        # Runs on completion and when the client disconnects (Django closes the
+        # generator). Shielded so a second cancellation cannot lose the usage.
+        await asyncio.shield(
+            finish_stream(upstream, run_token, family, model, metered, meter)
+        )
+
+
+async def finish_stream(
+    upstream: httpx.Response,
+    run_token: RunToken,
+    family: ModelFamily,
+    model: str,
+    metered: bool,
+    meter: StreamMeter,
+) -> None:
+    await upstream.aclose()
+    if metered:
+        await record_usage(run_token, family, model, upstream.status_code, meter.usage)

secondshift/gateway/tokens.py

@@ -0,0 +1,68 @@
+"""Per-run gateway credentials. Only a token's SHA-256 is ever stored."""
+
+import hashlib
+import secrets
+from datetime import datetime
+
+from django.http import HttpRequest
+
+from secondshift.connections.models import Connection
+from secondshift.gateway.dao import aactive_token, active_token, create_run_token
+from secondshift.gateway.models import RunToken
+
+TOKEN_PREFIX = "ssrt_"
+
+
+class GatewayUnavailable(Exception):
+    pass
+
+
+def token_digest(token: str) -> str:
+    return hashlib.sha256(token.encode()).hexdigest()
+
+
+def issue_token(
+    run_id: int,
+    workspace_id: int,
+    connection: Connection,
+    expires_at: datetime,
+    budget_tokens: int,
+) -> str:
+    """Create a gateway token for one run. The plaintext is returned once."""
+    if (
+        connection.kind != Connection.Kind.MODEL
+        or connection.workspace_id != workspace_id
+    ):
+        raise GatewayUnavailable("The model provider belongs to another workspace.")
+    token = TOKEN_PREFIX + secrets.token_urlsafe(32)
+    create_run_token(run_id, connection, token_digest(token), expires_at, budget_tokens)
+    return token
+
+
+def same_workspace(run_token: RunToken | None) -> RunToken | None:
+    if run_token is None:
+        return None
+    if run_token.connection.workspace_id != run_token.run.workspace_id:
+        return None
+    return run_token
+
+
+def authenticate(token: str) -> RunToken | None:
+    if not token.startswith(TOKEN_PREFIX):
+        return None
+    return same_workspace(active_token(token_digest(token)))
+
+
+async def aauthenticate(token: str) -> RunToken | None:
+    if not token.startswith(TOKEN_PREFIX):
+        return None
+    return same_workspace(await aactive_token(token_digest(token)))
+
+
+def request_token(request: HttpRequest) -> str:
+    # Anthropic SDKs send x-api-key. OpenAI SDKs send a bearer token.
+    if api_key := request.headers.get("x-api-key"):
+        return api_key.strip()
+    authorization = request.headers.get("authorization", "")
+    scheme, _, value = authorization.partition(" ")
+    return value.strip() if scheme.lower() == "bearer" else ""

secondshift/gateway/urls.py

@@ -0,0 +1,17 @@
+"""Internal gateway routes: model traffic, run-scoped agent tools, and health."""
+
+from django.urls import path
+from health_check.views import HealthCheckView
+
+from secondshift.config.constants import AGENT_BROKER_PATH
+from secondshift.gateway.views import model_gateway
+from secondshift.runs.execution.agent_broker import agent_mcp
+
+urlpatterns = [
+    path(AGENT_BROKER_PATH, agent_mcp),
+    path("gateway/<str:family>/v1/<path:path>", model_gateway),
+    path(
+        "-/health/",
+        HealthCheckView.as_view(checks=["health_check.checks.Database"]),
+    ),
+]

secondshift/gateway/views.py

@@ -0,0 +1,11 @@
+from django.http import HttpRequest, HttpResponse
+from django.views.decorators.csrf import csrf_exempt
+
+from secondshift.gateway.proxy import forward
+
+
+@csrf_exempt
+async def model_gateway(request: HttpRequest, family: str, path: str) -> HttpResponse:
+    # Authenticated by a per-run token, never a browser session, so CSRF does
+    # not apply.
+    return await forward(request, family, path)

secondshift/github/__init__.py

@@ -0,0 +1,6 @@
+"""GitHub access: App and token credentials, the REST API, and git transport.
+
+This package knows GitHub, not SecondShift's models. Callers pass plain
+values such as an installation ID or a repository ID. The Django apps turn
+their own rows into those values.
+"""

secondshift/github/app.py

@@ -0,0 +1,109 @@
+"""Server-owned GitHub App credentials and installation-scoped API access."""
+
+import base64
+import json
+import time
+from dataclasses import dataclass
+
+from cryptography.hazmat.primitives import hashes, serialization
+from cryptography.hazmat.primitives.asymmetric import padding
+from django.conf import settings
+
+from secondshift.config.constants import GITHUB_API_URL, GITHUB_BASE_URL
+from secondshift.github.errors import GitHubError
+from secondshift.github.rest import paginated, request_json
+
+
+def configured() -> bool:
+    return all(
+        (
+            settings.GITHUB_APP_ID,
+            settings.GITHUB_APP_SLUG,
+            settings.GITHUB_APP_PRIVATE_KEY,
+            settings.GITHUB_APP_CLIENT_ID,
+            settings.GITHUB_APP_CLIENT_SECRET,
+        )
+    )
+
+
+def encoded(value: bytes) -> str:
+    return base64.urlsafe_b64encode(value).rstrip(b"=").decode()
+
+
+def app_token() -> str:
+    now = int(time.time())
+    header = encoded(b'{"alg":"RS256","typ":"JWT"}')
+    payload = encoded(
+        json.dumps(
+            {"iat": now - 60, "exp": now + 540, "iss": settings.GITHUB_APP_ID}
+        ).encode()
+    )
+    message = f"{header}.{payload}"
+    try:
+        key = serialization.load_pem_private_key(
+            settings.GITHUB_APP_PRIVATE_KEY.encode(), password=None
+        )
+        signature = key.sign(message.encode(), padding.PKCS1v15(), hashes.SHA256())
+    except (ValueError, TypeError, AttributeError) as error:
+        raise GitHubError("The server's GitHub App key is invalid.") from error
+    return f"{message}.{encoded(signature)}"
+
+
+def installation_token(
+    installation_id: str,
+    repository_id: int | None = None,
+    permissions: dict[str, str] | None = None,
+) -> str:
+    """A token that expires within an hour. Read-only unless permissions say more."""
+    if not installation_id.isdecimal():
+        raise GitHubError(
+            "The saved GitHub App installation is not valid. Reconnect the GitHub App."
+        )
+    payload: dict[str, object] = {"permissions": permissions or {"contents": "read"}}
+    if repository_id is not None:
+        payload["repository_ids"] = [repository_id]
+    data = request_json(
+        "POST",
+        f"{GITHUB_API_URL}/app/installations/{installation_id}/access_tokens",
+        app_token(),
+        json=payload,
+    )
+    return data["token"]
+
+
+def installation_repositories(installation_id: str) -> list[dict]:
+    return paginated(
+        f"{GITHUB_API_URL}/installation/repositories",
+        installation_token(installation_id),
+        "repositories",
+        "Installation has too many repositories to list.",
+    )
+
+
+@dataclass(frozen=True)
+class InstallationDetails:
+    """What the installation's owner granted, and where they change it."""
+
+    permissions: dict[str, str]
+    settings_url: str
+
+
+def installation(installation_id: str) -> InstallationDetails:
+    if not installation_id.isdecimal():
+        raise GitHubError(
+            "The saved GitHub App installation is not valid. Reconnect the GitHub App."
+        )
+    data = request_json(
+        "GET", f"{GITHUB_API_URL}/app/installations/{installation_id}", app_token()
+    )
+    permissions = data.get("permissions")
+    url = data.get("html_url")
+    return InstallationDetails(
+        permissions if isinstance(permissions, dict) else {},
+        url if isinstance(url, str) and url.startswith(GITHUB_BASE_URL) else "",
+    )
+
+
+def installation_permissions(installation_id: str) -> dict[str, str]:
+    """What the installation's owner granted the App, such as {"actions": "read"}."""
+    return installation(installation_id).permissions

secondshift/github/authors.py

@@ -0,0 +1,56 @@
+"""Who SecondShift's commits say wrote them.
+
+A commit names the same GitHub account that pushes it and opens the pull
+request. GitHub links a commit to an account by its noreply address, which
+holds the account's ID and login. An address SecondShift made up could link
+the commit to a stranger.
+
+The name is the login, which only changes when the account is renamed. A
+retry then builds the same commit as the attempt before it.
+"""
+
+from dataclasses import dataclass
+from urllib.parse import quote
+
+from secondshift.config.constants import GITHUB_API_URL
+from secondshift.github import app
+from secondshift.github.errors import GitHubError
+from secondshift.github.rest import request_json
+
+NOREPLY_DOMAIN = "users.noreply.github.com"
+BOT_SUFFIX = "[bot]"
+
+
+@dataclass(frozen=True)
+class GitAuthor:
+    name: str
+    email: str
+
+
+def noreply_email(account_id: int, login: str) -> str:
+    return f"{account_id}+{login}@{NOREPLY_DOMAIN}"
+
+
+def account_author(data: dict) -> GitAuthor:
+    account_id = data.get("id")
+    login = data.get("login")
+    if not isinstance(account_id, int) or not isinstance(login, str) or not login:
+        raise GitHubError("GitHub did not say which account publishes the changes.")
+    return GitAuthor(login, noreply_email(account_id, login))
+
+
+def app_author(token: str) -> GitAuthor:
+    """The App's bot account, such as secondshift-acme[bot].
+
+    GitHub names the App, so renaming it never leaves a stale slug here.
+    """
+    slug = request_json("GET", f"{GITHUB_API_URL}/app", app.app_token()).get("slug")
+    if not isinstance(slug, str) or not slug:
+        raise GitHubError("GitHub did not say which account publishes the changes.")
+    login = quote(f"{slug}{BOT_SUFFIX}")
+    return account_author(request_json("GET", f"{GITHUB_API_URL}/users/{login}", token))
+
+
+def token_author(token: str) -> GitAuthor:
+    """The person who owns a personal access token."""
+    return account_author(request_json("GET", f"{GITHUB_API_URL}/user", token))