#!/usr/bin/env python3 """Original LockedIn FDE exercise: reason about model gateway admission. Python 3.11+, standard library. Synthetic identities, routes and token counts. No HTTP server, authentication, model call or vendor policy implementation. One sequential, in-memory exercise window; no distributed counter or persistence. """ from __future__ import annotations from dataclasses import dataclass, replace import argparse import json import unittest @dataclass(frozen=True) class Policy: token_budget: int allowed_routes: frozenset[str] allowed_locations: frozenset[str] @dataclass(frozen=True) class Route: location: str context_limit: int @dataclass(frozen=True) class Request: route: str = "private-model" input_tokens: int = 400 output_limit: int = 600 tenant_hint: str = "north" # Untrusted request-body field: never selects a budget. @dataclass class Reservation: tenant: str reserved: int maximum_usage: int actual_usage: int | None = None POLICIES = { "north": Policy(1500, frozenset({"private-model", "external-model"}), frozenset({"approved-zone"})), "south": Policy(5000, frozenset({"private-model"}), frozenset({"approved-zone"})), } ROUTES = { "private-model": Route("approved-zone", 4096), "external-model": Route("unapproved-zone", 4096), } def whole_number(value: int, minimum: int = 0) -> bool: return type(value) is int and value >= minimum class GatewayExercise: def __init__(self, candidate: str = "reserve-output"): if candidate not in {"reserve-output", "prompt-only"}: raise ValueError("Unknown teaching candidate") self.candidate = candidate self.reservations: dict[str, Reservation] = {} self.blocked_tenants: set[str] = set() def remaining(self, tenant: str) -> int: # Pending/unknown responses retain the whole reservation. Confirmed # usage replaces it only when a valid receipt is supplied below. committed = sum(r.reserved if r.actual_usage is None else r.actual_usage for r in self.reservations.values() if r.tenant == tenant) return POLICIES[tenant].token_budget - committed def admit(self, verified_tenant: str, request: Request) -> dict: """verified_tenant is a supplied trusted fixture, NOT authentication. A real application must derive it from verified server identity. We do not accept request.tenant_hint as identity, even if the name exists. """ policy = POLICIES.get(verified_tenant) if policy is None: return {"decision": "refuse", "reason": "unknown_identity"} if verified_tenant in self.blocked_tenants: return {"decision": "refuse", "reason": "usage_reconciliation"} if not whole_number(request.input_tokens) or not whole_number(request.output_limit, 1): return {"decision": "refuse", "reason": "invalid_token_count"} route = ROUTES.get(request.route) if route is None or request.route not in policy.allowed_routes: return {"decision": "refuse", "reason": "route_not_allowed"} if route.location not in policy.allowed_locations: return {"decision": "refuse", "reason": "location_not_allowed"} maximum = request.input_tokens + request.output_limit if maximum > route.context_limit: return {"decision": "refuse", "reason": "context_limit"} # prompt-only is intentionally defective: pending outputs disappear # from admission accounting. Context validation above still applies. reserved = request.input_tokens if self.candidate == "prompt-only" else maximum remaining = self.remaining(verified_tenant) if reserved > remaining: return {"decision": "refuse", "reason": "token_budget", "remaining": remaining, "needed": reserved} reservation_id = f"R{len(self.reservations) + 1}" self.reservations[reservation_id] = Reservation(verified_tenant, reserved, maximum) return {"decision": "admit", "reservation": reservation_id, "budget_owner": verified_tenant, "reserved": reserved, "maximum_usage": maximum, "remaining": self.remaining(verified_tenant)} def settle(self, verified_tenant: str, reservation_id: str, actual_usage: int | None) -> dict: reservation = self.reservations.get(reservation_id) if reservation is None or reservation.tenant != verified_tenant: raise ValueError("Reservation is not available to this identity") if verified_tenant in self.blocked_tenants: raise ValueError("Manual reconciliation required; this exercise has no release operation") if actual_usage is None: return {"state": "settled" if reservation.actual_usage is not None else "usage_unknown", "remaining": self.remaining(verified_tenant)} if not whole_number(actual_usage): raise ValueError("Invalid usage; reservation retained") if actual_usage > reservation.reserved: self.blocked_tenants.add(verified_tenant) raise ValueError("Usage needs reconciliation; reservation retained") if reservation.actual_usage is not None and reservation.actual_usage != actual_usage: self.blocked_tenants.add(verified_tenant) raise ValueError("A conflicting receipt cannot replace settled usage") reservation.actual_usage = actual_usage return {"state": "settled", "used": actual_usage, "remaining": self.remaining(verified_tenant)} def compare(candidate: str) -> dict: gateway = GatewayExercise(candidate) decisions = [gateway.admit("north", Request()), gateway.admit("north", Request())] potential = sum(r.maximum_usage for r in gateway.reservations.values()) return {"case": "two overlapping requests, before either response", "candidate": candidate, "token_budget": 1500, "decisions": decisions, "maximum_admitted_usage": potential, "within_budget_at_maximum": potential <= 1500} def demo() -> dict: gateway = GatewayExercise() first = gateway.admit("north", Request()) unknown = gateway.settle("north", first["reservation"], None) blocked = gateway.admit("north", Request(tenant_hint="south")) receipt = gateway.settle("north", first["reservation"], 550) smaller = gateway.admit("north", Request(output_limit=500)) return {"case": "identity, missing usage and approved routing", "first": first, "missing_usage": unknown, "forged_tenant_hint": blocked, "confirmed_receipt": receipt, "smaller_next_request": smaller, "unapproved_fallback": gateway.admit("north", Request(route="external-model")), "south_remaining": gateway.remaining("south"), "note": "Supplied observations only. No model or gateway was contacted."} class GatewayTests(unittest.TestCase): def test_overlapping_output_reservations(self): self.assertTrue(compare("reserve-output")["within_budget_at_maximum"]) broken = compare("prompt-only") self.assertFalse(broken["within_budget_at_maximum"]) self.assertEqual(broken["maximum_admitted_usage"], 2000) def test_identity_not_body_hint(self): gateway = GatewayExercise() gateway.admit("north", Request()) self.assertEqual(gateway.admit("north", Request(tenant_hint="south"))["reason"], "token_budget") self.assertEqual(gateway.remaining("south"), 5000) self.assertEqual(gateway.admit("unknown", Request())["reason"], "unknown_identity") def test_route_and_location_before_reservation(self): gateway = GatewayExercise() self.assertEqual(gateway.admit("north", Request(route="other"))["reason"], "route_not_allowed") self.assertEqual(gateway.admit("north", Request(route="external-model"))["reason"], "location_not_allowed") self.assertEqual(gateway.remaining("north"), 1500) def test_invalid_token_counts(self): for field, value in [("input_tokens", -1), ("input_tokens", True), ("input_tokens", 1.5), ("output_limit", 0), ("output_limit", -1), ("output_limit", "600")]: with self.subTest(field=field, value=value): gateway = GatewayExercise() self.assertEqual(gateway.admit("north", replace(Request(), **{field: value}))["reason"], "invalid_token_count") self.assertEqual(gateway.remaining("north"), 1500) def test_context_and_exact_budget_boundary(self): gateway = GatewayExercise() self.assertEqual(gateway.admit("south", Request(input_tokens=4096, output_limit=1))["reason"], "context_limit") self.assertEqual(gateway.admit("north", Request(input_tokens=900, output_limit=600))["remaining"], 0) self.assertEqual(gateway.admit("north", Request(input_tokens=0, output_limit=1))["reason"], "token_budget") def test_missing_usage_retains_reservation(self): gateway = GatewayExercise() reservation = gateway.admit("north", Request())["reservation"] self.assertEqual(gateway.settle("north", reservation, None)["state"], "usage_unknown") self.assertEqual(gateway.remaining("north"), 500) self.assertEqual(gateway.admit("north", Request())["decision"], "refuse") def test_confirmed_usage_and_idempotent_receipt(self): gateway = GatewayExercise() reservation = gateway.admit("north", Request())["reservation"] for _ in range(2): self.assertEqual(gateway.settle("north", reservation, 550)["remaining"], 950) self.assertEqual(gateway.admit("north", Request(output_limit=500))["remaining"], 50) with self.assertRaises(ValueError): gateway.settle("north", reservation, 400) self.assertEqual(gateway.remaining("north"), 50) self.assertEqual(gateway.admit("north", Request(input_tokens=0, output_limit=1))["reason"], "usage_reconciliation") def test_invalid_receipt_does_not_release_capacity(self): gateway = GatewayExercise() reservation = gateway.admit("north", Request())["reservation"] for usage in [-1, True]: with self.assertRaises(ValueError): gateway.settle("north", reservation, usage) with self.assertRaises(ValueError): gateway.settle("south", reservation, 0) self.assertEqual(gateway.remaining("north"), 500) self.assertEqual(gateway.settle("north", reservation, 0)["remaining"], 1500) def test_over_envelope_receipt_stops_new_work(self): gateway = GatewayExercise() reservation = gateway.admit("north", Request())["reservation"] with self.assertRaises(ValueError): gateway.settle("north", reservation, 1001) self.assertEqual(gateway.admit("north", Request(input_tokens=0, output_limit=1))["reason"], "usage_reconciliation") with self.assertRaises(ValueError): gateway.settle("north", reservation, 550) self.assertEqual(gateway.admit("south", Request())["decision"], "admit") def main() -> None: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("command", choices=["demo", "compare", "test"]) parser.add_argument("--candidate", choices=["reserve-output", "prompt-only"], default="reserve-output") args = parser.parse_args() if args.command == "test": result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(GatewayTests)) raise SystemExit(0 if result.wasSuccessful() else 1) print(json.dumps(demo() if args.command == "demo" else compare(args.candidate), indent=2)) if __name__ == "__main__": main()