Model providers

// npm install jose
import { createRemoteJWKSet, jwtVerify } from "jose";

// The JWKS is served with max-age=3600, and a live key is published 30 days before it signs, so
// refetch for an unknown kid at most once an hour (TLP-TOK-9).
const JWKS = createRemoteJWKSet(new URL("https://app.reflections.ai/.well-known/jwks.json"), {
  cacheMaxAge: 3_600_000,
  cooldownDuration: 3_600_000,
});
// The obligation types of TLP 1.0. Reject a token that carries one you do not handle.
const OBLIGATIONS = new Set([
  "generation_check", "hold_until_release", "draft_watermark", "generation_instructions",
  "disclosure", "attribution", "c2pa_manifest", "no_training", "delete_material", "usage_report",
]);

type Tlp = {
  ver: string;
  decision: string;
  max_generations: number;
  execution: { provider: string; model: string };
  obligations: string[];
};

// reserve(jti, n, max, exp) atomically adds n to the count for jti and returns false
// if that would exceed max. Implement it on your shared datastore.
type Reserve = (jti: string, n: number, max: number, exp: number) => Promise<boolean>;

export async function authorizeJob(
  token: string,
  customerAccount: string, // the acct_ ID stored for the calling customer
  model: string, // the model this job runs
  outputs: number, // outputs this job produces
  reserve: Reserve,
) {
  const { payload } = await jwtVerify<{ tlp: Tlp }>(token, JWKS, {
    issuer: "https://app.reflections.ai",
    audience: customerAccount,
    typ: "tlp-auth+jwt",
    algorithms: ["ES256"],
    clockTolerance: 60,
    maxTokenAge: "15m", // also rejects an iat in the future
    requiredClaims: ["exp", "iat", "jti", "sub"],
  });
  const { tlp } = payload;
  if (tlp.ver.split(".")[0] !== "1") throw new Error("unsupported_version");
  if (!tlp.obligations.every((type) => OBLIGATIONS.has(type))) throw new Error("unknown_obligation");
  if (tlp.decision !== "allow") throw new Error("not_allowed");
  if (tlp.execution.model !== model) throw new Error("model_mismatch");
  if (!(await reserve(payload.jti!, outputs, payload.tlp.max_generations, payload.exp!))) {
    throw new Error("max_generations_reached");
  }
  return payload;
}

from typing import Callable
import time
import jwt  # pip install "pyjwt[crypto]"

JWKS = jwt.PyJWKClient("https://app.reflections.ai/.well-known/jwks.json", lifespan=3600)
refetch_after = 0.0
# The obligation types of TLP 1.0. Reject a token that carries one you do not handle.
OBLIGATIONS = {
    "generation_check", "hold_until_release", "draft_watermark", "generation_instructions",
    "disclosure", "attribution", "c2pa_manifest", "no_training", "delete_material", "usage_report",
}

def signing_key(token: str):
    """The token's key. Refetches the JWKS for an unknown kid at most once an hour (TLP-TOK-9)."""
    global refetch_after
    kid = jwt.get_unverified_header(token).get("kid")
    keys = {k.key_id: k.key for k in JWKS.get_signing_keys()}
    if kid not in keys and time.time() >= refetch_after:
        refetch_after = time.time() + 3600
        keys = {k.key_id: k.key for k in JWKS.get_signing_keys(refresh=True)}
    if kid not in keys:
        raise PermissionError("unknown_kid")
    return keys[kid]

# reserve(jti, n, max, exp) atomically adds n to the count for jti and returns False
# if that would exceed max. Implement it on your shared datastore.
Reserve = Callable[[str, int, int, int], bool]

def authorize_job(token: str, customer_account: str, model: str, outputs: int, reserve: Reserve) -> dict:
    if jwt.get_unverified_header(token).get("typ") != "tlp-auth+jwt":
        raise PermissionError("wrong_type")
    claims = jwt.decode(
        token,
        signing_key(token),
        algorithms=["ES256"],
        audience=customer_account,
        issuer="https://app.reflections.ai",
        leeway=60,
        options={"require": ["exp", "iat", "jti", "sub"]},  # also rejects an iat in the future
    )
    tlp = claims["tlp"]
    if str(tlp["ver"]).split(".")[0] != "1":
        raise PermissionError("unsupported_version")
    if not set(tlp["obligations"]) <= OBLIGATIONS:
        raise PermissionError("unknown_obligation")
    if tlp["decision"] != "allow":
        raise PermissionError("not_allowed")
    if tlp["execution"]["model"] != model:
        raise PermissionError("model_mismatch")
    if not reserve(claims["jti"], outputs, tlp["max_generations"], claims["exp"]):
        raise PermissionError("max_generations_reached")
    return claims