From eb26efd94036522c101de51898e511c508e5bda7 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 22:05:57 +0200 Subject: [PATCH 01/16] wip(budget): checkpoint governed inference broker candidate Local integration checkpoint only. Rust and real Kind qualification remain pending; no publication or enforcement approval is implied. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .github/workflows/ci.yml | 33 +- .gitignore | 2 + cli/src/cli.ts | 4 +- cli/src/commands/budget.test.ts | 59 ++ cli/src/commands/budget.ts | 157 ++++ controller/Cargo.toml | 4 + controller/src/config_hash.rs | 5 + controller/src/crd.rs | 5 + controller/src/crd_validations.rs | 12 +- controller/src/inference_budget/account.rs | 42 + controller/src/inference_budget/admission.rs | 174 ++++ controller/src/inference_budget/auth.rs | 580 ++++++++++++++ controller/src/inference_budget/auth_tests.rs | 101 +++ controller/src/inference_budget/binding.rs | 509 ++++++++++++ controller/src/inference_budget/claim.rs | 111 +++ controller/src/inference_budget/config.rs | 160 ++++ controller/src/inference_budget/mod.rs | 41 + controller/src/inference_budget/pod.rs | 348 ++++++++ controller/src/inference_budget/pod_tests.rs | 128 +++ controller/src/inference_budget/recovery.rs | 198 +++++ controller/src/inference_budget/scope.rs | 60 ++ controller/src/inference_budget/service.rs | 258 ++++++ controller/src/inference_budget/store.rs | 310 ++++++++ .../src/inference_budget/store_tests.rs | 475 +++++++++++ controller/src/inference_budget/team.rs | 80 ++ controller/src/inference_budget/transport.rs | 98 +++ controller/src/kars_profile.rs | 1 + controller/src/kars_receipt_launch.rs | 1 + controller/src/kars_task.rs | 12 +- .../src/kars_task_authorization_tests.rs | 2 + controller/src/kars_task_execution.rs | 10 + controller/src/kars_task_reconciler.rs | 36 +- controller/src/kars_task_tests.rs | 8 + controller/src/kars_team.rs | 5 + controller/src/kars_team_reconciler.rs | 25 +- controller/src/kars_team_reconciler/specs.rs | 7 +- controller/src/kars_team_reconciler/tasks.rs | 4 +- controller/src/kars_team_reconciler/tests.rs | 2 + controller/src/main.rs | 4 + controller/src/providers/signing.rs | 57 +- controller/src/reconciler/mod.rs | 58 +- .../files/inference-budget-admission.json | 156 ++++ .../kars/templates/controller-deployment.yaml | 12 + .../kars/templates/crd-karsbudgetaccount.yaml | 115 +++ .../helm/kars/templates/crd-karsprofile.yaml | 4 + deploy/helm/kars/templates/crd-karstask.yaml | 21 +- deploy/helm/kars/templates/crd-karsteam.yaml | 22 + deploy/helm/kars/templates/crd.yaml | 4 + .../templates/inference-budget-admission.yaml | 24 + .../helm/kars/templates/inference-budget.yaml | 112 +++ .../operator-default-deny-networkpolicy.yaml | 12 + .../kars/tests/src/inference-budget.test.ts | 124 +++ deploy/helm/kars/values.yaml | 13 + .../governed-inference-budgets-2026-09-08.md | 56 ++ docs/governed-inference-budgets.md | 207 +++++ inference-router/src/blocklist.rs | 23 + inference-router/src/failover.rs | 1 + inference-router/src/forward_proxy.rs | 11 + .../src/inference_budget/client.rs | 322 ++++++++ .../src/inference_budget/client_tests.rs | 331 ++++++++ .../src/inference_budget/dispatch.rs | 115 +++ .../src/inference_budget/egress.rs | 43 + inference-router/src/inference_budget/mod.rs | 11 + .../src/inference_budget/readiness.rs | 24 + .../src/inference_budget/response.rs | 45 ++ .../src/inference_budget/usage.rs | 278 +++++++ inference-router/src/lib.rs | 3 + inference-router/src/providers/signing.rs | 6 + inference-router/src/proxy.rs | 38 +- .../src/routes/anthropic_messages.rs | 9 + .../src/routes/chat_completions.rs | 15 + inference-router/src/routes/handoff/mod.rs | 8 +- inference-router/src/routes/inference.rs | 20 + inference-router/src/routes/mod.rs | 29 + inference-router/src/routes/model_routing.rs | 1 + inference-router/src/routes/spawn_policy.rs | 6 + inference-router/src/spawn/mod.rs | 5 + .../tests/agt_governance_integration.rs | 1 + .../tests/anthropic_buffered_guardrail.rs | 1 + .../tests/chat_output_guardrail_nonjson.rs | 1 + .../tests/common/governed_services.rs | 1 + .../tests/egress_blocked_endpoint.rs | 1 + inference-router/tests/failover_walk.rs | 3 + inference-router/tests/foundry_route_guard.rs | 1 + .../tests/multi_provider_guardrails.rs | 2 + .../tests/policy_status_endpoint.rs | 1 + inference-router/tests/proxy_fake_upstream.rs | 3 + shared/inference_budget/catalog.rs | 270 +++++++ shared/inference_budget/ledger.rs | 748 ++++++++++++++++++ shared/inference_budget/ledger_lifecycle.rs | 113 +++ shared/inference_budget/ledger_tests.rs | 455 +++++++++++ shared/inference_budget/ledger_validation.rs | 156 ++++ shared/inference_budget/mod.rs | 12 + shared/inference_budget/tariff_tests.rs | 202 +++++ shared/inference_budget/tariffs.rs | 548 +++++++++++++ shared/inference_budget/types.rs | 361 +++++++++ tests/e2e/inference-budget-api.mjs | 176 +++++ tests/e2e/inference-budget-enforcement.mjs | 255 ++++++ tests/e2e/run.sh | 5 + 99 files changed, 9665 insertions(+), 63 deletions(-) create mode 100644 cli/src/commands/budget.test.ts create mode 100644 cli/src/commands/budget.ts create mode 100644 controller/src/inference_budget/account.rs create mode 100644 controller/src/inference_budget/admission.rs create mode 100644 controller/src/inference_budget/auth.rs create mode 100644 controller/src/inference_budget/auth_tests.rs create mode 100644 controller/src/inference_budget/binding.rs create mode 100644 controller/src/inference_budget/claim.rs create mode 100644 controller/src/inference_budget/config.rs create mode 100644 controller/src/inference_budget/mod.rs create mode 100644 controller/src/inference_budget/pod.rs create mode 100644 controller/src/inference_budget/pod_tests.rs create mode 100644 controller/src/inference_budget/recovery.rs create mode 100644 controller/src/inference_budget/scope.rs create mode 100644 controller/src/inference_budget/service.rs create mode 100644 controller/src/inference_budget/store.rs create mode 100644 controller/src/inference_budget/store_tests.rs create mode 100644 controller/src/inference_budget/team.rs create mode 100644 controller/src/inference_budget/transport.rs create mode 100644 deploy/helm/kars/files/inference-budget-admission.json create mode 100644 deploy/helm/kars/templates/crd-karsbudgetaccount.yaml create mode 100644 deploy/helm/kars/templates/inference-budget-admission.yaml create mode 100644 deploy/helm/kars/templates/inference-budget.yaml create mode 100644 deploy/helm/kars/tests/src/inference-budget.test.ts create mode 100644 docs/audits/governed-inference-budgets-2026-09-08.md create mode 100644 docs/governed-inference-budgets.md create mode 100644 inference-router/src/inference_budget/client.rs create mode 100644 inference-router/src/inference_budget/client_tests.rs create mode 100644 inference-router/src/inference_budget/dispatch.rs create mode 100644 inference-router/src/inference_budget/egress.rs create mode 100644 inference-router/src/inference_budget/mod.rs create mode 100644 inference-router/src/inference_budget/readiness.rs create mode 100644 inference-router/src/inference_budget/response.rs create mode 100644 inference-router/src/inference_budget/usage.rs create mode 100644 shared/inference_budget/catalog.rs create mode 100644 shared/inference_budget/ledger.rs create mode 100644 shared/inference_budget/ledger_lifecycle.rs create mode 100644 shared/inference_budget/ledger_tests.rs create mode 100644 shared/inference_budget/ledger_validation.rs create mode 100644 shared/inference_budget/mod.rs create mode 100644 shared/inference_budget/tariff_tests.rs create mode 100644 shared/inference_budget/tariffs.rs create mode 100644 shared/inference_budget/types.rs create mode 100644 tests/e2e/inference-budget-api.mjs create mode 100644 tests/e2e/inference-budget-enforcement.mjs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 716b086df..301ffc954 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -427,6 +427,37 @@ jobs: --config "$GITHUB_WORKSPACE/cli/vitest.config.ts" \ --root deploy/helm/kars/tests + inference-budget-api: + name: Governed inference API and CEL (Kind v1.31) + needs: changes + if: needs.changes.outputs.code == 'true' + runs-on: ubuntu-latest + timeout-minutes: 12 + permissions: + contents: read + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: "22" + - uses: azure/setup-helm@9bc31f4ebc9c6b171d7bfbaa5d006ae7abdb4310 # v5.0.1 + - uses: azure/setup-kubectl@829323503d1be3d00ca8346e5391ca0b07a9ab0d # v4 + with: + version: v1.30.5 + - uses: helm/kind-action@ef37e7f390d99f746eb8b610417061a60e82a6cc # v1.14.0 + with: + install_only: true + version: v0.24.0 + - name: Restore existing YAML and test dependencies + run: npm ci --prefix cli + - name: Create disposable supported apiserver before any Rust image build + run: kind create cluster --name kars-budget-api --image kindest/node:v1.31.0 --config tests/e2e/kind-config.yaml + - name: Validate real CRD/CEL and Pod audience identity + run: node tests/e2e/inference-budget-api.mjs + - name: Remove only the disposable API fixture cluster + if: always() + run: kind delete cluster --name kars-budget-api + security-scan: name: Security Scan runs-on: ubuntu-latest @@ -572,7 +603,7 @@ jobs: # Fetch enough history to diff. git fetch --no-tags --depth=50 origin "$base" "$head" 2>/dev/null || true if git diff --name-only "$base" "$head" 2>/dev/null \ - | grep -E '^(controller/|inference-router/|a2a-gateway/|kars-a2a-core/|deploy/helm/|sandbox-images/|tests/e2e/|Cargo\.toml|Cargo\.lock|Makefile)' >/dev/null; then + | grep -E '^(controller/|inference-router/|shared/|a2a-gateway/|kars-a2a-core/|deploy/helm/|sandbox-images/|tests/e2e/|Cargo\.toml|Cargo\.lock|Makefile)' >/dev/null; then echo "run=true" >> "$GITHUB_OUTPUT" else echo "run=false" >> "$GITHUB_OUTPUT" diff --git a/.gitignore b/.gitignore index 4f6ded393..f990cf503 100644 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,8 @@ *.userosscache *.sln.docstates *.env +.budget-kind-*/ +**/.budget-client-test-*/ # User-specific files (MonoDevelop/Xamarin Studio) *.userprefs diff --git a/cli/src/cli.ts b/cli/src/cli.ts index e724fdcb4..d5097a325 100644 --- a/cli/src/cli.ts +++ b/cli/src/cli.ts @@ -11,6 +11,7 @@ import { upgradeCommand } from "./commands/upgrade.js"; import { devCommand } from "./commands/dev.js"; import { addCommand } from "./commands/add.js"; import { credentialsCommand } from "./commands/credentials.js"; +import { budgetCommand } from "./commands/budget.js"; import { namespaceCommand } from "./commands/namespace.js"; import { configCommand } from "./commands/config.js"; import { connectCommand } from "./commands/connect.js"; @@ -73,6 +74,7 @@ export function createCli(): Command { // Configuration program.addCommand(credentialsCommand()); program.addCommand(namespaceCommand()); + program.addCommand(budgetCommand()); program.addCommand(configCommand()); program.addCommand(modelCommand()); program.addCommand(policyCommand()); @@ -116,7 +118,7 @@ export function createCli(): Command { Command groups: Lifecycle up, dev, add, push, destroy Operations connect, status, list, logs, inspect - Configuration credentials, model, policy, egress, config + Configuration credentials, budget, model, policy, egress, config Observability trace, eval, operator, audit, headlamp Agent mobility handoff, mesh, pair Interop convert, a2a, a2a-agent, migrate diff --git a/cli/src/commands/budget.test.ts b/cli/src/commands/budget.test.ts new file mode 100644 index 000000000..16133d9a5 --- /dev/null +++ b/cli/src/commands/budget.test.ts @@ -0,0 +1,59 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { describe, expect, it } from "vitest"; +import { budgetStatus, governedPlan } from "./budget.js"; + +const plan = { + apiVersion: "kars.azure.com/v1alpha1", kind: "KarsTask", + metadata: { name: "mission", namespace: "workspace" }, + spec: { objective: "Build", envelope: { tier: 3, authorityCeiling: 3, delegationDepth: 2, + budget: { tokens: 500, usdMicros: 0 } }, execution: { launch: false } }, +}; + +describe("explicit governed inference CLI", () => { + it("creates a scoped plan without mapping aggregate limits into daily budgets", () => { + const result = governedPlan(JSON.stringify(plan), "workspace"); + expect(result).toEqual({ ...plan, spec: { ...plan.spec, + envelope: { ...plan.spec.envelope, budget: { ...plan.spec.envelope.budget, scope: "GovernedInference" } } } }); + expect(JSON.stringify(result)).not.toContain("daily"); + }); + + it.each([-1, 0.5, Number.MAX_SAFE_INTEGER + 1, "10"])("rejects inexact or negative cap %s", (tokens) => { + const bad = structuredClone(plan) as Record; + bad.spec.envelope.budget.tokens = tokens; + expect(() => governedPlan(JSON.stringify(bad), "workspace")).toThrow(/exact nonnegative/); + }); + + it("rejects cross-workspace creation, existing UID imports and fully unbounded plans", () => { + expect(() => governedPlan(JSON.stringify(plan), "other")).toThrow(/workspace/); + expect(() => governedPlan(JSON.stringify({ ...plan, metadata: { ...plan.metadata, uid: "old" } }), "workspace")).toThrow(/new UID/); + const unbounded = structuredClone(plan); + unbounded.spec.envelope.budget.tokens = 0; + expect(() => governedPlan(JSON.stringify(unbounded), "workspace")).toThrow(/positive/); + }); + + it("reports only a pinned account and refuses replacement/missing ledger instead of zeroing", async () => { + const root = { kind: "KarsTask", resource: { namespace: "workspace", name: "mission", uid: "task-uid" }, + workspaceUid: "workspace-uid", clusterUid: "cluster-uid" }; + const account = { namespace: "accounting", name: "root-account", uid: "account-uid" }; + const task = { metadata: { ...plan.metadata, uid: "task-uid" }, status: { + phase: "Ready", inferenceBudget: { root, account, taskUid: "task-uid" }, + } }; + const ledger = { version: "governed-inference/v1", scope: "GovernedInference", phase: "Active", + accountUid: account.uid, limits: { tokens: 500 }, attempts: {}, meters: { + reserved: { tokens: 40, usdMicros: 0 }, settled: { tokens: 30, usdMicros: 0 }, + uncertain: { tokens: 20, usdMicros: 0 }, unpricedAttempts: 2, + } }; + const stored = { metadata: { uid: account.uid }, spec: { scope: "GovernedInference", root }, status: { ledger } }; + const execute = async (args: string[]) => JSON.stringify(args[1] === "karstask" ? task : stored); + const result = await budgetStatus(execute, "task", "mission", "workspace"); + expect(result.meters).toMatchObject({ uncertain: { tokens: "20", usdMicros: "0" } }); + expect(result.priceCoverage).toContain("unknown"); + stored.metadata.uid = "replacement"; + await expect(budgetStatus(execute, "task", "mission", "workspace")).rejects.toThrow(/replaced/); + stored.metadata.uid = account.uid; + delete (stored.status as Record).ledger; + await expect(budgetStatus(execute, "task", "mission", "workspace")).rejects.toThrow(/missing/); + }); +}); diff --git a/cli/src/commands/budget.ts b/cli/src/commands/budget.ts new file mode 100644 index 000000000..80abd1517 --- /dev/null +++ b/cli/src/commands/budget.ts @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { Command } from "commander"; +import { readFile } from "node:fs/promises"; +import { parseAllDocuments } from "yaml"; + +type ObjectValue = Record; +type Execute = (args: string[], input?: string) => Promise; +const scope = "GovernedInference"; + +function object(value: unknown, message: string): ObjectValue { + if (!value || typeof value !== "object" || Array.isArray(value)) throw new Error(message); + return value as ObjectValue; +} + +function name(value: unknown): string { + if (typeof value !== "string" || !/^[a-z0-9](?:[a-z0-9.-]{0,251}[a-z0-9])?$/.test(value)) { + throw new Error("A valid Kubernetes resource name is required"); + } + return value; +} + +function amount(value: unknown): number { + if (typeof value !== "number" || !Number.isSafeInteger(value) || value < 0) { + throw new Error("Budget amounts must be exact nonnegative safe integers; no floating prices"); + } + return value; +} + +/** Explicit opt-in on CREATE only: never silently convert a running legacy UID. */ +export function governedPlan(text: string, namespace: string): ObjectValue { + const documents = parseAllDocuments(text); + if (documents.length !== 1 || documents[0].errors.length) throw new Error("Expected one valid Task or Team manifest"); + const plan = object(documents[0].toJSON(), "Expected a Task or Team manifest"); + if (plan.apiVersion !== "kars.azure.com/v1alpha1" || !["KarsTask", "KarsTeam"].includes(String(plan.kind))) { + throw new Error("Governed inference plans must be KarsTask or KarsTeam v1alpha1"); + } + const metadata = object(plan.metadata, "Manifest metadata is required"); + name(metadata.name); + if (metadata.uid || metadata.resourceVersion || plan.status) throw new Error("Create a new UID; existing state cannot be imported or reset"); + if (metadata.namespace !== undefined && metadata.namespace !== namespace) { + throw new Error("Manifest workspace differs from --namespace"); + } + metadata.namespace = name(namespace); + const spec = object(plan.spec, "Manifest spec is required"); + const envelope = object(spec.envelope, "Manifest envelope is required"); + const budget = object(envelope.budget, "Set envelope.budget.tokens and/or usdMicros explicitly"); + if (budget.scope !== undefined && budget.scope !== scope) throw new Error("Unsupported budget scope"); + const tokens = amount(budget.tokens ?? 0); + const usdMicros = amount(budget.usdMicros ?? 0); + if (tokens === 0 && usdMicros === 0) throw new Error("At least one positive governed-inference cap is required; zero means unbounded"); + budget.scope = scope; + return plan; +} + +function resource(kind: string): string { + if (kind === "task") return "karstask"; + if (kind === "team") return "karsteam"; + throw new Error("--kind must be task or team"); +} + +function readObject(text: string): ObjectValue { + try { return object(JSON.parse(text), "Invalid API object"); } + catch { throw new Error("Kubernetes returned an invalid budget object"); } +} + +function metadata(value: ObjectValue): ObjectValue { return object(value.metadata, "Object identity is missing"); } + +export async function budgetStatus( + execute: Execute, kind: string, target: string, namespace: string, +): Promise { + const owner = readObject(await execute(["get", resource(kind), name(target), "-n", name(namespace), "-o", "json"])); + const ownerMeta = metadata(owner); + const status = object(owner.status, "Task/Team has no controller status yet"); + const binding = kind === "task" ? object(status.inferenceBudget, "Task has no governed-inference account binding") : undefined; + const reference = object(binding?.account ?? status.inferenceBudgetAccount, "No lifetime account has been bound"); + if (kind === "task" && binding?.taskUid !== ownerMeta.uid) throw new Error("Task UID binding is stale"); + const account = readObject(await execute(["get", "karsbudgetaccount", name(reference.name), + "-n", name(reference.namespace), "-o", "json"])); + if (metadata(account).uid !== reference.uid) throw new Error("Budget account was replaced; refusing to report a fresh zero balance"); + const spec = object(account.spec, "Account spec is missing"); + const root = object(spec.root, "Account root is missing"); + const identity = object(root.resource, "Account root identity is missing"); + if (identity.namespace !== namespace || spec.scope !== scope) throw new Error("Budget scope/workspace mismatch"); + if (kind === "team" && (root.kind !== "KarsTeam" || identity.uid !== ownerMeta.uid || identity.name !== ownerMeta.name)) { + throw new Error("Team lifetime UID binding is stale"); + } + if (binding && JSON.stringify(binding.root) !== JSON.stringify(root)) { + // JSON object ordering is not authority; compare exact identity fields. + const pinned = object(binding.root, "Task root binding is missing"); + const pinnedIdentity = object(pinned.resource, "Task root identity is missing"); + if (["kind", "workspaceUid", "clusterUid"].some((key) => pinned[key] !== root[key]) + || ["namespace", "name", "uid"].some((key) => pinnedIdentity[key] !== identity[key])) { + throw new Error("Task/account root UID binding differs"); + } + } + const accountStatus = object(account.status, "Budget account is uninitialized"); + const ledger = object(accountStatus.ledger, "Budget ledger is missing; no zero fallback"); + if (ledger.accountUid !== reference.uid || ledger.scope !== scope || ledger.version !== "governed-inference/v1") { + throw new Error("Ledger version or identity differs"); + } + const totals = object(ledger.meters, "Budget meters are missing"); + const meters: ObjectValue = {}; + for (const group of ["reserved", "settled", "uncertain"]) { + const values = object(totals[group], "Budget meters are incomplete"); + meters[group] = { tokens: String(amount(values.tokens)), usdMicros: String(amount(values.usdMicros)) }; + } + const unpriced = amount(totals.unpricedAttempts); + const pendingUnpriced = Object.values(object(ledger.attempts, "Budget attempt ledger is missing")) + .some((entry) => { + const attempt = object(entry, "Malformed budget attempt"); + return ["Reserved", "InFlight"].includes(String(attempt.phase)) + && object(attempt.quote, "Budget quote is missing").priceCovered === false; + }); + return { + scope, root, account: reference, accountPhase: ledger.phase, + taskOrTeamPhase: status.phase, limits: ledger.limits, meters, + unpricedAttempts: String(unpriced), + priceCoverage: unpriced > 0 || pendingUnpriced ? "incomplete — monetary cost is unknown" : "configured maxima only — not an invoice", + meaning: "Governed inference tokens and configured maximum prices only; compute, tools, storage and invoice costs are excluded", + dispatch: "Every provider send still requires live router identity, full task authorization, valid bounds/tariffs and an atomic broker grant", + }; +} + +async function kubectl(args: string[], input?: string): Promise { + const { execa } = await import("execa"); + try { + return (await execa("kubectl", [...args, "--request-timeout=20s"], { + input, stdio: ["pipe", "pipe", "pipe"], + })).stdout; + } catch { + throw new Error("Kubernetes budget operation failed; check context, API readiness and operator RBAC"); + } +} + +export function budgetCommand(): Command { + const command = new Command("budget") + .description("Governed inference token/configured-maximum-price accounts (not total task spend)"); + command.command("create") + .description("Opt in a NEW Task/Team manifest; controller configuration and launch gates still apply") + .requiredOption("-f, --file ", "One KarsTask or KarsTeam YAML/JSON manifest") + .option("-n, --namespace ", "Task/Team workspace", "kars-system") + .action(async (options: { file: string; namespace: string }) => { + const plan = governedPlan(await readFile(options.file, "utf8"), options.namespace); + await kubectl(["create", "-f", "-", "-o", "name"], JSON.stringify(plan)); + console.log("Governed-inference plan created. This is not an enforcement/readiness assertion; inspect controller status before launch."); + }); + command.command("status ") + .description("Read the pinned durable account; never substitute absent accounting with zero") + .option("--kind ", "Resource owning the binding", "task") + .option("-n, --namespace ", "Task/Team workspace", "kars-system") + .action(async (target: string, options: { kind: string; namespace: string }) => { + console.log(JSON.stringify(await budgetStatus(kubectl, options.kind, target, options.namespace), null, 2)); + }); + return command; +} diff --git a/controller/Cargo.toml b/controller/Cargo.toml index 3a359e94f..75aa226d7 100644 --- a/controller/Cargo.toml +++ b/controller/Cargo.toml @@ -56,6 +56,10 @@ jsonwebtoken.workspace = true # HTTP server (controller metrics endpoint) axum = "0.8" +# Existing locked TLS transport/parser packages; broker identity stays behind +# the signing wrapper and the combined privacy-issuance gate. +tokio-rustls = "0.26" +rustls-pemfile = "2" # Sigstore cosign verification + OCI fetch (S12.b: signed egress allowlist artifacts). # Status-only fetcher; actual NetworkPolicy still derives from inline allowedEndpoints. diff --git a/controller/src/config_hash.rs b/controller/src/config_hash.rs index b7ac69f1e..91e3415c1 100644 --- a/controller/src/config_hash.rs +++ b/controller/src/config_hash.rs @@ -37,6 +37,11 @@ use std::sync::LazyLock; /// Adding/removing entries from this list is itself a config-hash /// change and should be called out in the audit trail. pub const CONFIG_HASH_INPUTS: &[&str] = &[ + "KARS_INFERENCE_BUDGET_ENABLED", + "KARS_INFERENCE_BUDGET_CATALOG", + "KARS_INFERENCE_BUDGET_TLS_SECRET", + "KARS_INFERENCE_BUDGET_ROUTER_DIGEST", + "KARS_INFERENCE_BUDGET_ADDR", "KARS_DISABLE_ENTRA_AUTH", // Multi-provider endpoints (never secrets — matches the // AZURE_OPENAI_API_KEY precedent). diff --git a/controller/src/crd.rs b/controller/src/crd.rs index 36c89d02c..904d969ab 100644 --- a/controller/src/crd.rs +++ b/controller/src/crd.rs @@ -77,6 +77,11 @@ pub struct KarsSandboxSpec { #[serde(default, skip_serializing_if = "Option::is_none")] pub credentials_ref: Option, + /// Controller-generated governed-inference account binding. A task-owned + /// finite Sandbox cannot omit, repoint, or fall back around this binding. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inference_budget_ref: Option, + /// Network policy pub network_policy: Option, diff --git a/controller/src/crd_validations.rs b/controller/src/crd_validations.rs index 67b609d1b..60d8e847c 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -595,8 +595,8 @@ pub fn kars_task_validations() -> Vec { ..ValidationRule::default() }, ValidationRule { - rule: "!has(self.execution) || !self.execution.launch || !has(self.envelope.budget) || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0))".into(), - message: Some("UnsupportedLaunchBudget: total/subtree token and usdMicros ceilings are not enforced; bounded tasks may be planned but cannot launch".into()), + rule: crate::inference_budget::scope::LAUNCH_RULE.into(), + message: Some("UnsupportedLaunchBudget: positive launch budgets require explicit GovernedInference scope and a configured durable broker".into()), reason: Some("FieldValueForbidden".into()), ..ValidationRule::default() }, @@ -638,7 +638,9 @@ pub fn kars_task_validations() -> Vec { /// Panics only if kube-rs ever produces a CRD whose `spec` is missing. #[must_use] pub fn kars_task_crd() -> CustomResourceDefinition { - inject_spec_validations(KarsTask::crd(), kars_task_validations()) + let mut validations = kars_task_validations(); + validations.extend(crate::inference_budget::scope::validations()); + inject_spec_validations(KarsTask::crd(), validations) .expect("kube-rs derive must produce a spec property on KarsTask") } @@ -691,7 +693,9 @@ pub fn kars_team_validations() -> Vec { /// `KarsTeam` CRD — the standing-team / org primitive (design note §11). #[must_use] pub fn kars_team_crd() -> CustomResourceDefinition { - inject_spec_validations(KarsTeam::crd(), kars_team_validations()) + let mut validations = kars_team_validations(); + validations.extend(crate::inference_budget::scope::validations()); + inject_spec_validations(KarsTeam::crd(), validations) .expect("kube-rs derive must produce a spec property on KarsTeam") } diff --git a/controller/src/inference_budget/account.rs b/controller/src/inference_budget/account.rs new file mode 100644 index 000000000..462578755 --- /dev/null +++ b/controller/src/inference_budget/account.rs @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use kube::CustomResource; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +use crate::inference_budget_contract::{BudgetScope, Limits, RootIdentity, ledger::Ledger}; + +#[derive(CustomResource, Clone, Debug, Deserialize, Serialize, JsonSchema)] +#[kube( + group = "kars.azure.com", + version = "v1alpha1", + kind = "KarsBudgetAccount", + namespaced, + status = "KarsBudgetAccountStatus", + shortname = "kbudget", + printcolumn = r#"{"name":"Scope","type":"string","jsonPath":".spec.scope"}"#, + printcolumn = r#"{"name":"Phase","type":"string","jsonPath":".status.ledger.phase"}"# +)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct KarsBudgetAccountSpec { + pub scope: BudgetScope, + /// Root Task UID, or lifetime Team UID. This is never a display-name key. + pub root: RootIdentity, + pub limits: Limits, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct KarsBudgetAccountStatus { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ledger: Option, +} + +pub const BOOTSTRAP: &str = "kars.azure.com/inference-budget-bootstrap"; +pub const MANAGED_BY: &str = "app.kubernetes.io/managed-by"; +pub const OWNER: &str = "kars-inference-budget"; + +pub fn name_for_root(root: &RootIdentity) -> String { + format!("inference-budget-{}", root.resource.uid) +} diff --git a/controller/src/inference_budget/admission.rs b/controller/src/inference_budget/admission.rs new file mode 100644 index 000000000..1ee1ff996 --- /dev/null +++ b/controller/src/inference_budget/admission.rs @@ -0,0 +1,174 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::store::StoreError; +use crate::inference_budget_contract::BudgetError; +use k8s_openapi::api::admissionregistration::v1::{ + ValidatingAdmissionPolicy, ValidatingAdmissionPolicyBinding, +}; +use kube::{Api, Client, ResourceExt}; +use serde_json::{Value, json}; + +const BUNDLE: &str = + include_str!("../../../deploy/helm/kars/files/inference-budget-admission.json"); + +/// Erase only semantically empty/default selectors. Nonempty match conditions, +/// exclusions, namespace selectors, param refs or bypasses are never ignored. +fn normalized(value: &Value) -> Value { + let mut value = prune_empty(value); + for section in ["matchConstraints", "matchResources"] { + for field in ["resourceRules", "excludeResourceRules"] { + if let Some(rules) = value + .get_mut(section) + .and_then(|section| section.get_mut(field)) + .and_then(Value::as_array_mut) + { + for rule in rules { + if let Some(rule) = rule.as_object_mut() { + rule.entry("scope").or_insert(json!("*")); + } + } + } + } + } + value +} + +fn prune_empty(value: &Value) -> Value { + match value { + Value::Object(map) => Value::Object( + map.iter() + .filter_map(|(key, value)| { + let value = prune_empty(value); + if value.is_null() + || value.as_array().is_some_and(Vec::is_empty) + || value.as_object().is_some_and(serde_json::Map::is_empty) + { + None + } else { + Some((key.clone(), value)) + } + }) + .collect(), + ), + Value::Array(values) => Value::Array(values.iter().map(prune_empty).collect()), + value => value.clone(), + } +} + +fn denied() -> StoreError { + BudgetError::Authorization.into() +} + +fn api_error(error: kube::Error) -> StoreError { + StoreError::Api { + stage: "verify budget admission", + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +fn matches_policy(actual: &Value, expected: &Value) -> bool { + normalized(actual) == normalized(expected) +} + +/// A Helm flag or an object name is not an effective admission proof. Compare +/// the exact shared policy specification and an unrestricted Deny binding, +/// including current-generation CEL compilation status, on every issuance. +pub async fn verify(client: &Client, accounting_namespace: &str) -> Result<(), StoreError> { + let policies: Api = Api::all(client.clone()); + let bindings: Api = Api::all(client.clone()); + let bundle: Value = + serde_json::from_str(&BUNDLE.replace("__ACCOUNTING_NAMESPACE__", accounting_namespace)) + .map_err(|_| denied())?; + for expected in bundle + .get("items") + .and_then(Value::as_array) + .ok_or_else(denied)? + { + let name = expected + .get("name") + .and_then(Value::as_str) + .ok_or_else(denied)?; + let policy = policies.get(name).await.map_err(api_error)?; + let policy_value = serde_json::to_value(&policy).map_err(|_| denied())?; + let spec = policy_value.get("spec").ok_or_else(denied)?; + if policy.metadata.deletion_timestamp.is_some() + || policy.uid().is_none() + || policy.metadata.generation.is_none() + || policy_value + .pointer("/status/observedGeneration") + .and_then(Value::as_i64) + != policy.metadata.generation + || !policy_value + .pointer("/status/typeChecking/expressionWarnings") + .is_none_or(|warnings| warnings.as_array().is_some_and(Vec::is_empty)) + || !matches_policy(spec, expected.get("spec").ok_or_else(denied)?) + { + return Err(denied()); + } + let binding = bindings.get(name).await.map_err(api_error)?; + let binding_value = serde_json::to_value(&binding).map_err(|_| denied())?; + if binding.metadata.deletion_timestamp.is_some() + || binding.uid().is_none() + || !matches_policy( + binding_value.get("spec").ok_or_else(denied)?, + &json!({"policyName": name, "validationActions": ["Deny", "Audit"]}), + ) + { + return Err(denied()); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn api_empty_selectors_do_not_change_a_policy() { + let wanted = + json!({"failurePolicy":"Fail","matchConstraints":{"matchPolicy":"Equivalent"}}); + let defaulted = json!({"failurePolicy":"Fail","matchConstraints":{ + "matchPolicy":"Equivalent", "objectSelector":{}, "namespaceSelector":null + }, "matchConditions":[]}); + assert!(matches_policy(&wanted, &defaulted)); + } + + #[test] + fn names_modes_and_restrictive_bindings_are_not_proof() { + let wanted = json!({"policyName":"guard", "validationActions":["Deny","Audit"]}); + for modified in [ + json!({"policyName":"guard","validationActions":["Warn","Audit"]}), + json!({"policyName":"guard","validationActions":["Deny","Audit"], + "matchResources":{"namespaceSelector":{"matchLabels":{"never":"true"}}}}), + json!({"policyName":"other","validationActions":["Deny","Audit"]}), + ] { + assert!(!matches_policy(&wanted, &modified)); + } + let wanted = json!({"failurePolicy":"Fail","validations":[{"expression":"false"}]}); + let mut bypass = wanted.clone(); + bypass["matchConditions"] = json!([{"name":"skip","expression":"false"}]); + assert!(!matches_policy(&wanted, &bypass)); + let mut selector_bypass = wanted.clone(); + selector_bypass["matchConstraints"] = json!({ + "namespaceSelector":{"matchLabels":{"scope":"*"}} + }); + assert!(!matches_policy(&wanted, &selector_bypass)); + } + + #[test] + fn shared_bundle_has_no_duplicate_names_or_cel_namespace_accessors() { + let bundle: Value = serde_json::from_str(BUNDLE).unwrap(); + let mut seen = std::collections::BTreeSet::new(); + for policy in bundle["items"].as_array().unwrap() { + assert!(seen.insert(policy["name"].as_str().unwrap())); + assert_eq!(policy["spec"]["failurePolicy"], "Fail"); + assert!(!policy.to_string().contains(".namespace")); + } + assert_eq!(seen.len(), 8); + } +} diff --git a/controller/src/inference_budget/auth.rs b/controller/src/inference_budget/auth.rs new file mode 100644 index 000000000..0d11e3f37 --- /dev/null +++ b/controller/src/inference_budget/auth.rs @@ -0,0 +1,580 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Router-only broker authentication. Identity comes from a reviewed, +//! audience-bound Pod token and live UID relationships, never agent headers. + +use axum::http::HeaderMap; +use k8s_openapi::api::{ + apps::v1::{Deployment, ReplicaSet}, + authentication::v1::{TokenReview, TokenReviewSpec}, + authorization::v1::SubjectAccessReview, + core::v1::{Namespace, Pod}, +}; +use kube::{Api, Client, ResourceExt, api::PostParams}; +use serde_json::json; + +#[cfg(test)] +#[path = "auth_tests.rs"] +mod tests; + +use super::{ + account::KarsBudgetAccount, + config::{AUDIENCE, PRIVATE_MOUNT, TOKEN_VOLUME}, + store::StoreError, +}; +use crate::{ + crd::KarsSandbox, + inference_budget_contract::{ + BudgetError, ExecutionIdentity, RootKind, RouterBinding, ledger::Ledger, + }, + kars_task::KarsTask, + kars_team::KarsTeam, +}; + +fn api_error(stage: &'static str, error: kube::Error) -> StoreError { + StoreError::Api { + stage, + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +fn denied() -> StoreError { + BudgetError::Authorization.into() +} + +fn bearer(headers: &HeaderMap) -> Result<&str, StoreError> { + headers + .get("authorization") + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.strip_prefix("Bearer ")) + .filter(|value| !value.is_empty() && value.len() <= 16_384) + .ok_or_else(denied) +} + +fn token_extra<'a>( + extra: &'a Option>>, + key: &str, +) -> Option<&'a str> { + let values = extra.as_ref()?.get(key)?; + if values.len() != 1 { + return None; + } + values.first().map(String::as_str) +} + +fn has_router_only_projection(pod: &Pod) -> bool { + let Some(spec) = &pod.spec else { return false }; + let private = spec + .volumes + .iter() + .flatten() + .filter(|volume| { + volume.name == TOKEN_VOLUME + && volume.projected.as_ref().is_some_and(|projected| { + let sources = projected.sources.as_deref().unwrap_or_default(); + sources.len() == 1 + && sources[0] + .service_account_token + .as_ref() + .is_some_and(|token| { + token.audience.as_deref() == Some(AUDIENCE) + && token.path == "token" + && token + .expiration_seconds + .is_some_and(|seconds| (600..=3600).contains(&seconds)) + }) + }) + }) + .count() + == 1; + let duplicate_audience = spec + .volumes + .iter() + .flatten() + .filter(|volume| volume.name != TOKEN_VOLUME) + .any(|volume| { + volume.projected.as_ref().is_some_and(|projected| { + projected.sources.iter().flatten().any(|source| { + source + .service_account_token + .as_ref() + .is_some_and(|token| token.audience.as_deref() == Some(AUDIENCE)) + }) + }) + }); + let router = spec + .containers + .iter() + .filter(|container| container.name == "inference-router") + .filter(|container| { + container.security_context.as_ref().is_some_and(|security| { + security.run_as_user == Some(1001) + && security.allow_privilege_escalation == Some(false) + && security.read_only_root_filesystem == Some(true) + }) + }) + .filter(|container| { + container.volume_mounts.iter().flatten().any(|mount| { + mount.name == TOKEN_VOLUME + && mount.mount_path == PRIVATE_MOUNT + && mount.read_only == Some(true) + }) + }) + .count() + == 1; + let leaked = spec + .containers + .iter() + .filter(|container| container.name != "inference-router") + .any(|container| { + container + .volume_mounts + .iter() + .flatten() + .any(|mount| mount.name == TOKEN_VOLUME) + }) + || spec.init_containers.iter().flatten().any(|container| { + container + .volume_mounts + .iter() + .flatten() + .any(|mount| mount.name == TOKEN_VOLUME) + }) + || spec.ephemeral_containers.iter().flatten().any(|container| { + container + .volume_mounts + .iter() + .flatten() + .any(|mount| mount.name == TOKEN_VOLUME) + }); + private + && router + && !leaked + && !duplicate_audience + && spec.host_pid != Some(true) + && spec.host_network != Some(true) + && spec.host_ipc != Some(true) + && spec.share_process_namespace != Some(true) +} + +pub async fn authenticate( + client: &Client, + headers: &HeaderMap, + account: &KarsBudgetAccount, + identity: &ExecutionIdentity, + dispatch: bool, +) -> Result<(), StoreError> { + identity.validate()?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or_else(denied)?; + if dispatch && ledger.phase != crate::inference_budget_contract::AccountPhase::Active { + return Err(denied()); + } + let reviews: Api = Api::all(client.clone()); + let review = reviews + .create( + &PostParams::default(), + &TokenReview { + spec: TokenReviewSpec { + audiences: Some(vec![AUDIENCE.into()]), + token: Some(bearer(headers)?.into()), + }, + ..Default::default() + }, + ) + .await + .map_err(|error| api_error("verify budget audience identity", error))?; + let status = review.status.ok_or_else(denied)?; + let user = status.user.ok_or_else(denied)?; + let runtime_namespace = format!("kars-{}", identity.sandbox.name); + if status.authenticated != Some(true) + || status.error.is_some() + || status.audiences.as_deref() != Some([AUDIENCE.to_string()].as_slice()) + || user.username.as_deref() + != Some(format!("system:serviceaccount:{runtime_namespace}:sandbox").as_str()) + || token_extra(&user.extra, "authentication.kubernetes.io/pod-uid") + != Some(identity.pod_uid.as_str()) + || token_extra(&user.extra, "authentication.kubernetes.io/pod-name") + != Some(identity.pod_name.as_str()) + { + return Err(denied()); + } + + let namespaces: Api = Api::all(client.clone()); + let workspace = namespaces + .get(&ledger.root.resource.namespace) + .await + .map_err(|error| api_error("verify budget workspace incarnation", error))?; + let cluster = namespaces + .get("kube-system") + .await + .map_err(|error| api_error("verify budget cluster incarnation", error))?; + if workspace.metadata.uid.as_deref() != Some(ledger.root.workspace_uid.as_str()) + || workspace.metadata.deletion_timestamp.is_some() + || cluster.metadata.uid.as_deref() != Some(ledger.root.cluster_uid.as_str()) + { + return Err(denied()); + } + let namespace = namespaces + .get(&runtime_namespace) + .await + .map_err(|e| api_error("verify budget namespace", e))?; + if namespace.metadata.uid.as_deref() != Some(identity.runtime_namespace_uid.as_str()) + || namespace + .labels() + .get("kars.azure.com/inference-budget") + .map(String::as_str) + != Some("v1") + { + return Err(denied()); + } + let sandboxes: Api = Api::namespaced(client.clone(), &identity.sandbox.namespace); + let sandbox = sandboxes + .get(&identity.sandbox.name) + .await + .map_err(|e| api_error("verify budget Sandbox", e))?; + if sandbox.metadata.uid.as_deref() != Some(identity.sandbox.uid.as_str()) + || !crate::reconciler::namespace_ownership::claimed(&namespace, &sandbox) + .map_err(|_| denied())? + || (dispatch && sandbox.metadata.deletion_timestamp.is_some()) + { + return Err(denied()); + } + let pods: Api = Api::namespaced(client.clone(), &runtime_namespace); + let pod = pods + .get(&identity.pod_name) + .await + .map_err(|e| api_error("verify budget Pod", e))?; + if pod.metadata.uid.as_deref() != Some(identity.pod_uid.as_str()) + || pod + .spec + .as_ref() + .and_then(|spec| spec.service_account_name.as_deref()) + != Some("sandbox") + || (dispatch && pod.metadata.deletion_timestamp.is_some()) + || !has_router_only_projection(&pod) + { + return Err(denied()); + } + verify_workload_owner(client, &pod, &sandbox, &runtime_namespace, dispatch).await?; + + // The ordinary agent credential must not be able to mint a replacement + // budget-audience token for its shared Pod service account. + let access: Api = Api::all(client.clone()); + for (resource, subresource, name) in [ + ("serviceaccounts", "token", "sandbox"), + ("pods", "exec", identity.pod_name.as_str()), + ("pods", "attach", identity.pod_name.as_str()), + ] { + let request: SubjectAccessReview = serde_json::from_value(json!({ + "apiVersion": "authorization.k8s.io/v1", "kind": "SubjectAccessReview", + "spec": { + "user": user.username, "groups": user.groups, + "resourceAttributes": {"namespace": runtime_namespace, "verb": "create", + "group": "", "resource": resource, "subresource": subresource, "name": name} + } + })) + .map_err(|_| denied())?; + let result = access + .create(&PostParams::default(), &request) + .await + .map_err(|e| api_error("verify agent cannot access budget token", e))?; + let status = result.status.ok_or_else(denied)?; + if status.allowed + || status + .evaluation_error + .as_ref() + .is_some_and(|error| !error.is_empty()) + { + return Err(denied()); + } + } + + // This is the combined-privacy prerequisite, not an agent-admin-token + // fallback. It performs real shared Secret GET/LIST/WATCH denial checks and + // validates the current v2 registration when present. + let epoch = crate::sre_authority::privacy_epoch(client, &runtime_namespace) + .await + .map_err(|_| denied())?; + let encoded = pod + .spec + .as_ref() + .and_then(|spec| { + spec.containers + .iter() + .find(|container| container.name == "inference-router") + }) + .and_then(|container| container.env.as_ref()) + .and_then(|env| { + env.iter() + .find(|value| value.name == "KARS_INFERENCE_BUDGET_BINDING") + }) + .and_then(|value| value.value.as_deref()) + .ok_or_else(denied)?; + let binding: RouterBinding = serde_json::from_str(encoded).map_err(|_| denied())?; + if binding.privacy_epoch != epoch + || binding.sandbox != identity.sandbox + || binding.runtime_namespace != runtime_namespace + || binding.runtime_namespace_uid != identity.runtime_namespace_uid + || binding.task.task_uid != identity.task_uid + || binding.task.authorization_digest != identity.authorization_digest + || binding.task.account.uid != ledger.account_uid + || binding.task.root != ledger.root + { + return Err(denied()); + } + verify_task_binding(client, account, ledger, identity, &sandbox, dispatch).await +} + +async fn verify_workload_owner( + client: &Client, + pod: &Pod, + sandbox: &KarsSandbox, + namespace: &str, + dispatch: bool, +) -> Result<(), StoreError> { + let refs = pod.metadata.owner_references.as_deref().unwrap_or_default(); + let replica = refs + .iter() + .find(|owner| { + owner.controller == Some(true) + && owner.kind == "ReplicaSet" + && owner.api_version == "apps/v1" + }) + .ok_or_else(denied)?; + let replicas: Api = Api::namespaced(client.clone(), namespace); + let replica_set = replicas + .get(&replica.name) + .await + .map_err(|e| api_error("verify budget ReplicaSet", e))?; + if replica_set.metadata.uid.as_deref() != Some(replica.uid.as_str()) { + return Err(denied()); + } + let refs = replica_set + .metadata + .owner_references + .as_deref() + .unwrap_or_default(); + let owner = refs + .iter() + .find(|owner| { + owner.controller == Some(true) + && owner.kind == "Deployment" + && owner.api_version == "apps/v1" + && owner.name == sandbox.name_any() + }) + .ok_or_else(denied)?; + let deployments: Api = Api::namespaced(client.clone(), namespace); + let deployment = deployments + .get(&owner.name) + .await + .map_err(|e| api_error("verify budget Deployment", e))?; + if deployment.metadata.uid.as_deref() != Some(owner.uid.as_str()) + || deployment.labels().get("kars.azure.com/sandbox") != sandbox.metadata.name.as_ref() + || deployment.labels().get("kars.azure.com/parent-namespace") + != sandbox.metadata.namespace.as_ref() + { + return Err(denied()); + } + let pod_router = pod + .spec + .as_ref() + .and_then(|spec| { + spec.containers + .iter() + .find(|container| container.name == "inference-router") + }) + .ok_or_else(denied)?; + let template = if dispatch { + deployment + .spec + .as_ref() + .and_then(|spec| spec.template.spec.as_ref()) + } else { + replica_set + .spec + .as_ref() + .and_then(|spec| spec.template.as_ref()) + .and_then(|template| template.spec.as_ref()) + }; + let template_router = template + .and_then(|spec| { + spec.containers + .iter() + .find(|container| container.name == "inference-router") + }) + .ok_or_else(denied)?; + if pod_router.image != template_router.image + || pod_router.command != template_router.command + || pod_router.args != template_router.args + || !router_env_matches(pod_router, template_router) + || pod_router.env_from != template_router.env_from + || pod_router.security_context != template_router.security_context + || !router_mounts_match(pod_router, template_router) + { + return Err(denied()); + } + if dispatch { + let settings = super::config::Settings::from_env()?.ok_or_else(denied)?; + if pod_router + .image + .as_ref() + .is_none_or(|image| !image.ends_with(&format!("@{}", settings.router_image_digest))) + { + return Err(denied()); + } + } + Ok(()) +} + +fn router_env_matches( + actual: &k8s_openapi::api::core::v1::Container, + template: &k8s_openapi::api::core::v1::Container, +) -> bool { + let actual = actual.env.as_deref().unwrap_or_default(); + let template = template.env.as_deref().unwrap_or_default(); + let mut names = std::collections::BTreeSet::new(); + if actual.iter().any(|entry| !names.insert(&entry.name)) + || template.iter().any(|entry| !actual.contains(entry)) + { + return false; + } + actual + .iter() + .filter(|entry| !template.contains(entry)) + .all(|entry| { + entry.value_from.is_none() + && entry + .value + .as_ref() + .is_some_and(|value| match entry.name.as_str() { + "AZURE_CLIENT_ID" | "AZURE_TENANT_ID" | "AZURE_AUTHORITY_HOST" => { + !value.is_empty() + } + "AZURE_FEDERATED_TOKEN_FILE" => { + value == "/var/run/secrets/azure/tokens/azure-identity-token" + } + _ => false, + }) + }) +} + +fn router_mounts_match( + actual: &k8s_openapi::api::core::v1::Container, + template: &k8s_openapi::api::core::v1::Container, +) -> bool { + let actual = actual.volume_mounts.as_deref().unwrap_or_default(); + let template = template.volume_mounts.as_deref().unwrap_or_default(); + template.iter().all(|mount| actual.contains(mount)) + && actual + .iter() + .filter(|mount| !template.contains(mount)) + .all(|mount| { + mount.read_only == Some(true) + && mount.sub_path.is_none() + && mount.sub_path_expr.is_none() + && mount.mount_propagation.is_none() + && ((mount.name.starts_with("kube-api-access-") + && mount.mount_path == "/var/run/secrets/kubernetes.io/serviceaccount") + || (mount.name == "azure-identity-token" + && mount.mount_path == "/var/run/secrets/azure/tokens")) + }) +} + +async fn verify_task_binding( + client: &Client, + account: &KarsBudgetAccount, + ledger: &Ledger, + identity: &ExecutionIdentity, + sandbox: &KarsSandbox, + dispatch: bool, +) -> Result<(), StoreError> { + let node = ledger.nodes.get(&identity.task_uid).ok_or_else(denied)?; + let refs = sandbox + .metadata + .owner_references + .as_deref() + .unwrap_or_default(); + if !refs.iter().any(|owner| { + owner.controller == Some(true) + && owner.kind == "KarsTask" + && owner.api_version == "kars.azure.com/v1alpha1" + && owner.uid == identity.task_uid + && owner.name == node.authority.task.name + }) { + return Err(denied()); + } + if !dispatch { + let session = ledger.sessions.get(&identity.pod_uid).ok_or_else(denied)?; + return if session.identity == *identity { + Ok(()) + } else { + Err(denied()) + }; + } + let tasks: Api = Api::namespaced(client.clone(), &ledger.root.resource.namespace); + for uid in ledger.ancestors(&identity.task_uid)? { + let node = &ledger.nodes[&uid]; + let task = tasks + .get(&node.authority.task.name) + .await + .map_err(|e| api_error("verify current task authority", e))?; + let status = task.status.as_ref().ok_or_else(denied)?; + if task.metadata.uid.as_deref() != Some(uid.as_str()) + || !crate::kars_task_reconciler::task_is_ready(&task) + || task.envelope_digest() != node.authority.authorization_digest + || (uid == identity.task_uid + && !task + .spec + .execution + .as_ref() + .is_some_and(|execution| execution.launch)) + { + return Err(denied()); + } + let binding = status.inference_budget.as_ref().ok_or_else(denied)?; + if binding.account.uid != ledger.account_uid + || binding.account.name != account.name_any() + || Some(binding.account.namespace.as_str()) != account.metadata.namespace.as_deref() + || binding.root != ledger.root + || binding.task_uid != uid + || binding.parent_task_uid != node.authority.parent_uid + || binding.root_task_uid != node.authority.root_task_uid + || binding.authorization_digest != node.authority.authorization_digest + { + return Err(denied()); + } + } + if ledger.root.kind == RootKind::KarsTeam { + let teams: Api = Api::namespaced(client.clone(), &ledger.root.resource.namespace); + let team = teams + .get(&ledger.root.resource.name) + .await + .map_err(|e| api_error("verify team lifetime budget", e))?; + if team.metadata.uid.as_deref() != Some(ledger.root.resource.uid.as_str()) + || team.metadata.deletion_timestamp.is_some() + || team.spec.paused + || super::binding::limits(&team.spec.envelope)? != ledger.limits + || team + .status + .as_ref() + .and_then(|status| status.inference_budget_account.as_ref()) + .is_none_or(|reference| { + reference.uid != ledger.account_uid + || reference.name != account.name_any() + || Some(reference.namespace.as_str()) + != account.metadata.namespace.as_deref() + }) + { + return Err(denied()); + } + } + Ok(()) +} diff --git a/controller/src/inference_budget/auth_tests.rs b/controller/src/inference_budget/auth_tests.rs new file mode 100644 index 000000000..2e495ca9b --- /dev/null +++ b/controller/src/inference_budget/auth_tests.rs @@ -0,0 +1,101 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use serde_json::{Value, json}; + +fn pod() -> Value { + json!({ + "apiVersion":"v1", "kind":"Pod", + "metadata":{"name":"task-pod","namespace":"kars-task","uid":"pod-uid"}, + "spec":{ + "serviceAccountName":"sandbox", + "volumes":[{"name":TOKEN_VOLUME,"projected":{"sources":[{"serviceAccountToken":{ + "audience":AUDIENCE,"path":"token","expirationSeconds":600 + }}]}}], + "containers":[ + {"name":"agent","image":"agent:latest"}, + {"name":"inference-router","image":"router:latest", + "securityContext":{"runAsUser":1001,"allowPrivilegeEscalation":false,"readOnlyRootFilesystem":true}, + "volumeMounts":[{"name":TOKEN_VOLUME,"mountPath":PRIVATE_MOUNT,"readOnly":true}]} + ] + } + }) +} + +fn valid(value: Value) -> bool { + has_router_only_projection(&serde_json::from_value::(value).unwrap()) +} + +#[test] +fn only_the_secure_router_may_mount_the_budget_audience_token() { + assert!(valid(pod())); + let mut leaked = pod(); + leaked["spec"]["containers"][0]["volumeMounts"] = + leaked["spec"]["containers"][1]["volumeMounts"].clone(); + assert!(!valid(leaked)); + let mut wrong_user = pod(); + wrong_user["spec"]["containers"][1]["securityContext"]["runAsUser"] = json!(1000); + assert!(!valid(wrong_user)); +} + +#[test] +fn alternate_named_budget_projection_and_shared_process_namespaces_fail_closed() { + let mut alternate = pod(); + let mut volume = alternate["spec"]["volumes"][0].clone(); + volume["name"] = json!("disguised"); + alternate["spec"]["volumes"] + .as_array_mut() + .unwrap() + .push(volume); + alternate["spec"]["containers"][0]["volumeMounts"] = + json!([{"name":"disguised","mountPath":"/agent-token"}]); + assert!(!valid(alternate)); + for field in ["hostPID", "hostNetwork", "hostIPC", "shareProcessNamespace"] { + let mut shared = pod(); + shared["spec"][field] = json!(true); + assert!(!valid(shared)); + } +} + +#[test] +fn standard_admission_injections_are_not_mistaken_for_replaced_router_authority() { + use k8s_openapi::api::core::v1::Container; + let template: Container = + serde_json::from_value(pod()["spec"]["containers"][1].clone()).unwrap(); + let mut actual = template.clone(); + actual.env = Some(serde_json::from_value(json!([ + {"name":"AZURE_CLIENT_ID","value":"operator-client"}, + {"name":"AZURE_TENANT_ID","value":"operator-tenant"}, + {"name":"AZURE_FEDERATED_TOKEN_FILE","value":"/var/run/secrets/azure/tokens/azure-identity-token"} + ])).unwrap()); + actual.volume_mounts.as_mut().unwrap().extend( + serde_json::from_value::>(json!([ + {"name":"kube-api-access-abcde","mountPath":"/var/run/secrets/kubernetes.io/serviceaccount","readOnly":true}, + {"name":"azure-identity-token","mountPath":"/var/run/secrets/azure/tokens","readOnly":true} + ])).unwrap() + ); + assert!(router_env_matches(&actual, &template)); + assert!(router_mounts_match(&actual, &template)); + actual.env.as_mut().unwrap().push( + serde_json::from_value(json!({ + "name":"KARS_INFERENCE_BUDGET_REQUIRED","value":"false" + })) + .unwrap(), + ); + assert!(!router_env_matches(&actual, &template)); + actual.volume_mounts.as_mut().unwrap()[1].mount_path = PRIVATE_MOUNT.into(); + assert!(!router_mounts_match(&actual, &template)); +} + +#[test] +fn audience_or_expiry_substitution_does_not_authorize_a_router() { + let mut wrong = pod(); + wrong["spec"]["volumes"][0]["projected"]["sources"][0]["serviceAccountToken"]["audience"] = + json!("api"); + assert!(!valid(wrong)); + let mut persistent = pod(); + persistent["spec"]["volumes"][0]["projected"]["sources"][0]["serviceAccountToken"]["expirationSeconds"] = + json!(86400); + assert!(!valid(persistent)); +} diff --git a/controller/src/inference_budget/binding.rs b/controller/src/inference_budget/binding.rs new file mode 100644 index 000000000..b3b62c30d --- /dev/null +++ b/controller/src/inference_budget/binding.rs @@ -0,0 +1,509 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use k8s_openapi::api::core::v1::Namespace; +use kube::{ + Api, Client, ResourceExt, + api::{Patch, PatchParams}, +}; +use serde_json::json; +use std::collections::BTreeSet; + +use super::{ + account::{KarsBudgetAccountSpec, name_for_root}, + config::Settings, + store::{Store, StoreError}, +}; +use crate::{ + inference_budget_contract::{ + AccountReference, BudgetError, BudgetScope, Limits, ResourceIdentity, RootIdentity, + RootKind, TaskAuthority, TaskBudgetBinding, + }, + kars_task::{KarsTask, TaskEnvelope}, + kars_team::KarsTeam, +}; + +fn api_error(stage: &'static str, error: kube::Error) -> StoreError { + StoreError::Api { + stage, + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +pub fn limits(envelope: &TaskEnvelope) -> Result { + let Some(budget) = &envelope.budget else { + return Ok(Limits::default()); + }; + let convert = |value: Option| -> Result, StoreError> { + value + .map(|value| { + u64::try_from(value).map_err(|_| StoreError::Ledger(BudgetError::Authorization)) + }) + .transpose() + .map(|value| value.filter(|value| *value > 0)) + }; + let limits = Limits { + tokens: convert(budget.tokens)?, + usd_micros: convert(budget.usd_micros)?, + }; + if limits.finite() && budget.scope != Some(BudgetScope::GovernedInference) { + return Err(BudgetError::Authorization.into()); + } + Ok(limits) +} + +pub fn has_finite(envelope: &TaskEnvelope) -> bool { + envelope.budget.as_ref().is_some_and(|budget| { + budget.tokens.is_some_and(|value| value > 0) + || budget.usd_micros.is_some_and(|value| value > 0) + }) +} + +fn resource(task: &KarsTask) -> Result { + let resource = ResourceIdentity { + namespace: task.namespace().ok_or(BudgetError::Identity)?, + name: task.name_any(), + uid: task.uid().ok_or(BudgetError::Identity)?, + }; + resource.validate()?; + Ok(resource) +} + +async fn chain(client: &Client, task: &KarsTask) -> Result, StoreError> { + let namespace = task.namespace().ok_or(BudgetError::Identity)?; + let api: Api = Api::namespaced(client.clone(), &namespace); + let mut nodes = Vec::new(); + let mut current = task.clone(); + let mut seen = BTreeSet::new(); + loop { + let id = resource(¤t)?; + if !seen.insert(id.uid) + || nodes.len() >= crate::inference_budget_contract::MAX_NODES + || current.metadata.deletion_timestamp.is_some() + { + return Err(BudgetError::Authorization.into()); + } + limits(¤t.spec.envelope)?; + let parent = current + .spec + .parent_ref + .as_ref() + .map(|reference| reference.name.clone()); + nodes.push(current.clone()); + let Some(name) = parent else { break }; + let parent = api + .get(&name) + .await + .map_err(|e| api_error("resolve budget parent UID", e))?; + if current + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + .and_then(|binding| binding.parent_task_uid.as_deref()) + .is_some_and(|uid| parent.metadata.uid.as_deref() != Some(uid)) + || !crate::kars_task::spec_attenuation_violations(¤t.spec, &parent.spec) + .is_empty() + { + return Err(BudgetError::Authorization.into()); + } + current = parent; + } + nodes.reverse(); + Ok(nodes) +} + +pub(super) async fn needs_account(client: &Client, task: &KarsTask) -> Result { + let nodes = chain(client, task).await?; + if nodes.iter().any(|task| { + has_finite(&task.spec.envelope) + || task + .status + .as_ref() + .is_some_and(|status| status.inference_budget.is_some()) + }) { + return Ok(true); + } + let first = nodes.first().ok_or(BudgetError::Identity)?; + Ok(team_for_root(client, first).await?.is_some_and(|team| { + has_finite(&team.spec.envelope) + || team + .status + .as_ref() + .is_some_and(|status| status.inference_budget_account.is_some()) + })) +} + +pub async fn close_task(client: &Client, task: &KarsTask) -> Result<(), StoreError> { + let Some(binding) = task + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + else { + return Ok(()); + }; + if task.metadata.uid.as_deref() != Some(binding.task_uid.as_str()) { + return Err(BudgetError::Identity.into()); + } + let store = Store::new(client.clone(), &binding.account.namespace); + // A crash may occur after the protected UID pin but before bootstrap + // initialization/task registration. Seal that exact pending anchor only; + // initialize rejects missing/replaced/sealed-corrupt prior accounting. + store + .initialize(&binding.root, &binding.account.uid) + .await?; + store + .transact(&binding.root, &binding.account.uid, |ledger| { + if binding.root.kind == RootKind::KarsTask + && binding.root.resource.uid == binding.task_uid + { + ledger.close_account() + } else { + ledger.close_registered_task(&binding.task_uid) + } + }) + .await +} + +pub async fn prepare_task(client: &Client, task: &KarsTask) -> Result { + if !has_finite(&task.spec.envelope) + && task + .status + .as_ref() + .is_none_or(|status| status.inference_budget.is_none()) + { + return Ok(task.clone()); + } + if !needs_account(client, task).await? { + return Ok(task.clone()); + } + let settings = Settings::from_env()?.ok_or(BudgetError::Contract)?; + let binding = ensure_task(client, task, &settings).await?; + let api: Api = Api::namespaced( + client.clone(), + &task.namespace().ok_or(BudgetError::Identity)?, + ); + let live = api + .get(&task.name_any()) + .await + .map_err(|error| api_error("read prepared task", error))?; + if live.metadata.uid != task.metadata.uid + || live.metadata.generation != task.metadata.generation + || live + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + != Some(&binding) + { + return Err(BudgetError::Identity.into()); + } + Ok(live) +} + +fn task_binding( + task: &KarsTask, + parent: Option<&KarsTask>, + root: &RootIdentity, + root_task_uid: &str, + account: &AccountReference, +) -> Result { + Ok(TaskBudgetBinding { + scope: BudgetScope::GovernedInference, + account: account.clone(), + root: root.clone(), + task_uid: task.uid().ok_or(BudgetError::Identity)?, + parent_task_uid: parent.and_then(ResourceExt::uid), + root_task_uid: root_task_uid.into(), + authorization_digest: task.envelope_digest(), + }) +} + +async fn pin_task( + client: &Client, + task: &KarsTask, + binding: &TaskBudgetBinding, +) -> Result { + let api: Api = Api::namespaced( + client.clone(), + &task.namespace().ok_or(BudgetError::Identity)?, + ); + let live = api + .get(&task.name_any()) + .await + .map_err(|e| api_error("read task budget binding", e))?; + if live.metadata.uid != task.metadata.uid + || live.envelope_digest() != task.envelope_digest() + || live.metadata.deletion_timestamp.is_some() + { + return Err(BudgetError::Identity.into()); + } + if let Some(old) = live + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + { + if old.account != binding.account + || old.root != binding.root + || old.task_uid != binding.task_uid + || old.parent_task_uid != binding.parent_task_uid + || old.root_task_uid != binding.root_task_uid + { + return Err(BudgetError::Identity.into()); + } + if old == binding { + return Ok(live); + } + } + let uid = live.uid().ok_or(BudgetError::Identity)?; + let rv = live.resource_version().ok_or(BudgetError::Identity)?; + let updated = api.patch_status(&live.name_any(), &PatchParams::default(), &Patch::Merge(json!({ + "metadata": {"uid": uid, "resourceVersion": rv}, "status": {"inferenceBudget": binding} + }))).await.map_err(|e| api_error("pin task inference account", e))?; + if updated + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + != Some(binding) + { + return Err(StoreError::Missing); + } + Ok(updated) +} + +async fn team_for_root(client: &Client, task: &KarsTask) -> Result, StoreError> { + let owners = task + .metadata + .owner_references + .as_deref() + .unwrap_or_default(); + let Some(owner) = owners.iter().find(|owner| { + owner.controller == Some(true) + && owner.kind == "KarsTeam" + && owner.api_version == "kars.azure.com/v1alpha1" + }) else { + return Ok(None); + }; + let api: Api = Api::namespaced( + client.clone(), + &task.namespace().ok_or(BudgetError::Identity)?, + ); + let team = api + .get(&owner.name) + .await + .map_err(|e| api_error("resolve lifetime Team UID", e))?; + if team.metadata.uid.as_deref() != Some(owner.uid.as_str()) + || team.metadata.deletion_timestamp.is_some() + { + return Err(BudgetError::Identity.into()); + } + Ok(Some(team)) +} + +/// Bind a complete immutable UID ancestry and current full authorization before +/// execution. A missing/corrupt pinned account is never recreated as a balance +/// of zero. This routine does not relax any launch gate by itself. +pub async fn ensure_task( + client: &Client, + task: &KarsTask, + settings: &Settings, +) -> Result { + settings + .catalog(client, chrono::Utc::now().timestamp()) + .await?; + let nodes = chain(client, task).await?; + let first = nodes.first().ok_or(BudgetError::Identity)?; + let first_id = resource(first)?; + let team = team_for_root(client, first).await?; + let namespaces: Api = Api::all(client.clone()); + let workspace = namespaces + .get(&first_id.namespace) + .await + .map_err(|e| api_error("resolve budget workspace UID", e))?; + let cluster = namespaces + .get("kube-system") + .await + .map_err(|e| api_error("resolve budget cluster UID", e))?; + let root = RootIdentity { + kind: if team.is_some() { + RootKind::KarsTeam + } else { + RootKind::KarsTask + }, + resource: match &team { + Some(team) => ResourceIdentity { + namespace: first_id.namespace.clone(), + name: team.name_any(), + uid: team.uid().ok_or(BudgetError::Identity)?, + }, + None => first_id.clone(), + }, + workspace_uid: workspace.uid().ok_or(BudgetError::Identity)?, + cluster_uid: cluster.uid().ok_or(BudgetError::Identity)?, + }; + root.validate()?; + let root_limits = limits( + team.as_ref() + .map(|team| &team.spec.envelope) + .unwrap_or(&first.spec.envelope), + )?; + let store = Store::new(client.clone(), &settings.accounting_namespace); + let existing = match &team { + Some(team) => team + .status + .as_ref() + .and_then(|status| status.inference_budget_account.clone()), + None => first + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + .map(|binding| binding.account.clone()), + }; + let account = if let Some(reference) = existing { + if reference.namespace != settings.accounting_namespace + || reference.name != name_for_root(&root) + { + return Err(BudgetError::Identity.into()); + } + reference + } else { + crate::sre_authority::privacy_epoch(client, &settings.accounting_namespace) + .await + .map_err(|_| BudgetError::Authorization)?; + let signer = crate::providers::signing::load_existing(client) + .await + .map_err(|_| BudgetError::Authorization)?; + let anchor = store + .create_anchor( + KarsBudgetAccountSpec { + scope: BudgetScope::GovernedInference, + root: root.clone(), + limits: root_limits, + }, + &signer, + ) + .await?; + let reference = AccountReference { + namespace: settings.accounting_namespace.clone(), + name: anchor.name_any(), + uid: anchor.uid().ok_or(BudgetError::Identity)?, + }; + if let Some(team) = &team { + let api: Api = Api::namespaced(client.clone(), &root.resource.namespace); + let updated = api + .patch_status( + &team.name_any(), + &PatchParams::default(), + &Patch::Merge(json!({ + "metadata": {"uid": team.uid(), "resourceVersion": team.resource_version()}, + "status": {"inferenceBudgetAccount": reference} + })), + ) + .await + .map_err(|e| api_error("pin lifetime Team account", e))?; + if updated + .status + .as_ref() + .and_then(|status| status.inference_budget_account.as_ref()) + != Some(&reference) + { + return Err(StoreError::Missing); + } + } else { + pin_task( + client, + first, + &task_binding(first, None, &root, &first_id.uid, &reference)?, + ) + .await?; + } + reference + }; + let initialized = store.initialize(&root, &account.uid).await?; + if initialized + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .is_none_or(|ledger| ledger.phase != crate::inference_budget_contract::AccountPhase::Active) + { + return Err(BudgetError::Closed.into()); + } + let current_limits = initialized + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)? + .limits; + if current_limits != root_limits { + store + .transact(&root, &account.uid, |ledger| { + ledger.narrow_root_limits(root_limits) + }) + .await?; + } + let default_model = crate::kars_task::blueprint::controller_default_model(); + let mut output = None; + for (index, node) in nodes.iter().enumerate() { + let authority = TaskAuthority { + task: resource(node)?, + parent_uid: index.checked_sub(1).and_then(|index| nodes[index].uid()), + root_task_uid: first_id.uid.clone(), + authorization_digest: node.spec.authorization_digest_with_model(&default_model), + effective_authorization: node + .spec + .authorization_configuration_with_model(&default_model), + limits: limits(&node.spec.envelope)?, + }; + let current = store.read(&root, &account.uid).await?; + let old = current + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .and_then(|ledger| ledger.nodes.get(&authority.task.uid)); + match old { + Some(old) if old.authority != authority => { + store + .transact(&root, &account.uid, |ledger| { + ledger.close_subtree(&authority.task.uid) + }) + .await?; + store + .transact(&root, &account.uid, |ledger| { + ledger.update_authority(authority.clone()) + }) + .await?; + store + .transact(&root, &account.uid, |ledger| { + ledger.resume_task(&authority.task.uid, &authority.authorization_digest) + }) + .await?; + } + Some(old) if !old.active => { + store + .transact(&root, &account.uid, |ledger| { + ledger.resume_task(&authority.task.uid, &authority.authorization_digest) + }) + .await?; + } + Some(_) => {} + None => { + store + .transact(&root, &account.uid, |ledger| { + ledger.register_task(authority.clone()) + }) + .await?; + } + } + let binding = task_binding( + node, + index.checked_sub(1).map(|index| &nodes[index]), + &root, + &first_id.uid, + &account, + )?; + pin_task(client, node, &binding).await?; + output = Some(binding); + } + output.ok_or_else(|| BudgetError::Identity.into()) +} diff --git a/controller/src/inference_budget/claim.rs b/controller/src/inference_budget/claim.rs new file mode 100644 index 000000000..dc078c2db --- /dev/null +++ b/controller/src/inference_budget/claim.rs @@ -0,0 +1,111 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + account::{KarsBudgetAccount, KarsBudgetAccountSpec}, + store::StoreError, +}; +use crate::{inference_budget_contract::BudgetError, providers::signing::ReceiptSigner}; +use kube::ResourceExt; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +pub const ANNOTATION: &str = "kars.azure.com/inference-budget-authority"; + +#[derive(Serialize, Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct Claim { + key_id: String, + nonce: String, + signature: String, +} + +fn note(spec: &KarsBudgetAccountSpec, namespace: &str, nonce: &str) -> Result, StoreError> { + let mut value = json!({ + "domain":"kars.azure.com/inference-budget-bootstrap/v1", + "accountingNamespace":namespace, "name":super::account::name_for_root(&spec.root), + "grant":spec, "nonce":nonce, + }); + value.sort_all_objects(); + serde_json::to_vec(&value).map_err(|_| BudgetError::Corrupt.into()) +} + +/// Authorship is signed atomically with CREATE, not inferred from labels on a +/// pre-existing object. Root status subsequently pins the API-generated UID. +pub fn issue( + spec: &KarsBudgetAccountSpec, + namespace: &str, + signer: &ReceiptSigner, +) -> Result { + let nonce = crate::providers::signing::generate_service_token(); + let signature = signer.sign_note(¬e(spec, namespace, &nonce)?); + serde_json::to_string(&Claim { + key_id: signer.key_id.clone(), + nonce, + signature, + }) + .map_err(|_| BudgetError::Corrupt.into()) +} + +pub fn verify( + account: &KarsBudgetAccount, + namespace: &str, + signer: &ReceiptSigner, +) -> Result<(), StoreError> { + let claim: Claim = serde_json::from_str( + account + .annotations() + .get(ANNOTATION) + .ok_or(BudgetError::Identity)?, + ) + .map_err(|_| BudgetError::Identity)?; + if claim.key_id != signer.key_id + || claim.nonce.len() != 64 + || !signer.verify_note( + ¬e(&account.spec, namespace, &claim.nonce)?, + &claim.signature, + ) + { + return Err(BudgetError::Identity.into()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::inference_budget_contract::{ + BudgetScope, Limits, ResourceIdentity, RootIdentity, RootKind, + }; + + #[test] + fn forged_labels_or_a_copied_grant_from_another_workspace_do_not_prove_authorship() { + let signer = ReceiptSigner::from_bytes(&[7; 32]); + let spec = KarsBudgetAccountSpec { + scope: BudgetScope::GovernedInference, + limits: Limits::default(), + root: RootIdentity { + kind: RootKind::KarsTask, + resource: ResourceIdentity { + namespace: "workspace".into(), + name: "root".into(), + uid: "root-uid".into(), + }, + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }, + }; + let mut account = + KarsBudgetAccount::new(&super::super::account::name_for_root(&spec.root), spec); + assert!(verify(&account, "accounting", &signer).is_err()); + account.metadata.annotations = Some(std::collections::BTreeMap::from([( + ANNOTATION.into(), + issue(&account.spec, "accounting", &signer).unwrap(), + )])); + assert!(verify(&account, "accounting", &signer).is_ok()); + assert!(verify(&account, "foreign-accounting", &signer).is_err()); + assert!(verify(&account, "accounting", &ReceiptSigner::from_bytes(&[8; 32])).is_err()); + account.spec.root.resource.uid = "recreated-root".into(); + assert!(verify(&account, "accounting", &signer).is_err()); + } +} diff --git a/controller/src/inference_budget/config.rs b/controller/src/inference_budget/config.rs new file mode 100644 index 000000000..066ee51b8 --- /dev/null +++ b/controller/src/inference_budget/config.rs @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::store::StoreError; +use crate::inference_budget_contract::{BudgetError, catalog::Catalog}; +use k8s_openapi::api::core::v1::ConfigMap; +use kube::{Api, Client, ResourceExt}; + +pub const AUDIENCE: &str = "kars.azure.com/governed-inference-budget"; +pub const PRIVATE_MOUNT: &str = "/var/run/kars/inference-budget"; +pub const TOKEN_VOLUME: &str = "kars-inference-budget-token"; +pub const CATALOG_KEY: &str = "contracts.json"; + +#[derive(Clone, Debug)] +pub struct Settings { + pub accounting_namespace: String, + pub catalog_name: String, + pub address: String, + pub tls_secret: String, + pub router_image_digest: String, +} + +impl Settings { + /// Absent/false preserves standalone operation and the existing launch + /// rejection for finite budgets. Enabling is not proof of readiness. + pub fn from_env() -> Result, StoreError> { + let enabled = std::env::var("KARS_INFERENCE_BUDGET_ENABLED").unwrap_or_default(); + match enabled.as_str() { + "" | "false" => return Ok(None), + "true" => {} + _ => return Err(BudgetError::Contract.into()), + } + let settings = Self { + accounting_namespace: crate::providers::signing::receipt_namespace(), + catalog_name: std::env::var("KARS_INFERENCE_BUDGET_CATALOG") + .unwrap_or_else(|_| "kars-inference-budget-contracts".into()), + address: std::env::var("KARS_INFERENCE_BUDGET_ADDR") + .unwrap_or_else(|_| "0.0.0.0:9447".into()), + tls_secret: std::env::var("KARS_INFERENCE_BUDGET_TLS_SECRET") + .map_err(|_| BudgetError::Contract)?, + router_image_digest: std::env::var("KARS_INFERENCE_BUDGET_ROUTER_DIGEST") + .map_err(|_| BudgetError::Contract)?, + }; + if !crate::inference_budget_contract::valid_label(&settings.accounting_namespace) + || !crate::inference_budget_contract::valid_name(&settings.catalog_name) + || !crate::inference_budget_contract::valid_name(&settings.tls_secret) + || settings.address.parse::().is_err() + || !crate::inference_budget_contract::valid_digest(&settings.router_image_digest) + { + return Err(BudgetError::Contract.into()); + } + Ok(Some(settings)) + } + + pub async fn catalog( + &self, + client: &Client, + now: i64, + ) -> Result { + super::admission::verify(client, &self.accounting_namespace).await?; + let api: Api = Api::namespaced(client.clone(), &self.accounting_namespace); + let value = api + .get(&self.catalog_name) + .await + .map_err(|error| StoreError::Api { + stage: "read operator inference contracts", + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + })?; + if value.metadata.uid.as_deref().is_none_or(str::is_empty) + || value + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || value.metadata.deletion_timestamp.is_some() + || value + .annotations() + .get("kars.azure.com/inference-budget-contracts") + .map(String::as_str) + != Some("v1") + { + return Err(BudgetError::Contract.into()); + } + let encoded = value + .data + .as_ref() + .and_then(|data| data.get(CATALOG_KEY)) + .ok_or(BudgetError::Contract)?; + if encoded.len() > 262_144 { + return Err(BudgetError::Capacity.into()); + } + let catalog: Catalog = serde_json::from_str(encoded).map_err(|_| BudgetError::Contract)?; + catalog.validate(now)?; + let mut canonical = serde_json::to_value(&catalog).map_err(|_| BudgetError::Contract)?; + canonical.sort_all_objects(); + let digest = crate::providers::signing::sha256_hex( + &serde_json::to_vec(&canonical).map_err(|_| BudgetError::Contract)?, + ); + Ok(ConfiguredCatalog { + catalog, + uid: value.metadata.uid.ok_or(BudgetError::Identity)?, + resource_version: value + .metadata + .resource_version + .ok_or(BudgetError::Identity)?, + digest, + }) + } +} + +#[derive(Clone, Debug)] +pub struct ConfiguredCatalog { + pub catalog: Catalog, + pub uid: String, + pub resource_version: String, + pub digest: String, +} + +impl Settings { + pub async fn public_ca(&self, client: &Client) -> Result<(String, String), StoreError> { + let api: Api = Api::namespaced(client.clone(), &self.accounting_namespace); + let map = api + .get(&self.catalog_name) + .await + .map_err(|error| StoreError::Api { + stage: "read budget public CA", + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + })?; + if map.metadata.uid.is_none() + || map.metadata.resource_version.is_none() + || map.metadata.deletion_timestamp.is_some() + || map + .annotations() + .get("kars.azure.com/inference-budget-contracts") + .map(String::as_str) + != Some("v1") + { + return Err(BudgetError::Contract.into()); + } + let ca = map + .data + .as_ref() + .and_then(|data| data.get("ca.crt")) + .ok_or(BudgetError::Contract)?; + if !ca.contains("-----BEGIN CERTIFICATE-----") + || ca.contains("PRIVATE KEY") + || ca.len() > 65_536 + { + return Err(BudgetError::Contract.into()); + } + let digest = crate::providers::signing::sha256_hex(ca.as_bytes()); + Ok((ca.clone(), digest)) + } +} diff --git a/controller/src/inference_budget/mod.rs b/controller/src/inference_budget/mod.rs new file mode 100644 index 000000000..7c7b26df5 --- /dev/null +++ b/controller/src/inference_budget/mod.rs @@ -0,0 +1,41 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +pub mod account; +pub mod admission; +pub mod auth; +pub mod binding; +pub mod claim; +pub mod config; +pub mod pod; +pub mod recovery; +pub mod scope; +pub mod service; +pub mod store; +pub mod team; +pub mod transport; + +pub fn start(client: kube::Client) { + match config::Settings::from_env() { + Ok(None) => {} + Err(_) => tracing::error!( + "Governed inference broker configuration is invalid; finite inference remains unavailable" + ), + Ok(Some(settings)) => { + recovery::start(client.clone(), settings.clone()); + tokio::spawn(async move { + loop { + if transport::run(client.clone(), settings.clone()) + .await + .is_err() + { + tracing::error!( + "Governed inference broker unavailable; finite inference remains fail-closed" + ); + } + tokio::time::sleep(std::time::Duration::from_secs(30)).await; + } + }); + } + } +} diff --git a/controller/src/inference_budget/pod.rs b/controller/src/inference_budget/pod.rs new file mode 100644 index 000000000..389a25660 --- /dev/null +++ b/controller/src/inference_budget/pod.rs @@ -0,0 +1,348 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + binding, + config::{AUDIENCE, PRIVATE_MOUNT, Settings, TOKEN_VOLUME}, + store::{Store, StoreError}, +}; +use crate::{ + crd::KarsSandbox, + inference_budget_contract::{BudgetError, ResourceIdentity, RouterBinding}, + kars_task::KarsTask, +}; +use k8s_openapi::api::core::v1::{ConfigMap, Namespace}; +use kube::{ + Api, Client, ResourceExt, + api::{Patch, PatchParams, PostParams}, +}; +use serde_json::{Value, json}; + +const CA_NAME: &str = "kars-inference-budget-ca"; +const CA_VOLUME: &str = "kars-inference-budget-ca"; + +#[cfg(test)] +#[path = "pod_tests.rs"] +mod tests; + +pub fn egress(sandbox: &KarsSandbox) -> Result, StoreError> { + let Some(reference) = &sandbox.spec.inference_budget_ref else { + return Ok(None); + }; + let settings = Settings::from_env()?.ok_or(BudgetError::Contract)?; + if reference.account.namespace != settings.accounting_namespace { + return Err(BudgetError::Identity.into()); + } + Ok(Some(json!({ + "to": [{ + "namespaceSelector": {"matchLabels": {"kubernetes.io/metadata.name": settings.accounting_namespace}}, + "podSelector": {"matchLabels": { + "app.kubernetes.io/name": "kars", "app.kubernetes.io/component": "controller" + }} + }], + "ports": [{"protocol": "TCP", "port": 9447}] + }))) +} + +pub struct Plan { + pub binding: RouterBinding, + pub ca_version: String, + pub endpoint: String, + pub router_image_digest: String, +} + +fn api_error(stage: &'static str, error: kube::Error) -> StoreError { + StoreError::Api { + stage, + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +pub async fn prepare( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result, StoreError> { + let owners = sandbox + .metadata + .owner_references + .as_deref() + .unwrap_or_default(); + let Some(owner) = owners.iter().find(|owner| { + owner.controller == Some(true) + && owner.kind == "KarsTask" + && owner.api_version == "kars.azure.com/v1alpha1" + }) else { + if sandbox.spec.inference_budget_ref.is_some() { + return Err(BudgetError::Authorization.into()); + } + return Ok(None); + }; + let workspace = sandbox.namespace().ok_or(BudgetError::Identity)?; + let tasks: Api = Api::namespaced(client.clone(), &workspace); + let task = tasks + .get(&owner.name) + .await + .map_err(|error| api_error("read sandbox budget authority", error))?; + if task.metadata.uid.as_deref() != Some(owner.uid.as_str()) + || task.name_any() != sandbox.name_any() + { + return Err(BudgetError::Identity.into()); + } + if !binding::needs_account(client, &task).await? { + return Ok(None); + } + let settings = Settings::from_env()?.ok_or(BudgetError::Contract)?; + super::admission::verify(client, &settings.accounting_namespace).await?; + let bound = task + .status + .as_ref() + .and_then(|status| status.inference_budget.clone()) + .ok_or(StoreError::Missing)?; + if sandbox.spec.inference_budget_ref.as_ref() != Some(&bound) + || bound.task_uid != owner.uid + || bound.authorization_digest != task.envelope_digest() + || bound.account.namespace != settings.accounting_namespace + { + return Err(BudgetError::Authorization.into()); + } + let store = Store::new(client.clone(), &settings.accounting_namespace); + let account = store.read(&bound.root, &bound.account.uid).await?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + if !ledger.requires_enforcement(&bound.task_uid)? { + return Ok(None); + } + let node = ledger + .nodes + .get(&bound.task_uid) + .ok_or(BudgetError::Authorization)?; + if !node.active + || node.authority.authorization_digest != bound.authorization_digest + || !task + .spec + .execution + .as_ref() + .is_some_and(|execution| execution.launch) + { + return Err(BudgetError::Authorization.into()); + } + crate::reconciler::namespace_ownership::recheck(client, sandbox, namespace) + .await + .map_err(|_| BudgetError::Identity)?; + fence_namespace(client, sandbox, namespace).await?; + let epoch = crate::sre_authority::privacy_epoch(client, &namespace.name_any()) + .await + .map_err(|_| BudgetError::Authorization)?; + let (ca, ca_version) = settings.public_ca(client).await?; + mirror_public_ca(client, sandbox, namespace, ca).await?; + Ok(Some(Plan { + binding: RouterBinding { + task: bound, + sandbox: ResourceIdentity { + namespace: workspace, + name: sandbox.name_any(), + uid: sandbox.uid().ok_or(BudgetError::Identity)?, + }, + runtime_namespace: namespace.name_any(), + runtime_namespace_uid: namespace.uid().ok_or(BudgetError::Identity)?, + privacy_epoch: epoch, + }, + ca_version, + endpoint: format!( + "https://kars-inference-budget.{}.svc:9447", + settings.accounting_namespace + ), + router_image_digest: settings.router_image_digest, + })) +} + +async fn fence_namespace( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result<(), StoreError> { + let live = crate::reconciler::namespace_ownership::recheck(client, sandbox, namespace) + .await + .map_err(|_| BudgetError::Identity)?; + let key = "kars.azure.com/inference-budget"; + if let Some(value) = live.labels().get(key) { + return if value == "v1" { + Ok(()) + } else { + Err(BudgetError::Identity.into()) + }; + } + let api: Api = Api::all(client.clone()); + let updated = api + .patch( + &live.name_any(), + &PatchParams::default(), + &Patch::Merge(json!({ + "metadata": {"uid": live.uid(), "resourceVersion": live.resource_version(), + "labels": {key: "v1"}} + })), + ) + .await + .map_err(|error| api_error("activate budget namespace fence", error))?; + if updated.metadata.uid != live.metadata.uid + || updated.labels().get(key).map(String::as_str) != Some("v1") + { + return Err(BudgetError::Identity.into()); + } + Ok(()) +} + +async fn mirror_public_ca( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, + ca: String, +) -> Result<(), StoreError> { + let api: Api = Api::namespaced(client.clone(), &namespace.name_any()); + let existing = api + .get_opt(CA_NAME) + .await + .map_err(|error| api_error("read budget CA projection", error))?; + if let Some(existing) = existing { + if existing + .annotations() + .get("kars.azure.com/budget-sandbox-uid") + != sandbox.metadata.uid.as_ref() + || existing + .annotations() + .get("kars.azure.com/budget-namespace-uid") + != namespace.metadata.uid.as_ref() + || existing.metadata.uid.is_none() + || existing.metadata.resource_version.is_none() + { + return Err(BudgetError::Identity.into()); + } + if existing.data.as_ref().and_then(|data| data.get("ca.crt")) == Some(&ca) { + return Ok(()); + } + api.patch( + CA_NAME, + &PatchParams::default(), + &Patch::Merge(json!({ + "metadata": {"uid": existing.uid(), "resourceVersion": existing.resource_version()}, + "data": {"ca.crt": ca} + })), + ) + .await + .map_err(|error| api_error("update budget CA projection", error))?; + } else { + let map: ConfigMap = serde_json::from_value(json!({ + "apiVersion": "v1", "kind": "ConfigMap", + "metadata": { + "name": CA_NAME, "namespace": namespace.name_any(), + "annotations": { + "kars.azure.com/budget-sandbox-uid": sandbox.uid(), + "kars.azure.com/budget-namespace-uid": namespace.uid() + }, + "ownerReferences": [{ + "apiVersion": "v1", "kind": "Namespace", "name": namespace.name_any(), + "uid": namespace.uid(), "controller": true, "blockOwnerDeletion": false + }] + }, "data": {"ca.crt": ca} + })) + .map_err(|_| BudgetError::Corrupt)?; + api.create(&PostParams::default(), &map) + .await + .map_err(|error| api_error("create budget CA projection", error))?; + } + Ok(()) +} + +impl Plan { + pub fn apply( + &self, + pod: &mut Value, + annotations: &mut std::collections::BTreeMap, + ) -> Result<(), StoreError> { + let volumes = pod + .get_mut("volumes") + .and_then(Value::as_array_mut) + .ok_or(BudgetError::Corrupt)?; + if volumes.iter().any(|volume| { + matches!( + volume.get("name").and_then(Value::as_str), + Some(TOKEN_VOLUME | CA_VOLUME) + ) + }) { + return Err(BudgetError::Corrupt.into()); + } + + pub async fn decorate( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, + pod: &mut Value, + ) -> Result, StoreError> { + let mut annotations = std::collections::BTreeMap::new(); + if let Some(plan) = prepare(client, sandbox, namespace).await? { + plan.apply(pod, &mut annotations)?; + } + Ok(annotations) + } + volumes.push(json!({ + "name": TOKEN_VOLUME, + "projected": {"sources": [{"serviceAccountToken": {"audience": AUDIENCE, "expirationSeconds": 600, "path": "token"}}]} + })); + volumes.push(json!({"name": CA_VOLUME, "configMap": {"name": CA_NAME}})); + let containers = pod + .get_mut("containers") + .and_then(Value::as_array_mut) + .ok_or(BudgetError::Corrupt)?; + let router = containers + .iter_mut() + .find(|container| { + container.get("name").and_then(Value::as_str) == Some("inference-router") + }) + .ok_or(BudgetError::Corrupt)?; + let env = router + .get_mut("env") + .and_then(Value::as_array_mut) + .ok_or(BudgetError::Corrupt)?; + env.extend([ + json!({"name": "KARS_INFERENCE_BUDGET_REQUIRED", "value": "true"}), + json!({"name": "KARS_INFERENCE_BUDGET_BINDING", "value": serde_json::to_string(&self.binding).map_err(|_| BudgetError::Corrupt)?}), + json!({"name": "KARS_INFERENCE_BUDGET_ENDPOINT", "value": self.endpoint}), + json!({"name": "KARS_INFERENCE_BUDGET_CA", "value": "/etc/kars/inference-budget-ca/ca.crt"}), + json!({"name": "POD_NAME", "valueFrom": {"fieldRef": {"fieldPath": "metadata.name"}}}), + json!({"name": "POD_UID", "valueFrom": {"fieldRef": {"fieldPath": "metadata.uid"}}}), + ]); + let mounts = router + .get_mut("volumeMounts") + .and_then(Value::as_array_mut) + .ok_or(BudgetError::Corrupt)?; + mounts.push(json!({"name": TOKEN_VOLUME, "mountPath": PRIVATE_MOUNT, "readOnly": true})); + mounts.push(json!({"name": CA_VOLUME, "mountPath": "/etc/kars/inference-budget-ca", "readOnly": true})); + router["readinessProbe"] = json!({ + "httpGet": {"path": "/readyz", "port": "inference"}, "initialDelaySeconds": 3, "periodSeconds": 5 + }); + let image = router + .get("image") + .and_then(Value::as_str) + .ok_or(BudgetError::Contract)?; + let base = image.split('@').next().ok_or(BudgetError::Contract)?; + router["image"] = json!(format!("{base}@{}", self.router_image_digest)); + annotations.insert( + "kars.azure.com/inference-budget-ca-version".into(), + self.ca_version.clone(), + ); + annotations.insert( + "kars.azure.com/inference-budget-binding".into(), + crate::providers::signing::sha256_hex( + &serde_json::to_vec(&self.binding).map_err(|_| BudgetError::Corrupt)?, + ), + ); + Ok(()) + } +} diff --git a/controller/src/inference_budget/pod_tests.rs b/controller/src/inference_budget/pod_tests.rs new file mode 100644 index 000000000..b9cab598e --- /dev/null +++ b/controller/src/inference_budget/pod_tests.rs @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::inference_budget_contract::{ + AccountReference, BudgetScope, RootIdentity, RootKind, TaskBudgetBinding, +}; + +fn identity(name: &str) -> ResourceIdentity { + ResourceIdentity { + namespace: "workspace".into(), + name: name.into(), + uid: format!("{name}-uid"), + } +} + +fn plan() -> Plan { + Plan { + binding: RouterBinding { + task: TaskBudgetBinding { + scope: BudgetScope::GovernedInference, + account: AccountReference { + namespace: "accounting".into(), + name: "root".into(), + uid: "account-uid".into(), + }, + root: RootIdentity { + kind: RootKind::KarsTask, + resource: identity("task"), + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }, + task_uid: "task-uid".into(), + parent_task_uid: None, + root_task_uid: "task-uid".into(), + authorization_digest: format!("sha256:{}", "a".repeat(64)), + }, + sandbox: identity("task"), + runtime_namespace: "kars-task".into(), + runtime_namespace_uid: "namespace-uid".into(), + privacy_epoch: None, + }, + ca_version: "public-ca-version".into(), + endpoint: "https://kars-inference-budget.accounting.svc:9447".into(), + router_image_digest: format!("sha256:{}", "7".repeat(64)), + } +} + +#[test] +fn every_runtime_receives_only_a_router_private_token_mount() { + for runtime in [ + "openclaw", + "hermes", + "openai-agents", + "maf-python", + "anthropic", + "pydantic-ai", + "langgraph", + "byo", + ] { + let agent = if runtime == "openclaw" { + "openclaw" + } else { + "agent" + }; + let mut pod = json!({ + "volumes": [], "serviceAccountName": "sandbox", + "containers": [ + {"name": agent, "env": [{"name":"KARS_RUNTIME_KIND","value":runtime}], "volumeMounts": []}, + {"name":"inference-router","image":"router:latest","env": [], "volumeMounts":[]} + ] + }); + let original_agent = pod["containers"][0].clone(); + let mut annotations = std::collections::BTreeMap::new(); + plan().apply(&mut pod, &mut annotations).unwrap(); + assert_eq!(pod["containers"][0], original_agent, "{runtime}"); + let router = &pod["containers"][1]; + assert_eq!(router["volumeMounts"][0]["mountPath"], PRIVATE_MOUNT); + assert_eq!(router["volumeMounts"][0]["readOnly"], true); + assert_eq!( + pod["volumes"][0]["projected"]["sources"][0]["serviceAccountToken"]["audience"], + AUDIENCE + ); + assert_eq!(router["readinessProbe"]["httpGet"]["path"], "/readyz"); + assert_eq!( + router["image"], + format!("router:latest@sha256:{}", "7".repeat(64)) + ); + assert_eq!( + annotations["kars.azure.com/inference-budget-ca-version"], + "public-ca-version" + ); + } +} + +#[tokio::test] +async fn no_task_owner_no_reference_is_a_byte_safe_legacy_noop_without_api_access() { + let server = wiremock::MockServer::start().await; + let config = kube::Config::new(server.uri().parse().unwrap()); + let client = Client::try_from(config).unwrap(); + let sandbox: KarsSandbox = serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1", "kind":"KarsSandbox", + "metadata":{"name":"legacy","namespace":"workspace","uid":"legacy-uid"}, "spec":{} + })) + .unwrap(); + let namespace: Namespace = serde_json::from_value(json!({ + "apiVersion":"v1","kind":"Namespace","metadata":{"name":"kars-legacy","uid":"namespace-uid"} + })) + .unwrap(); + let mut pod = json!({"containers":[{"name":"agent","envFrom":[{"secretRef":{"name":"legacy-credentials"}}]}]}); + let before = pod.clone(); + assert!( + decorate(&client, &sandbox, &namespace, &mut pod) + .await + .unwrap() + .is_empty() + ); + assert_eq!(pod, before); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn collision_never_adds_a_second_private_volume() { + let mut pod = json!({"volumes":[{"name":TOKEN_VOLUME}], "containers":[]}); + let before = pod.clone(); + assert!(plan().apply(&mut pod, &mut Default::default()).is_err()); + assert_eq!(before, pod); +} diff --git a/controller/src/inference_budget/recovery.rs b/controller/src/inference_budget/recovery.rs new file mode 100644 index 000000000..445933085 --- /dev/null +++ b/controller/src/inference_budget/recovery.rs @@ -0,0 +1,198 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + account::{KarsBudgetAccount, MANAGED_BY, OWNER}, + config::Settings, + store::{Store, StoreError}, +}; +use crate::{ + inference_budget_contract::{BudgetError, RootKind}, + kars_task::KarsTask, + kars_team::KarsTeam, +}; +use k8s_openapi::api::core::v1::{Namespace, Pod}; +use kube::{Api, Client, ResourceExt, api::ListParams}; + +fn api_error(error: kube::Error) -> StoreError { + StoreError::Api { + stage: "reconcile inference liabilities", + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +pub fn start(client: Client, settings: Settings) { + tokio::spawn(async move { + loop { + if scan(&client, &settings).await.is_err() { + tracing::error!( + "Governed inference recovery unavailable; outstanding liabilities remain funded" + ); + } + tokio::time::sleep(std::time::Duration::from_secs(15)).await; + } + }); +} + +async fn scan(client: &Client, settings: &Settings) -> Result<(), StoreError> { + let api: Api = + Api::namespaced(client.clone(), &settings.accounting_namespace); + let mut params = ListParams::default() + .labels(&format!("{MANAGED_BY}={OWNER}")) + .limit(50); + let store = Store::new(client.clone(), &settings.accounting_namespace); + loop { + let page = api.list(¶ms).await.map_err(api_error)?; + for account in page.items { + if account.status.is_none() { + continue; + } + if recover(client, &store, &account).await.is_err() { + tracing::error!(account = %account.name_any(), + "Governed inference account recovery failed closed; no balances reset"); + } + } + let Some(token) = page.metadata.continue_.filter(|token| !token.is_empty()) else { + break; + }; + params = params.continue_token(&token); + } + Ok(()) +} + +async fn recover( + client: &Client, + store: &Store, + account: &KarsBudgetAccount, +) -> Result<(), StoreError> { + let root = &account.spec.root; + let uid = account.uid().ok_or(BudgetError::Identity)?; + let now = chrono::Utc::now().timestamp(); + store + .transact(root, &uid, |ledger| ledger.expire_undispatched(now)) + .await?; + // Provider transport deadline is 600s. A vanished router's liability is + // committed after that deadline plus a margin, not refunded on TTL. + store + .transact(root, &uid, |ledger| { + ledger.commit_uncertain_before(now - 660) + }) + .await?; + let namespaces: Api = Api::all(client.clone()); + let workspace = namespaces + .get_opt(&root.resource.namespace) + .await + .map_err(api_error)?; + if workspace.is_none_or(|namespace| { + namespace.metadata.uid.as_deref() != Some(root.workspace_uid.as_str()) + || namespace.metadata.deletion_timestamp.is_some() + }) { + return store + .transact(root, &uid, |ledger| ledger.close_account()) + .await; + } + let tasks: Api = Api::namespaced(client.clone(), &root.resource.namespace); + if root.kind == RootKind::KarsTeam { + let teams: Api = Api::namespaced(client.clone(), &root.resource.namespace); + let team = teams + .get_opt(&root.resource.name) + .await + .map_err(api_error)?; + if team.as_ref().is_none_or(|team| { + team.metadata.uid.as_deref() != Some(root.resource.uid.as_str()) + || team.metadata.deletion_timestamp.is_some() + }) { + return store + .transact(root, &uid, |ledger| ledger.close_account()) + .await; + } + if team.is_some_and(|team| team.spec.paused) { + let current = store.read(root, &uid).await?; + let ledger = current + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + for node in ledger + .nodes + .values() + .filter(|node| node.authority.parent_uid.is_none()) + { + store + .transact(root, &uid, |ledger| { + ledger.close_subtree(&node.authority.task.uid) + }) + .await?; + } + return Ok(()); + } + } else { + let task = tasks + .get_opt(&root.resource.name) + .await + .map_err(api_error)?; + if task.is_none_or(|task| { + task.metadata.uid.as_deref() != Some(root.resource.uid.as_str()) + || task.metadata.deletion_timestamp.is_some() + }) { + return store + .transact(root, &uid, |ledger| ledger.close_account()) + .await; + } + } + let current = store.read(root, &uid).await?; + let ledger = current + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + for node in ledger.nodes.values().filter(|node| node.active) { + let task = tasks + .get_opt(&node.authority.task.name) + .await + .map_err(api_error)?; + if task.as_ref().is_none_or(|task| { + task.metadata.uid.as_deref() != Some(node.authority.task.uid.as_str()) + || task.metadata.deletion_timestamp.is_some() + || task.envelope_digest() != node.authority.authorization_digest + }) { + store + .transact(root, &uid, |ledger| { + ledger.close_subtree(&node.authority.task.uid) + }) + .await?; + continue; + } + let launched = task.is_some_and(|task| { + task.spec + .execution + .is_some_and(|execution| execution.launch) + }); + for session in ledger.sessions.values().filter(|session| { + !session.closed && session.identity.task_uid == node.authority.task.uid + }) { + let namespace = format!("kars-{}", session.identity.sandbox.name); + let pods: Api = Api::namespaced(client.clone(), &namespace); + let pod = pods + .get_opt(&session.identity.pod_name) + .await + .map_err(api_error)?; + if !launched + || pod.is_none_or(|pod| { + pod.metadata.uid.as_deref() != Some(session.identity.pod_uid.as_str()) + || pod.metadata.deletion_timestamp.is_some() + }) + { + store + .transact(root, &uid, |ledger| { + ledger.close_session(&session.identity.pod_uid) + }) + .await?; + } + } + } + Ok(()) +} diff --git a/controller/src/inference_budget/scope.rs b/controller/src/inference_budget/scope.rs new file mode 100644 index 000000000..dddbab909 --- /dev/null +++ b/controller/src/inference_budget/scope.rs @@ -0,0 +1,60 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::{inference_budget_contract::BudgetScope, kars_task::TaskEnvelope}; + +pub fn unsupported(envelope: &TaskEnvelope) -> bool { + super::binding::has_finite(envelope) + && envelope + .budget + .as_ref() + .is_none_or(|budget| budget.scope != Some(BudgetScope::GovernedInference)) +} + +pub const LAUNCH_RULE: &str = "!has(self.execution) || !self.execution.launch || !has(self.envelope.budget) || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference')"; +pub const OPT_IN_RULE: &str = "!has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0)))"; +pub const RETAIN_SCOPE_RULE: &str = "!has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference')"; + +pub fn validations() +-> Vec { + use k8s_openapi::apiextensions_apiserver::pkg::apis::apiextensions::v1::ValidationRule; + [ + (OPT_IN_RULE, "First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted"), + (RETAIN_SCOPE_RULE, "GovernedInference scope cannot be removed or changed for an existing UID"), + ].into_iter().map(|(rule, message)| ValidationRule { + rule:rule.into(), message:Some(message.into()), reason:Some("FieldValueForbidden".into()), + ..Default::default() + }).collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::kars_task::{KarsTaskSpec, TaskBudget, TaskExecution}; + + #[test] + fn explicit_scope_is_syntax_not_permission_to_bypass_the_async_broker_gate() { + let mut spec = KarsTaskSpec { + envelope: TaskEnvelope { + budget: Some(TaskBudget { + tokens: Some(100), + ..Default::default() + }), + ..Default::default() + }, + execution: Some(TaskExecution { + launch: true, + runtime: None, + }), + ..Default::default() + }; + assert!(unsupported(&spec.envelope)); + assert!(crate::kars_task::validate_execution_contract(&spec).is_err()); + spec.envelope.budget.as_mut().unwrap().scope = Some(BudgetScope::GovernedInference); + assert!(!unsupported(&spec.envelope)); + assert!(crate::kars_task::validate_execution_contract(&spec).is_ok()); + // Actual materialization still calls binding::prepare_task and requires + // a current, privacy-qualified broker account; the pure validator cannot + // create that authority. + } +} diff --git a/controller/src/inference_budget/service.rs b/controller/src/inference_budget/service.rs new file mode 100644 index 000000000..01339ddb6 --- /dev/null +++ b/controller/src/inference_budget/service.rs @@ -0,0 +1,258 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use axum::{ + Json, Router, + extract::{DefaultBodyLimit, State}, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + routing::post, +}; +use kube::Client; +use serde_json::json; +use std::sync::Arc; + +use super::{ + auth, + config::Settings, + store::{Store, StoreError}, +}; +use crate::inference_budget_contract::{ + AttemptCommand, BrokerRequest, BudgetError, ReserveRequest, SessionRequest, Settlement, +}; + +#[derive(Clone)] +pub struct Broker { + pub client: Client, + pub settings: Settings, + pub store: Store, +} + +pub fn router(broker: Broker) -> Router { + Router::new() + .route("/v1/catalog", post(catalog)) + .route("/v1/session", post(session)) + .route("/v1/reserve", post(reserve)) + .route("/v1/begin", post(begin)) + .route("/v1/settle", post(settle)) + .layer(DefaultBodyLimit::max(32_768)) + .with_state(Arc::new(broker)) +} + +struct Failure(StoreError); + +impl From for Failure { + fn from(error: StoreError) -> Self { + Self(error) + } +} + +impl From for Failure { + fn from(error: BudgetError) -> Self { + Self(error.into()) + } +} + +impl IntoResponse for Failure { + fn into_response(self) -> Response { + let (status, code) = match &self.0 { + StoreError::Ledger(BudgetError::Exhausted) => { + (StatusCode::TOO_MANY_REQUESTS, "inference_budget_exhausted") + } + StoreError::Ledger(BudgetError::Capacity) => { + (StatusCode::SERVICE_UNAVAILABLE, "inference_budget_capacity") + } + StoreError::Ledger(BudgetError::Identity | BudgetError::Authorization) => { + (StatusCode::FORBIDDEN, "inference_budget_authority") + } + StoreError::Ledger( + BudgetError::AlreadyDispatched | BudgetError::Sequence | BudgetError::Expired, + ) => (StatusCode::CONFLICT, "inference_budget_attempt_state"), + StoreError::Ledger(BudgetError::Contract) => { + (StatusCode::SERVICE_UNAVAILABLE, "inference_budget_contract") + } + StoreError::Ledger(BudgetError::Closed) => { + (StatusCode::FORBIDDEN, "inference_budget_closed") + } + StoreError::Ledger( + BudgetError::Breach | BudgetError::Corrupt | BudgetError::Overflow, + ) => (StatusCode::SERVICE_UNAVAILABLE, "inference_budget_frozen"), + _ => ( + StatusCode::SERVICE_UNAVAILABLE, + "inference_budget_unavailable", + ), + }; + // Deliberately neither the token, wire request, nor API response body. + ( + status, + Json(json!({"error": {"code": code, "message": self.0.to_string()}})), + ) + .into_response() + } +} + +async fn catalog( + State(broker): State>, + headers: HeaderMap, + Json(request): Json>, +) -> Result, Failure> { + let account = broker + .store + .read(&request.root, &request.payload.account_uid) + .await?; + auth::authenticate( + &broker.client, + &headers, + &account, + &request.payload.identity, + true, + ) + .await?; + let catalog = broker + .settings + .catalog(&broker.client, chrono::Utc::now().timestamp()) + .await?; + Ok(Json(json!({ + "catalog": catalog.catalog, "uid": catalog.uid, "resourceVersion": catalog.resource_version, + "digest": catalog.digest, "scope": "GovernedInference", + "moneyRequired": account.status.as_ref().and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?.requires_price(&request.payload.identity.task_uid)?, + }))) +} + +async fn session( + State(broker): State>, + headers: HeaderMap, + Json(request): Json>, +) -> Result, Failure> { + let account = broker + .store + .read(&request.root, &request.payload.account_uid) + .await?; + auth::authenticate( + &broker.client, + &headers, + &account, + &request.payload.identity, + true, + ) + .await?; + let result = broker + .store + .transact(&request.root, &request.payload.account_uid, |ledger| { + ledger.register_session(request.payload.identity.clone()) + }) + .await?; + Ok(Json(result)) +} + +async fn reserve( + State(broker): State>, + headers: HeaderMap, + Json(request): Json>, +) -> Result, Failure> { + let account = broker + .store + .read(&request.root, &request.payload.account_uid) + .await?; + auth::authenticate( + &broker.client, + &headers, + &account, + &request.payload.identity, + true, + ) + .await?; + let now = chrono::Utc::now().timestamp(); + let catalog = broker.settings.catalog(&broker.client, now).await?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + catalog.catalog.accepts_quote( + &request.payload.quote, + now, + ledger.requires_price(&request.payload.identity.task_uid)?, + )?; + broker + .store + .transact(&request.root, &request.payload.account_uid, |ledger| { + ledger.expire_undispatched(now) + }) + .await?; + let result = broker + .store + .transact(&request.root, &request.payload.account_uid, |ledger| { + ledger.reserve(&request.payload, chrono::Utc::now().timestamp()) + }) + .await?; + Ok(Json(result)) +} + +async fn begin( + State(broker): State>, + headers: HeaderMap, + Json(request): Json>, +) -> Result, Failure> { + let account = broker + .store + .read(&request.root, &request.payload.account_uid) + .await?; + auth::authenticate( + &broker.client, + &headers, + &account, + &request.payload.identity, + true, + ) + .await?; + let now = chrono::Utc::now().timestamp(); + let catalog = broker.settings.catalog(&broker.client, now).await?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + let attempt = ledger + .attempts + .get(&request.payload.key.storage_key()) + .ok_or(BudgetError::Sequence)?; + catalog.catalog.accepts_quote( + &attempt.quote, + now, + ledger.requires_price(&request.payload.identity.task_uid)?, + )?; + let result = broker + .store + .transact(&request.root, &request.payload.account_uid, |ledger| { + ledger.begin_dispatch(&request.payload, chrono::Utc::now().timestamp()) + }) + .await?; + Ok(Json(result)) +} + +async fn settle( + State(broker): State>, + headers: HeaderMap, + Json(request): Json>, +) -> Result, Failure> { + let command = &request.payload.attempt; + let account = broker + .store + .read(&request.root, &command.account_uid) + .await?; + auth::authenticate(&broker.client, &headers, &account, &command.identity, false).await?; + // Settlement uses the accepted contract snapshot. A later tariff expiry or + // policy change cannot erase already accepted liability. + let result = broker + .store + .transact(&request.root, &command.account_uid, |ledger| { + ledger.settle(&request.payload) + }) + .await?; + if result.breach { + tracing::error!(account_uid = %command.account_uid, "Governed inference contract bound breached; account frozen"); + } + Ok(Json(result)) +} diff --git a/controller/src/inference_budget/store.rs b/controller/src/inference_budget/store.rs new file mode 100644 index 000000000..fb2879f12 --- /dev/null +++ b/controller/src/inference_budget/store.rs @@ -0,0 +1,310 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Durable single-object accounting. No informer balance and no in-memory +//! authority. A grant is returned only after its UID/RV-fenced write succeeds. + +use super::account::*; +use crate::inference_budget_contract::{ + BudgetError, RootIdentity, + ledger::{Ledger, Mutation}, +}; +use kube::{ + Api, Client, ResourceExt, + api::{Patch, PatchParams, PostParams}, +}; +use serde_json::json; + +const RETRIES: usize = 8; +const MAX_OPERATION_SECONDS: u64 = 10; + +#[derive(Debug, thiserror::Error)] +pub enum StoreError { + #[error(transparent)] + Ledger(#[from] BudgetError), + #[error("Inference budget API {stage} failed (status {code:?})")] + Api { + stage: &'static str, + code: Option, + }, + #[error("Inference budget account is missing, replaced, or uninitialized")] + Missing, + #[error("Inference budget account CAS contention/deadline exceeded")] + Contention, +} + +fn api_error(stage: &'static str, error: kube::Error) -> StoreError { + StoreError::Api { + stage, + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +#[derive(Clone)] +pub struct Store { + accounts: Api, + namespace: String, +} + +impl Store { + pub fn new(client: Client, accounting_namespace: &str) -> Self { + Self { + accounts: Api::namespaced(client, accounting_namespace), + namespace: accounting_namespace.into(), + } + } + + fn validate_identity( + account: &KarsBudgetAccount, + root: &RootIdentity, + uid: &str, + ) -> Result<(), StoreError> { + if account.metadata.uid.as_deref() != Some(uid) + || account + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || account.metadata.name.as_deref() != Some(name_for_root(root).as_str()) + || account.metadata.deletion_timestamp.is_some() + || account.spec.root != *root + || account.labels().get(MANAGED_BY).map(String::as_str) != Some(OWNER) + || account + .metadata + .owner_references + .as_ref() + .is_some_and(|refs| !refs.is_empty()) + { + return Err(StoreError::Missing); + } + Ok(()) + } + + fn ledger(account: &KarsBudgetAccount) -> Result<&Ledger, StoreError> { + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + if account.annotations().get(BOOTSTRAP).map(String::as_str) != Some("sealed") + || account.metadata.uid.as_deref() != Some(ledger.account_uid.as_str()) + || account.spec.root != ledger.root + || account.spec.scope != ledger.scope + || !ledger.limits.attenuates(account.spec.limits) + { + return Err(StoreError::Missing); + } + ledger.validate()?; + for node in ledger.nodes.values() { + let mut effective = node.authority.effective_authorization.clone(); + effective.sort_all_objects(); + let bytes = serde_json::to_vec(&effective).map_err(|_| BudgetError::Corrupt)?; + let digest = format!("sha256:{}", crate::providers::signing::sha256_hex(&bytes)); + if digest != node.authority.authorization_digest { + return Err(BudgetError::Corrupt.into()); + } + } + Ok(ledger) + } + + /// Reserve an account anchor. The owning root's protected status must pin + /// the returned UID BEFORE `initialize` is called. Existing callers with a + /// pin use `read`, never this bootstrap path. + pub async fn create_anchor( + &self, + spec: KarsBudgetAccountSpec, + signer: &crate::providers::signing::ReceiptSigner, + ) -> Result { + spec.root.validate()?; + let name = name_for_root(&spec.root); + let mut account = KarsBudgetAccount::new(&name, spec); + let authority = super::claim::issue(&account.spec, &self.namespace, signer)?; + account.metadata.labels = Some(std::collections::BTreeMap::from([( + MANAGED_BY.into(), + OWNER.into(), + )])); + account.metadata.annotations = Some(std::collections::BTreeMap::from([ + (BOOTSTRAP.into(), "pending".into()), + (super::claim::ANNOTATION.into(), authority), + ])); + match self.accounts.create(&PostParams::default(), &account).await { + Ok(created) => Ok(created), + Err(kube::Error::Api(status)) if status.code == 409 => { + let existing = self + .accounts + .get(&name) + .await + .map_err(|e| api_error("read account anchor", e))?; + let uid = existing + .metadata + .uid + .as_deref() + .ok_or(StoreError::Missing)?; + Self::validate_identity(&existing, &account.spec.root, uid)?; + super::claim::verify(&existing, &self.namespace, signer)?; + if existing.spec.limits.normalized() != account.spec.limits.normalized() { + return Err(BudgetError::Authorization.into()); + } + if existing.annotations().get(BOOTSTRAP).map(String::as_str) == Some("sealed") { + Self::ledger(&existing)?; + } else if existing.annotations().get(BOOTSTRAP).map(String::as_str) + != Some("pending") + { + return Err(StoreError::Missing); + } + Ok(existing) + } + Err(error) => Err(api_error("create account anchor", error)), + } + } + + pub async fn initialize( + &self, + root: &RootIdentity, + pinned_uid: &str, + ) -> Result { + root.validate()?; + if !crate::inference_budget_contract::valid_uid(pinned_uid) { + return Err(BudgetError::Identity.into()); + } + let name = name_for_root(root); + for _ in 0..RETRIES { + let account = self + .accounts + .get(&name) + .await + .map_err(|e| api_error("read account bootstrap", e))?; + Self::validate_identity(&account, root, pinned_uid)?; + if account.annotations().get(BOOTSTRAP).map(String::as_str) == Some("sealed") { + Self::ledger(&account)?; + return Ok(account); + } + if account.annotations().get(BOOTSTRAP).map(String::as_str) != Some("pending") { + return Err(StoreError::Missing); + } + if let Some(ledger) = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + { + // A crash after status initialization must validate the existing + // ledger, not reset it. Pending accounts cannot dispatch. + ledger.validate()?; + if ledger.account_uid != pinned_uid + || ledger.root != *root + || ledger.limits.normalized() != account.spec.limits.normalized() + || !ledger.nodes.is_empty() + || !ledger.sessions.is_empty() + || !ledger.attempts.is_empty() + || ledger.meters != Default::default() + { + return Err(BudgetError::Corrupt.into()); + } + match self.accounts.patch(&name, &PatchParams::default(), &Patch::Merge(json!({ + "metadata": {"uid": pinned_uid, "resourceVersion": account.resource_version(), + "annotations": {BOOTSTRAP: "sealed"}} + }))).await { + Ok(sealed) => { Self::ledger(&sealed)?; return Ok(sealed); } + Err(kube::Error::Api(status)) if status.code == 409 => continue, + Err(error) => return Err(api_error("seal account bootstrap", error)), + } + } + let ledger = Ledger::new(pinned_uid.into(), root.clone(), account.spec.limits)?; + match self.accounts.patch_status(&name, &PatchParams::default(), &Patch::Merge(json!({ + "metadata": {"uid": pinned_uid, "resourceVersion": account.resource_version()}, + "status": {"ledger": ledger} + }))).await { + Ok(_) => {} + Err(kube::Error::Api(status)) if status.code == 409 => continue, + Err(error) => return Err(api_error("initialize account ledger", error)), + } + } + Err(StoreError::Contention) + } + + pub async fn read( + &self, + root: &RootIdentity, + account_uid: &str, + ) -> Result { + root.validate()?; + if !crate::inference_budget_contract::valid_uid(account_uid) { + return Err(BudgetError::Identity.into()); + } + let account = self + .accounts + .get(&name_for_root(root)) + .await + .map_err(|e| api_error("read account", e))?; + Self::validate_identity(&account, root, account_uid)?; + Self::ledger(&account)?; + Ok(account) + } + + pub async fn transact( + &self, + root: &RootIdentity, + account_uid: &str, + operation: impl Fn(&Ledger) -> Result, BudgetError>, + ) -> Result { + let deadline = + tokio::time::Instant::now() + std::time::Duration::from_secs(MAX_OPERATION_SECONDS); + for _ in 0..RETRIES { + if tokio::time::Instant::now() >= deadline { + return Err(StoreError::Contention); + } + let account = tokio::time::timeout_at(deadline, self.read(root, account_uid)) + .await + .map_err(|_| StoreError::Contention)??; + let mutation = operation(Self::ledger(&account)?)?; + if !mutation.changed { + return Ok(mutation.value); + } + mutation.next.validate()?; + if mutation.next.account_uid != account_uid || mutation.next.root != *root { + return Err(BudgetError::Identity.into()); + } + if tokio::time::Instant::now() >= deadline { + return Err(StoreError::Contention); + } + // Merge-patching maps would retain compacted attempt rows. Replace + // the status as one value using PUT with UID/RV preconditions. + let mut next = account.clone(); + next.status = Some(KarsBudgetAccountStatus { + ledger: Some(mutation.next), + }); + let bytes = serde_json::to_vec(&next).map_err(|_| BudgetError::Corrupt)?; + let committed = tokio::time::timeout_at( + deadline, + self.accounts + .replace_status(&name_for_root(root), &PostParams::default(), bytes), + ) + .await + .map_err(|_| StoreError::Contention)?; + match committed { + Ok(stored) => { + Self::validate_identity(&stored, root, account_uid)?; + if Some(Self::ledger(&stored)?) + != next + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + { + return Err(BudgetError::Corrupt.into()); + } + return Ok(mutation.value); + } + Err(kube::Error::Api(status)) if status.code == 409 => continue, + Err(error) => return Err(api_error("commit ledger transition", error)), + } + } + Err(StoreError::Contention) + } +} + +#[cfg(test)] +#[path = "store_tests.rs"] +mod tests; diff --git a/controller/src/inference_budget/store_tests.rs b/controller/src/inference_budget/store_tests.rs new file mode 100644 index 000000000..cfac368b0 --- /dev/null +++ b/controller/src/inference_budget/store_tests.rs @@ -0,0 +1,475 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::inference_budget_contract::{ + Amounts, AttemptCommand, AttemptKey, BudgetScope, ExecutionIdentity, Limits, ReserveRequest, + ResourceIdentity, RootKind, TaskAuthority, + ledger::Ledger, + tariffs::{MaximumPrice, ModelContract, Operation, OutputField}, +}; +use std::sync::{Arc, Mutex}; +use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; + +const COLLECTION: &str = "/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karsbudgetaccounts"; +const OBJECT: &str = + "/apis/kars.azure.com/v1alpha1/namespaces/kars-system/karsbudgetaccounts/inference-budget-root"; + +fn root() -> RootIdentity { + RootIdentity { + kind: RootKind::KarsTask, + resource: ResourceIdentity { + namespace: "workspace".into(), + name: "root".into(), + uid: "root".into(), + }, + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + } +} + +fn identity(pod: &str) -> ExecutionIdentity { + ExecutionIdentity { + task_uid: "root".into(), + authorization_digest: authorization().authorization_digest, + sandbox: ResourceIdentity { + namespace: "workspace".into(), + name: "runtime".into(), + uid: "sandbox-uid".into(), + }, + runtime_namespace_uid: "runtime-namespace-uid".into(), + pod_name: format!("pod-{pod}"), + pod_uid: pod.into(), + } +} + +fn authorization() -> TaskAuthority { + let mut effective = json!({"domain": "test", "taskUid": "root"}); + effective.sort_all_objects(); + let digest = format!( + "sha256:{}", + crate::providers::signing::sha256_hex(&serde_json::to_vec(&effective).unwrap()) + ); + TaskAuthority { + task: root().resource, + parent_uid: None, + root_task_uid: "root".into(), + authorization_digest: digest, + effective_authorization: effective, + limits: Limits { + tokens: Some(50), + usd_micros: Some(10), + }, + } +} + +fn account() -> KarsBudgetAccount { + let spec = KarsBudgetAccountSpec { + scope: BudgetScope::GovernedInference, + root: root(), + limits: Limits { + tokens: Some(50), + usd_micros: Some(10), + }, + }; + let mut account = KarsBudgetAccount::new("inference-budget-root", spec); + account.metadata.namespace = Some("kars-system".into()); + account.metadata.uid = Some("account-uid".into()); + account.metadata.resource_version = Some("1".into()); + account.labels_mut().insert(MANAGED_BY.into(), OWNER.into()); + account + .annotations_mut() + .insert(BOOTSTRAP.into(), "sealed".into()); + let ledger = Ledger::new("account-uid".into(), root(), account.spec.limits).unwrap(); + let ledger = ledger.register_task(authorization()).unwrap().next; + let ledger = ledger.register_session(identity("a")).unwrap().next; + let ledger = ledger.register_session(identity("b")).unwrap().next; + account.status = Some(KarsBudgetAccountStatus { + ledger: Some(ledger), + }); + account +} + +fn reserve(pod: &str) -> ReserveRequest { + let contract = ModelContract { + id: "model".into(), + version: "v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "configured".into(), + endpoint: "https://configured.example".into(), + model: "model".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxTokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 5 }), + }; + let (_, quote) = contract + .normalize( + br#"{"messages":[{"role":"user","content":"hello"}]}"#, + 100, + true, + ) + .unwrap(); + ReserveRequest { + account_uid: "account-uid".into(), + identity: identity(pod), + sequence: 1, + wire_digest: format!("sha256:{}", "a".repeat(64)), + quote, + } +} + +#[derive(Clone, Copy)] +enum Fault { + None, + ConflictOnce, + CommitThenFail, + RecreateOnWrite, +} + +struct State { + account: Option, + fault: Fault, + version: u64, + writes: usize, +} + +#[derive(Clone)] +struct Server(Arc>); + +fn failure(code: u16) -> ResponseTemplate { + ResponseTemplate::new(code).set_body_json(json!({ + "apiVersion": "v1", "kind": "Status", "status": "Failure", + "reason": if code == 404 { "NotFound" } else { "Conflict" }, "code": code, + "message": "private-request-body-must-not-appear" + })) +} + +fn response(value: &KarsBudgetAccount) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(value) +} + +impl Respond for Server { + fn respond(&self, request: &Request) -> ResponseTemplate { + let mut state = self.0.lock().unwrap(); + match (request.method.as_str(), request.url.path()) { + ("GET", OBJECT) => state + .account + .as_ref() + .map_or_else(|| failure(404), response), + ("POST", COLLECTION) => { + if state.account.is_some() { + return failure(409); + } + let mut account: KarsBudgetAccount = serde_json::from_slice(&request.body).unwrap(); + // Kubernetes ignores status on the CRD main-resource CREATE. + account.status = None; + account.metadata.uid = Some("account-uid".into()); + account.metadata.namespace = Some("kars-system".into()); + account.metadata.resource_version = Some(state.version.to_string()); + let output = response(&account); + state.account = Some(account); + state.writes += 1; + output + } + ("PUT", path) if path == format!("{OBJECT}/status") => { + let mut incoming: KarsBudgetAccount = + serde_json::from_slice(&request.body).unwrap(); + match state.fault { + Fault::ConflictOnce => { + state.fault = Fault::None; + state.version += 1; + let version = state.version.to_string(); + state.account.as_mut().unwrap().metadata.resource_version = Some(version); + return failure(409); + } + Fault::RecreateOnWrite => { + state.account.as_mut().unwrap().metadata.uid = Some("replacement".into()); + } + _ => {} + } + let Some(existing) = state.account.as_ref() else { + return failure(404); + }; + if incoming.metadata.uid != existing.metadata.uid + || incoming.metadata.resource_version != existing.metadata.resource_version + { + return failure(409); + } + state.version += 1; + state.writes += 1; + incoming.metadata.resource_version = Some(state.version.to_string()); + state.account = Some(incoming.clone()); + if matches!(state.fault, Fault::CommitThenFail) { + state.fault = Fault::None; + failure(500) + } else { + response(&incoming) + } + } + ("PATCH", path) if path == OBJECT || path == format!("{OBJECT}/status") => { + let patch: serde_json::Value = serde_json::from_slice(&request.body).unwrap(); + let Some(existing) = state.account.as_ref() else { + return failure(404); + }; + if patch["metadata"]["uid"] != json!(existing.metadata.uid) + || patch["metadata"]["resourceVersion"] + != json!(existing.metadata.resource_version) + { + return failure(409); + } + state.version += 1; + state.writes += 1; + let version = state.version.to_string(); + let account = state.account.as_mut().unwrap(); + if path.ends_with("/status") { + account.status = Some(serde_json::from_value(patch["status"].clone()).unwrap()); + } else { + account.annotations_mut().insert( + BOOTSTRAP.into(), + patch["metadata"]["annotations"][BOOTSTRAP] + .as_str() + .unwrap() + .into(), + ); + } + account.metadata.resource_version = Some(version); + response(account) + } + _ => failure(500), + } + } +} + +async fn setup( + account: Option, + fault: Fault, +) -> (MockServer, Store, Arc>) { + let server = MockServer::start().await; + let state = Arc::new(Mutex::new(State { + account, + fault, + version: 100, + writes: 0, + })); + Mock::given(wiremock::matchers::any()) + .respond_with(Server(state.clone())) + .mount(&server) + .await; + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + (server, Store::new(client, "kars-system"), state) +} + +#[tokio::test] +async fn concurrent_siblings_cannot_both_reserve_past_a_shared_ceiling() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + let a = reserve("a"); + let b = reserve("b"); + let (first, second) = tokio::join!( + store.transact(&root(), "account-uid", |ledger| ledger.reserve(&a, 100)), + store.transact(&root(), "account-uid", |ledger| ledger.reserve(&b, 100)), + ); + assert_eq!(usize::from(first.is_ok()) + usize::from(second.is_ok()), 1); + assert_eq!( + state + .lock() + .unwrap() + .account + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .ledger + .as_ref() + .unwrap() + .meters + .reserved + .tokens, + 30 + ); +} + +#[tokio::test] +async fn resource_version_conflicts_retry_from_authoritative_state() { + let (_server, store, state) = setup(Some(account()), Fault::ConflictOnce).await; + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + assert_eq!(state.lock().unwrap().writes, 1); +} + +#[tokio::test] +async fn lost_reserve_ack_recovers_without_allocating_twice() { + let (_server, store, state) = setup(Some(account()), Fault::CommitThenFail).await; + let request = reserve("a"); + assert!( + store + .transact(&root(), "account-uid", |ledger| ledger + .reserve(&request, 100)) + .await + .is_err() + ); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + assert_eq!(state.lock().unwrap().writes, 1); +} + +#[tokio::test] +async fn lost_begin_ack_never_reissues_permission_to_send() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + state.lock().unwrap().fault = Fault::CommitThenFail; + let command = AttemptCommand { + account_uid: "account-uid".into(), + key: AttemptKey { + pod_uid: "a".into(), + sequence: 1, + }, + identity: request.identity, + wire_digest: request.wire_digest, + }; + assert!( + store + .transact(&root(), "account-uid", |ledger| ledger + .begin_dispatch(&command, 101)) + .await + .is_err() + ); + assert!(matches!( + store + .transact(&root(), "account-uid", |ledger| ledger + .begin_dispatch(&command, 101)) + .await, + Err(StoreError::Ledger(BudgetError::AlreadyDispatched)) + )); + assert_eq!( + state + .lock() + .unwrap() + .account + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .ledger + .as_ref() + .unwrap() + .meters + .reserved, + Amounts { + tokens: 30, + usd_micros: 5 + } + ); +} + +#[tokio::test] +async fn missing_recreated_and_corrupt_accounts_are_never_reinitialized_by_requests() { + for case in ["missing", "uid", "status", "counters"] { + let mut account = account(); + if case == "uid" { + account.metadata.uid = Some("replacement".into()); + } + if case == "status" { + account.status = None; + } + if case == "counters" { + account + .status + .as_mut() + .unwrap() + .ledger + .as_mut() + .unwrap() + .meters + .reserved + .tokens = 1; + } + let (server, store, state) = setup( + if case == "missing" { + None + } else { + Some(account) + }, + Fault::None, + ) + .await; + assert!(store.read(&root(), "account-uid").await.is_err(), "{case}"); + assert_eq!(state.lock().unwrap().writes, 0); + assert!( + server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.method == "GET") + ); + } +} + +#[tokio::test] +async fn bootstrap_is_metadata_first_and_sealed_only_after_uid_bound_status() { + let (_server, store, _) = setup(None, Fault::None).await; + let spec = account().spec; + let anchor = store + .create_anchor( + spec, + &crate::providers::signing::ReceiptSigner::from_bytes(&[7; 32]), + ) + .await + .unwrap(); + assert!(store.read(&root(), "account-uid").await.is_err()); + let initialized = store + .initialize(&root(), anchor.metadata.uid.as_deref().unwrap()) + .await + .unwrap(); + assert_eq!( + initialized.annotations().get(BOOTSTRAP).map(String::as_str), + Some("sealed") + ); + assert!( + initialized + .status + .as_ref() + .unwrap() + .ledger + .as_ref() + .unwrap() + .nodes + .is_empty() + ); +} + +#[tokio::test] +async fn replacement_between_read_and_commit_cannot_receive_a_grant() { + let (_server, store, state) = setup(Some(account()), Fault::RecreateOnWrite).await; + let request = reserve("a"); + assert!( + store + .transact(&root(), "account-uid", |ledger| ledger + .reserve(&request, 100)) + .await + .is_err() + ); + assert_eq!(state.lock().unwrap().writes, 0); +} diff --git a/controller/src/inference_budget/team.rs b/controller/src/inference_budget/team.rs new file mode 100644 index 000000000..46e881092 --- /dev/null +++ b/controller/src/inference_budget/team.rs @@ -0,0 +1,80 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + binding, + config::Settings, + store::{Store, StoreError}, +}; +use crate::{ + inference_budget_contract::{AccountPhase, BudgetError, RootKind}, + kars_task::KarsTask, + kars_team::KarsTeam, +}; +use kube::{Api, Client, ResourceExt}; + +/// Cadence must wait for the principal's protected immutable account pin. +/// This checks accounting readiness, not completion of a business operation. +pub async fn ready(client: &Client, team: &KarsTeam, principal: &str) -> Result<(), StoreError> { + if !binding::has_finite(&team.spec.envelope) { + return Ok(()); + } + let limits = binding::limits(&team.spec.envelope)?; + let settings = Settings::from_env()?.ok_or(BudgetError::Contract)?; + settings + .catalog(client, chrono::Utc::now().timestamp()) + .await?; + let tasks: Api = Api::namespaced( + client.clone(), + &team.namespace().ok_or(BudgetError::Identity)?, + ); + let principal = tasks + .get(principal) + .await + .map_err(|error| StoreError::Api { + stage: "read finite Team principal", + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + })?; + let binding = principal + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + .ok_or(StoreError::Missing)?; + if !crate::kars_task_reconciler::task_is_ready(&principal) + || binding.root.kind != RootKind::KarsTeam + || team.metadata.uid.as_deref() != Some(binding.root.resource.uid.as_str()) + || team.metadata.namespace.as_deref() != Some(binding.root.resource.namespace.as_str()) + || team.metadata.name.as_deref() != Some(binding.root.resource.name.as_str()) + || binding.account.namespace != settings.accounting_namespace + { + return Err(BudgetError::Authorization.into()); + } + let account = Store::new(client.clone(), &settings.accounting_namespace) + .read(&binding.root, &binding.account.uid) + .await?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + if ledger.phase != AccountPhase::Active || ledger.limits != limits { + return Err(BudgetError::Closed.into()); + } + let used = ledger.meters.total()?; + if limits.tokens.is_some_and(|limit| used.tokens >= limit) + || limits + .usd_micros + .is_some_and(|limit| used.usd_micros >= limit) + { + return Err(BudgetError::Exhausted.into()); + } + if ledger.nodes.len() >= crate::inference_budget_contract::MAX_NODES + || ledger.sessions.len() >= crate::inference_budget_contract::MAX_SESSIONS + { + return Err(BudgetError::Capacity.into()); + } + Ok(()) +} diff --git a/controller/src/inference_budget/transport.rs b/controller/src/inference_budget/transport.rs new file mode 100644 index 000000000..4aab858de --- /dev/null +++ b/controller/src/inference_budget/transport.rs @@ -0,0 +1,98 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + config::Settings, + service::{self, Broker}, + store::Store, +}; +use k8s_openapi::api::core::v1::Secret; +use kube::{Api, Client, ResourceExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio_rustls::{TlsAcceptor, server::TlsStream}; + +struct Listener { + listener: TcpListener, + tls: TlsAcceptor, +} + +impl axum::serve::Listener for Listener { + type Io = TlsStream; + type Addr = std::net::SocketAddr; + + async fn accept(&mut self) -> (Self::Io, Self::Addr) { + loop { + let Ok((stream, address)) = self.listener.accept().await else { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + continue; + }; + if let Ok(Ok(tls)) = + tokio::time::timeout(std::time::Duration::from_secs(10), self.tls.accept(stream)) + .await + { + return (tls, address); + } + tracing::warn!("Inference budget TLS connection rejected"); + } + } + + fn local_addr(&self) -> std::io::Result { + self.listener.local_addr() + } +} + +pub async fn server_identity( + client: &Client, + settings: &Settings, +) -> anyhow::Result> { + crate::sre_authority::privacy_epoch(client, &settings.accounting_namespace) + .await + .map_err(|_| anyhow::anyhow!("Inference budget TLS privacy proof is unavailable"))?; + let secrets: Api = Api::namespaced(client.clone(), &settings.accounting_namespace); + let secret = secrets + .get(&settings.tls_secret) + .await + .map_err(|_| anyhow::anyhow!("Inference budget TLS Secret is unavailable"))?; + if secret.type_.as_deref() != Some("kubernetes.io/tls") + || secret.metadata.uid.as_deref().is_none_or(str::is_empty) + || secret.metadata.deletion_timestamp.is_some() + || secret + .annotations() + .get("kars.azure.com/inference-budget-tls") + .map(String::as_str) + != Some("v1") + { + anyhow::bail!("Inference budget TLS Secret identity is invalid"); + } + let data = secret + .data + .as_ref() + .ok_or_else(|| anyhow::anyhow!("Inference budget TLS material is absent"))?; + let cert = data + .get("tls.crt") + .ok_or_else(|| anyhow::anyhow!("Inference budget certificate is absent"))?; + let key = data + .get("tls.key") + .ok_or_else(|| anyhow::anyhow!("Inference budget key is absent"))?; + crate::providers::signing::inference_budget_tls_config(&cert.0, &key.0) +} + +pub async fn run(client: Client, settings: Settings) -> anyhow::Result<()> { + settings + .catalog(&client, chrono::Utc::now().timestamp()) + .await?; + let tls = server_identity(&client, &settings).await?; + let socket = TcpListener::bind(&settings.address).await?; + let listener = Listener { + listener: socket, + tls: TlsAcceptor::from(tls), + }; + let store = Store::new(client.clone(), &settings.accounting_namespace); + let router = service::router(Broker { + client, + settings, + store, + }); + axum::serve(listener, router).await?; + Ok(()) +} diff --git a/controller/src/kars_profile.rs b/controller/src/kars_profile.rs index 25985f459..749945337 100644 --- a/controller/src/kars_profile.rs +++ b/controller/src/kars_profile.rs @@ -216,6 +216,7 @@ mod tests { p.spec.default_envelope.budget = Some(crate::kars_task::TaskBudget { tokens: Some(1000), usd_micros: Some(2000), + ..Default::default() }); p.spec.default_envelope.tool_policy_ref = Some(crate::mcp_server::LocalObjectRef { name: "bounded-tools".into(), diff --git a/controller/src/kars_receipt_launch.rs b/controller/src/kars_receipt_launch.rs index b0be01fc6..c3097b8bf 100644 --- a/controller/src/kars_receipt_launch.rs +++ b/controller/src/kars_receipt_launch.rs @@ -130,6 +130,7 @@ mod tests { budget: Some(TaskBudget { tokens: Some(1000), usd_micros: Some(2000), + ..Default::default() }), tool_policy_ref: Some(LocalObjectRef { name: "bounded-tools".into(), diff --git a/controller/src/kars_task.rs b/controller/src/kars_task.rs index 6a7ed088b..cd6760b5d 100644 --- a/controller/src/kars_task.rs +++ b/controller/src/kars_task.rs @@ -646,8 +646,10 @@ pub fn validate_execution_contract(spec: &KarsTaskSpec) -> Result<(), String> { } if spec.execution.as_ref().is_some_and(|e| e.launch) && (budget.tokens.is_some_and(|n| n > 0) || budget.usd_micros.is_some_and(|n| n > 0)) + && budget.scope + != Some(crate::inference_budget_contract::BudgetScope::GovernedInference) { - return Err("UnsupportedLaunchBudget: total/subtree token and usdMicros ceilings are not enforced by this foundation; bounded tasks may be planned but cannot launch".into()); + return Err("UnsupportedLaunchBudget: legacy total/subtree budgets are planning-only; a new GovernedInference-scoped task requires the configured durable broker before materialization".into()); } } Ok(()) @@ -702,6 +704,10 @@ pub fn spec_attenuation_violations( #[derive(Debug, Serialize, Deserialize, Default, Clone, JsonSchema)] #[serde(rename_all = "camelCase")] pub struct TaskBudget { + /// Explicit opt-in to governed-inference tokens/configured maximum prices. + /// An absent scope retains the foundation's planning-only interpretation. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub scope: Option, /// Maximum total tokens the task subtree may consume. `0`/absent means /// "no token cap declared". Positive ceilings are planning declarations: /// launch is rejected until durable total/subtree enforcement is available. @@ -719,6 +725,10 @@ pub struct TaskBudget { #[derive(Debug, Serialize, Deserialize, Default, Clone, JsonSchema)] #[serde(rename_all = "camelCase")] pub struct KarsTaskStatus { + /// Controller-owned immutable account/UID ancestry binding. Not a task-name + /// metering label and never supplied by an agent header. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inference_budget: Option, /// One of: `Pending`, `Ready`, `Degraded`. #[serde(default, skip_serializing_if = "Option::is_none")] pub phase: Option, diff --git a/controller/src/kars_task_authorization_tests.rs b/controller/src/kars_task_authorization_tests.rs index 1e2a793bd..8025822c9 100644 --- a/controller/src/kars_task_authorization_tests.rs +++ b/controller/src/kars_task_authorization_tests.rs @@ -205,6 +205,7 @@ fn defaults_aliases_and_runtime_precedence_have_one_canonical_digest() { explicit.envelope.budget = Some(TaskBudget { tokens: Some(0), usd_micros: Some(0), + ..Default::default() }); assert_eq!(explicit.authorization_digest_with_model(&model()), digest); explicit.blueprint.as_mut().unwrap().runtime = Some("MAF".into()); @@ -225,6 +226,7 @@ fn shared_authorization_snapshot_exposes_the_exact_effective_digest_input() { task.envelope.budget = Some(TaskBudget { tokens: Some(0), usd_micros: Some(0), + ..Default::default() }); let configuration = task.authorization_configuration_with_model(&model()); assert_eq!( diff --git a/controller/src/kars_task_execution.rs b/controller/src/kars_task_execution.rs index 9e44e0f59..ab55ff1c8 100644 --- a/controller/src/kars_task_execution.rs +++ b/controller/src/kars_task_execution.rs @@ -137,6 +137,9 @@ pub async fn materialize( let envelope = &task.spec.envelope; let blueprint = crate::kars_task::blueprint::effective_blueprint(&task.spec); let runtime = runtime_spec(task)?; + let prepared = crate::inference_budget::binding::prepare_task(client, task) + .await + .map_err(|error| contract_error(error.to_string()))?; // 1. InferencePolicy scoped to this sandbox. Model: blueprint wins, else // the controller default (required — without it the sandbox degrades). @@ -160,6 +163,13 @@ pub async fn materialize( "sandbox": { "isolation": blueprint.isolation }, "networkPolicy": network_policy(&blueprint), }); + if let Some(binding) = prepared + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + { + sandbox_spec["inferenceBudgetRef"] = json!(binding); + } // Agent instructions (the system prompt) — combine the objective with any // standing instructions the blueprint carries, so the agent knows both diff --git a/controller/src/kars_task_reconciler.rs b/controller/src/kars_task_reconciler.rs index 3a721f5d5..95be75c9e 100644 --- a/controller/src/kars_task_reconciler.rs +++ b/controller/src/kars_task_reconciler.rs @@ -44,6 +44,8 @@ const REQUEUE_PENDING: Duration = Duration::from_secs(10); #[derive(Debug, thiserror::Error)] enum ReconcileError { + #[error(transparent)] + InferenceBudget(#[from] crate::inference_budget::store::StoreError), #[error("Kubernetes API error: {0}")] Kube(#[from] kube::Error), #[error("JSON serialization error: {0}")] @@ -53,6 +55,7 @@ enum ReconcileError { impl ReconcileError { fn class(&self) -> &'static str { match self { + ReconcileError::InferenceBudget(_) => "inference_budget", ReconcileError::Kube(_) => "kube_api", ReconcileError::SerdeJson(_) => "serde", } @@ -108,6 +111,7 @@ struct Ctx { } async fn reconcile(task: Arc, ctx: Arc) -> Result { + let mut task = task; let name = task.name_any(); let ns = task.namespace().unwrap_or_else(|| "default".into()); let tasks: Api = Api::namespaced(ctx.client.clone(), &ns); @@ -115,6 +119,7 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result task = Arc::new(prepared), + Err(error) => { + new_status = degraded_status( + prior_ready, + generation, + &format!("Governed inference unavailable: {error}"), + new_status.lineage.clone(), + ); + // Preparation can pin an account before a later CAS failure. + // Preserve the pin and use the current UID/RV for status only. + let live = tasks.get(&name).await?; + if live.metadata.uid != task.metadata.uid + || live.metadata.generation != task.metadata.generation + { + return Ok(Action::requeue(REQUEUE_PENDING)); + } + task = Arc::new(live); + } + } + } + new_status.inference_budget = task + .status + .as_ref() + .and_then(|status| status.inference_budget.clone()); + // Execution bridge (§20 launch gate). Only a governance-Ready task may // execute. Launch materializes a governed sandbox; un-launch tears it down. // Any execution error is surfaced (Degraded) but never fails the whole @@ -263,7 +295,9 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result TaskEnvelope { budget: Some(TaskBudget { tokens: Some(100_000), usd_micros: Some(5_000_000), + ..Default::default() }), tool_policy_ref: Some(LocalObjectRef { name: "default-tools".into(), @@ -50,10 +51,12 @@ fn bounded_launch_fails_closed_but_planning_remains_available() { TaskBudget { tokens: Some(100), usd_micros: None, + ..Default::default() }, TaskBudget { tokens: None, usd_micros: Some(100), + ..Default::default() }, ] { let mut spec = KarsTaskSpec::default(); @@ -80,6 +83,7 @@ fn bounded_launch_fails_closed_but_planning_remains_available() { spec.envelope.budget = Some(TaskBudget { tokens: Some(0), usd_micros: Some(0), + ..Default::default() }); assert!(validate_execution_contract(&spec).is_ok()); } @@ -178,6 +182,7 @@ fn parent_envelope() -> TaskEnvelope { budget: Some(TaskBudget { tokens: Some(1_000_000), usd_micros: Some(50_000_000), + ..Default::default() }), tool_policy_ref: Some(LocalObjectRef { name: "strict-tools".into(), @@ -196,6 +201,7 @@ fn valid_child_attenuates_on_every_axis() { budget: Some(TaskBudget { tokens: Some(100_000), usd_micros: Some(5_000_000), + ..Default::default() }), tool_policy_ref: Some(LocalObjectRef { name: "strict-tools".into(), @@ -294,6 +300,7 @@ fn child_budget_over_parent_cap_is_amplification() { child.budget = Some(TaskBudget { tokens: Some(2_000_000), usd_micros: Some(1_000_000), + ..Default::default() }); assert!( child @@ -420,6 +427,7 @@ fn child_envelope() -> TaskEnvelope { budget: Some(TaskBudget { tokens: Some(100_000), usd_micros: Some(5_000_000), + ..Default::default() }), tool_policy_ref: Some(LocalObjectRef { name: "strict-tools".into(), diff --git a/controller/src/kars_team.rs b/controller/src/kars_team.rs index 9d4d94162..87475c51a 100644 --- a/controller/src/kars_team.rs +++ b/controller/src/kars_team.rs @@ -185,6 +185,9 @@ pub struct TeamCadence { #[derive(Debug, Serialize, Deserialize, Default, Clone, JsonSchema)] #[serde(rename_all = "camelCase")] pub struct KarsTeamStatus { + /// Lifetime Team-UID inference account; cadence runs never reset its balance. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub inference_budget_account: Option, /// Lifecycle phase: `Forming` (validating + materializing), `Active` /// (running, cadence ticking), `Hibernating` (paused/idle), `Degraded` /// (envelope invalid — no authority to operate), `Retired`. @@ -370,6 +373,7 @@ mod tests { budget: Some(TaskBudget { tokens: Some(1_000_000), usd_micros: None, + ..Default::default() }), tool_policy_ref: None, egress_allowlist_ref: None, @@ -413,6 +417,7 @@ mod tests { budget: Some(TaskBudget { tokens: Some(100_000), usd_micros: None, + ..Default::default() }), tool_policy_ref: None, egress_allowlist_ref: None, diff --git a/controller/src/kars_team_reconciler.rs b/controller/src/kars_team_reconciler.rs index ea9240933..8b33ded5a 100644 --- a/controller/src/kars_team_reconciler.rs +++ b/controller/src/kars_team_reconciler.rs @@ -212,7 +212,19 @@ async fn reconcile_valid( .cadence .as_ref() .and_then(|cadence| cadence.every_minutes); - let bounded_plan = specs::has_positive_budget(&team.spec.envelope); + let finite = specs::has_positive_budget(&team.spec.envelope); + let unsupported = specs::unsupported_budget(&team.spec.envelope); + let budget_error = if finite && !unsupported { + crate::inference_budget::team::ready(client, team, &principal_name) + .await + .err() + } else { + None + }; + if budget_error.is_some() { + tasks::revoke_all(tasks_api, team).await?; + } + let bounded_plan = unsupported || budget_error.is_some(); let cadence_blocked = bounded_plan && every.is_some() && !team.spec.paused; let mut generated = prior.generated_task_count; let mut last_generated = prior.last_generated_task.clone(); @@ -307,8 +319,16 @@ async fn reconcile_valid( } let detail = if team.spec.paused { "Team hibernating — members and runs governed-but-idle; charter loop paused.".into() - } else if bounded_plan { + } else if let Some(error) = &budget_error { + format!( + "Governed inference accounting unavailable: {error}. No cadence or finite execution is permitted." + ) + } else if unsupported { "UnsupportedLaunchBudget: Team and member plans are governed-but-idle. Finite total/subtree token or monetary budgets require durable enforcement; cadence and bounded execution are unavailable.".into() + } else if finite { + format!( + "Governed inference account bound to Team UID for its lifetime. Standing operation {health}; compute, tool, storage and invoice costs are excluded." + ) } else { format!( "Standing operation {health}: {generated} run(s), {} delivered, {entries} knowledge entries.", @@ -347,6 +367,7 @@ async fn reconcile_valid( commons_entry_count: Some(entries), last_success_at, last_digest_at, + inference_budget_account: prior.inference_budget_account.clone(), ..Default::default() }, ) diff --git a/controller/src/kars_team_reconciler/specs.rs b/controller/src/kars_team_reconciler/specs.rs index 4faf99ee5..b232d7a24 100644 --- a/controller/src/kars_team_reconciler/specs.rs +++ b/controller/src/kars_team_reconciler/specs.rs @@ -10,8 +10,7 @@ use crate::mcp_server::LocalObjectRef; /// Matches the existing KarsTask objective CEL rule; covered by a drift test. pub(crate) const MAX_OBJECTIVE_CHARS: usize = 4096; -/// The foundation cannot enforce durable total/subtree or monetary budgets. -/// Finite budgets are valid plans, but must never become running tasks. +/// Positive ceilings require the explicit governed-inference capability. pub(crate) fn has_positive_budget(envelope: &TaskEnvelope) -> bool { envelope.budget.as_ref().is_some_and(|budget| { budget.tokens.is_some_and(|value| value > 0) @@ -19,6 +18,10 @@ pub(crate) fn has_positive_budget(envelope: &TaskEnvelope) -> bool { }) } +pub(crate) fn unsupported_budget(envelope: &TaskEnvelope) -> bool { + crate::inference_budget::scope::unsupported(envelope) +} + pub(crate) fn default_member_envelope(parent: &TaskEnvelope) -> TaskEnvelope { let tier = parent .tier diff --git a/controller/src/kars_team_reconciler/tasks.rs b/controller/src/kars_team_reconciler/tasks.rs index 305e55420..4415b63a9 100644 --- a/controller/src/kars_team_reconciler/tasks.rs +++ b/controller/src/kars_team_reconciler/tasks.rs @@ -151,7 +151,7 @@ pub(super) async fn reconcile_revocations( if !authorized { retire(tasks, task).await?; } else if team.spec.paused - || specs::has_positive_budget(&task.spec.envelope) + || specs::unsupported_budget(&task.spec.envelope) || (role == Some("principal") && !within_seat(&task.spec, &principal)) { idle(tasks, task).await?; @@ -179,7 +179,7 @@ pub(super) async fn apply_task( mut spec: KarsTaskSpec, role: &str, ) -> Result { - if role == "taskforce" && specs::has_positive_budget(&spec.envelope) { + if role == "taskforce" && specs::unsupported_budget(&spec.envelope) { return Err(ReconcileError::Invalid( "UnsupportedLaunchBudget: finite total/subtree and monetary budgets are planning-only until durable enforcement is available".into(), )); diff --git a/controller/src/kars_team_reconciler/tests.rs b/controller/src/kars_team_reconciler/tests.rs index bdb7552f4..c12144159 100644 --- a/controller/src/kars_team_reconciler/tests.rs +++ b/controller/src/kars_team_reconciler/tests.rs @@ -70,6 +70,7 @@ pub(super) fn team() -> KarsTeam { budget: Some(TaskBudget { tokens: Some(1_000), usd_micros: Some(2_000), + ..Default::default() }), ..Default::default() }, @@ -407,6 +408,7 @@ fn only_finite_positive_budgets_block_execution_not_planning() { team.spec.envelope.budget = Some(TaskBudget { tokens: Some(0), usd_micros: None, + ..Default::default() }); assert!(!specs::has_positive_budget(&team.spec.envelope)); team.spec.envelope.budget.as_mut().unwrap().usd_micros = Some(1); diff --git a/controller/src/main.rs b/controller/src/main.rs index 6a8a9bb4f..dea825dc8 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -36,6 +36,9 @@ mod fedcred; mod fedcred_reaper; mod field_managers; mod helm_drift; +mod inference_budget; +#[path = "../../shared/inference_budget/mod.rs"] +mod inference_budget_contract; mod inference_policy; mod inference_policy_compile; mod inference_policy_reconciler; @@ -133,6 +136,7 @@ async fn main() -> Result<()> { ); let client = Client::try_default().await?; + inference_budget::start(client.clone()); // S7.E: Prometheus + health server. Default ON; opt out via // `CONTROLLER_METRICS_ADDR=disabled` (or empty). Failures here are diff --git a/controller/src/providers/signing.rs b/controller/src/providers/signing.rs index 5225bc718..750aaffd5 100644 --- a/controller/src/providers/signing.rs +++ b/controller/src/providers/signing.rs @@ -38,7 +38,7 @@ use anyhow::{Context, Result}; use base64::{Engine as _, engine::general_purpose::STANDARD as BASE64}; -use ed25519_dalek::{Signer, SigningKey, VerifyingKey}; +use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey}; use k8s_openapi::api::core::v1::{ConfigMap, Secret}; use kube::{ Client, @@ -153,6 +153,19 @@ impl ReceiptSigner { pub fn sign_note(&self, note: &[u8]) -> String { BASE64.encode(self.signing_key.sign(note).to_bytes()) } + + pub fn verify_note(&self, note: &[u8], signature: &str) -> bool { + let Ok(bytes) = BASE64.decode(signature) else { + return false; + }; + let Ok(signature) = Signature::from_slice(&bytes) else { + return false; + }; + self.signing_key + .verifying_key() + .verify(note, &signature) + .is_ok() + } } /// Hex SHA-256 fingerprint of an Ed25519 public key. @@ -172,6 +185,48 @@ pub fn sha256_hex(bytes: &[u8]) -> String { out } +/// Standard rustls server configuration for operator-provided broker TLS +/// material. No private key or PEM parse error is returned to callers/logs. +pub fn inference_budget_tls_config( + cert: &[u8], + key: &[u8], +) -> Result> { + let certificates = rustls_pemfile::certs(&mut std::io::Cursor::new(cert)) + .collect::, _>>() + .map_err(|_| anyhow::anyhow!("Inference budget TLS certificate is invalid"))?; + if certificates.is_empty() { + anyhow::bail!("Inference budget TLS certificate is empty"); + } + let private = rustls_pemfile::private_key(&mut std::io::Cursor::new(key)) + .map_err(|_| anyhow::anyhow!("Inference budget TLS key is invalid"))? + .ok_or_else(|| anyhow::anyhow!("Inference budget TLS key is missing"))?; + let config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certificates, private) + .map_err(|_| anyhow::anyhow!("Inference budget TLS identity is invalid"))?; + Ok(std::sync::Arc::new(config)) +} + +/// Read the existing controller identity without creating, rotating or +/// publishing anything. Callers must establish their current privacy proof. +pub async fn load_existing(client: &Client) -> Result { + let secrets: Api = Api::namespaced(client.clone(), &receipt_namespace()); + let secret = secrets + .get(IDENTITY_SECRET_NAME) + .await + .map_err(|_| anyhow::anyhow!("Controller signing identity unavailable"))?; + if secret.metadata.deletion_timestamp.is_some() || secret.metadata.uid.is_none() { + anyhow::bail!("Controller signing identity is not live"); + } + let bytes = secret + .data + .as_ref() + .and_then(|data| data.get("signing_key")) + .and_then(|bytes| <[u8; 32]>::try_from(bytes.0.as_slice()).ok()) + .ok_or_else(|| anyhow::anyhow!("Controller signing identity malformed"))?; + Ok(ReceiptSigner::from_bytes(&bytes)) +} + /// Operator-only service credential, generated by the standard CSPRNG. pub fn generate_service_token() -> String { use rand::distr::{Alphanumeric, SampleString}; diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index d7dabe3cc..4efd6dc12 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -63,6 +63,8 @@ fn sandbox_node_selector(default_pool: &str) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result Duration { // Transient kube API errors (throttling, connection reset, 5xx): // retry soon so we don't starve legitimate work. ReconcileError::Kube(_) + | ReconcileError::InferenceBudget(_) | ReconcileError::NamespaceOwnership(_) | ReconcileError::Credentials(_) => 30, // Serde errors are deterministic — the same body will fail again. @@ -3170,6 +3165,7 @@ fn error_policy(sandbox: Arc, error: &ReconcileError, _ctx: Arc "configuration", ReconcileError::NamespaceOwnership(_) => "namespace_ownership", ReconcileError::Credentials(_) => "credentials", + ReconcileError::InferenceBudget(_) => "inference_budget", }; crate::metrics::record_reconcile_error("KarsSandbox", class); tracing::error!( diff --git a/deploy/helm/kars/files/inference-budget-admission.json b/deploy/helm/kars/files/inference-budget-admission.json new file mode 100644 index 000000000..b2a01b42e --- /dev/null +++ b/deploy/helm/kars/files/inference-budget-admission.json @@ -0,0 +1,156 @@ +{ + "items": [ + { + "name": "kars-inference-budget-public-ca", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": [""], "apiVersions": ["v1"], + "operations": ["CREATE", "UPDATE"], "resources": ["configmaps"] + }] + }, + "validations": [{ + "expression": "request.name != 'kars-inference-budget-ca' || request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller'", + "message": "The router's budget CA projection is controller-owned; replacing it could forge broker grants" + }] + } + }, + { + "name": "kars-inference-budget-store", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": ["kars.azure.com"], "apiVersions": ["v1alpha1"], + "operations": ["CREATE", "UPDATE", "DELETE"], + "resources": ["karsbudgetaccounts", "karsbudgetaccounts/status"] + }] + }, + "validations": [{ + "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller' || (request.operation == 'DELETE' && request.userInfo.username in ['system:kube-controller-manager', 'system:serviceaccount:kube-system:namespace-controller'])", + "message": "Governed inference ledgers are controller-owned and cannot be reset, deleted or rewritten by routers or agents" + }] + } + }, + { + "name": "kars-inference-budget-status", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": ["kars.azure.com"], "apiVersions": ["v1alpha1"], + "operations": ["UPDATE"], "resources": ["karstasks/status", "karsteams/status"] + }] + }, + "validations": [{ + "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller'", + "message": "Only the controller may publish UID-bound budget and task authority" + }] + } + }, + { + "name": "kars-inference-budget-runtime-binding", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": ["kars.azure.com"], "apiVersions": ["v1alpha1"], + "operations": ["CREATE", "UPDATE"], "resources": ["karssandboxes"] + }] + }, + "validations": [{ + "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller' || (!has(object.spec.inferenceBudgetRef) && (oldObject == null || !has(oldObject.spec.inferenceBudgetRef)))", + "message": "Finite runtime authority and pod configuration are controller-only; edit the owning Task instead" + }] + } + }, + { + "name": "kars-inference-budget-namespace", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": [""], "apiVersions": ["v1"], + "operations": ["CREATE", "UPDATE"], "resources": ["namespaces"] + }] + }, + "variables": [{ + "name": "oldBound", + "expression": "oldObject != null && has(oldObject.metadata.labels) && 'kars.azure.com/inference-budget' in oldObject.metadata.labels" + }, { + "name": "newBound", + "expression": "has(object.metadata.labels) && 'kars.azure.com/inference-budget' in object.metadata.labels" + }], + "validations": [{ + "expression": "!variables.oldBound || (variables.newBound && object.metadata.labels['kars.azure.com/inference-budget'] == oldObject.metadata.labels['kars.azure.com/inference-budget'])", + "message": "An enforced runtime namespace cannot lose its private-budget admission fence" + }, { + "expression": "!variables.newBound || variables.oldBound || request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller'", + "message": "Only the core controller can activate the budget namespace fence" + }] + } + }, + { + "name": "kars-inference-budget-workloads", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "namespaceSelector": {"matchLabels": {"kars.azure.com/inference-budget": "v1"}}, + "resourceRules": [{ + "apiGroups": [""], "apiVersions": ["v1"], + "operations": ["CREATE", "UPDATE"], "resources": ["pods", "pods/ephemeralcontainers"] + }, { + "apiGroups": ["apps"], "apiVersions": ["v1"], + "operations": ["CREATE", "UPDATE"], "resources": ["deployments", "replicasets"] + }] + }, + "validations": [{ + "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller' || (request.resource.resource != 'deployments' && request.subResource != 'ephemeralcontainers' && request.userInfo.username in ['system:kube-controller-manager', 'system:serviceaccount:kube-system:deployment-controller', 'system:serviceaccount:kube-system:replicaset-controller'])", + "message": "Budget runtime workloads may be produced only by the core controller and Kubernetes Deployment/ReplicaSet controllers" + }] + } + }, + { + "name": "kars-inference-budget-token", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "resourceRules": [{ + "apiGroups": [""], "apiVersions": ["v1"], + "operations": ["CREATE"], "resources": ["serviceaccounts/token"] + }] + }, + "validations": [{ + "expression": "!object.spec.audiences.exists(a, a == 'kars.azure.com/governed-inference-budget') || ('system:nodes' in request.userInfo.groups && has(object.spec.boundObjectRef) && object.spec.boundObjectRef.kind == 'Pod' && has(object.spec.boundObjectRef.uid))", + "message": "Only kubelet node identities may obtain a Pod-bound governed-inference audience token" + }] + } + }, + { + "name": "kars-inference-budget-private-connect", + "spec": { + "failurePolicy": "Fail", + "matchConstraints": { + "matchPolicy": "Equivalent", + "namespaceSelector": {"matchLabels": {"kars.azure.com/inference-budget": "v1"}}, + "resourceRules": [{ + "apiGroups": [""], "apiVersions": ["v1"], + "operations": ["CONNECT"], "resources": ["pods/exec", "pods/attach"] + }] + }, + "validations": [{ + "expression": "false", + "message": "Exec/attach would expose the private budget audience token; use authenticated router APIs or port-forward" + }] + } + } + ] +} diff --git a/deploy/helm/kars/templates/controller-deployment.yaml b/deploy/helm/kars/templates/controller-deployment.yaml index f7337d57f..085969041 100644 --- a/deploy/helm/kars/templates/controller-deployment.yaml +++ b/deploy/helm/kars/templates/controller-deployment.yaml @@ -37,6 +37,18 @@ spec: env: - name: RUST_LOG value: "info,kars_controller=debug" + {{- if ((.Values.inferenceBudget | default dict).enabled | default false) }} + - name: KARS_INFERENCE_BUDGET_ENABLED + value: "true" + - name: KARS_INFERENCE_BUDGET_CATALOG + value: "kars-inference-budget-contracts" + - name: KARS_INFERENCE_BUDGET_TLS_SECRET + value: {{ required "inferenceBudget.tlsSecretName is required" .Values.inferenceBudget.tlsSecretName | quote }} + - name: KARS_INFERENCE_BUDGET_ROUTER_DIGEST + value: {{ required "inferenceBudget.routerImageDigest is required" .Values.inferenceBudget.routerImageDigest | quote }} + - name: KARS_INFERENCE_BUDGET_ADDR + value: "0.0.0.0:9447" + {{- end }} # Downward API — used by leader-election + S12.d # SignerPolicy watcher to scope to the controller's own # namespace without cluster-wide watch RBAC. diff --git a/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml b/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml new file mode 100644 index 000000000..7fd263daf --- /dev/null +++ b/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml @@ -0,0 +1,115 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. +apiVersion: apiextensions.k8s.io/v1 +kind: CustomResourceDefinition +metadata: + name: karsbudgetaccounts.kars.azure.com + labels: + app.kubernetes.io/name: kars + app.kubernetes.io/component: crd + annotations: + helm.sh/resource-policy: keep +spec: + group: kars.azure.com + scope: Namespaced + names: + kind: KarsBudgetAccount + plural: karsbudgetaccounts + singular: karsbudgetaccount + shortNames: [kbudget] + versions: + - name: v1alpha1 + served: true + storage: true + subresources: + status: {} + additionalPrinterColumns: + - name: Scope + type: string + jsonPath: .spec.scope + - name: Phase + type: string + jsonPath: .status.ledger.phase + schema: + openAPIV3Schema: + type: object + required: [spec] + properties: + apiVersion: {type: string} + kind: {type: string} + metadata: {type: object} + spec: + type: object + required: [scope, root, limits] + x-kubernetes-validations: + - rule: "self == oldSelf" + message: "Budget account grant/root identity is immutable; narrowing is an authoritative ledger transition" + properties: + scope: + type: string + enum: [GovernedInference] + description: "Tokens and configured maximum prices for governed inference only, not all-in task, compute, tool, or invoice costs" + root: + type: object + required: [kind, resource, workspaceUid, clusterUid] + properties: + kind: {type: string, enum: [KarsTask, KarsTeam]} + workspaceUid: {type: string, minLength: 1, maxLength: 128} + clusterUid: {type: string, minLength: 1, maxLength: 128} + resource: + type: object + required: [namespace, name, uid] + properties: + namespace: {type: string, minLength: 1, maxLength: 63} + name: {type: string, minLength: 1, maxLength: 253} + uid: {type: string, minLength: 1, maxLength: 128} + limits: + type: object + properties: + tokens: {type: integer, format: int64, minimum: 0} + usdMicros: + type: integer + format: int64 + minimum: 0 + description: "Configured maximum governed-inference price units; zero/absent is unbounded" + status: + type: object + properties: + ledger: + type: object + x-kubernetes-preserve-unknown-fields: true + required: [version, scope, accountUid, root, limits, phase, meters, nodes, sessions, attempts] + properties: + version: {type: string, enum: [governed-inference/v1]} + scope: {type: string, enum: [GovernedInference]} + accountUid: {type: string, minLength: 1, maxLength: 128} + phase: {type: string, enum: [Active, Closing, Closed, Frozen]} + root: + type: object + x-kubernetes-preserve-unknown-fields: true + limits: + type: object + properties: + tokens: {type: integer, format: int64, minimum: 0} + usdMicros: {type: integer, format: int64, minimum: 0} + meters: + type: object + x-kubernetes-preserve-unknown-fields: true + nodes: + type: object + maxProperties: 128 + additionalProperties: + type: object + x-kubernetes-preserve-unknown-fields: true + sessions: + type: object + maxProperties: 256 + additionalProperties: + type: object + x-kubernetes-preserve-unknown-fields: true + attempts: + type: object + maxProperties: 256 + additionalProperties: + type: object + x-kubernetes-preserve-unknown-fields: true diff --git a/deploy/helm/kars/templates/crd-karsprofile.yaml b/deploy/helm/kars/templates/crd-karsprofile.yaml index 43863b884..bd2022b78 100644 --- a/deploy/helm/kars/templates/crd-karsprofile.yaml +++ b/deploy/helm/kars/templates/crd-karsprofile.yaml @@ -57,6 +57,10 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + enum: [GovernedInference] + description: "Governed inference only; no all-in task cost claim" tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means diff --git a/deploy/helm/kars/templates/crd-karstask.yaml b/deploy/helm/kars/templates/crd-karstask.yaml index 78fd122f2..ba8276bfb 100644 --- a/deploy/helm/kars/templates/crd-karstask.yaml +++ b/deploy/helm/kars/templates/crd-karstask.yaml @@ -181,6 +181,10 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + enum: [GovernedInference] + description: "Explicit governed-inference token/configured maximum-price scope; other costs are excluded" tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -327,9 +331,18 @@ spec: - message: envelope.egressAllowlistRef is unsupported by this foundation; use blueprint.egress for enforced Strict destinations reason: FieldValueInvalid rule: '!has(self.envelope.egressAllowlistRef)' - - message: 'UnsupportedLaunchBudget: total/subtree token and usdMicros ceilings are not enforced; bounded tasks may be planned but cannot launch' + - message: 'UnsupportedLaunchBudget: positive launch budgets require explicit GovernedInference scope and a configured durable broker' + reason: FieldValueForbidden + rule: >- + !has(self.execution) || !self.execution.launch || !has(self.envelope.budget) || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') + - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted reason: FieldValueForbidden - rule: '!has(self.execution) || !self.execution.launch || !has(self.envelope.budget) || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0))' + rule: >- + !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) + - message: GovernedInference scope cannot be removed or changed for an existing UID + reason: FieldValueForbidden + rule: >- + !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') - message: spec.blueprint.isolation must be one of standard, enhanced, confidential reason: FieldValueInvalid rule: '!has(self.blueprint) || !has(self.blueprint.isolation) || self.blueprint.isolation in [''standard'',''enhanced'',''confidential'']' @@ -349,6 +362,10 @@ spec: description: '`KarsTask.status`.' nullable: true properties: + inferenceBudget: + type: object + x-kubernetes-preserve-unknown-fields: true + description: "Controller-owned immutable budget account and Task UID ancestry binding" conditions: description: |- Standard K8s conditions. `Ready` is set `True` once the envelope has diff --git a/deploy/helm/kars/templates/crd-karsteam.yaml b/deploy/helm/kars/templates/crd-karsteam.yaml index bbd68ceb7..070c88c11 100644 --- a/deploy/helm/kars/templates/crd-karsteam.yaml +++ b/deploy/helm/kars/templates/crd-karsteam.yaml @@ -218,6 +218,10 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + enum: [GovernedInference] + description: "Governed inference over the lifetime of this Team UID; other costs are excluded" tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -460,6 +464,9 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + enum: [GovernedInference] tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -546,6 +553,14 @@ spec: - envelope type: object x-kubernetes-validations: + - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted + reason: FieldValueForbidden + rule: >- + !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) + - message: GovernedInference scope cannot be removed or changed for an existing UID + reason: FieldValueForbidden + rule: >- + !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') - message: spec.charter must be 1-8192 characters (or empty when spec.profileRef is set, to inherit the profile's charter) reason: FieldValueInvalid rule: (has(self.profileRef) && size(self.charter) == 0) || (size(self.charter) > 0 && size(self.charter) <= 8192) @@ -568,6 +583,13 @@ spec: description: '`KarsTeam.status` — the controller is the sole writer.' nullable: true properties: + inferenceBudgetAccount: + type: object + required: [namespace, name, uid] + properties: + namespace: {type: string, minLength: 1, maxLength: 63} + name: {type: string, minLength: 1, maxLength: 253} + uid: {type: string, minLength: 1, maxLength: 128} commonsEntryCount: description: Number of entries in the team's knowledge commons (shared memory size). format: int64 diff --git a/deploy/helm/kars/templates/crd.yaml b/deploy/helm/kars/templates/crd.yaml index e146f7e24..5d49be0d1 100644 --- a/deploy/helm/kars/templates/crd.yaml +++ b/deploy/helm/kars/templates/crd.yaml @@ -70,6 +70,10 @@ spec: maxLength: 253 aiConformanceReference: type: boolean + inferenceBudgetRef: + type: object + x-kubernetes-preserve-unknown-fields: true + description: "Controller-generated Task/account UID binding; validated against live task authority before private routing" credentialsRef: type: object description: "Explicit agent credential collection in this Sandbox's workspace; replaces legacy credentials while set. Missing/replaced sources fail closed." diff --git a/deploy/helm/kars/templates/inference-budget-admission.yaml b/deploy/helm/kars/templates/inference-budget-admission.yaml new file mode 100644 index 000000000..0c7a6146b --- /dev/null +++ b/deploy/helm/kars/templates/inference-budget-admission.yaml @@ -0,0 +1,24 @@ +{{- $budget := .Values.inferenceBudget | default dict -}} +{{- if ($budget.enabled | default false) -}} +{{- $bundle := (.Files.Get "files/inference-budget-admission.json" | replace "__ACCOUNTING_NAMESPACE__" .Release.Namespace | fromJson) -}} +{{- range $policy := $bundle.items }} +--- +apiVersion: admissionregistration.k8s.io/v1 +kind: ValidatingAdmissionPolicy +metadata: + name: {{ $policy.name }} + labels: + app.kubernetes.io/name: kars + app.kubernetes.io/component: inference-budget +spec: + {{- $policy.spec | toYaml | nindent 2 }} +--- +apiVersion: admissionregistration.k8s.io/v1 +kind: ValidatingAdmissionPolicyBinding +metadata: + name: {{ $policy.name }} +spec: + policyName: {{ $policy.name }} + validationActions: [Deny, Audit] +{{- end }} +{{- end }} diff --git a/deploy/helm/kars/templates/inference-budget.yaml b/deploy/helm/kars/templates/inference-budget.yaml new file mode 100644 index 000000000..20f97f9ef --- /dev/null +++ b/deploy/helm/kars/templates/inference-budget.yaml @@ -0,0 +1,112 @@ +{{- /* +Governed inference only. No prices, model limits, free providers, or TLS +identities are guessed. Old/reused values without this block render no broker. +*/ -}} +{{- $budget := .Values.inferenceBudget | default dict -}} +{{- if and (hasKey $budget "enabled") (not (kindIs "bool" $budget.enabled)) (not (kindIs "invalid" $budget.enabled)) -}} +{{- fail "inferenceBudget.enabled must be a boolean" -}} +{{- end -}} +{{- if ($budget.enabled | default false) -}} +{{- $version := required "inferenceBudget.catalogVersion is required" $budget.catalogVersion -}} +{{- $contracts := required "inferenceBudget.contracts is required" $budget.contracts -}} +{{- $ca := required "inferenceBudget.caBundle is required" $budget.caBundle -}} +{{- $tls := required "inferenceBudget.tlsSecretName is required" $budget.tlsSecretName -}} +{{- $routerDigest := required "inferenceBudget.routerImageDigest is required" $budget.routerImageDigest -}} +{{- if not (regexMatch "^sha256:[a-f0-9]{64}$" $routerDigest) -}} +{{- fail "inferenceBudget.routerImageDigest must be a qualified sha256 manifest digest" -}} +{{- end -}} +{{- if contains "PRIVATE KEY" $ca -}}{{- fail "inferenceBudget.caBundle must contain public certificates only" -}}{{- end -}} +apiVersion: v1 +kind: ConfigMap +metadata: + name: kars-inference-budget-contracts + namespace: {{ .Release.Namespace }} + annotations: + kars.azure.com/inference-budget-contracts: v1 + labels: + app.kubernetes.io/name: kars + app.kubernetes.io/component: inference-budget +data: + contracts.json: {{ dict "version" $version "contracts" $contracts "nonInferenceEgressHosts" ($budget.nonInferenceEgressHosts | default list) | toJson | quote }} + ca.crt: | + {{- $ca | nindent 4 }} +--- +apiVersion: v1 +kind: Service +metadata: + name: kars-inference-budget + namespace: {{ .Release.Namespace }} +spec: + type: ClusterIP + selector: + app.kubernetes.io/name: kars + app.kubernetes.io/component: controller + ports: + - name: budget-tls + port: 9447 + targetPort: 9447 + protocol: TCP +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: Role +metadata: + name: kars-inference-budget-store + namespace: {{ .Release.Namespace }} +rules: + - apiGroups: ["kars.azure.com"] + resources: ["karsbudgetaccounts", "karsbudgetaccounts/status"] + verbs: ["get", "list", "watch", "create", "update", "patch"] + - apiGroups: [""] + resources: ["configmaps"] + resourceNames: ["kars-inference-budget-contracts"] + verbs: ["get"] + - apiGroups: [""] + resources: ["secrets"] + resourceNames: [{{ $tls | quote }}] + verbs: ["get"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: RoleBinding +metadata: + name: kars-inference-budget-store + namespace: {{ .Release.Namespace }} +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: Role + name: kars-inference-budget-store +subjects: + - kind: ServiceAccount + name: kars-controller + namespace: {{ .Release.Namespace }} +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRole +metadata: + name: kars-inference-budget-reviewer +rules: + - apiGroups: ["authentication.k8s.io"] + resources: ["tokenreviews"] + verbs: ["create"] + - apiGroups: ["authorization.k8s.io"] + resources: ["subjectaccessreviews"] + verbs: ["create"] + - apiGroups: ["apps"] + resources: ["replicasets"] + verbs: ["get"] + - apiGroups: ["admissionregistration.k8s.io"] + resources: ["validatingadmissionpolicies", "validatingadmissionpolicybindings"] + verbs: ["get"] +--- +apiVersion: rbac.authorization.k8s.io/v1 +kind: ClusterRoleBinding +metadata: + name: kars-inference-budget-reviewer +roleRef: + apiGroup: rbac.authorization.k8s.io + kind: ClusterRole + name: kars-inference-budget-reviewer +subjects: + - kind: ServiceAccount + name: kars-controller + namespace: {{ .Release.Namespace }} +{{- end }} diff --git a/deploy/helm/kars/templates/operator-default-deny-networkpolicy.yaml b/deploy/helm/kars/templates/operator-default-deny-networkpolicy.yaml index 19e60327c..a9df93cfd 100644 --- a/deploy/helm/kars/templates/operator-default-deny-networkpolicy.yaml +++ b/deploy/helm/kars/templates/operator-default-deny-networkpolicy.yaml @@ -28,6 +28,18 @@ spec: - Ingress - Egress ingress: + {{- if ((.Values.inferenceBudget | default dict).enabled | default false) }} + - from: + - namespaceSelector: + matchLabels: + kars.azure.com/inference-budget: v1 + podSelector: + matchLabels: + kars.azure.com/component: sandbox + ports: + - protocol: TCP + port: 9447 + {{- end }} # Allow Prometheus / port-forward scrapes of /metrics on :9091. - from: [] ports: diff --git a/deploy/helm/kars/tests/src/inference-budget.test.ts b/deploy/helm/kars/tests/src/inference-budget.test.ts new file mode 100644 index 000000000..f2ee1f9c2 --- /dev/null +++ b/deploy/helm/kars/tests/src/inference-budget.test.ts @@ -0,0 +1,124 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { execFileSync } from "node:child_process"; +import { cpSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; + +const require = createRequire(new URL("../../../../../cli/package.json", import.meta.url)); +const { parse, parseAllDocuments } = require("yaml"); +const chart = fileURLToPath(new URL("../../", import.meta.url)); +let counter = 0; + +function render(inferenceBudget?: unknown) { + const fixture = join(chart, "tests", `.budget-fixture-${process.pid}-${++counter}`); + mkdirSync(fixture); + try { + cpSync(join(chart, "templates"), join(fixture, "templates"), { recursive: true }); + cpSync(join(chart, "files"), join(fixture, "files"), { recursive: true }); + cpSync(join(chart, "Chart.yaml"), join(fixture, "Chart.yaml")); + const values = parse(readFileSync(join(chart, "values.yaml"), "utf8")); + delete values.inferenceBudget; + if (inferenceBudget !== undefined) values.inferenceBudget = inferenceBudget; + values.controller.replicas = 3; + values.controller.extraEnv = [{ name: "CUSTOMER_SETTING", value: "preserved" }]; + writeFileSync(join(fixture, "values.yaml"), JSON.stringify(values)); + const result = execFileSync("helm", [ + "template", "retained-release", fixture, "--namespace", "customer-system", + ], { encoding: "utf8", timeout: 20_000, stdio: ["ignore", "pipe", "pipe"] }); + return parseAllDocuments(result).map((document: { toJSON(): unknown; errors: unknown[] }) => { + if (document.errors.length) throw new Error("Invalid rendered YAML"); + return document.toJSON(); + }).filter(Boolean); + } finally { + rmSync(fixture, { recursive: true, force: true }); + } +} + +const configured = { + enabled: true, + catalogVersion: "test-v1", + tlsSecretName: "operator-provided-budget-tls", + routerImageDigest: `sha256:${"7".repeat(64)}`, + caBundle: "-----BEGIN CERTIFICATE-----\npublic-test-fixture\n-----END CERTIFICATE-----", + nonInferenceEgressHosts: ["api.github.com"], + contracts: [{ + id: "bounded-chat", version: "v1", validUntil: "2030-01-01T00:00:00Z", + providerId: "private", endpoint: "https://models.example.test", model: "bounded", + operation: "ChatCompletions", outputField: "MaxTokens", + maximumInputTokens: 100, maximumOutputTokens: 20, maximumWireBytes: 4096, + outputBoundIncludesReasoning: true, + maximumPrice: { kind: "perRequest", maximumMicros: 5 }, + }], +}; + +describe("governed inference budget Helm contract", () => { + it("keeps generated Task/Team CEL in sync with the Rust scope contract", () => { + const source = readFileSync(join(chart, "../../../controller/src/inference_budget/scope.rs"), "utf8"); + const rule = (name: string): string => { + const value = source.match(new RegExp(`pub const ${name}: &str =\\s*"([^"]+)"`)); + if (!value) throw new Error(`Missing canonical scope rule ${name}`); + return value[1]; + }; + const documents = render(); + for (const kind of ["KarsTask", "KarsTeam"]) { + const definition = documents.find((doc: { kind: string; spec?: { names?: { kind?: string } } }) => + doc.kind === "CustomResourceDefinition" && doc.spec?.names?.kind === kind); + const rules = definition.spec.versions[0].schema.openAPIV3Schema.properties.spec["x-kubernetes-validations"] + .map((validation: { rule: string }) => validation.rule); + expect(rules).toContain(rule("OPT_IN_RULE")); + expect(rules).toContain(rule("RETAIN_SCOPE_RULE")); + if (kind === "KarsTask") expect(rules).toContain(rule("LAUNCH_RULE")); + } + }); + + it.each([undefined, null, {}, { enabled: false }])( + "leaves existing customers and replica count intact with %s", (value) => { + const documents = render(value); + expect(documents.some((doc: { kind: string; metadata: { name: string } }) => + doc.kind === "Service" && doc.metadata.name === "kars-inference-budget")).toBe(false); + const controller = documents.find((doc: { kind: string; metadata: { name: string } }) => + doc.kind === "Deployment" && doc.metadata.name === "kars-controller"); + expect(controller.spec.replicas).toBe(3); + const env = controller.spec.template.spec.containers[0].env; + expect(env).toContainEqual({ name: "CUSTOMER_SETTING", value: "preserved" }); + expect(env.some((item: { name: string }) => item.name.startsWith("KARS_INFERENCE_BUDGET"))).toBe(false); + }, + ); + + it("uses the exact runtime-verified admission bundle and custom controller namespace", () => { + const documents = render(configured); + const bundle = JSON.parse(readFileSync(join(chart, "files/inference-budget-admission.json"), "utf8") + .replaceAll("__ACCOUNTING_NAMESPACE__", "customer-system")); + for (const policy of bundle.items) { + const actual = documents.find((doc: { kind: string; metadata: { name: string } }) => + doc.kind === "ValidatingAdmissionPolicy" && doc.metadata.name === policy.name); + expect(actual.spec).toEqual(policy.spec); + const binding = documents.find((doc: { kind: string; metadata: { name: string } }) => + doc.kind === "ValidatingAdmissionPolicyBinding" && doc.metadata.name === policy.name); + expect(binding.spec).toEqual({ policyName: policy.name, validationActions: ["Deny", "Audit"] }); + } + const catalog = documents.find((doc: { kind: string; metadata: { name: string } }) => + doc.kind === "ConfigMap" && doc.metadata.name === "kars-inference-budget-contracts"); + expect(JSON.parse(catalog.data["contracts.json"])).toEqual({ + version: configured.catalogVersion, contracts: configured.contracts, + nonInferenceEgressHosts: configured.nonInferenceEgressHosts, + }); + expect(documents.some((doc: { kind: string }) => doc.kind === "Secret")).toBe(false); + }); + + it.each(["catalogVersion", "contracts", "caBundle", "tlsSecretName", "routerImageDigest"])( + "refuses enabling without operator-supplied %s", (field) => { + const incomplete: Record = { ...configured }; + delete incomplete[field]; + expect(() => render(incomplete)).toThrow(/inferenceBudget/); + }, + ); + + it("never places private TLS keys in a public CA ConfigMap", () => { + expect(() => render({ ...configured, caBundle: "-----BEGIN PRIVATE KEY-----" })) + .toThrow(/public certificates only/); + }); +}); diff --git a/deploy/helm/kars/values.yaml b/deploy/helm/kars/values.yaml index 08bf3e63a..dfc86023a 100644 --- a/deploy/helm/kars/values.yaml +++ b/deploy/helm/kars/values.yaml @@ -9,6 +9,19 @@ localInference: # namespace, nonempty matchLabels and destination pod TCP ports. targets: [] +inferenceBudget: + # Required for finite mode: operator-qualified image containing this broker + # client. A floating legacy :latest sidecar cannot silently ignore enforcement. + routerImageDigest: "" + # Opt-in governed inference only, not an all-in invoice/compute/tool budget. + enabled: false + catalogVersion: "" + contracts: [] + nonInferenceEgressHosts: [] + # Explicit operator-owned TLS identity. No key or pricing is generated here. + tlsSecretName: "" + caBundle: "" + # Controller configuration controller: image: diff --git a/docs/audits/governed-inference-budgets-2026-09-08.md b/docs/audits/governed-inference-budgets-2026-09-08.md new file mode 100644 index 000000000..aa2aadf9d --- /dev/null +++ b/docs/audits/governed-inference-budgets-2026-09-08.md @@ -0,0 +1,56 @@ +# Governed inference budget capability audit + +Date: **2026-09-08 UTC** +Status: **Implementation/integration candidate — publication not approved** + +## Claim being evaluated + +Durable token ceilings and operator-configured maximum-price caps for governed +inference only. No claim covers compute, GPU/VM, tool/MCP, storage, networking, +all-in task spend, invoice accuracy, taxes, or exchange rates. + +## Current evidence + +| Area | Evidence | Status | +|---|---|---| +| Shared arithmetic/state engine | UID ancestry, ancestor reservation, BeginDispatch, settlement, cancellation, replay and breach tests written | Rust execution pending | +| Kubernetes store | UID/RV PUT/CAS, concurrent sibling, lost acknowledgement, corruption/replacement, bootstrap tests written | Rust execution pending | +| Router dispatch | Buffered/stream actual-send integration and conservative usage tests written | Rust execution pending | +| Private authority | Shared exact admission bundle, Pod projection and audience checks; real core privacy helper referenced | Prerequisite merge and live qualification pending | +| Helm/default/reused values | 12 budget/schema cases plus 6 local-inference compatibility cases | **18 passed locally** | +| CLI | Scoped CREATE and pinned-account reporting | **7 passed locally** | +| CLI static validation | Existing TypeScript typecheck and targeted oxlint | **Passed locally** | +| Public API/CEL | Independent pinned Kind v0.24 / Kubernetes v1.31 preflight added, no Rust image dependency | Not executed | +| Complete broker Kind integration | Real loaded router digest, TLS broker, sibling token/price caps, route closure and cancellation scenario added to the existing E2E runner | Not executed | +| Affected-crate strict Clippy/static guards | Required, no waivers | Pending | +| Launch/Team cadence gates | Explicit scope syntax, first-opt-in transition constraints, and mandatory asynchronous broker/account checks | Source implemented; Rust/API qualification pending | +| Router code identity | Finite mode requires an operator-qualified immutable router manifest digest | Helm configuration tested; runtime qualification pending | + +No Cargo command, dependency installation, local Docker test, deployment, customer +mutation, H100 operation, main-branch change, image publication, or budget commit/ +push was performed for this evidence. Existing authorized cached CLI dependencies +were used after the local runner was found missing. + +## Remaining release decisions/gates + +1. Integrate the real privacy issuer prerequisite; no fallback implementation. +2. Qualify Team lifecycle/selective launch integration, first-opt-in transition + constraints, and source/route/cancellation regressions, including the actual + broker-in-Kind scenario and separate schema/identity preflight. +3. Run affected Rust tests, strict Clippy, schema/drift/LOC and complete disposable + Kind enforcement tests, including actual API/CEL evidence. +4. Independently review the bootstrap signature, authoritative ancestor CAS, + provider maximum-bound assumptions, every actual-send path, private token + accessibility, and conservative uncertainty/capacity behavior. +5. Obtain genuine required human audit signoffs before protected publication. + +## Signoffs + +- Implementation author: changes under active development; not a signoff. +- Independent technical reviewer: **pending**. +- Security/privacy reviewer: **pending**. +- Financial-scope/model-contract owner: **pending**. +- Required human approval/signatures: **pending**. + +No reviewer identity, email, approval, waiver, or signature is inferred or +fabricated. Prior feature waivers do not apply to this capability. diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md new file mode 100644 index 000000000..4fe9ee08b --- /dev/null +++ b/docs/governed-inference-budgets.md @@ -0,0 +1,207 @@ +# Governed inference budgets — v1 contract + +**Implementation candidate, not yet qualified for publication.** Only the new +explicit scope can pass syntactic launch validation; materialization and cadence +still require the actual configured broker, immutable account and privacy proof. +Rust qualification and the disposable Kubernetes API gate have not yet run. +See the [capability audit](audits/governed-inference-budgets-2026-09-08.md). + +## What the limits mean + +`spec.envelope.budget.scope: GovernedInference` explicitly selects accounting for +inference dispatched through the Kars router. The two ceilings are: + +- `tokens`: provider input plus output tokens, including cache input and reasoning + categories covered by the configured model contract. +- `usdMicros`: **configured maximum inference-price units**, not measured invoices. + One million units represents one US dollar in the operator's maximum tariff. + These are deliberately not estimates of total task cost. + +Compute, GPU/VM time, tools and MCP services, storage, network, external services, +taxes, currency conversion, and invoice adjustments are outside this scope. +They are **not asserted to be free**. An omitted or zero currency limit remains +unbounded. Token-only accounts require trustworthy token bounds, not an invented +price. No Copilot, local model, or provider is assigned a guessed zero price. + +Legacy tasks without the new scope retain the foundation's planning-only +interpretation of positive aggregate budgets. Ordinary standalone sandboxes +without a governed Task binding keep their existing daily/monthly tracker, +credentials, runtime environment, and inference behavior. Aggregate limits are +never copied into a per-sandbox daily allowance. + +## Root lifetime and immutable identity + +A standalone task tree has one account for its **root Task UID**. A Team and all +its principals, members, and cadence runs share one account for the **Team UID's +lifetime**, not a fresh allowance per run or per day. + +The controller pins a `KarsBudgetAccount` name and UID in protected Task/Team +status. Accounts live in the controller namespace, outside the task workspace, +so deleting the workspace cannot silently remove its accounting first. +Bindings include cluster and workspace namespace UIDs, root kind/name/UID, Task +and parent/root Task UIDs, full effective authorization and digest, Sandbox UID, +runtime namespace UID, and Pod UID. + +Account bootstrap is metadata-first. A domain-separated note signed through the +existing controller signing provider proves authorship even after a lost CREATE +acknowledgement; labels alone cannot authorize adoption of a pre-existing ledger. +The root pins the API-generated account UID before initialization. An initialized +ledger is sealed before dispatch. Missing, replaced, malformed, or uninitialized +pinned accounting never becomes a new zero balance. + +The public fields are: + +- `KarsTask.status.inferenceBudget`: scope, account reference, immutable root and + ancestry, and the full effective authorization digest. +- `KarsTeam.status.inferenceBudgetAccount`: the lifetime account reference. +- `KarsSandbox.spec.inferenceBudgetRef`: the controller-owned Task binding. + +Workspace producers create Tasks/Teams; they do not create accounts, write +balances, manufacture status, or inject this Sandbox field. Root UID recreation +is a new grant identity, not continuation of an old account. + +## Reserve, authorize one send, settle + +All ancestor meters and the root meter are updated in **one Kubernetes object +status PUT**, with the current UID/resourceVersion and bounded conflict retries. +No router-local cache is balance authority. + +1. **Reserve:** hold the complete configured maximum input/context plus the + request's enforced output maximum, and the corresponding maximum price. +2. **BeginDispatch:** atomically consume the one-time dispatch permission. Only + the first transition returns permission. A lost acknowledgement does not + authorize retransmission of that attempt. +3. **Settle:** trustworthy complete provider usage can reduce the reservation. + Missing, malformed, incomplete, disconnected, cancelled, or uncertain accepted + work is charged at the full reserved maximum. + +The boundary is every **actual final provider/model/transformed-wire send**, +including compatibility translations and each retry/fallback. A broker denial +is not a provider health failure or permission to try another provider. +Output filtering does not refund already-performed inference. + +Only provably **undispatched Reserved** attempts expire/refund (30-second maximum). +An InFlight attempt never refunds merely because its TTL, router, Pod, or Task +disappears. Recovery commits stale InFlight maxima after the transport deadline +and a margin. Cancellation closes sessions/subtrees without erasing spend. +Normal terminal rows compact behind durable Pod sequence high-water marks. +Uncertain attempts retain their contract as tombstones: late over-bound usage +freezes the account without granting another refund. + +Arithmetic is checked, uses signed-Kubernetes-compatible integer ranges, and +rounds input/output tariff categories upward separately. An observed provider +bound breach freezes the account and preserves the observation. Such a breach +invalidates the configured contract's guarantee; it is not reported as successful +hard-ceiling enforcement. + +The full-context reservation is deliberately conservative. A token ceiling +smaller than the model's hard context/output maximum can refuse even a short +prompt; concurrency needs room for all simultaneous maxima. This version does +not substitute a character heuristic for exact pre-dispatch token evidence. + +## Operator contracts and unavoidable configuration + +Enable the optional Helm `inferenceBudget` section only after supplying: + +- A versioned catalog and expiry for each supported exact + provider authentication identity, endpoint, model, and operation. +- The provider-enforced complete input/context bound, including framing and tool + schemas; Kars does not estimate this from characters. +- A maximum output field that bounds **all** output, including reasoning/thinking. +- For monetary ceilings, a per-request maximum or maximum input/output + micro-unit rates per million tokens plus any fixed maximum charge. +- A TLS Secret and public CA for + `kars-inference-budget..svc:9447`. +- `routerImageDigest`, the operator-qualified SHA-256 manifest digest of a router + implementing this contract. Finite pods keep the configured router repository + and `:latest` tag but pin that immutable digest. A floating legacy sidecar that + ignores the new budget environment must not be mistaken for enforcement. + +The TLS Secret has type `kubernetes.io/tls`, `tls.crt`/`tls.key`, and annotation +`kars.azure.com/inference-budget-tls: v1`. It is operator-provided, not generated +with a development certificate or mounted into agents. The CA ConfigMap contains +public certificates only. Certificate issuance and accurate, maintained provider +bounds/tariffs are operator responsibilities, not external database requirements. +TLS/certificate rotation must preserve trustworthy CA overlap during rollout. + +Finite accounts fail closed for missing/expired/unknown bounds, missing prices +when any ancestor needs a currency cap, or unavailable admission/privacy proof. +Adding a monetary cap after unpriced historical or InFlight work is rejected; +that history is not retrospectively priced at zero. + +The initial closed operation families are text Chat Completions, Anthropic +Messages, and Responses. Contracts reject hosted tools, multimodal generation, +async/background work, stateful prior responses, multiple output candidates, +unknown request options, and incompatible output-limit fields. Embeddings, +legacy completions, image/audio/video generation, Foundry internal generation, +fine-tuning, agent runs, evaluations, memory generation, and generic Foundry +proxy operations are not authorized by these text contracts. +Standalone moderation API stages also lack a v1 bounded contract: when a policy +requires one, the entire finite request is rejected before that API call. +The stage is never silently disabled or asserted to have zero cost. Native +provider guardrail annotations remain part of the supported governed send. + +Opaque CONNECT, redirected TLS and raw HTTP tunnels are unsupported for finite +accounts: hostname checks cannot prove encrypted HTTP authority or prevent +domain-fronting/coalescing. Native SDKs needing such tunnels must use supported +mediated router operations instead; this does not affect unbounded sandboxes. +Mediated HTTP egress additionally requires an exact operator-declared +`nonInferenceEgressHosts` entry **and** the existing Kars egress policy. This +cannot include configured/catalog model hosts. The declaration is an exclusion +of non-inference service costs, not permission to tunnel model traffic. +Existing content safety and governance checks still apply. + +Plain Sandbox spawn and handoff cannot carry authoritative accounting ancestry. +They fail closed in finite mode; use controller-enrolled Task delegation and +retry the same Task UID with a new Pod, preserving its account. + +## Authentication and authority prerequisites + +Routers receive a short-lived, Pod-bound projected ServiceAccount token with the +exclusive audience `kars.azure.com/governed-inference-budget`. Only the secure +router mounts it. Broker requests use TokenReview and live UID/ownership checks, +not agent headers, task names, or the legacy shared admin token. + +The enabled feature installs a shared, runtime-verified admission bundle: +controller-only accounting/status and finite Sandbox authority; namespace +fencing; trusted Deployment/ReplicaSet production; kubelet-only audience minting; +protected public CA projection; and no exec/attach into finite runtime namespaces. +The broker compares exact policies and unrestricted Deny bindings with current +CEL compilation status, not their names or Helm flags. + +Issuance also requires the core's real privacy-epoch helper: actual shared Secret +GET/LIST/WATCH denial checks and, when present, current v2 SRE registration proof. +There is no agent-readable fallback credential. + +## Standalone CLI and upgrade workflow + +`kars budget create -f task.yaml -n ` explicitly scopes a **new** +Task/Team manifest's existing positive budget to governed inference. It uses +Kubernetes CREATE, never force/adoption, and does not imply launch readiness. +`kars budget status --kind task|team -n ` reads the pinned +account and reports reserved, settled, uncertain, and unpriced amounts. +CLI inputs/output arithmetic use exact JavaScript-safe integers; larger account +values are rejected rather than rounded. + +Existing `kars up`, upgrade, and old/reused Helm values omit the broker by default. +Do not introduce finite enforcement retroactively into a running unbounded UID: +create a reviewed new scoped Task/Team plan. Do not clear pins, reparent funded +tasks, replace accounting objects, or delete data to “fix” a budget denial. +Inspect conditions, account phase, catalog validity, TLS, UID conflicts, and +admission/privacy status instead. + +Disabling/uninstalling the optional broker makes finite sends unavailable; it +does not turn finite runtimes into legacy unbounded routers. Workspace or Team +deletion closes authority and conservatively accounts for accepted work. +Persist/backup account objects and root pins together. Restoring missing +accounting requires operator recovery of the original authoritative state, +not a new zero account. + +## Explicit v1 capacity + +One account has at most 128 retained Task identities, 256 retained Pod sessions, +256 detailed attempts, a 64-attempt replay window per Pod, and a 512-KiB ledger. +An effective authorization snapshot is capped at 64 KiB. These are deliberate +fail-closed lifetime/capacity limits, including for long-running Teams: this +version does not promise unlimited cadence history. Capacity exhaustion requires +operator planning; there is no automatic balance reset or replay-fence eviction. diff --git a/inference-router/src/blocklist.rs b/inference-router/src/blocklist.rs index da852bab4..36aa802fd 100644 --- a/inference-router/src/blocklist.rs +++ b/inference-router/src/blocklist.rs @@ -62,9 +62,27 @@ pub struct Blocklist { /// surface) was removed alongside this slice — see /// `docs/internal/crd-well-oiled-machine/slice-5-egress-polish-and-observability.md`. allowlist: Arc>>, + inference_budget: Arc>, } impl Blocklist { + pub fn opaque_proxy_allowed(&self) -> bool { + self.inference_budget.get().is_none() + } + + pub fn bind_inference_budget( + &self, + client: Arc, + model_hosts: Vec, + ) -> anyhow::Result<()> { + self.inference_budget + .set(crate::inference_budget::egress::Fence { + client, + model_hosts, + }) + .map_err(|_| anyhow::anyhow!("inference budget egress fence already bound")) + } + /// Create a new empty blocklist (disabled mode — passes everything). pub fn disabled() -> Self { Self { @@ -75,6 +93,7 @@ impl Blocklist { learn_mode: Arc::new(AtomicBool::new(false)), learned_domains: Arc::new(RwLock::new(HashSet::new())), allowlist: Arc::new(RwLock::new(HashSet::new())), + inference_budget: Arc::default(), } } @@ -107,6 +126,7 @@ impl Blocklist { learn_mode: Arc::new(AtomicBool::new(false)), learned_domains: Arc::new(RwLock::new(HashSet::new())), allowlist: Arc::new(RwLock::new(HashSet::new())), + inference_budget: Arc::default(), } } @@ -325,6 +345,9 @@ impl Blocklist { /// `BlockedBuffer` observability surface in `forward_proxy`) can /// continue to attribute denials. pub async fn check_egress(&self, url: &str, _sandbox: &str) -> Result<(), String> { + if let Some(fence) = self.inference_budget.get() { + fence.check(extract_domain(url)).await?; + } // 1. Blocklist: hard deny let block_result = self.is_blocked(url).await; if block_result.is_blocked() { diff --git a/inference-router/src/failover.rs b/inference-router/src/failover.rs index ecb522045..6c3188026 100644 --- a/inference-router/src/failover.rs +++ b/inference-router/src/failover.rs @@ -467,6 +467,7 @@ mod tests { fn upstream(dep: &str) -> UpstreamConfig { UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: "https://example.openai.azure.com".into(), deployment: dep.to_string(), sandbox_name: "sbx".into(), diff --git a/inference-router/src/forward_proxy.rs b/inference-router/src/forward_proxy.rs index 7e0541dd1..08b618535 100644 --- a/inference-router/src/forward_proxy.rs +++ b/inference-router/src/forward_proxy.rs @@ -240,6 +240,17 @@ async fn handle_connection( sandbox: &str, blocked_egress: &BlockedBuffer, ) -> anyhow::Result<()> { + if !blocklist.opaque_proxy_allowed() { + // Host-only CONNECT/SNI checks cannot prove encrypted HTTP authority + // (domain fronting/coalescing). Finite mode uses mediated router APIs. + send_response( + &mut stream, + 403, + "Opaque tunnels are unsupported for governed inference; use mediated router APIs", + ) + .await?; + return Ok(()); + } // Read the initial request. For TLS ClientHello, we may need multiple reads // if the handshake is fragmented across TCP segments (rare, but possible). let mut buf = vec![0u8; 16384]; diff --git a/inference-router/src/inference_budget/client.rs b/inference-router/src/inference_budget/client.rs new file mode 100644 index 000000000..e34fef15d --- /dev/null +++ b/inference-router/src/inference_budget/client.rs @@ -0,0 +1,322 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::inference_budget_contract::{ + AttemptCommand, BrokerRequest, ExecutionIdentity, ReserveRequest, RouterBinding, + SessionRequest, Settlement, Usage, + catalog::Catalog, + ledger::{Reservation, Session, SettlementResult}, + tariffs::{Operation, Quote}, +}; +use serde::{Deserialize, Serialize}; +use std::{ + path::PathBuf, + sync::{ + Arc, + atomic::{AtomicBool, Ordering}, + }, +}; +use tokio::sync::Mutex; + +const TOKEN_PATH: &str = "/var/run/kars/inference-budget/token"; + +#[cfg(test)] +#[path = "client_tests.rs"] +mod tests; + +#[derive(Debug, thiserror::Error)] +#[error( + "Governed inference budget {stage} unavailable/denied (HTTP {status:?}); no provider fallback" +)] +pub struct Error { + pub stage: &'static str, + pub status: Option, +} + +fn failure(stage: &'static str) -> Error { + Error { + stage, + status: None, + } +} + +pub struct Client { + pub binding: RouterBinding, + pub identity: ExecutionIdentity, + endpoint: String, + token_path: PathBuf, + http: reqwest::Client, + reservation_lock: Mutex<()>, +} + +#[derive(Deserialize)] +struct CatalogResponse { + catalog: Catalog, + #[serde(rename = "moneyRequired")] + money_required: bool, +} + +impl Client { + pub fn from_env() -> Result>, Error> { + match std::env::var("KARS_INFERENCE_BUDGET_REQUIRED") + .unwrap_or_default() + .as_str() + { + "" | "false" => return Ok(None), + "true" => {} + _ => return Err(failure("required-mode configuration")), + } + let binding: RouterBinding = serde_json::from_str( + &std::env::var("KARS_INFERENCE_BUDGET_BINDING").map_err(|_| failure("binding"))?, + ) + .map_err(|_| failure("binding"))?; + let identity = ExecutionIdentity { + task_uid: binding.task.task_uid.clone(), + authorization_digest: binding.task.authorization_digest.clone(), + sandbox: binding.sandbox.clone(), + runtime_namespace_uid: binding.runtime_namespace_uid.clone(), + pod_name: std::env::var("POD_NAME").map_err(|_| failure("Pod name"))?, + pod_uid: std::env::var("POD_UID").map_err(|_| failure("Pod UID"))?, + }; + identity.validate().map_err(|_| failure("identity"))?; + let endpoint = + std::env::var("KARS_INFERENCE_BUDGET_ENDPOINT").map_err(|_| failure("endpoint"))?; + let url = reqwest::Url::parse(&endpoint).map_err(|_| failure("endpoint"))?; + if url.scheme() != "https" + || !url.username().is_empty() + || url.password().is_some() + || url.query().is_some() + || url.fragment().is_some() + { + return Err(failure("TLS endpoint")); + } + let ca_path = + std::env::var("KARS_INFERENCE_BUDGET_CA").map_err(|_| failure("CA configuration"))?; + let ca = std::fs::read(ca_path).map_err(|_| failure("CA material"))?; + let certificate = + reqwest::Certificate::from_pem(&ca).map_err(|_| failure("CA material"))?; + let http = reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .tls_built_in_root_certs(false) + .add_root_certificate(certificate) + .timeout(std::time::Duration::from_secs(15)) + .build() + .map_err(|_| failure("TLS client"))?; + Ok(Some(Arc::new(Self { + binding, + identity, + endpoint: endpoint.trim_end_matches('/').into(), + token_path: TOKEN_PATH.into(), + http, + reservation_lock: Mutex::new(()), + }))) + } + + async fn rpc( + &self, + path: &'static str, + payload: T, + ) -> Result { + let token = tokio::fs::read_to_string(&self.token_path) + .await + .map_err(|_| failure("private audience token"))?; + if token.trim().is_empty() { + return Err(failure("private audience token")); + } + let request = BrokerRequest { + root: self.binding.task.root.clone(), + payload, + }; + let response = self + .http + .post(format!("{}{path}", self.endpoint)) + .bearer_auth(token.trim()) + .json(&request) + .send() + .await + .map_err(|_| failure(path))?; + let status = response.status(); + if !status.is_success() { + return Err(Error { + stage: path, + status: Some(status.as_u16()), + }); + } + response + .json() + .await + .map_err(|_| failure("broker response")) + } + + fn session_request(&self) -> SessionRequest { + SessionRequest { + account_uid: self.binding.task.account.uid.clone(), + identity: self.identity.clone(), + } + } + + pub async fn catalog(&self) -> Result { + let response: CatalogResponse = self.rpc("/v1/catalog", self.session_request()).await?; + response + .catalog + .validate(chrono::Utc::now().timestamp()) + .map_err(|_| failure("operator contracts"))?; + Ok(response.catalog) + } + + pub async fn ready_for(&self, upstream: &crate::proxy::UpstreamConfig) -> Result<(), Error> { + let response: CatalogResponse = self.rpc("/v1/catalog", self.session_request()).await?; + let now = chrono::Utc::now().timestamp(); + response + .catalog + .validate(now) + .map_err(|_| failure("operator contracts"))?; + let applicable = response.catalog.contracts.iter().any(|contract| { + contract.provider_id == upstream.telemetry_provider() + && contract.endpoint.trim_end_matches('/') + == upstream.endpoint.trim_end_matches('/') + && contract.model == upstream.deployment + && contract.validate(now).is_ok() + && (!response.money_required || contract.maximum_price.is_some()) + }); + if !applicable { + return Err(failure("selected model bounds/maximum prices")); + } + Ok(()) + } + + /// Called for each actual final provider/model/wire SEND. Resolving provider + /// credentials/configuration happens before this boundary, not after begin. + pub async fn begin( + self: &Arc, + provider_id: &str, + endpoint: &str, + model: &str, + operation: Operation, + wire: &[u8], + ) -> Result<(bytes::Bytes, AttemptGuard), Error> { + let catalog = self.catalog().await?; + let contract = catalog + .select( + provider_id, + endpoint, + model, + operation, + chrono::Utc::now().timestamp(), + ) + .map_err(|_| failure("provider/operation contract"))?; + // The broker independently checks every ancestor's price requirement. + // This normalization never invents a price for a token-only contract. + let (body, quote) = contract + .normalize(wire, chrono::Utc::now().timestamp(), false) + .map_err(|_| failure("wire request bounds"))?; + let wire_digest = format!("sha256:{}", crate::providers::signing::sha256_hex(&body)); + let lock = self.reservation_lock.lock().await; + // Refresh the durable sequence rather than resetting a process-local + // counter after restart or a lost reserve acknowledgement. + let session: Session = self.rpc("/v1/session", self.session_request()).await?; + if session.identity != self.identity || session.closed { + return Err(failure("session identity")); + } + let reserve = ReserveRequest { + account_uid: self.binding.task.account.uid.clone(), + identity: self.identity.clone(), + sequence: session.next_sequence, + wire_digest: wire_digest.clone(), + quote: quote.clone(), + }; + let reserved: Reservation = self.rpc("/v1/reserve", reserve).await?; + if reserved.key.pod_uid != self.identity.pod_uid + || reserved.key.sequence != session.next_sequence + || reserved.maximum != quote.maximum + || reserved.phase != crate::inference_budget_contract::AttemptPhase::Reserved + { + return Err(failure("reservation identity")); + } + drop(lock); + let command = AttemptCommand { + account_uid: self.binding.task.account.uid.clone(), + key: reserved.key, + identity: self.identity.clone(), + wire_digest, + }; + let begun: Reservation = self.rpc("/v1/begin", command.clone()).await?; + if begun.key != command.key + || begun.maximum != quote.maximum + || begun.phase != crate::inference_budget_contract::AttemptPhase::InFlight + || begun.expires_at != reserved.expires_at + || begun.expires_at <= chrono::Utc::now().timestamp() + || quote + .contract + .validate(chrono::Utc::now().timestamp()) + .is_err() + { + return Err(failure("dispatch identity")); + } + Ok(( + body.into(), + AttemptGuard { + client: self.clone(), + command, + quote, + finalized: AtomicBool::new(false), + }, + )) + } + + async fn settle( + &self, + command: AttemptCommand, + usage: Option, + ) -> Result { + self.rpc( + "/v1/settle", + Settlement { + attempt: command, + usage, + }, + ) + .await + } +} + +pub struct AttemptGuard { + client: Arc, + pub command: AttemptCommand, + pub quote: Quote, + finalized: AtomicBool, +} + +impl AttemptGuard { + pub async fn finish(&self, usage: Option) -> Result { + if self.finalized.swap(true, Ordering::AcqRel) { + return Err(failure("duplicate local settlement")); + } + let result = self.client.settle(self.command.clone(), usage).await; + if result.is_err() { + self.finalized.store(false, Ordering::Release); + } + result + } +} + +impl Drop for AttemptGuard { + fn drop(&mut self) { + if self.finalized.swap(true, Ordering::AcqRel) { + return; + } + if let Ok(runtime) = tokio::runtime::Handle::try_current() { + let client = self.client.clone(); + let command = self.command.clone(); + runtime.spawn(async move { + if client.settle(command, None).await.is_err() { + // The durable InFlight reservation stays charged/held. + // The broker's recovery loop conservatively finalizes it. + tracing::warn!( + "Governed inference settlement pending; reservation remains funded" + ); + } + }); + } + } +} diff --git a/inference-router/src/inference_budget/client_tests.rs b/inference-router/src/inference_budget/client_tests.rs new file mode 100644 index 000000000..85866b5d7 --- /dev/null +++ b/inference-router/src/inference_budget/client_tests.rs @@ -0,0 +1,331 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::inference_budget_contract::{ + AccountReference, BudgetScope, Limits, ResourceIdentity, RootIdentity, RootKind, TaskAuthority, + TaskBudgetBinding, + ledger::Ledger, + tariffs::{MaximumPrice, ModelContract, OutputField}, +}; +use crate::{auth::WorkloadIdentityAuth, provider::ProviderKind, proxy::UpstreamConfig}; +use axum::http::{HeaderMap, Method}; +use futures::TryStreamExt; +use serde_json::json; +use std::sync::{Mutex as StdMutex, atomic::AtomicUsize}; +use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; + +static NEXT: AtomicUsize = AtomicUsize::new(0); + +#[tokio::test] +async fn finite_fence_cannot_be_bypassed_by_disabling_blocklist_or_enabling_learning() { + let fixture = Fixture::new(100, false).await; + let blocklist = crate::blocklist::Blocklist::disabled(); + assert!(blocklist.opaque_proxy_allowed()); + blocklist + .bind_inference_budget(fixture.client.clone(), vec![]) + .unwrap(); + assert!(!blocklist.opaque_proxy_allowed()); + blocklist.set_learn_mode(true); + assert!(!blocklist.opaque_proxy_allowed()); + assert!( + blocklist + .check_egress("unknown-model.example", "forged-agent-header") + .await + .is_err() + ); + assert!( + fixture + .provider + .received_requests() + .await + .unwrap() + .is_empty() + ); +} + +#[derive(Clone)] +struct BrokerFixture { + ledger: Arc>, + catalog: Catalog, + lose_begin_ack: bool, +} + +impl Respond for BrokerFixture { + fn respond(&self, request: &Request) -> ResponseTemplate { + assert_eq!( + request.headers.get("authorization").unwrap(), + "Bearer private-test-token" + ); + let now = chrono::Utc::now().timestamp(); + let mut ledger = self.ledger.lock().unwrap(); + let result = match request.url.path() { + "/v1/catalog" => Ok(json!({"catalog":self.catalog,"moneyRequired":true})), + "/v1/session" => { + let request: BrokerRequest = + serde_json::from_slice(&request.body).unwrap(); + ledger + .register_session(request.payload.identity) + .map(|change| { + *ledger = change.next; + json!(change.value) + }) + } + "/v1/reserve" => { + let request: BrokerRequest = + serde_json::from_slice(&request.body).unwrap(); + self.catalog + .accepts_quote(&request.payload.quote, now, true) + .unwrap(); + ledger.reserve(&request.payload, now).map(|change| { + *ledger = change.next; + json!(change.value) + }) + } + "/v1/begin" => { + let request: BrokerRequest = + serde_json::from_slice(&request.body).unwrap(); + ledger.begin_dispatch(&request.payload, now).map(|change| { + *ledger = change.next; + json!(change.value) + }) + } + "/v1/settle" => { + let request: BrokerRequest = + serde_json::from_slice(&request.body).unwrap(); + ledger.settle(&request.payload).map(|change| { + *ledger = change.next; + json!(change.value) + }) + } + _ => panic!("unexpected broker route"), + }; + if self.lose_begin_ack && request.url.path() == "/v1/begin" { + return ResponseTemplate::new(503); + } + match result { + Ok(response) => ResponseTemplate::new(200).set_body_json(response), + Err(_) => ResponseTemplate::new(429), + } + } +} + +struct Fixture { + provider: MockServer, + _broker: MockServer, + client: Arc, + ledger: Arc>, + directory: PathBuf, +} + +impl Drop for Fixture { + fn drop(&mut self) { + std::fs::remove_dir_all(&self.directory).unwrap(); + } +} + +impl Fixture { + async fn new(limit: u64, lose_begin_ack: bool) -> Self { + let provider = MockServer::start().await; + let broker = MockServer::start().await; + let task = ResourceIdentity { + namespace: "workspace".into(), + name: "task".into(), + uid: "task-uid".into(), + }; + let root = RootIdentity { + kind: RootKind::KarsTask, + resource: task.clone(), + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }; + let digest = format!("sha256:{}", "a".repeat(64)); + let authority = TaskAuthority { + task, + parent_uid: None, + root_task_uid: "task-uid".into(), + authorization_digest: digest.clone(), + effective_authorization: json!({"test":true}), + limits: Limits { + tokens: Some(limit), + usd_micros: Some(1000), + }, + }; + let ledger = Ledger::new("account-uid".into(), root.clone(), authority.limits) + .unwrap() + .register_task(authority) + .unwrap() + .next; + let ledger = Arc::new(StdMutex::new(ledger)); + let model = ModelContract { + id: "model".into(), + version: "v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "ollama".into(), + endpoint: provider.uri(), + model: "model".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxTokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 5 }), + }; + Mock::given(wiremock::matchers::method("POST")) + .respond_with(BrokerFixture { + ledger: ledger.clone(), + catalog: Catalog { + version: "v1".into(), + contracts: vec![model], + non_inference_egress_hosts: vec![], + }, + lose_begin_ack, + }) + .mount(&broker) + .await; + let identity = ExecutionIdentity { + task_uid: "task-uid".into(), + authorization_digest: digest.clone(), + sandbox: ResourceIdentity { + namespace: "workspace".into(), + name: "task".into(), + uid: "sandbox-uid".into(), + }, + runtime_namespace_uid: "runtime-uid".into(), + pod_name: "pod".into(), + pod_uid: "pod-uid".into(), + }; + let directory = std::env::current_dir().unwrap().join(format!( + ".budget-client-test-{}-{}", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); + std::fs::create_dir(&directory).unwrap(); + let token_path = directory.join("token"); + std::fs::write(&token_path, "private-test-token").unwrap(); + let client = Arc::new(Client { + binding: RouterBinding { + task: TaskBudgetBinding { + scope: BudgetScope::GovernedInference, + account: AccountReference { + namespace: "accounting".into(), + name: "account".into(), + uid: "account-uid".into(), + }, + root, + task_uid: "task-uid".into(), + parent_task_uid: None, + root_task_uid: "task-uid".into(), + authorization_digest: digest, + }, + sandbox: identity.sandbox.clone(), + runtime_namespace: "kars-task".into(), + runtime_namespace_uid: "runtime-uid".into(), + privacy_epoch: None, + }, + identity, + endpoint: broker.uri(), + token_path, + http: reqwest::Client::new(), + reservation_lock: Mutex::new(()), + }); + Self { + provider, + _broker: broker, + client, + ledger, + directory, + } + } + + fn upstream(&self) -> UpstreamConfig { + let mut upstream = + UpstreamConfig::azure(self.provider.uri(), "model".into(), "task".into()); + upstream.provider = ProviderKind::Ollama; + upstream.inference_budget = Some(self.client.clone()); + upstream + } + + async fn buffered(&self) -> anyhow::Result<(axum::http::StatusCode, HeaderMap, bytes::Bytes)> { + crate::proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &self.upstream(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + bytes::Bytes::from_static(br#"{"messages":[{"role":"user","content":"text"}]}"#), + ) + .await + } +} + +#[tokio::test] +async fn actual_buffered_send_injects_output_bound_and_settles_before_returning() { + let fixture = Fixture::new(100, false).await; + Mock::given(wiremock::matchers::path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8} + }))) + .mount(&fixture.provider) + .await; + assert!(fixture.buffered().await.unwrap().0.is_success()); + let sent = fixture.provider.received_requests().await.unwrap(); + assert_eq!(sent.len(), 1); + let body: serde_json::Value = serde_json::from_slice(&sent[0].body).unwrap(); + assert_eq!(body["max_tokens"], 20); + let ledger = fixture.ledger.lock().unwrap(); + assert_eq!(ledger.meters.reserved.tokens, 0); + assert_eq!(ledger.meters.settled.tokens, 8); + assert_eq!(ledger.meters.settled.usd_micros, 5); +} + +#[tokio::test] +async fn unknown_usage_consumes_full_maximum_and_second_send_is_denied_without_health_failover() { + let fixture = Fixture::new(50, false).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[]}))) + .mount(&fixture.provider) + .await; + fixture.buffered().await.unwrap(); + assert_eq!(fixture.ledger.lock().unwrap().meters.uncertain.tokens, 30); + let error = fixture.buffered().await.err().unwrap(); + assert!(!crate::proxy::failure::retryable_failure(&error)); + assert_eq!(fixture.provider.received_requests().await.unwrap().len(), 1); +} + +#[tokio::test] +async fn lost_begin_ack_funds_uncertain_work_but_never_sends_or_regrants_that_attempt() { + let fixture = Fixture::new(50, true).await; + assert!(fixture.buffered().await.is_err()); + assert_eq!(fixture.provider.received_requests().await.unwrap().len(), 0); + assert_eq!(fixture.ledger.lock().unwrap().meters.reserved.tokens, 30); + assert!(fixture.buffered().await.is_err()); + assert_eq!(fixture.provider.received_requests().await.unwrap().len(), 0); +} + +#[tokio::test] +async fn actual_stream_eof_requires_terminal_usage_and_charges_output_once() { + let fixture = Fixture::new(100, false).await; + Mock::given(wiremock::matchers::method("POST")).respond_with( + ResponseTemplate::new(200).insert_header("content-type", "text/event-stream") + .set_body_string("data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5}}\n\ndata: [DONE]\n\n"), + ).mount(&fixture.provider).await; + let (_, _, stream) = crate::proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::new(), + fixture.upstream(), + "chat/completions", + HeaderMap::new(), + bytes::Bytes::from_static( + br#"{"stream":true,"messages":[{"role":"user","content":"text"}]}"#, + ), + ) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(fixture.ledger.lock().unwrap().meters.settled.tokens, 8); +} diff --git a/inference-router/src/inference_budget/dispatch.rs b/inference-router/src/inference_budget/dispatch.rs new file mode 100644 index 000000000..3bb42242e --- /dev/null +++ b/inference-router/src/inference_budget/dispatch.rs @@ -0,0 +1,115 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + client::{AttemptGuard, Error}, + usage, +}; +use crate::{inference_budget_contract::catalog, proxy::UpstreamConfig}; +use axum::http::{Method, StatusCode}; +use bytes::Bytes; +use futures::{StreamExt, stream::BoxStream}; + +/// Runs after provider resolution and all wire transformations, immediately +/// before the actual transport send. Each retry needs a separate grant. +pub async fn begin( + upstream: &UpstreamConfig, + method: &Method, + path: &str, + body: Bytes, +) -> anyhow::Result<(Bytes, Option)> { + let Some(client) = &upstream.inference_budget else { + return Ok((body, None)); + }; + let denied = || Error { + stage: "unsupported inference operation", + status: None, + }; + if method != Method::POST { + return Err(denied().into()); + } + let operation = catalog::operation(path).ok_or_else(denied)?; + let (wire, guard) = client + .begin( + upstream.telemetry_provider(), + &upstream.endpoint, + &upstream.deployment, + operation, + &body, + ) + .await?; + Ok((wire, Some(guard))) +} + +pub async fn finish( + guard: Option, + status: StatusCode, + body: &[u8], +) -> anyhow::Result<()> { + if let Some(guard) = guard { + let usage = status + .is_success() + .then(|| usage::buffered(body, guard.quote.contract.operation)) + .flatten(); + let result = guard.finish(usage).await?; + if result.breach { + return Err(Error { + stage: "provider contract breached; account frozen", + status: None, + } + .into()); + } + } + Ok(()) +} + +/// EOF settles complete usage before finishing the client stream. Transport +/// errors and downstream cancellation drop the guard and keep the full maximum. +/// A settlement outage also leaves the durable reservation funded. +pub fn stream( + inner: BoxStream<'static, Result>, + guard: Option, + is_sse: bool, +) -> BoxStream<'static, Result> { + let Some(guard) = guard else { return inner }; + let usage = usage::StreamUsage::new(guard.quote.contract.operation); + futures::stream::unfold( + (inner, Some(guard), Some(usage), false), + move |(mut inner, mut guard, mut usage, mut failed)| async move { + match inner.next().await { + Some(chunk) => { + match &chunk { + Ok(bytes) if is_sse && !failed => { + if let Some(usage) = &mut usage { + usage.push(bytes); + } + } + Err(_) => failed = true, + _ => {} + } + Some((chunk, (inner, guard, usage, failed))) + } + None => { + if let Some(guard) = guard.take() { + let evidence = if is_sse && !failed { + usage.take().and_then(usage::StreamUsage::finish) + } else { + None + }; + match guard.finish(evidence).await { + Ok(result) if result.breach => tracing::error!( + "Governed inference provider bound breached; account frozen" + ), + Err(_) => tracing::warn!( + "Governed inference settlement pending; reservation remains funded" + ), + _ => {} + } + } + None + } + } + }, + ) + .boxed() +} diff --git a/inference-router/src/inference_budget/egress.rs b/inference-router/src/inference_budget/egress.rs new file mode 100644 index 000000000..69695d038 --- /dev/null +++ b/inference-router/src/inference_budget/egress.rs @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::Client; +use std::sync::Arc; + +pub struct Fence { + pub client: Arc, + pub model_hosts: Vec, +} + +impl Fence { + pub async fn check(&self, host: &str) -> Result<(), String> { + let catalog = self.client.catalog().await.map_err(|_| { + "Governed inference budget unavailable; mediated egress is closed".to_owned() + })?; + if catalog.allows_mediated_non_inference(host, &self.model_hosts) { + Ok(()) + } else { + Err("Governed inference requires a brokered model route; this destination is not an operator-declared non-inference exclusion".into()) + } + } +} + +pub fn model_hosts(config: &crate::config::Config) -> Vec { + [ + config.azure_openai_endpoint.as_deref(), + config.foundry_endpoint.as_deref(), + config.foundry_project_endpoint.as_deref(), + Some(config.anthropic_endpoint.as_str()), + config.ollama_endpoint.as_deref(), + ] + .into_iter() + .flatten() + .chain( + config + .providers + .values() + .map(|provider| provider.endpoint.as_str()), + ) + .filter_map(crate::proxy::endpoint_host) + .collect() +} diff --git a/inference-router/src/inference_budget/mod.rs b/inference-router/src/inference_budget/mod.rs new file mode 100644 index 000000000..c1cbb0126 --- /dev/null +++ b/inference-router/src/inference_budget/mod.rs @@ -0,0 +1,11 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +pub mod client; +pub mod dispatch; +pub mod egress; +pub mod readiness; +pub mod response; +pub mod usage; + +pub use client::Client; diff --git a/inference-router/src/inference_budget/readiness.rs b/inference-router/src/inference_budget/readiness.rs new file mode 100644 index 000000000..9fcd755d9 --- /dev/null +++ b/inference-router/src/inference_budget/readiness.rs @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::{proxy::UpstreamConfig, routes::AppState}; + +pub async fn ready(state: &AppState, mut upstream: UpstreamConfig) -> anyhow::Result<()> { + let client = state + .inference_budget + .as_ref() + .ok_or_else(|| anyhow::anyhow!("budget binding absent"))?; + let policy = crate::inference_policy_loader::current_snapshot(&state.inference_policy).await; + if !policy.guardrails.is_empty() { + anyhow::bail!("Mandatory standalone moderation needs a supported bounded contract"); + } + crate::routes::apply_provider_resolution(state, &mut upstream, &policy)?; + let candidates = crate::failover::candidates_for_request(&upstream, &policy, b"{}"); + let candidate = candidates + .first() + .ok_or_else(|| anyhow::anyhow!("model route absent"))?; + let target = crate::failover::resolve_candidate(&upstream, &state.config, candidate)?; + client.ready_for(&target).await?; + crate::proxy::credential_for_upstream(&state.auth, Some(&state.copilot), &target).await?; + Ok(()) +} diff --git a/inference-router/src/inference_budget/response.rs b/inference-router/src/inference_budget/response.rs new file mode 100644 index 000000000..f03041516 --- /dev/null +++ b/inference-router/src/inference_budget/response.rs @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use axum::{ + http::StatusCode, + response::{IntoResponse, Response}, +}; + +pub fn denial(error: &anyhow::Error) -> Option { + let error = error.downcast_ref::()?; + let (status, code) = match error.status { + Some(429) => (StatusCode::TOO_MANY_REQUESTS, "inference_budget_exhausted"), + Some(403) => (StatusCode::FORBIDDEN, "inference_budget_authority"), + Some(409) => (StatusCode::CONFLICT, "inference_budget_attempt_state"), + _ => ( + StatusCode::SERVICE_UNAVAILABLE, + "inference_budget_unavailable", + ), + }; + Some(crate::errors::openai_coded(status, error.to_string(), code, code).into_response()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn budget_errors_are_not_reported_as_provider_bad_gateways() { + for (code, expected) in [ + (Some(429), 429), + (Some(403), 403), + (Some(409), 409), + (None, 503), + ] { + let error: anyhow::Error = super::super::client::Error { + stage: "/v1/reserve", + status: code, + } + .into(); + assert_eq!(denial(&error).unwrap().status().as_u16(), expected); + assert!(!crate::proxy::failure::retryable_failure(&error)); + } + assert!(denial(&anyhow::anyhow!("ordinary provider error")).is_none()); + } +} diff --git a/inference-router/src/inference_budget/usage.rs b/inference-router/src/inference_budget/usage.rs new file mode 100644 index 000000000..a1d170da0 --- /dev/null +++ b/inference-router/src/inference_budget/usage.rs @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Usage is settlement evidence, never pre-dispatch authority. Missing, +//! incomplete, malformed, or interrupted evidence commits the full reservation. + +use crate::inference_budget_contract::{Usage, tariffs::Operation}; +use serde_json::Value; + +fn number(value: &Value, key: &str) -> Option { + value.get(key)?.as_u64() +} + +fn optional(value: &Value, key: &str) -> Option { + match value.get(key) { + None => Some(0), + Some(value) => value.as_u64(), + } +} + +fn detail(value: &Value, object: &str, key: &str) -> Option { + match value.get(object) { + None => Some(0), + Some(value) if value.is_object() => optional(value, key), + _ => None, + } +} + +pub fn buffered(body: &[u8], operation: Operation) -> Option { + let response: Value = serde_json::from_slice(body).ok()?; + parse(response.get("usage")?, operation) +} + +fn parse(usage: &Value, operation: Operation) -> Option { + match operation { + Operation::ChatCompletions => { + let input = number(usage, "prompt_tokens")?; + let output = number(usage, "completion_tokens")?; + if let Some(total) = usage.get("total_tokens") { + if total.as_u64()? != input.checked_add(output)? { + return None; + } + } + Some(Usage { + input_tokens: input, + output_tokens: output, + cached_input_tokens: detail(usage, "prompt_tokens_details", "cached_tokens")?, + cache_creation_input_tokens: 0, + reasoning_output_tokens: detail( + usage, + "completion_tokens_details", + "reasoning_tokens", + )?, + }) + } + Operation::Responses => { + let input = number(usage, "input_tokens")?; + let output = number(usage, "output_tokens")?; + if let Some(total) = usage.get("total_tokens") { + if total.as_u64()? != input.checked_add(output)? { + return None; + } + } + Some(Usage { + input_tokens: input, + output_tokens: output, + cached_input_tokens: detail(usage, "input_tokens_details", "cached_tokens")?, + cache_creation_input_tokens: 0, + reasoning_output_tokens: detail( + usage, + "output_tokens_details", + "reasoning_tokens", + )?, + }) + } + Operation::AnthropicMessages => { + let fresh = number(usage, "input_tokens")?; + let cached = optional(usage, "cache_read_input_tokens")?; + let creation = optional(usage, "cache_creation_input_tokens")?; + Some(Usage { + input_tokens: fresh.checked_add(cached)?.checked_add(creation)?, + output_tokens: number(usage, "output_tokens")?, + cached_input_tokens: cached, + cache_creation_input_tokens: creation, + // Anthropic output_tokens already includes thinking output. + reasoning_output_tokens: 0, + }) + } + } +} + +pub struct StreamUsage { + operation: Operation, + buffer: Vec, + usage: Option, + terminal: bool, + invalid: bool, +} + +impl StreamUsage { + pub fn new(operation: Operation) -> Self { + Self { + operation, + buffer: Vec::new(), + usage: None, + terminal: false, + invalid: false, + } + } + + pub fn push(&mut self, bytes: &[u8]) { + if self.invalid { + return; + } + self.buffer.extend_from_slice(bytes); + if self.buffer.len() > 262_144 { + self.invalid = true; + self.buffer.clear(); + return; + } + while let Some(end) = self.buffer.iter().position(|byte| *byte == b'\n') { + let mut line: Vec = self.buffer.drain(..=end).collect(); + while line + .last() + .is_some_and(|byte| matches!(*byte, b'\n' | b'\r')) + { + line.pop(); + } + let Some(data) = line.strip_prefix(b"data:") else { + continue; + }; + let data = data.strip_prefix(b" ").unwrap_or(data); + if data == b"[DONE]" { + if self.operation == Operation::ChatCompletions { + self.terminal = true; + } + continue; + } + if data.is_empty() { + continue; + } + let Ok(event) = serde_json::from_slice::(data) else { + self.invalid = true; + continue; + }; + self.event(&event); + } + } + + fn event(&mut self, event: &Value) { + if self.terminal { + self.invalid = true; + return; + } + match self.operation { + Operation::ChatCompletions => { + if let Some(usage) = event.get("usage").filter(|value| !value.is_null()) { + let parsed = parse(usage, self.operation); + if parsed.is_none() + || self + .usage + .as_ref() + .is_some_and(|old| Some(old) != parsed.as_ref()) + { + self.invalid = true; + } + self.usage = parsed; + } + if event.get("error").is_some() { + self.invalid = true; + } + } + Operation::Responses => match event.get("type").and_then(Value::as_str) { + Some("response.completed") => { + self.usage = event + .pointer("/response/usage") + .and_then(|usage| parse(usage, self.operation)); + self.terminal = true; + } + Some("response.failed" | "error") => self.invalid = true, + _ => {} + }, + Operation::AnthropicMessages => match event.get("type").and_then(Value::as_str) { + Some("message_start") => { + self.usage = event + .pointer("/message/usage") + .and_then(|usage| parse(usage, self.operation)); + if self.usage.is_none() { + self.invalid = true; + } + } + Some("message_delta") => { + match ( + self.usage.as_mut(), + event + .get("usage") + .and_then(|usage| number(usage, "output_tokens")), + ) { + (Some(usage), Some(output)) if output >= usage.output_tokens => { + usage.output_tokens = output + } + _ => self.invalid = true, + } + } + Some("message_stop") => self.terminal = true, + Some("error") => self.invalid = true, + _ => {} + }, + } + } + + /// Only a cleanly completed transport with a terminal provider event is + /// eligible for a refund. Client disconnect/drop uses None instead. + pub fn finish(self) -> Option { + if self.invalid || !self.terminal || !self.buffer.iter().all(u8::is_ascii_whitespace) { + return None; + } + self.usage + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn native_cache_tokens_are_included_once_in_the_input_total() { + let usage = buffered(br#"{"usage":{"input_tokens":3,"cache_read_input_tokens":5,"cache_creation_input_tokens":7,"output_tokens":11}}"#, + Operation::AnthropicMessages).unwrap(); + assert_eq!(usage.input_tokens, 15); + assert_eq!(usage.cached_input_tokens, 5); + assert_eq!(usage.cache_creation_input_tokens, 7); + } + + #[test] + fn chat_stream_usage_survives_every_byte_boundary_without_counting_chunks_as_tokens() { + let wire = b"data: {\"choices\":[]}\n\ndata: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":3,\"total_tokens\":8}}\n\ndata: [DONE]\n\n"; + for size in 1..wire.len() { + let mut stream = StreamUsage::new(Operation::ChatCompletions); + for chunk in wire.chunks(size) { + stream.push(chunk); + } + assert_eq!(stream.finish().unwrap().output_tokens, 3, "{size}"); + } + } + + #[test] + fn partial_missing_or_malformed_usage_does_not_prove_a_refund() { + for wire in [ + &b"data: {\"choices\":[]}\n\ndata: [DONE]\n\n"[..], + &b"data: {\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":3}}\n\n"[..], + &b"data: {\"usage\":{\"prompt_tokens\":-1,\"completion_tokens\":3}}\n\ndata: [DONE]\n\n"[..], + ] { + let mut stream = StreamUsage::new(Operation::ChatCompletions); + stream.push(wire); + assert!(stream.finish().is_none()); + } + } + + #[test] + fn native_stream_requires_start_usage_and_terminal_stop() { + let mut stream = StreamUsage::new(Operation::AnthropicMessages); + stream.push(b"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":0}}}\n\n"); + stream.push(b"data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":7}}\n\n"); + stream.push(b"data: {\"type\":\"message_stop\"}\n\n"); + assert_eq!(stream.finish().unwrap().output_tokens, 7); + } + + #[test] + fn responses_only_settles_complete_final_usage() { + let mut stream = StreamUsage::new(Operation::Responses); + stream.push(b"data: {\"type\":\"response.completed\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3}}}\n\n"); + assert_eq!(stream.finish().unwrap().input_tokens, 2); + let mut incomplete = StreamUsage::new(Operation::Responses); + incomplete.push(b"data: {\"type\":\"response.incomplete\",\"response\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":3}}}\n\n"); + assert!(incomplete.finish().is_none()); + } +} diff --git a/inference-router/src/lib.rs b/inference-router/src/lib.rs index 9257c75db..d686056c8 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -36,6 +36,9 @@ pub mod governance; pub mod governed_services; pub mod guardrails; pub mod handoff; +pub mod inference_budget; +#[path = "../../shared/inference_budget/mod.rs"] +pub mod inference_budget_contract; pub mod inference_policy_loader; pub mod mcp; pub mod memory_binding_loader; diff --git a/inference-router/src/providers/signing.rs b/inference-router/src/providers/signing.rs index a8461ab1b..e9d722dfc 100644 --- a/inference-router/src/providers/signing.rs +++ b/inference-router/src/providers/signing.rs @@ -33,6 +33,12 @@ pub struct KeyRef(pub String); #[derive(Debug, Clone, PartialEq, Eq)] pub struct Signature(pub Vec); +/// Standard content digest for immutable wire-request authorization bindings. +pub fn sha256_hex(payload: &[u8]) -> String { + use sha2::{Digest, Sha256}; + hex::encode(Sha256::digest(payload)) +} + #[derive(Debug, thiserror::Error)] pub enum SigningError { #[error("unknown key ref: {0:?}")] diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 122ff9605..2f48a90fa 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -35,6 +35,7 @@ mod authentication_tests; #[derive(Clone)] pub struct UpstreamConfig { pub telemetry: Option>, + pub inference_budget: Option>, pub endpoint: String, pub deployment: String, pub sandbox_name: String, @@ -52,7 +53,7 @@ pub struct UpstreamConfig { } impl UpstreamConfig { - fn telemetry_provider(&self) -> &str { + pub(crate) fn telemetry_provider(&self) -> &str { match &self.authentication { AuthenticationProvenance::Named { provider_id } => provider_id, AuthenticationProvenance::LegacyDefault => self.provider.as_tag(), @@ -64,6 +65,7 @@ impl UpstreamConfig { pub fn azure(endpoint: String, deployment: String, sandbox_name: String) -> Self { Self { telemetry: None, + inference_budget: None, endpoint, deployment, sandbox_name, @@ -314,15 +316,14 @@ pub async fn forward( tracing::info!(sandbox = %upstream.sandbox_name, body_len = body.len(), "Sending upstream request"); - let retryable = is_idempotent(&method, path); - let response = send_with_retry( + let (response, budget_attempt) = send_with_retry( client, &method, &upstream_url, &headers, body, - retryable, - &upstream.sandbox_name, + upstream, + path, ) .await .inspect_err(|_| { @@ -358,6 +359,7 @@ pub async fn forward( ForwardFailure::response_body(status, error) })?; let latency = start.elapsed(); + crate::inference_budget::dispatch::finish(budget_attempt, status, &response_body).await?; record_metrics(upstream, status, latency, &response_body); if let Some(observation) = observation.as_mut() { @@ -438,22 +440,29 @@ async fn send_with_retry( url: &str, headers: &HeaderMap, body: Bytes, - retryable: bool, - sandbox_name: &str, -) -> Result { + upstream: &UpstreamConfig, + path: &str, +) -> Result<( + reqwest::Response, + Option, +)> { const MAX_ATTEMPTS: u32 = 3; const BACKOFF_MS: [u64; 2] = [250, 750]; // after attempts 1 and 2 + let retryable = is_idempotent(method, path); let attempts = if retryable { MAX_ATTEMPTS } else { 1 }; let mut last_err: Option = None; + let sandbox_name = &upstream.sandbox_name; for attempt in 1..=attempts { + let (wire, budget_attempt) = + crate::inference_budget::dispatch::begin(upstream, method, path, body.clone()).await?; // RequestBuilder is not Clone, so rebuild per attempt. Body is a // Bytes (cheap ref-counted clone) — no real cost. let response = client .request(method.clone(), url) .headers(headers.clone()) - .body(body.clone()) + .body(wire) .timeout(INFERENCE_REQUEST_TIMEOUT) .send() .await; @@ -477,7 +486,7 @@ async fn send_with_retry( .await; continue; } - return Ok(resp); + return Ok((resp, budget_attempt)); } Err(err) => { if retryable && is_retryable_error(&err) && attempt < attempts { @@ -579,6 +588,8 @@ pub async fn forward_stream( })?; let start = Instant::now(); + let (body, budget_attempt) = + crate::inference_budget::dispatch::begin(&upstream, &Method::POST, path, body).await?; let response = client .post(&upstream_url) @@ -644,6 +655,7 @@ pub async fn forward_stream( if let Some(observation) = observation.as_mut() { observation.buffered(status.as_u16(), &body_bytes); } + crate::inference_budget::dispatch::finish(budget_attempt, status, &body_bytes).await?; tracing::warn!( sandbox = %upstream.sandbox_name, status = %status.as_u16(), @@ -716,7 +728,11 @@ pub async fn forward_stream( Ok(( status, response_headers, - crate::task_telemetry::observe::wrap_stream(metered.boxed(), observation, is_sse), + crate::task_telemetry::observe::wrap_stream( + crate::inference_budget::dispatch::stream(metered.boxed(), budget_attempt, is_sse), + observation, + is_sse, + ), )) } diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index de4799019..30d3f9fac 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -562,6 +562,9 @@ pub(super) async fn anthropic_messages( } Err(e) => { tracing::warn!(sandbox = %sandbox_name, error = %e, "Anthropic upstream call failed"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } deny_response( StatusCode::BAD_GATEWAY, &format!("Upstream error: {e}"), @@ -842,6 +845,9 @@ async fn forward_anthropic_passthrough( } Err(e) => { tracing::warn!(sandbox = %sandbox_name, error = %e, "Copilot Anthropic stream failed"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } deny_response( StatusCode::BAD_GATEWAY, &format!("Upstream error: {e}"), @@ -956,6 +962,9 @@ async fn forward_anthropic_passthrough( } Err(e) => { tracing::warn!(sandbox = %sandbox_name, error = %e, "Copilot Anthropic call failed"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } deny_response( StatusCode::BAD_GATEWAY, &format!("Upstream error: {e}"), diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 89ac9a0f2..99b2277f6 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -115,6 +115,12 @@ pub(super) fn build_guardrail_pipeline( if policy.guardrails.is_empty() { return Ok(None); } + if state.inference_budget.is_some() { + return Err(GuardrailError::Config { + provider: "governed-inference".into(), + reason: "Mandatory standalone moderation has no bounded inference contract in v1; the request is blocked, not sent unmetered or without its guardrail".into(), + }); + } GuardrailPipeline::from_stages(&policy.guardrails, &state.config, &state.client) .map(|p| Some(Arc::new(p))) } @@ -966,6 +972,9 @@ pub(super) async fn chat_completions( } Err(e) => { tracing::error!(sandbox = %sandbox_name, "Stream proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } errors::openai( StatusCode::BAD_GATEWAY, "Failed to reach inference backend", @@ -1046,6 +1055,9 @@ pub(super) async fn chat_completions( } Err(e) => { tracing::error!(sandbox = %sandbox_name, "Responses fallback proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } errors::openai( StatusCode::BAD_GATEWAY, "Failed to reach inference backend", @@ -1330,6 +1342,9 @@ pub(super) async fn chat_completions( } Err(e) => { tracing::error!(sandbox = %sandbox_name, "Proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } ( StatusCode::BAD_GATEWAY, Json(serde_json::json!({ diff --git a/inference-router/src/routes/handoff/mod.rs b/inference-router/src/routes/handoff/mod.rs index 7cd757b53..0ead1d943 100644 --- a/inference-router/src/routes/handoff/mod.rs +++ b/inference-router/src/routes/handoff/mod.rs @@ -2,10 +2,6 @@ // Licensed under the MIT License. //! handoff route handlers and router builders. -//! -//! Extracted from `routes/mod.rs` as part of the Q1 split. -//! Function bodies are byte-identical to the originals (verified by -//! `item-manifest` drift-check). use axum::Json; use axum::Router; @@ -151,6 +147,10 @@ async fn handoff_init_handler( State(state): State, Json(body): Json, ) -> axum::response::Response { + if state.inference_budget.is_some() { + return errors::flat(StatusCode::FORBIDDEN, + "Governed inference account lineage is not transferable by handoff; retry the same Task UID through the controller").into_response(); + } // ── Registry mode guard ────────────────────────────────────────────────── // Handoff requires a global registry — both agents must be in the same // registry for identity succession to work. diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index d7dce2e2e..3075f6293 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -272,6 +272,9 @@ async fn completions( Ok((status, _, resp_body)) => (status, Body::from(resp_body)).into_response(), Err(e) => { tracing::error!("Proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } StatusCode::BAD_GATEWAY.into_response() } } @@ -414,6 +417,9 @@ async fn responses( } Err(e) => { tracing::error!(sandbox = %sandbox_name, "Responses proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } errors::openai( StatusCode::BAD_GATEWAY, "Failed to reach inference backend", @@ -464,6 +470,9 @@ async fn embeddings( Ok((status, _, resp_body)) => (status, Body::from(resp_body)).into_response(), Err(e) => { tracing::error!("Proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } StatusCode::BAD_GATEWAY.into_response() } } @@ -544,6 +553,9 @@ async fn images_generations( } Err(e) => { tracing::error!(deployment = %deployment, "Image generation proxy error: {e:#}"); + if let Some(response) = crate::inference_budget::response::denial(&e) { + return response; + } ( StatusCode::BAD_GATEWAY, Json(serde_json::json!({"error": {"message": format!("Image generation proxy error: {e}")}})), @@ -801,6 +813,14 @@ async fn foundry_proxy( headers: HeaderMap, body: Bytes, ) -> impl IntoResponse { + if state.inference_budget.is_some() { + return errors::openai_coded( + StatusCode::FORBIDDEN, + "This Foundry operation is outside the closed governed-inference contract; use a supported inference route", + "unsupported_budget_operation", + "unsupported_budget_operation", + ).into_response(); + } let sandbox_name = resolve_sandbox_name(&headers); // The raw path is concatenated into the upstream URL below, and diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index a2eb5a6ea..33f24408c 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -84,6 +84,7 @@ pub struct AppState { pub client: reqwest::Client, pub config: Arc, pub budget: TokenBudgetTracker, + pub inference_budget: Option>, pub governance: Arc, /// Four-seam policy contract view of `governance`. Today it's the same /// `Arc` coerced to `Arc` — the @@ -319,6 +320,13 @@ impl AppState { )); let blocked_egress = Arc::new(BlockedBuffer::with_defaults()); blocked_egress.bind_services(&services); + let inference_budget = crate::inference_budget::Client::from_env()?; + if let Some(client) = &inference_budget { + blocklist.bind_inference_budget( + client.clone(), + crate::inference_budget::egress::model_hosts(&config), + )?; + } Ok(Self { services, auth: Arc::new(WorkloadIdentityAuth::new()), @@ -326,6 +334,7 @@ impl AppState { client: client.clone(), config: Arc::new(config), budget, + inference_budget, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, @@ -378,6 +387,7 @@ impl AppState { .unwrap_or_else(|| self.config.default_model.clone()); let mut upstream = UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()); + upstream.inference_budget = self.inference_budget.clone(); if self.services.identity_valid { upstream.telemetry = Some(self.services.telemetry.clone()); } @@ -522,6 +532,25 @@ async fn healthz() -> &'static str { } async fn readyz(State(state): State) -> impl IntoResponse { + if state.inference_budget.is_some() { + return match crate::inference_budget::readiness::ready( + &state, + state.upstream_config(&state.sandbox_name), + ) + .await + { + Ok(()) => ( + StatusCode::OK, + "governed inference authority and model contracts available", + ) + .into_response(), + Err(_) => ( + StatusCode::SERVICE_UNAVAILABLE, + "not ready — governed inference authority, provider or contracts unavailable", + ) + .into_response(), + }; + } // Check that we can acquire a token (validates Workload Identity / IMDS setup) let audience = if state .config diff --git a/inference-router/src/routes/model_routing.rs b/inference-router/src/routes/model_routing.rs index 9ecd35dc0..a71c19ea8 100644 --- a/inference-router/src/routes/model_routing.rs +++ b/inference-router/src/routes/model_routing.rs @@ -382,6 +382,7 @@ mod tests { client: reqwest::Client::new(), config: Arc::new(config), budget: crate::budget::TokenBudgetTracker::new(0, 0), + inference_budget: None, policy_provider: governance.clone(), audit_sink: governance.clone(), signing_provider: governance.clone(), diff --git a/inference-router/src/routes/spawn_policy.rs b/inference-router/src/routes/spawn_policy.rs index c61926e12..6be93e2cb 100644 --- a/inference-router/src/routes/spawn_policy.rs +++ b/inference-router/src/routes/spawn_policy.rs @@ -24,6 +24,12 @@ pub async fn check_sandbox_spawn( parent: &str, child: &str, ) -> Result<(), Box> { + if state.inference_budget.is_some() { + return Err(Box::new(errors::flat( + StatusCode::FORBIDDEN, + "Finite governed inference requires controller-enrolled Task UID delegation; plain Sandbox spawn is unsupported", + ).into_response())); + } let request = PolicyRequest { principal: parent.to_string(), tool: format!("spawn:create:{child}"), diff --git a/inference-router/src/spawn/mod.rs b/inference-router/src/spawn/mod.rs index 9114012fd..3d9aed3d6 100644 --- a/inference-router/src/spawn/mod.rs +++ b/inference-router/src/spawn/mod.rs @@ -129,6 +129,11 @@ pub async fn create_sandbox( parent_name: &str, req: &SpawnRequest, ) -> Result { + if std::env::var("KARS_INFERENCE_BUDGET_REQUIRED") + .is_ok_and(|value| value != "false" && !value.is_empty()) + { + return Err("Governed inference cannot delegate through plain Sandbox spawn; use UID-bound Task delegation".into()); + } // Validate name: must be DNS-safe if req.agent_id.is_empty() || req.agent_id.len() > 63 { return Err("name must be 1-63 characters".into()); diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index 016ec6386..d0875fd91 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -65,6 +65,7 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/anthropic_buffered_guardrail.rs b/inference-router/tests/anthropic_buffered_guardrail.rs index 2151fddc9..c727cb10c 100644 --- a/inference-router/tests/anthropic_buffered_guardrail.rs +++ b/inference-router/tests/anthropic_buffered_guardrail.rs @@ -76,6 +76,7 @@ fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppSta providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/chat_output_guardrail_nonjson.rs b/inference-router/tests/chat_output_guardrail_nonjson.rs index ad80ae9cf..5c5cd6946 100644 --- a/inference-router/tests/chat_output_guardrail_nonjson.rs +++ b/inference-router/tests/chat_output_guardrail_nonjson.rs @@ -74,6 +74,7 @@ fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/common/governed_services.rs b/inference-router/tests/common/governed_services.rs index a59804578..9be102b79 100644 --- a/inference-router/tests/common/governed_services.rs +++ b/inference-router/tests/common/governed_services.rs @@ -82,6 +82,7 @@ pub fn state(workspace: &str, uid: &str) -> AppState { providers: Default::default(), }), budget: TokenBudgetTracker::new(1000, 100), + inference_budget: None, policy_provider: governance.clone(), audit_sink: governance.clone(), signing_provider: governance.clone(), diff --git a/inference-router/tests/egress_blocked_endpoint.rs b/inference-router/tests/egress_blocked_endpoint.rs index 6b8c8bfaf..81048db4b 100644 --- a/inference-router/tests/egress_blocked_endpoint.rs +++ b/inference-router/tests/egress_blocked_endpoint.rs @@ -63,6 +63,7 @@ fn test_state() -> AppState { providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/failover_walk.rs b/inference-router/tests/failover_walk.rs index 41fe902a5..8cf13d68e 100644 --- a/inference-router/tests/failover_walk.rs +++ b/inference-router/tests/failover_walk.rs @@ -132,6 +132,7 @@ async fn primary_503_falls_through_to_fallback_200() { let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -200,6 +201,7 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -260,6 +262,7 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: base, deployment: "primary-down".into(), sandbox_name: "sbx".into(), diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index b5913fecb..cff2878b8 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -84,6 +84,7 @@ fn test_state() -> AppState { providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/multi_provider_guardrails.rs b/inference-router/tests/multi_provider_guardrails.rs index 6f18201d8..0b3edb03e 100644 --- a/inference-router/tests/multi_provider_guardrails.rs +++ b/inference-router/tests/multi_provider_guardrails.rs @@ -62,6 +62,7 @@ async fn ollama_provider_forwards_openai_compat_without_auth() { let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: server.uri(), deployment: "llama3.1".into(), sandbox_name: "test-sandbox".into(), @@ -120,6 +121,7 @@ async fn anthropic_provider_forwards_messages_with_router_held_key() { let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint: server.uri(), deployment: "claude-sonnet-4-5".into(), sandbox_name: "test-sandbox".into(), diff --git a/inference-router/tests/policy_status_endpoint.rs b/inference-router/tests/policy_status_endpoint.rs index 8c1699b68..57bbe1894 100644 --- a/inference-router/tests/policy_status_endpoint.rs +++ b/inference-router/tests/policy_status_endpoint.rs @@ -74,6 +74,7 @@ fn test_state() -> (AppState, Arc) { providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), + inference_budget: None, policy_provider: Arc::clone(&governance) as Arc, audit_sink: Arc::clone(&governance) as Arc, signing_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/proxy_fake_upstream.rs b/inference-router/tests/proxy_fake_upstream.rs index 4475036d1..bbaa6144a 100644 --- a/inference-router/tests/proxy_fake_upstream.rs +++ b/inference-router/tests/proxy_fake_upstream.rs @@ -83,6 +83,7 @@ async fn api_key_mode_proxies_chat_completion_with_filter_results() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox".to_string(), @@ -165,6 +166,7 @@ async fn wi_mode_falls_back_to_imds_and_proxies_embeddings() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint, deployment: "text-embedding-3-small".to_string(), sandbox_name: "test-sandbox-wi".to_string(), @@ -239,6 +241,7 @@ async fn upstream_error_status_is_propagated() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { telemetry: None, + inference_budget: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox-429".to_string(), diff --git a/shared/inference_budget/catalog.rs b/shared/inference_budget/catalog.rs new file mode 100644 index 000000000..84aed096d --- /dev/null +++ b/shared/inference_budget/catalog.rs @@ -0,0 +1,270 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + tariffs::{ModelContract, Operation, Quote}, + types::*, +}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::BTreeSet; + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Catalog { + pub version: String, + pub contracts: Vec, + /// Operator-declared non-inference destinations for mediated HTTP egress. This is + /// not a free-cost assertion: those costs are outside governed inference. + #[serde(default)] + pub non_inference_egress_hosts: Vec, +} + +impl Catalog { + pub fn validate(&self, _now: i64) -> Result<(), BudgetError> { + if !valid_uid(&self.version) + || self.contracts.is_empty() + || self.contracts.len() > 64 + || self.non_inference_egress_hosts.len() > 64 + { + return Err(BudgetError::Contract); + } + let mut routes = BTreeSet::new(); + let mut versions = BTreeSet::new(); + for contract in &self.contracts { + contract.validate(i64::MIN)?; + let route = format!( + "{}|{}|{}|{:?}", + contract.provider_id, + contract.endpoint.trim_end_matches('/'), + contract.model, + contract.operation + ); + if !routes.insert(route) + || !versions.insert((contract.id.clone(), contract.version.clone())) + { + return Err(BudgetError::Contract); + } + } + for host in &self.non_inference_egress_hosts { + if !valid_name(host) + || host != &host.to_ascii_lowercase() + || self + .contracts + .iter() + .any(|contract| endpoint_host(&contract.endpoint) == Some(host.as_str())) + { + return Err(BudgetError::Contract); + } + } + Ok(()) + } + + pub fn select( + &self, + provider_id: &str, + endpoint: &str, + model: &str, + operation: Operation, + now: i64, + ) -> Result<&ModelContract, BudgetError> { + self.validate(now)?; + let selected = self + .contracts + .iter() + .find(|contract| { + contract.provider_id == provider_id + && contract.endpoint.trim_end_matches('/') == endpoint.trim_end_matches('/') + && contract.model == model + && contract.operation == operation + }) + .ok_or(BudgetError::Contract)?; + selected.validate(now)?; + Ok(selected) + } + + /// The broker validates quotes against its own current operator catalog. + /// A caller cannot supply its own maximum rates, expiry, or token bounds. + pub fn accepts_quote( + &self, + quote: &Quote, + now: i64, + money_required: bool, + ) -> Result<(), BudgetError> { + let stored = self.select( + "e.contract.provider_id, + "e.contract.endpoint, + "e.contract.model, + quote.contract.operation, + now, + )?; + if stored != "e.contract { + return Err(BudgetError::Contract); + } + quote.validate(now, money_required) + } + + pub fn allows_mediated_non_inference( + &self, + host: &str, + configured_model_hosts: &[String], + ) -> bool { + let host = host.trim_end_matches('.').to_ascii_lowercase(); + self.non_inference_egress_hosts + .iter() + .any(|allowed| allowed == &host) + && !configured_model_hosts + .iter() + .any(|model| model.eq_ignore_ascii_case(&host)) + && !self + .contracts + .iter() + .any(|contract| endpoint_host(&contract.endpoint) == Some(host.as_str())) + } +} + +pub fn endpoint_host(endpoint: &str) -> Option<&str> { + let (_, rest) = endpoint.split_once("://")?; + let authority = rest.split('/').next()?; + if authority.is_empty() || authority.contains('@') || authority.starts_with('[') { + return None; + } + Some(authority.split(':').next()?) +} + +/// Exact, closed final-dispatch path classification. No substring such as +/// "completion" grants access to an unimplemented provider operation. +pub fn operation(path: &str) -> Option { + if path.contains(['?', '#']) { + return None; + } + let path = path.trim_matches('/'); + match path { + "chat/completions" | "v1/chat/completions" | "openai/v1/chat/completions" => { + Some(Operation::ChatCompletions) + } + "messages" | "v1/messages" | "anthropic/v1/messages" => Some(Operation::AnthropicMessages), + "responses" | "v1/responses" | "openai/v1/responses" => Some(Operation::Responses), + _ => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::inference_budget_contract::tariffs::{MaximumPrice, OutputField}; + + fn catalog() -> Catalog { + Catalog { + version: "catalog-v1".into(), + contracts: vec![ModelContract { + id: "chat".into(), + version: "v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "provider".into(), + endpoint: "https://models.example".into(), + model: "model".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxTokens, + maximum_input_tokens: 100, + maximum_output_tokens: 50, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 10 }), + }], + non_inference_egress_hosts: vec!["api.github.com".into()], + } + } + + #[test] + fn no_provider_model_or_operation_fallback_implicitly_acquires_a_contract() { + let catalog = catalog(); + assert!( + catalog + .select( + "provider", + "https://models.example", + "model", + Operation::ChatCompletions, + 1 + ) + .is_ok() + ); + assert!( + catalog + .select( + "other", + "https://models.example", + "model", + Operation::ChatCompletions, + 1 + ) + .is_err() + ); + assert!( + catalog + .select( + "provider", + "https://other.example", + "model", + Operation::ChatCompletions, + 1 + ) + .is_err() + ); + assert!( + catalog + .select( + "provider", + "https://models.example", + "other-model", + Operation::ChatCompletions, + 1 + ) + .is_err() + ); + assert!( + catalog + .select( + "provider", + "https://models.example", + "model", + Operation::Responses, + 1 + ) + .is_err() + ); + } + + #[test] + fn opaque_model_tunnels_are_not_non_inference_cost_exclusions() { + let catalog = catalog(); + assert!(catalog.allows_mediated_non_inference("api.github.com", &[])); + assert!(!catalog.allows_mediated_non_inference("models.example", &[])); + assert!(!catalog.allows_mediated_non_inference("unknown.example", &[])); + assert!( + !catalog.allows_mediated_non_inference("api.github.com", &["api.github.com".into()]) + ); + } + + #[test] + fn unsupported_hidden_or_async_generation_never_matches_a_text_contract() { + for path in [ + "embeddings", + "completions", + "images/generations", + "agents/run", + "openai/fine-tuning/jobs", + "memory_stores/search", + "openai/containers", + "knowledgebases/retrieve", + ] { + assert!(operation(path).is_none(), "{path}"); + } + assert_eq!(operation("/v1/responses"), Some(Operation::Responses)); + assert_eq!( + operation("/anthropic/v1/messages"), + Some(Operation::AnthropicMessages) + ); + } +} diff --git a/shared/inference_budget/ledger.rs b/shared/inference_budget/ledger.rs new file mode 100644 index 000000000..bb3f69695 --- /dev/null +++ b/shared/inference_budget/ledger.rs @@ -0,0 +1,748 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Single-account transitions. A caller must persist `next` with the account +//! UID/resourceVersion CAS before returning a grant or performing a send. + +use super::{tariffs::Quote, types::*}; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use std::collections::{BTreeMap, BTreeSet}; + +#[path = "ledger_lifecycle.rs"] +mod lifecycle; +#[path = "ledger_validation.rs"] +mod validation; + +#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Meters { + pub reserved: Amounts, + pub settled: Amounts, + pub uncertain: Amounts, + pub unpriced_attempts: u64, +} + +impl Meters { + pub fn total(&self) -> Result { + self.reserved.add(self.settled)?.add(self.uncertain) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Node { + pub authority: TaskAuthority, + pub active: bool, + pub meters: Meters, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Session { + pub identity: ExecutionIdentity, + pub closed: bool, + pub next_sequence: u64, + /// Terminal attempts through this sequence are never reissued, even after + /// their detailed rows have been compacted. + pub closed_through: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Attempt { + pub key: AttemptKey, + pub identity: ExecutionIdentity, + pub wire_digest: String, + pub quote: Quote, + pub phase: AttemptPhase, + pub expires_at: i64, + pub charged: Amounts, + #[serde(default, skip_serializing_if = "Option::is_none")] + /// Decimal JSON text preserves even absurd over-bound provider integers + /// without coercing them through Kubernetes' signed integer representation. + pub observed_breach: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Ledger { + pub version: String, + pub scope: BudgetScope, + pub account_uid: String, + pub root: RootIdentity, + pub limits: Limits, + pub phase: AccountPhase, + pub meters: Meters, + pub nodes: BTreeMap, + /// Retain closed session identities and high-water marks. Capacity is + /// explicit: never discard a replay fence to make room for new work. + pub sessions: BTreeMap, + pub attempts: BTreeMap, +} + +#[derive(Clone, Debug)] +pub struct Mutation { + pub next: Ledger, + pub value: T, + pub changed: bool, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct Reservation { + pub key: AttemptKey, + pub maximum: Amounts, + pub expires_at: i64, + pub phase: AttemptPhase, +} + +#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct SettlementResult { + pub finalized: bool, + pub uncertain: bool, + pub breach: bool, + /// None means the idempotent terminal row was already compacted. It is not + /// a zero charge, and never authorizes another refund. + pub charged: Option, +} + +impl Ledger { + pub fn new( + account_uid: String, + root: RootIdentity, + limits: Limits, + ) -> Result { + root.validate()?; + if !valid_uid(&account_uid) { + return Err(BudgetError::Identity); + } + Ok(Self { + version: CONTRACT_VERSION.into(), + scope: BudgetScope::GovernedInference, + account_uid, + root, + limits: limits.normalized(), + phase: AccountPhase::Active, + meters: Meters::default(), + nodes: BTreeMap::new(), + sessions: BTreeMap::new(), + attempts: BTreeMap::new(), + }) + } + + fn mutation(&self, next: Self, value: T) -> Result, BudgetError> { + next.validate()?; + let changed = *self != next; + Ok(Mutation { + next, + value, + changed, + }) + } + + fn accepting(&self) -> Result<(), BudgetError> { + if self.phase != AccountPhase::Active { + return Err(BudgetError::Closed); + } + Ok(()) + } + + pub fn ancestors(&self, task_uid: &str) -> Result, BudgetError> { + let mut path = Vec::new(); + let mut seen = BTreeSet::new(); + let mut current = Some(task_uid); + while let Some(uid) = current { + if path.len() >= MAX_NODES || !seen.insert(uid) { + return Err(BudgetError::Corrupt); + } + let node = self.nodes.get(uid).ok_or(BudgetError::Authorization)?; + path.push(uid.to_string()); + current = node.authority.parent_uid.as_deref(); + } + Ok(path) + } + + fn authorized(&self, identity: &ExecutionIdentity) -> Result, BudgetError> { + identity.validate()?; + let path = self.ancestors(&identity.task_uid)?; + let node = self + .nodes + .get(&identity.task_uid) + .ok_or(BudgetError::Authorization)?; + if identity.authorization_digest != node.authority.authorization_digest + || identity.sandbox.namespace != self.root.resource.namespace + || path.iter().any(|uid| !self.nodes[uid].active) + { + return Err(BudgetError::Authorization); + } + Ok(path) + } + + pub fn requires_price(&self, task_uid: &str) -> Result { + Ok(self.limits.normalized().usd_micros.is_some() + || self.ancestors(task_uid)?.iter().any(|uid| { + self.nodes[uid] + .authority + .limits + .normalized() + .usd_micros + .is_some() + })) + } + + pub fn requires_enforcement(&self, task_uid: &str) -> Result { + Ok(self.limits.finite() + || self + .ancestors(task_uid)? + .iter() + .any(|uid| self.nodes[uid].authority.limits.finite())) + } + + pub fn register_task(&self, mut authority: TaskAuthority) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + authority.task.validate()?; + authority.limits = authority.limits.normalized(); + if authority.task.namespace != self.root.resource.namespace + || !valid_uid(&authority.root_task_uid) + || !valid_digest(&authority.authorization_digest) + || !authority.effective_authorization.is_object() + || serde_json::to_vec(&authority.effective_authorization) + .map_err(|_| BudgetError::Corrupt)? + .len() + > 65_536 + { + return Err(BudgetError::Authorization); + } + if let Some(existing) = self.nodes.get(&authority.task.uid) { + if existing.authority != authority || !existing.active { + return Err(BudgetError::Authorization); + } + return self.mutation(self.clone(), ()); + } + if self.nodes.len() >= MAX_NODES { + return Err(BudgetError::Capacity); + } + if let Some(parent_uid) = &authority.parent_uid { + let parent = self + .nodes + .get(parent_uid) + .ok_or(BudgetError::Authorization)?; + if !parent.active + || parent.authority.root_task_uid != authority.root_task_uid + || !authority.limits.attenuates(parent.authority.limits) + { + return Err(BudgetError::Authorization); + } + } else if authority.root_task_uid != authority.task.uid + || !authority.limits.attenuates(self.limits) + || (self.root.kind == RootKind::KarsTask && authority.task != self.root.resource) + { + return Err(BudgetError::Authorization); + } + let mut next = self.clone(); + next.nodes.insert( + authority.task.uid.clone(), + Node { + authority, + active: true, + meters: Meters::default(), + }, + ); + self.mutation(next, ()) + } + + /// Controller-only attenuation. Immutable UID ancestry and prior charges + /// survive a new full authorization snapshot; old router identities cannot + /// reserve/begin under the new digest. + pub fn update_authority(&self, authority: TaskAuthority) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + let old = self + .nodes + .get(&authority.task.uid) + .ok_or(BudgetError::Authorization)?; + if authority.task != old.authority.task + || authority.parent_uid != old.authority.parent_uid + || authority.root_task_uid != old.authority.root_task_uid + || !authority.limits.attenuates(old.authority.limits) + || !authority.limits.allows(old.meters.total()?) + || !valid_digest(&authority.authorization_digest) + || !authority.effective_authorization.is_object() + { + return Err(BudgetError::Authorization); + } + if authority.limits.normalized().usd_micros.is_some() + && (old.meters.unpriced_attempts > 0 + || self.attempts.values().any(|attempt| { + !attempt.phase.terminal() + && !attempt.quote.price_covered + && self + .ancestors(&attempt.identity.task_uid) + .is_ok_and(|path| path.contains(&authority.task.uid)) + })) + { + return Err(BudgetError::Contract); + } + let mut next = self.clone(); + next.nodes + .get_mut(&authority.task.uid) + .ok_or(BudgetError::Corrupt)? + .authority = authority; + self.mutation(next, ()) + } + + pub fn register_session( + &self, + identity: ExecutionIdentity, + ) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + self.authorized(&identity)?; + if let Some(session) = self.sessions.get(&identity.pod_uid) { + if session.identity != identity || session.closed { + return Err(BudgetError::Identity); + } + return self.mutation(self.clone(), session.clone()); + } + if self.sessions.len() >= MAX_SESSIONS { + return Err(BudgetError::Capacity); + } + let session = Session { + identity: identity.clone(), + closed: false, + next_sequence: 1, + closed_through: 0, + }; + let mut next = self.clone(); + next.sessions + .insert(identity.pod_uid.clone(), session.clone()); + self.mutation(next, session) + } + + pub fn reserve( + &self, + request: &ReserveRequest, + now: i64, + ) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + if request.account_uid != self.account_uid || !valid_digest(&request.wire_digest) { + return Err(BudgetError::Identity); + } + let path = self.authorized(&request.identity)?; + let session = self + .sessions + .get(&request.identity.pod_uid) + .ok_or(BudgetError::Identity)?; + if session.identity != request.identity || session.closed { + return Err(BudgetError::Identity); + } + if request.sequence <= session.closed_through { + return Err(BudgetError::Sequence); + } + let key = AttemptKey { + pod_uid: request.identity.pod_uid.clone(), + sequence: request.sequence, + }; + if let Some(existing) = self.attempts.get(&key.storage_key()) { + if existing.identity != request.identity + || existing.wire_digest != request.wire_digest + || existing.quote != request.quote + { + return Err(BudgetError::Identity); + } + return self.mutation( + self.clone(), + Reservation { + key, + maximum: existing.quote.maximum, + expires_at: existing.expires_at, + phase: existing.phase, + }, + ); + } + if request.sequence != session.next_sequence { + return Err(BudgetError::Sequence); + } + if self.attempts.len() >= MAX_ATTEMPTS + || session + .next_sequence + .checked_sub(session.closed_through) + .ok_or(BudgetError::Corrupt)? + > MAX_REPLAY_WINDOW + { + return Err(BudgetError::Capacity); + } + request + .quote + .validate(now, self.requires_price(&request.identity.task_uid)?)?; + let maximum = request.quote.maximum; + if !self.limits.allows(self.meters.total()?.add(maximum)?) + || path.iter().any(|uid| { + let node = &self.nodes[uid]; + node.meters + .total() + .and_then(|total| total.add(maximum)) + .map_or(true, |total| !node.authority.limits.allows(total)) + }) + { + return Err(BudgetError::Exhausted); + } + let expires_at = now + .checked_add(RESERVATION_TTL_SECONDS) + .ok_or(BudgetError::Overflow)? + .min( + chrono::DateTime::parse_from_rfc3339(&request.quote.contract.valid_until) + .map_err(|_| BudgetError::Contract)? + .timestamp(), + ); + let mut next = self.clone(); + next.meters.reserved = next.meters.reserved.add(maximum)?; + for uid in path { + let node = next.nodes.get_mut(&uid).ok_or(BudgetError::Corrupt)?; + node.meters.reserved = node.meters.reserved.add(maximum)?; + } + next.sessions + .get_mut(&key.pod_uid) + .ok_or(BudgetError::Corrupt)? + .next_sequence = request + .sequence + .checked_add(1) + .ok_or(BudgetError::Overflow)?; + next.attempts.insert( + key.storage_key(), + Attempt { + key: key.clone(), + identity: request.identity.clone(), + wire_digest: request.wire_digest.clone(), + quote: request.quote.clone(), + phase: AttemptPhase::Reserved, + expires_at, + charged: Amounts::default(), + observed_breach: None, + }, + ); + self.mutation( + next, + Reservation { + key, + maximum, + expires_at, + phase: AttemptPhase::Reserved, + }, + ) + } + + fn attempt(&self, command: &AttemptCommand) -> Result<&Attempt, BudgetError> { + if command.account_uid != self.account_uid { + return Err(BudgetError::Identity); + } + let attempt = self + .attempts + .get(&command.key.storage_key()) + .ok_or(BudgetError::Sequence)?; + if attempt.key != command.key + || attempt.identity != command.identity + || attempt.wire_digest != command.wire_digest + { + return Err(BudgetError::Identity); + } + Ok(attempt) + } + + pub fn begin_dispatch( + &self, + command: &AttemptCommand, + now: i64, + ) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + self.authorized(&command.identity)?; + let session = self + .sessions + .get(&command.key.pod_uid) + .ok_or(BudgetError::Identity)?; + if session.closed || session.identity != command.identity { + return Err(BudgetError::Identity); + } + let attempt = self.attempt(command)?; + if attempt.phase != AttemptPhase::Reserved { + return Err(BudgetError::AlreadyDispatched); + } + if now >= attempt.expires_at { + return Err(BudgetError::Expired); + } + attempt + .quote + .validate(now, self.requires_price(&attempt.identity.task_uid)?)?; + let response = Reservation { + key: attempt.key.clone(), + maximum: attempt.quote.maximum, + expires_at: attempt.expires_at, + phase: AttemptPhase::InFlight, + }; + let mut next = self.clone(); + next.attempts + .get_mut(&command.key.storage_key()) + .ok_or(BudgetError::Corrupt)? + .phase = AttemptPhase::InFlight; + self.mutation(next, response) + } + + fn charge( + &mut self, + key: &AttemptKey, + phase: AttemptPhase, + charged: Amounts, + ) -> Result<(), BudgetError> { + let attempt = self + .attempts + .get(&key.storage_key()) + .ok_or(BudgetError::Corrupt)? + .clone(); + if attempt.phase.terminal() { + return Err(BudgetError::Corrupt); + } + let path = self.ancestors(&attempt.identity.task_uid)?; + let apply = |meters: &mut Meters| -> Result<(), BudgetError> { + meters.reserved = meters.reserved.subtract(attempt.quote.maximum)?; + if phase == AttemptPhase::Uncertain { + meters.uncertain = meters.uncertain.add(charged)?; + } else { + meters.settled = meters.settled.add(charged)?; + } + if phase != AttemptPhase::Expired && !attempt.quote.price_covered { + meters.unpriced_attempts = meters + .unpriced_attempts + .checked_add(1) + .ok_or(BudgetError::Overflow)?; + } + Ok(()) + }; + apply(&mut self.meters)?; + for uid in path { + apply(&mut self.nodes.get_mut(&uid).ok_or(BudgetError::Corrupt)?.meters)?; + } + let attempt = self + .attempts + .get_mut(&key.storage_key()) + .ok_or(BudgetError::Corrupt)?; + attempt.phase = phase; + attempt.charged = charged; + Ok(()) + } + + pub fn settle( + &self, + settlement: &Settlement, + ) -> Result, BudgetError> { + self.validate()?; + let command = &settlement.attempt; + if command.account_uid != self.account_uid { + return Err(BudgetError::Identity); + } + let session = self + .sessions + .get(&command.key.pod_uid) + .ok_or(BudgetError::Identity)?; + if session.identity != command.identity { + return Err(BudgetError::Identity); + } + if command.key.sequence <= session.closed_through { + return self.mutation( + self.clone(), + SettlementResult { + finalized: true, + uncertain: false, + breach: false, + charged: None, + }, + ); + } + let attempt = self.attempt(command)?; + if attempt.phase.terminal() { + if attempt.phase == AttemptPhase::Uncertain + && let Some(usage) = &settlement.usage + && attempt.quote.usage(usage).is_err() + { + let mut next = self.clone(); + next.phase = AccountPhase::Frozen; + next.attempts + .get_mut(&command.key.storage_key()) + .ok_or(BudgetError::Corrupt)? + .observed_breach = + Some(serde_json::to_string(usage).map_err(|_| BudgetError::Corrupt)?); + return self.mutation( + next, + SettlementResult { + finalized: true, + uncertain: true, + breach: true, + charged: Some(attempt.charged), + }, + ); + } + return self.mutation( + self.clone(), + SettlementResult { + finalized: true, + uncertain: attempt.phase == AttemptPhase::Uncertain, + breach: attempt.observed_breach.is_some(), + charged: Some(attempt.charged), + }, + ); + } + if attempt.phase != AttemptPhase::InFlight { + return Err(BudgetError::AlreadyDispatched); + } + let (phase, charge, breach) = match &settlement.usage { + None => (AttemptPhase::Uncertain, attempt.quote.maximum, false), + Some(usage) => match attempt.quote.usage(usage) { + Ok(amount) if amount.within(attempt.quote.maximum) => { + (AttemptPhase::Settled, amount, false) + } + _ => (AttemptPhase::Uncertain, attempt.quote.maximum, true), + }, + }; + let mut next = self.clone(); + next.charge(&command.key, phase, charge)?; + if breach { + next.phase = AccountPhase::Frozen; + next.attempts + .get_mut(&command.key.storage_key()) + .ok_or(BudgetError::Corrupt)? + .observed_breach = settlement + .usage + .as_ref() + .map(serde_json::to_string) + .transpose() + .map_err(|_| BudgetError::Corrupt)?; + } + next.compact()?; + self.mutation( + next, + SettlementResult { + finalized: true, + uncertain: phase == AttemptPhase::Uncertain, + breach, + charged: Some(charge), + }, + ) + } + + /// Only Reserved attempts expire/refund. InFlight is potentially billed, + /// including when BeginDispatch's acknowledgement was lost. + pub fn expire_undispatched(&self, now: i64) -> Result, BudgetError> { + self.validate()?; + let expired: Vec<_> = self + .attempts + .values() + .filter(|attempt| attempt.phase == AttemptPhase::Reserved && now >= attempt.expires_at) + .map(|attempt| attempt.key.clone()) + .collect(); + let mut next = self.clone(); + for key in &expired { + next.charge(key, AttemptPhase::Expired, Amounts::default())?; + } + next.compact()?; + self.mutation(next, expired.len()) + } + + /// Cancellation is an atomic subtree fence. Undispatched work is released; + /// already authorized attempts are conservatively charged, never refunded. + pub fn close_subtree(&self, task_uid: &str) -> Result, BudgetError> { + self.validate()?; + if !self.nodes.contains_key(task_uid) { + return Err(BudgetError::Identity); + } + let affected: BTreeSet<_> = self + .nodes + .keys() + .filter(|uid| { + self.ancestors(uid) + .is_ok_and(|path| path.iter().any(|parent| parent == task_uid)) + }) + .cloned() + .collect(); + let mut next = self.clone(); + for uid in &affected { + next.nodes.get_mut(uid).ok_or(BudgetError::Corrupt)?.active = false; + } + for session in next + .sessions + .values_mut() + .filter(|session| affected.contains(&session.identity.task_uid)) + { + session.closed = true; + } + let attempts: Vec<_> = next + .attempts + .values() + .filter(|attempt| { + affected.contains(&attempt.identity.task_uid) && !attempt.phase.terminal() + }) + .map(|attempt| (attempt.key.clone(), attempt.phase, attempt.quote.maximum)) + .collect(); + for (key, phase, maximum) in attempts { + if phase == AttemptPhase::Reserved { + next.charge(&key, AttemptPhase::Expired, Amounts::default())?; + } else { + next.charge(&key, AttemptPhase::Uncertain, maximum)?; + } + } + next.compact()?; + self.mutation(next, ()) + } + + pub fn close_account(&self) -> Result, BudgetError> { + self.validate()?; + let mut next = self.clone(); + next.phase = AccountPhase::Closing; + let roots: Vec<_> = next + .nodes + .values() + .filter(|node| node.authority.parent_uid.is_none()) + .map(|node| node.authority.task.uid.clone()) + .collect(); + for root in roots { + next = next.close_subtree(&root)?.next; + } + next.phase = AccountPhase::Closed; + self.mutation(next, ()) + } + + fn compact(&mut self) -> Result<(), BudgetError> { + for (pod_uid, session) in &mut self.sessions { + loop { + let sequence = session + .closed_through + .checked_add(1) + .ok_or(BudgetError::Overflow)?; + let key = AttemptKey { + pod_uid: pod_uid.clone(), + sequence, + } + .storage_key(); + let Some(attempt) = self.attempts.get(&key) else { + break; + }; + if !attempt.phase.terminal() + || attempt.phase == AttemptPhase::Uncertain + || attempt.observed_breach.is_some() + { + break; + } + self.attempts.remove(&key); + session.closed_through = sequence; + } + } + Ok(()) + } +} + +#[cfg(test)] +#[path = "ledger_tests.rs"] +mod tests; diff --git a/shared/inference_budget/ledger_lifecycle.rs b/shared/inference_budget/ledger_lifecycle.rs new file mode 100644 index 000000000..f0bfebfd7 --- /dev/null +++ b/shared/inference_budget/ledger_lifecycle.rs @@ -0,0 +1,113 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; + +impl Ledger { + /// An interrupted initial registration has no session or reservation for + /// this UID. Ledger validation proves that a missing node cannot carry + /// outstanding attempts; do not strand a deleting, never-enrolled Task. + pub fn close_registered_task(&self, task_uid: &str) -> Result, BudgetError> { + self.validate()?; + if !valid_uid(task_uid) { + return Err(BudgetError::Identity); + } + if self.nodes.contains_key(task_uid) { + return self.close_subtree(task_uid); + } + self.mutation(self.clone(), ()) + } + + /// Controller-verified relaunch of the same Task UID keeps all spending. + /// Closed Pod sessions remain closed; a new runtime Pod UID must enroll. + pub fn resume_task(&self, task_uid: &str, digest: &str) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + let node = self.nodes.get(task_uid).ok_or(BudgetError::Authorization)?; + if node.authority.authorization_digest != digest + || self + .ancestors(task_uid)? + .iter() + .skip(1) + .any(|uid| !self.nodes[uid].active) + { + return Err(BudgetError::Authorization); + } + let mut next = self.clone(); + next.nodes + .get_mut(task_uid) + .ok_or(BudgetError::Corrupt)? + .active = true; + self.mutation(next, ()) + } + + pub fn close_session(&self, pod_uid: &str) -> Result, BudgetError> { + self.validate()?; + if !self.sessions.contains_key(pod_uid) { + return Err(BudgetError::Identity); + } + let mut next = self.clone(); + next.sessions + .get_mut(pod_uid) + .ok_or(BudgetError::Corrupt)? + .closed = true; + let attempts: Vec<_> = next + .attempts + .values() + .filter(|attempt| attempt.key.pod_uid == pod_uid && !attempt.phase.terminal()) + .map(|attempt| (attempt.key.clone(), attempt.phase, attempt.quote.maximum)) + .collect(); + for (key, phase, maximum) in attempts { + if phase == AttemptPhase::Reserved { + next.charge(&key, AttemptPhase::Expired, Amounts::default())?; + } else { + next.charge(&key, AttemptPhase::Uncertain, maximum)?; + } + } + next.compact()?; + self.mutation(next, ()) + } + + /// A stalled accepted attempt is permanently charged at its full maximum. + /// Its contract remains as a bounded tombstone so late over-bound evidence + /// can still freeze the account; capacity pressure never erases uncertainty. + pub fn commit_uncertain_before(&self, cutoff: i64) -> Result, BudgetError> { + self.validate()?; + let attempts: Vec<_> = self + .attempts + .values() + .filter(|attempt| { + attempt.phase == AttemptPhase::InFlight && attempt.expires_at <= cutoff + }) + .map(|attempt| (attempt.key.clone(), attempt.quote.maximum)) + .collect(); + let mut next = self.clone(); + for (key, maximum) in &attempts { + next.charge(key, AttemptPhase::Uncertain, *maximum)?; + } + next.compact()?; + self.mutation(next, attempts.len()) + } + + /// Account grant limits are immutable; the authoritative effective ledger + /// ceiling may narrow without resetting reservations or previous charges. + pub fn narrow_root_limits(&self, limits: Limits) -> Result, BudgetError> { + self.validate()?; + self.accepting()?; + if !limits.attenuates(self.limits) || !limits.allows(self.meters.total()?) { + return Err(BudgetError::Authorization); + } + if limits.normalized().usd_micros.is_some() + && (self.meters.unpriced_attempts > 0 + || self + .attempts + .values() + .any(|attempt| !attempt.phase.terminal() && !attempt.quote.price_covered)) + { + return Err(BudgetError::Contract); + } + let mut next = self.clone(); + next.limits = limits.normalized(); + self.mutation(next, ()) + } +} diff --git a/shared/inference_budget/ledger_tests.rs b/shared/inference_budget/ledger_tests.rs new file mode 100644 index 000000000..000d97ebf --- /dev/null +++ b/shared/inference_budget/ledger_tests.rs @@ -0,0 +1,455 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::inference_budget_contract::tariffs::{ + MaximumPrice, ModelContract, Operation, OutputField, +}; +use serde_json::json; + +fn resource(name: &str, uid: &str) -> ResourceIdentity { + ResourceIdentity { + namespace: "workspace".into(), + name: name.into(), + uid: uid.into(), + } +} + +fn digest() -> String { + format!("sha256:{}", "a".repeat(64)) +} + +fn task(uid: &str, parent: Option<&str>, limit: u64) -> TaskAuthority { + TaskAuthority { + task: resource(uid, uid), + parent_uid: parent.map(str::to_string), + root_task_uid: "root".into(), + authorization_digest: digest(), + effective_authorization: json!({"domain": "test", "uid": uid}), + limits: Limits { + tokens: Some(limit), + usd_micros: Some(1000), + }, + } +} + +fn identity(task: &str, pod: &str) -> ExecutionIdentity { + ExecutionIdentity { + task_uid: task.into(), + authorization_digest: digest(), + sandbox: resource(&format!("sandbox-{task}"), &format!("sandbox-{task}")), + runtime_namespace_uid: format!("namespace-{task}"), + pod_name: format!("pod-{pod}"), + pod_uid: pod.into(), + } +} + +fn ledger(limit: u64) -> Ledger { + let root = RootIdentity { + kind: RootKind::KarsTask, + resource: resource("root", "root"), + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }; + let ledger = Ledger::new( + "account-uid".into(), + root, + Limits { + tokens: Some(limit), + usd_micros: Some(1000), + }, + ) + .unwrap(); + ledger + .register_task(task("root", None, limit)) + .unwrap() + .next +} + +fn request(task: &str, pod: &str, sequence: u64) -> ReserveRequest { + let contract = ModelContract { + id: "model".into(), + version: "v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "provider".into(), + endpoint: "https://provider.example".into(), + model: "model".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxTokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 5 }), + }; + let (_, quote) = contract + .normalize( + br#"{"model":"model","messages":[{"role":"user","content":"hello"}]}"#, + 100, + true, + ) + .unwrap(); + ReserveRequest { + account_uid: "account-uid".into(), + identity: identity(task, pod), + sequence, + wire_digest: digest(), + quote, + } +} + +fn command(request: &ReserveRequest) -> AttemptCommand { + AttemptCommand { + account_uid: request.account_uid.clone(), + key: AttemptKey { + pod_uid: request.identity.pod_uid.clone(), + sequence: request.sequence, + }, + identity: request.identity.clone(), + wire_digest: request.wire_digest.clone(), + } +} + +fn usage() -> Usage { + Usage { + input_tokens: 3, + output_tokens: 5, + cached_input_tokens: 0, + cache_creation_input_tokens: 0, + reasoning_output_tokens: 0, + } +} + +#[test] +fn uncertain_tombstones_retain_contracts_and_freeze_on_late_over_bound_usage() { + let request = request("root", "pod-a", 1); + let mut ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + ledger = ledger.reserve(&request, 100).unwrap().next; + ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + ledger = ledger.commit_uncertain_before(200).unwrap().next; + assert_eq!(ledger.meters.uncertain.tokens, 30); + assert_eq!(ledger.attempts.len(), 1); + let settled = ledger + .settle(&Settlement { + attempt: command(&request), + usage: Some(Usage { + output_tokens: 21, + ..usage() + }), + }) + .unwrap(); + assert!(settled.value.breach); + assert_eq!(settled.next.phase, AccountPhase::Frozen); + assert_eq!(settled.next.meters.uncertain.tokens, 30); + assert_eq!(settled.next.sessions["pod-a"].next_sequence, 2); +} + +#[test] +fn ordinary_completed_attempts_compact_without_reissuing_their_sequences() { + let request = request("root", "pod-a", 1); + let mut ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + ledger = ledger.reserve(&request, 100).unwrap().next; + ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + ledger = ledger + .settle(&Settlement { + attempt: command(&request), + usage: Some(usage()), + }) + .unwrap() + .next; + assert!(ledger.attempts.is_empty()); + assert_eq!(ledger.sessions["pod-a"].closed_through, 1); + assert!(ledger.reserve(&request, 102).is_err()); + assert_eq!(ledger.meters.settled.tokens, 8); +} + +#[test] +fn sibling_reservations_share_one_root_ceiling_instead_of_getting_daily_copies() { + let ledger = ledger(50); + let ledger = ledger + .register_task(task("left", Some("root"), 50)) + .unwrap() + .next; + let ledger = ledger + .register_task(task("right", Some("root"), 50)) + .unwrap() + .next; + let ledger = ledger + .register_session(identity("left", "pod-left")) + .unwrap() + .next; + let ledger = ledger + .register_session(identity("right", "pod-right")) + .unwrap() + .next; + let ledger = ledger + .reserve(&request("left", "pod-left", 1), 100) + .unwrap() + .next; + assert_eq!(ledger.meters.reserved.tokens, 30); + assert!(matches!( + ledger.reserve(&request("right", "pod-right", 1), 100), + Err(BudgetError::Exhausted) + )); +} + +#[test] +fn all_ancestors_are_reserved_atomically_and_a_narrow_subtree_binds_grandchildren() { + let ledger = ledger(100) + .register_task(task("branch", Some("root"), 40)) + .unwrap() + .next; + let ledger = ledger + .register_task(task("leaf", Some("branch"), 40)) + .unwrap() + .next; + let ledger = ledger + .register_session(identity("leaf", "pod-leaf")) + .unwrap() + .next; + let ledger = ledger + .reserve(&request("leaf", "pod-leaf", 1), 100) + .unwrap() + .next; + for uid in ["root", "branch", "leaf"] { + assert_eq!(ledger.nodes[uid].meters.reserved.tokens, 30); + } + assert!(matches!( + ledger.reserve(&request("leaf", "pod-leaf", 2), 100), + Err(BudgetError::Exhausted) + )); +} + +#[test] +fn reserve_replay_is_idempotent_but_begin_dispatch_never_regrants() { + let request = request("root", "pod", 1); + let ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + let ledger = ledger.reserve(&request, 100).unwrap().next; + assert!(!ledger.reserve(&request, 101).unwrap().changed); + let ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + assert!(matches!( + ledger.begin_dispatch(&command(&request), 101), + Err(BudgetError::AlreadyDispatched) + )); + assert_eq!(ledger.meters.reserved.tokens, 30); +} + +#[test] +fn only_undispatched_reservations_expire_and_lost_begin_ack_remains_funded() { + let one = request("root", "pod", 1); + let two = request("root", "pod", 2); + let ledger = ledger(100) + .register_session(one.identity.clone()) + .unwrap() + .next; + let ledger = ledger.reserve(&one, 100).unwrap().next; + let ledger = ledger.reserve(&two, 100).unwrap().next; + let ledger = ledger.begin_dispatch(&command(&two), 101).unwrap().next; + let expired = ledger.expire_undispatched(1000).unwrap(); + assert_eq!(expired.value, 1); + assert_eq!(expired.next.meters.reserved.tokens, 30); + assert!( + expired + .next + .attempts + .values() + .any(|attempt| attempt.phase == AttemptPhase::InFlight) + ); +} + +#[test] +fn settlement_refunds_only_unused_bound_once_and_compaction_retains_replay_fences() { + let request = request("root", "pod", 1); + let ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + let ledger = ledger.reserve(&request, 100).unwrap().next; + let ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + let settlement = Settlement { + attempt: command(&request), + usage: Some(usage()), + }; + let next = ledger.settle(&settlement).unwrap(); + assert_eq!(next.next.meters.settled.tokens, 8); + assert_eq!(next.next.meters.reserved.tokens, 0); + assert_eq!(next.next.sessions["pod"].closed_through, 1); + assert!(next.next.attempts.is_empty()); + assert!(!next.next.settle(&settlement).unwrap().changed); + assert!(matches!( + next.next.reserve(&request, 102), + Err(BudgetError::Sequence) + )); +} + +#[test] +fn cancellation_charges_inflight_at_maximum_and_never_refunds_late_usage() { + let request = request("root", "pod", 1); + let ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + let ledger = ledger.reserve(&request, 100).unwrap().next; + let ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + let ledger = ledger.close_subtree("root").unwrap().next; + assert_eq!(ledger.meters.uncertain.tokens, 30); + assert_eq!(ledger.meters.reserved.tokens, 0); + assert!( + !ledger + .settle(&Settlement { + attempt: command(&request), + usage: Some(usage()) + }) + .unwrap() + .changed + ); + assert!( + ledger + .register_session(identity("root", "new-pod")) + .is_err() + ); +} + +#[test] +fn missing_usage_commits_the_full_bound_and_overbound_usage_freezes_the_account() { + for observed in [ + None, + Some(Usage { + input_tokens: 999, + ..usage() + }), + ] { + let request = request("root", "pod", 1); + let ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + let ledger = ledger.reserve(&request, 100).unwrap().next; + let ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + let breach = observed.is_some(); + let result = ledger + .settle(&Settlement { + attempt: command(&request), + usage: observed, + }) + .unwrap(); + assert_eq!(result.next.meters.uncertain.tokens, 30); + assert_eq!(result.value.breach, breach); + assert_eq!(result.next.phase == AccountPhase::Frozen, breach); + } +} + +#[test] +fn uid_reparenting_and_authorization_replays_are_not_new_budget_accounts() { + let ledger = ledger(100); + let mut different = task("root", None, 100); + different.task.name = "same-uid-different-name".into(); + assert!(ledger.register_task(different).is_err()); + let mut child = task("child", Some("root"), 100); + child.root_task_uid = "other-root".into(); + assert!(ledger.register_task(child).is_err()); + let request = request("root", "pod", 1); + let ledger = ledger + .register_session(request.identity.clone()) + .unwrap() + .next; + let mut replay = request.clone(); + replay.account_uid = "new-account".into(); + assert!(ledger.reserve(&replay, 100).is_err()); + replay = request; + replay.identity.authorization_digest = format!("sha256:{}", "b".repeat(64)); + assert!(ledger.reserve(&replay, 100).is_err()); +} + +#[test] +fn monetary_axis_is_separate_and_zero_remains_unbounded() { + assert_eq!( + Limits { + tokens: Some(0), + usd_micros: Some(0) + } + .normalized(), + Limits::default() + ); + let mut ledger = ledger(100); + ledger.limits.usd_micros = Some(4); + ledger + .nodes + .get_mut("root") + .unwrap() + .authority + .limits + .usd_micros = Some(4); + let request = request("root", "pod", 1); + let ledger = ledger + .register_session(request.identity.clone()) + .unwrap() + .next; + assert!(matches!( + ledger.reserve(&request, 100), + Err(BudgetError::Exhausted) + )); +} + +#[test] +fn corrupted_balances_or_missing_active_rows_fail_instead_of_resetting() { + let request = request("root", "pod", 1); + let ledger = ledger(100) + .register_session(request.identity.clone()) + .unwrap() + .next; + let mut ledger = ledger.reserve(&request, 100).unwrap().next; + ledger.meters.reserved.tokens = 0; + assert!(ledger.validate().is_err()); + ledger.meters.reserved.tokens = 30; + ledger.attempts.clear(); + assert!(ledger.validate().is_err()); +} + +#[test] +fn team_uid_is_the_lifetime_account_even_for_new_principal_task_uids() { + let root = RootIdentity { + kind: RootKind::KarsTeam, + resource: resource("team", "team-uid"), + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }; + let mut ledger = Ledger::new( + "account-uid".into(), + root, + Limits { + tokens: Some(50), + usd_micros: Some(1000), + }, + ) + .unwrap(); + for uid in ["principal-one", "principal-two"] { + let mut authority = task(uid, None, 50); + authority.root_task_uid = uid.into(); + ledger = ledger.register_task(authority).unwrap().next; + let request = request(uid, uid, 1); + ledger = ledger + .register_session(request.identity.clone()) + .unwrap() + .next; + if uid == "principal-one" { + ledger = ledger.reserve(&request, 100).unwrap().next; + ledger = ledger.begin_dispatch(&command(&request), 101).unwrap().next; + ledger = ledger.close_subtree(uid).unwrap().next; + } else { + assert!(matches!( + ledger.reserve(&request, 100), + Err(BudgetError::Exhausted) + )); + } + } +} diff --git a/shared/inference_budget/ledger_validation.rs b/shared/inference_budget/ledger_validation.rs new file mode 100644 index 000000000..7ee7db829 --- /dev/null +++ b/shared/inference_budget/ledger_validation.rs @@ -0,0 +1,156 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; + +impl Ledger { + pub fn validate(&self) -> Result<(), BudgetError> { + self.root.validate()?; + if self.version != CONTRACT_VERSION + || !valid_uid(&self.account_uid) + || self.nodes.len() > MAX_NODES + || self.sessions.len() > MAX_SESSIONS + || self.attempts.len() > MAX_ATTEMPTS + { + return Err(BudgetError::Corrupt); + } + let serialized = serde_json::to_value(self).map_err(|_| BudgetError::Corrupt)?; + if !signed_integer_range(&serialized) { + return Err(BudgetError::Overflow); + } + if serde_json::to_vec(self) + .map_err(|_| BudgetError::Corrupt)? + .len() + > MAX_LEDGER_BYTES + { + return Err(BudgetError::Capacity); + } + if !self.limits.allows(self.meters.total()?) + || (self.limits.normalized().usd_micros.is_some() && self.meters.unpriced_attempts > 0) + { + return Err(BudgetError::Corrupt); + } + let mut root_reserved = Amounts::default(); + let mut node_reserved = BTreeMap::::new(); + let mut children_used = BTreeMap::::new(); + let mut root_nodes_used = Amounts::default(); + for (uid, node) in &self.nodes { + node.authority.task.validate()?; + if uid != &node.authority.task.uid + || node.authority.task.namespace != self.root.resource.namespace + || !valid_digest(&node.authority.authorization_digest) + || !node.authority.effective_authorization.is_object() + || serde_json::to_vec(&node.authority.effective_authorization) + .map_err(|_| BudgetError::Corrupt)? + .len() + > 65_536 + || !node.authority.limits.allows(node.meters.total()?) + || (node.authority.limits.normalized().usd_micros.is_some() + && node.meters.unpriced_attempts > 0) + { + return Err(BudgetError::Corrupt); + } + let path = self.ancestors(uid)?; + let root_uid = path.last().ok_or(BudgetError::Corrupt)?; + if *root_uid != node.authority.root_task_uid { + return Err(BudgetError::Corrupt); + } + let used = node.meters.settled.add(node.meters.uncertain)?; + if let Some(parent) = &node.authority.parent_uid { + let prior = children_used.get(parent).copied().unwrap_or_default(); + children_used.insert(parent.clone(), prior.add(used)?); + } else { + if self.root.kind == RootKind::KarsTask && node.authority.task != self.root.resource + { + return Err(BudgetError::Corrupt); + } + root_nodes_used = root_nodes_used.add(used)?; + } + } + if !root_nodes_used.within(self.meters.settled.add(self.meters.uncertain)?) { + return Err(BudgetError::Corrupt); + } + for (uid, used) in children_used { + let node = self.nodes.get(&uid).ok_or(BudgetError::Corrupt)?; + if !used.within(node.meters.settled.add(node.meters.uncertain)?) { + return Err(BudgetError::Corrupt); + } + } + for (uid, session) in &self.sessions { + session.identity.validate()?; + if uid != &session.identity.pod_uid + || session.next_sequence == 0 + || session.closed_through >= session.next_sequence + || session.next_sequence - session.closed_through > MAX_REPLAY_WINDOW + 1 + || !self.nodes.contains_key(&session.identity.task_uid) + { + return Err(BudgetError::Corrupt); + } + for sequence in session.closed_through + 1..session.next_sequence { + if !self.attempts.contains_key( + &AttemptKey { + pod_uid: uid.clone(), + sequence, + } + .storage_key(), + ) { + return Err(BudgetError::Corrupt); + } + } + } + for (key, attempt) in &self.attempts { + let session = self + .sessions + .get(&attempt.key.pod_uid) + .ok_or(BudgetError::Corrupt)?; + if key != &attempt.key.storage_key() + || attempt.identity != session.identity + || attempt.key.sequence <= session.closed_through + || attempt.key.sequence >= session.next_sequence + || !valid_digest(&attempt.wire_digest) + || !attempt.charged.within(attempt.quote.maximum) + || attempt.observed_breach.as_ref().is_some_and(|observation| { + serde_json::from_str::(observation).is_err() + || self.phase == AccountPhase::Active + }) + { + return Err(BudgetError::Corrupt); + } + + fn signed_integer_range(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::Number(number) => number + .as_u64() + .is_none_or(|number| number <= MAX_LEDGER_INTEGER), + serde_json::Value::Array(values) => values.iter().all(signed_integer_range), + serde_json::Value::Object(values) => values.values().all(signed_integer_range), + _ => true, + } + } + // Expired contracts remain accounting evidence; validate their + // shape and arithmetic without making old charges disappear. + attempt + .quote + .validate(i64::MIN, attempt.quote.price_covered)?; + if !attempt.phase.terminal() { + if attempt.charged != Amounts::default() { + return Err(BudgetError::Corrupt); + } + root_reserved = root_reserved.add(attempt.quote.maximum)?; + for uid in self.ancestors(&attempt.identity.task_uid)? { + let previous = node_reserved.get(&uid).copied().unwrap_or_default(); + node_reserved.insert(uid, previous.add(attempt.quote.maximum)?); + } + } + } + if self.meters.reserved != root_reserved { + return Err(BudgetError::Corrupt); + } + for (uid, node) in &self.nodes { + if node.meters.reserved != node_reserved.get(uid).copied().unwrap_or_default() { + return Err(BudgetError::Corrupt); + } + } + Ok(()) + } +} diff --git a/shared/inference_budget/mod.rs b/shared/inference_budget/mod.rs new file mode 100644 index 000000000..da3a8e5d7 --- /dev/null +++ b/shared/inference_budget/mod.rs @@ -0,0 +1,12 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Governed-inference accounting contract, shared by controller and router. +//! This accounts for tokens and configured maximum inference prices only. + +pub mod catalog; +pub mod ledger; +pub mod tariffs; +pub mod types; + +pub use types::*; diff --git a/shared/inference_budget/tariff_tests.rs b/shared/inference_budget/tariff_tests.rs new file mode 100644 index 000000000..b0aefdcdb --- /dev/null +++ b/shared/inference_budget/tariff_tests.rs @@ -0,0 +1,202 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use serde_json::json; + +pub(super) fn contract() -> ModelContract { + ModelContract { + id: "text-model".into(), + version: "operator-v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "named-provider".into(), + endpoint: "https://operator.example/inference".into(), + model: "model-revision-1".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxCompletionTokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::TokenRates { + input_micros_per_million: 1_000_001, + output_micros_per_million: 2_000_001, + fixed_micros: 3, + }), + } +} + +fn request() -> Vec { + serde_json::to_vec( + &json!({"model": "model-revision-1", "messages": [{"role": "user", "content": "hello"}]}), + ) + .unwrap() +} + +#[test] +fn missing_maximum_is_injected_and_input_uses_operator_bound_not_character_estimate() { + let (wire, quote) = contract().normalize(&request(), 100, true).unwrap(); + assert_eq!( + serde_json::from_slice::(&wire).unwrap()["max_completion_tokens"], + 20 + ); + assert_eq!(quote.maximum.tokens, 30); + assert_eq!(quote.maximum.usd_micros, 55); + assert!(quote.price_covered); +} + +#[test] +fn no_price_is_allowed_only_when_no_monetary_axis_requires_it() { + let mut contract = contract(); + contract.maximum_price = None; + assert!(contract.normalize(&request(), 100, true).is_err()); + let (_, quote) = contract.normalize(&request(), 100, false).unwrap(); + assert!(!quote.price_covered); + assert_eq!(quote.maximum.tokens, 30); +} + +#[test] +fn expiry_unknown_bounds_or_uncovered_reasoning_are_not_zero_cost_fallbacks() { + for case in [ + "expired", + "no-input", + "no-output", + "reasoning", + "zero-price", + ] { + let mut contract = contract(); + match case { + "expired" => contract.valid_until = "1970-01-01T00:00:01Z".into(), + "no-input" => contract.maximum_input_tokens = 0, + "no-output" => contract.maximum_output_tokens = 0, + "reasoning" => contract.output_bound_includes_reasoning = false, + _ => contract.maximum_price = Some(MaximumPrice::PerRequest { maximum_micros: 0 }), + } + assert!(contract.normalize(&request(), 100, true).is_err(), "{case}"); + } +} + +#[test] +fn money_math_is_checked_and_rounds_each_category_up() { + let price = MaximumPrice::TokenRates { + input_micros_per_million: 1, + output_micros_per_million: 1, + fixed_micros: 0, + }; + assert_eq!(price.price(1, 1).unwrap(), 2); + let price = MaximumPrice::TokenRates { + input_micros_per_million: u64::MAX, + output_micros_per_million: u64::MAX, + fixed_micros: u64::MAX, + }; + assert!(price.price(u64::MAX, u64::MAX).is_err()); +} + +#[test] +fn unsupported_async_multimodal_server_tools_and_multiplicity_fail_closed() { + for extra in [ + json!({"n": 2}), + json!({"background": true}), + json!({"store": true}), + json!({"web_search_options": {}}), + json!({"tools": [{"type": "web_search"}]}), + json!({"messages": [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://example.test/a"}}]}]}), + json!({"max_completion_tokens": 21}), + json!({"max_completion_tokens": "20"}), + json!({"max_tokens": 20}), + ] { + let mut value: Value = serde_json::from_slice(&request()).unwrap(); + value + .as_object_mut() + .unwrap() + .extend(extra.as_object().unwrap().clone()); + assert!( + contract() + .normalize(&serde_json::to_vec(&value).unwrap(), 100, true) + .is_err() + ); + } +} + +#[test] +fn client_function_tools_remain_available_without_enabling_hosted_generation() { + let mut value: Value = serde_json::from_slice(&request()).unwrap(); + value["tools"] = json!([{"type": "function", "function": {"name": "read_file", "parameters": {"type": "object"}}}]); + assert!( + contract() + .normalize(&serde_json::to_vec(&value).unwrap(), 100, true) + .is_ok() + ); +} + +#[test] +fn native_and_responses_shapes_have_explicit_distinct_output_contracts() { + let mut contract = contract(); + contract.operation = Operation::AnthropicMessages; + contract.output_field = OutputField::MaxTokens; + let body = json!({"model": contract.model, "messages": [{"role": "user", "content": "hello"}], + "tools": [{"name": "read_file", "input_schema": {"type": "object"}}]}); + let (wire, _) = contract + .normalize(&serde_json::to_vec(&body).unwrap(), 100, true) + .unwrap(); + assert_eq!( + serde_json::from_slice::(&wire).unwrap()["max_tokens"], + 20 + ); + contract.operation = Operation::Responses; + contract.output_field = OutputField::MaxOutputTokens; + let body = + json!({"model": contract.model, "input": "hello", "store": false, "background": false}); + let (wire, _) = contract + .normalize(&serde_json::to_vec(&body).unwrap(), 100, true) + .unwrap(); + assert_eq!( + serde_json::from_slice::(&wire).unwrap()["max_output_tokens"], + 20 + ); +} + +#[test] +fn trustworthy_usage_is_bounded_and_cache_reasoning_are_subsets_not_extra_refunds() { + let (_, quote) = contract().normalize(&request(), 100, true).unwrap(); + let usage = Usage { + input_tokens: 3, + output_tokens: 5, + cached_input_tokens: 1, + cache_creation_input_tokens: 1, + reasoning_output_tokens: 2, + }; + assert_eq!(quote.usage(&usage).unwrap().tokens, 8); + assert!( + quote + .usage(&Usage { + input_tokens: 11, + ..usage.clone() + }) + .is_err() + ); + assert!( + quote + .usage(&Usage { + output_tokens: 21, + ..usage.clone() + }) + .is_err() + ); + assert!( + quote + .usage(&Usage { + cached_input_tokens: 3, + ..usage.clone() + }) + .is_err() + ); + assert!( + quote + .usage(&Usage { + reasoning_output_tokens: 6, + ..usage + }) + .is_err() + ); +} diff --git a/shared/inference_budget/tariffs.rs b/shared/inference_budget/tariffs.rs new file mode 100644 index 000000000..7c06303cf --- /dev/null +++ b/shared/inference_budget/tariffs.rs @@ -0,0 +1,548 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Operator-declared maximum bounds, never guessed token counts or prices. + +use super::types::*; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum Operation { + ChatCompletions, + AnthropicMessages, + Responses, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum OutputField { + MaxTokens, + MaxCompletionTokens, + MaxOutputTokens, +} + +impl OutputField { + pub fn key(self) -> &'static str { + match self { + Self::MaxTokens => "max_tokens", + Self::MaxCompletionTokens => "max_completion_tokens", + Self::MaxOutputTokens => "max_output_tokens", + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde( + tag = "kind", + rename_all = "camelCase", + rename_all_fields = "camelCase", + deny_unknown_fields +)] +pub enum MaximumPrice { + TokenRates { + /// Maximum rate across fresh/cache-read/cache-creation input categories. + input_micros_per_million: u64, + /// Maximum rate across visible and reasoning output categories. + output_micros_per_million: u64, + fixed_micros: u64, + }, + PerRequest { + maximum_micros: u64, + }, +} + +impl MaximumPrice { + pub fn price(&self, input: u64, output: u64) -> Result { + let amount = match self { + Self::PerRequest { maximum_micros } => return Ok(*maximum_micros), + Self::TokenRates { + input_micros_per_million, + output_micros_per_million, + fixed_micros, + } => { + let input = u128::from(input) + .checked_mul(u128::from(*input_micros_per_million)) + .ok_or(BudgetError::Overflow)?; + let output = u128::from(output) + .checked_mul(u128::from(*output_micros_per_million)) + .ok_or(BudgetError::Overflow)?; + let input = input.checked_add(999_999).ok_or(BudgetError::Overflow)? / 1_000_000; + let output = output.checked_add(999_999).ok_or(BudgetError::Overflow)? / 1_000_000; + input + .checked_add(output) + .ok_or(BudgetError::Overflow)? + .checked_add(u128::from(*fixed_micros)) + .ok_or(BudgetError::Overflow)? + } + }; + u64::try_from(amount).map_err(|_| BudgetError::Overflow) + } + + fn valid(&self) -> bool { + match self { + Self::PerRequest { maximum_micros } => { + *maximum_micros > 0 && *maximum_micros <= MAX_LEDGER_INTEGER + } + Self::TokenRates { + input_micros_per_million, + output_micros_per_million, + fixed_micros, + } => { + (*input_micros_per_million > 0 + || *output_micros_per_million > 0 + || *fixed_micros > 0) + && [ + *input_micros_per_million, + *output_micros_per_million, + *fixed_micros, + ] + .iter() + .all(|value| *value <= MAX_LEDGER_INTEGER) + } + } + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ModelContract { + pub id: String, + pub version: String, + pub valid_until: String, + /// Exact final named provider/authentication identity, never request headers. + pub provider_id: String, + pub endpoint: String, + pub model: String, + pub operation: Operation, + pub output_field: OutputField, + /// The provider-enforced maximum complete input/context bound, including + /// hidden framing, cache input, and function schemas. Not bytes/4. + pub maximum_input_tokens: u64, + pub maximum_output_tokens: u64, + pub maximum_wire_bytes: u64, + /// Operator attestation that the output field bounds ALL output, including + /// reasoning/thinking. Contracts without this guarantee cannot dispatch. + pub output_bound_includes_reasoning: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub maximum_price: Option, +} + +impl ModelContract { + pub fn validate(&self, now: i64) -> Result<(), BudgetError> { + let endpoint = reqwest::Url::parse(&self.endpoint).map_err(|_| BudgetError::Contract)?; + if !matches!(endpoint.scheme(), "http" | "https") || endpoint.host_str().is_none() { + return Err(BudgetError::Contract); + } + let expiry = chrono::DateTime::parse_from_rfc3339(&self.valid_until) + .map_err(|_| BudgetError::Contract)? + .timestamp(); + let field_matches = matches!( + (self.operation, self.output_field), + ( + Operation::ChatCompletions, + OutputField::MaxTokens | OutputField::MaxCompletionTokens + ) | (Operation::AnthropicMessages, OutputField::MaxTokens) + | (Operation::Responses, OutputField::MaxOutputTokens) + ); + if !valid_name(&self.id) + || !valid_uid(&self.version) + || self.provider_id.is_empty() + || self.provider_id.len() > 253 + || self.model.is_empty() + || self.model.len() > 253 + || !(self.endpoint.starts_with("https://") || self.endpoint.starts_with("http://")) + || self.endpoint.contains(['@', '?', '#']) + || self.endpoint.bytes().any(|b| b.is_ascii_control()) + || expiry <= now + || !field_matches + || !self.output_bound_includes_reasoning + || self.maximum_input_tokens == 0 + || self.maximum_output_tokens == 0 + || self.maximum_wire_bytes == 0 + || self.maximum_wire_bytes > 1_048_576 + || self + .maximum_price + .as_ref() + .is_some_and(|price| !price.valid()) + { + return Err(BudgetError::Contract); + } + let tokens = self + .maximum_input_tokens + .checked_add(self.maximum_output_tokens) + .ok_or(BudgetError::Overflow)?; + if tokens > MAX_LEDGER_INTEGER { + return Err(BudgetError::Overflow); + } + if let Some(price) = &self.maximum_price { + if price.price(self.maximum_input_tokens, self.maximum_output_tokens)? + > MAX_LEDGER_INTEGER + { + return Err(BudgetError::Overflow); + } + } + Ok(()) + } + + pub fn normalize( + &self, + body: &[u8], + now: i64, + money_required: bool, + ) -> Result<(Vec, Quote), BudgetError> { + self.validate(now)?; + if body.len() as u64 > self.maximum_wire_bytes + || (money_required && self.maximum_price.is_none()) + { + return Err(BudgetError::Contract); + } + let mut value: Value = serde_json::from_slice(body).map_err(|_| BudgetError::Contract)?; + validate_shape(&value, self.operation)?; + let map = value.as_object_mut().ok_or(BudgetError::Contract)?; + if map + .get("model") + .is_some_and(|model| model.as_str() != Some(self.model.as_str())) + { + return Err(BudgetError::Contract); + } + map.insert("model".into(), self.model.clone().into()); + let key = self.output_field.key(); + for other in ["max_tokens", "max_completion_tokens", "max_output_tokens"] { + if other != key && map.contains_key(other) { + return Err(BudgetError::Contract); + } + } + let output = match map.get(key) { + None => self.maximum_output_tokens, + Some(value) => value + .as_u64() + .filter(|value| *value > 0 && *value <= self.maximum_output_tokens) + .ok_or(BudgetError::Contract)?, + }; + map.insert(key.into(), output.into()); + let bytes = serde_json::to_vec(&value).map_err(|_| BudgetError::Contract)?; + if bytes.len() as u64 > self.maximum_wire_bytes { + return Err(BudgetError::Contract); + } + let maximum = Amounts { + tokens: self + .maximum_input_tokens + .checked_add(output) + .ok_or(BudgetError::Overflow)?, + usd_micros: self + .maximum_price + .as_ref() + .map(|price| price.price(self.maximum_input_tokens, output)) + .transpose()? + .unwrap_or(0), + }; + Ok(( + bytes, + Quote { + contract: self.clone(), + output_tokens: output, + maximum, + price_covered: self.maximum_price.is_some(), + }, + )) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Quote { + pub contract: ModelContract, + pub output_tokens: u64, + pub maximum: Amounts, + pub price_covered: bool, +} + +impl Quote { + pub fn validate(&self, now: i64, money_required: bool) -> Result<(), BudgetError> { + self.contract.validate(now)?; + if self.output_tokens == 0 + || self.output_tokens > self.contract.maximum_output_tokens + || self.price_covered != self.contract.maximum_price.is_some() + || (money_required && !self.price_covered) + { + return Err(BudgetError::Contract); + } + let expected = Amounts { + tokens: self + .contract + .maximum_input_tokens + .checked_add(self.output_tokens) + .ok_or(BudgetError::Overflow)?, + usd_micros: self + .contract + .maximum_price + .as_ref() + .map(|price| price.price(self.contract.maximum_input_tokens, self.output_tokens)) + .transpose()? + .unwrap_or(0), + }; + if expected != self.maximum { + return Err(BudgetError::Contract); + } + Ok(()) + } + + pub fn usage(&self, usage: &Usage) -> Result { + if usage.input_tokens > self.contract.maximum_input_tokens + || usage.output_tokens > self.output_tokens + || usage + .cached_input_tokens + .checked_add(usage.cache_creation_input_tokens) + .is_none_or(|cached| cached > usage.input_tokens) + || usage.reasoning_output_tokens > usage.output_tokens + { + return Err(BudgetError::Breach); + } + Ok(Amounts { + tokens: usage + .input_tokens + .checked_add(usage.output_tokens) + .ok_or(BudgetError::Overflow)?, + usd_micros: self + .contract + .maximum_price + .as_ref() + .map(|price| price.price(usage.input_tokens, usage.output_tokens)) + .transpose()? + .unwrap_or(0), + }) + } +} + +fn only_keys(value: &Value, keys: &[&str]) -> bool { + value + .as_object() + .is_some_and(|object| object.keys().all(|key| keys.contains(&key.as_str()))) +} + +fn cache_control(value: &Value) -> bool { + value.get("cache_control").is_none_or(|cache| { + only_keys(cache, &["type", "ttl"]) + && cache.get("type").and_then(Value::as_str) == Some("ephemeral") + && cache + .get("ttl") + .is_none_or(|ttl| matches!(ttl.as_str(), Some("5m" | "1h"))) + }) +} + +fn text_content(value: &Value, anthropic: bool) -> bool { + if value.is_null() || value.is_string() { + return true; + } + value.as_array().is_some_and(|blocks| { + blocks.len() <= 256 + && blocks + .iter() + .all(|block| match block.get("type").and_then(Value::as_str) { + Some("text" | "input_text" | "output_text") => { + block.get("text").is_some_and(Value::is_string) + && only_keys(block, &["type", "text", "cache_control"]) + && cache_control(block) + } + Some("tool_use") if anthropic => { + only_keys(block, &["type", "id", "name", "input", "cache_control"]) + && cache_control(block) + && block.get("name").is_some_and(Value::is_string) + && block.get("id").is_some_and(Value::is_string) + && block.get("input").is_some_and(Value::is_object) + } + Some("tool_result") if anthropic => { + only_keys( + block, + &[ + "type", + "tool_use_id", + "content", + "is_error", + "cache_control", + ], + ) && cache_control(block) + && block.get("tool_use_id").is_some_and(Value::is_string) + && block + .get("content") + .is_some_and(|value| text_content(value, false)) + } + _ => false, + }) + }) +} + +fn messages(value: &Value, anthropic: bool) -> bool { + value.as_array().is_some_and(|items| { + !items.is_empty() + && items.len() <= 256 + && items.iter().all(|item| { + let keys: &[&str] = if anthropic { + &["role", "content"] + } else { + &[ + "role", + "content", + "name", + "tool_calls", + "tool_call_id", + "refusal", + ] + }; + let role = item.get("role").and_then(Value::as_str); + only_keys(item, keys) + && matches!( + role, + Some("system" | "developer" | "user" | "assistant" | "tool") + ) + && item + .get("content") + .is_some_and(|content| text_content(content, anthropic)) + && item.get("tool_calls").is_none_or(|tools| { + tools.as_array().is_some_and(|items| { + items.iter().all(|tool| { + tool.get("type").and_then(Value::as_str) == Some("function") + }) + }) + }) + }) + }) +} + +fn validate_shape(value: &Value, operation: Operation) -> Result<(), BudgetError> { + let map = value.as_object().ok_or(BudgetError::Contract)?; + let allowed: &[&str] = match operation { + Operation::ChatCompletions => &[ + "model", + "messages", + "max_tokens", + "max_completion_tokens", + "stream", + "stream_options", + "temperature", + "top_p", + "frequency_penalty", + "presence_penalty", + "stop", + "seed", + "tools", + "tool_choice", + "parallel_tool_calls", + "response_format", + "user", + "metadata", + "n", + "reasoning_effort", + "store", + ], + Operation::AnthropicMessages => &[ + "model", + "messages", + "system", + "max_tokens", + "stream", + "temperature", + "top_p", + "top_k", + "stop_sequences", + "tools", + "tool_choice", + "thinking", + "metadata", + ], + Operation::Responses => &[ + "model", + "input", + "instructions", + "max_output_tokens", + "stream", + "tools", + "tool_choice", + "parallel_tool_calls", + "text", + "reasoning", + "store", + "background", + "metadata", + "user", + ], + }; + if map.keys().any(|key| !allowed.contains(&key.as_str())) + || map.get("n").is_some_and(|value| value.as_u64() != Some(1)) + || map + .get("background") + .is_some_and(|value| value.as_bool() != Some(false)) + || map + .get("store") + .is_some_and(|value| value.as_bool() != Some(false)) + || map.get("stream").is_some_and(|value| !value.is_boolean()) + { + return Err(BudgetError::Contract); + } + let valid_input = match operation { + Operation::ChatCompletions => map + .get("messages") + .is_some_and(|value| messages(value, false)), + Operation::AnthropicMessages => { + map.get("messages") + .is_some_and(|value| messages(value, true)) + && map + .get("system") + .is_none_or(|value| text_content(value, false)) + } + Operation::Responses => map.get("input").is_some_and(|input| { + input.is_string() + || input.as_array().is_some_and(|items| { + items.len() <= 256 + && items + .iter() + .all(|item| match item.get("type").and_then(Value::as_str) { + Some("function_call") => { + only_keys( + item, + &["type", "id", "call_id", "name", "arguments", "status"], + ) && item.get("name").is_some_and(Value::is_string) + && item.get("arguments").is_some_and(Value::is_string) + } + Some("function_call_output") => { + only_keys(item, &["type", "id", "call_id", "output", "status"]) + && item.get("call_id").is_some_and(Value::is_string) + && item.get("output").is_some_and(Value::is_string) + } + None | Some("message") => { + only_keys(item, &["type", "role", "content", "id", "status"]) + && item + .get("content") + .is_some_and(|content| text_content(content, false)) + } + _ => false, + }) + }) + }), + }; + if !valid_input { + return Err(BudgetError::Contract); + } + if let Some(tools) = map.get("tools") { + let tools = tools.as_array().ok_or(BudgetError::Contract)?; + if tools.len() > 128 + || tools.iter().any(|tool| { + if operation == Operation::AnthropicMessages { + tool.get("type").is_some() + || !tool.get("name").is_some_and(Value::is_string) + || !tool.get("input_schema").is_some_and(Value::is_object) + } else { + tool.get("type").and_then(Value::as_str) != Some("function") + } + }) + { + return Err(BudgetError::Contract); + } + } + Ok(()) +} + +#[cfg(test)] +#[path = "tariff_tests.rs"] +mod tests; diff --git a/shared/inference_budget/types.rs b/shared/inference_budget/types.rs new file mode 100644 index 000000000..efe7b139f --- /dev/null +++ b/shared/inference_budget/types.rs @@ -0,0 +1,361 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; + +pub const CONTRACT_VERSION: &str = "governed-inference/v1"; +pub const MAX_NODES: usize = 128; +pub const MAX_SESSIONS: usize = 256; +pub const MAX_ATTEMPTS: usize = 256; +pub const MAX_REPLAY_WINDOW: u64 = 64; +pub const MAX_LEDGER_BYTES: usize = 524_288; +pub const RESERVATION_TTL_SECONDS: i64 = 30; +pub const MAX_LEDGER_INTEGER: u64 = i64::MAX as u64; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum BudgetScope { + #[default] + GovernedInference, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum RootKind { + KarsTask, + KarsTeam, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ResourceIdentity { + pub namespace: String, + pub name: String, + pub uid: String, +} + +impl ResourceIdentity { + pub fn validate(&self) -> Result<(), BudgetError> { + if !valid_label(&self.namespace) || !valid_name(&self.name) || !valid_uid(&self.uid) { + return Err(BudgetError::Identity); + } + Ok(()) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct RootIdentity { + pub kind: RootKind, + pub resource: ResourceIdentity, + pub workspace_uid: String, + pub cluster_uid: String, +} + +impl RootIdentity { + pub fn validate(&self) -> Result<(), BudgetError> { + self.resource.validate()?; + if !valid_uid(&self.workspace_uid) || !valid_uid(&self.cluster_uid) { + return Err(BudgetError::Identity); + } + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Limits { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub tokens: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub usd_micros: Option, +} + +impl Limits { + pub fn normalized(self) -> Self { + Self { + tokens: self.tokens.filter(|n| *n > 0), + usd_micros: self.usd_micros.filter(|n| *n > 0), + } + } + + pub fn finite(self) -> bool { + let normalized = self.normalized(); + normalized.tokens.is_some() || normalized.usd_micros.is_some() + } + + pub fn allows(self, amounts: Amounts) -> bool { + let limits = self.normalized(); + limits.tokens.is_none_or(|limit| amounts.tokens <= limit) + && limits + .usd_micros + .is_none_or(|limit| amounts.usd_micros <= limit) + } + + pub fn attenuates(self, parent: Self) -> bool { + let child = self.normalized(); + let parent = parent.normalized(); + parent + .tokens + .is_none_or(|limit| child.tokens.is_some_and(|value| value <= limit)) + && parent + .usd_micros + .is_none_or(|limit| child.usd_micros.is_some_and(|value| value <= limit)) + } +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Amounts { + pub tokens: u64, + /// Maximum configured inference-price units, not an invoice or all-in cost. + pub usd_micros: u64, +} + +impl Amounts { + pub fn add(self, other: Self) -> Result { + let amount = Self { + tokens: self + .tokens + .checked_add(other.tokens) + .ok_or(BudgetError::Overflow)?, + usd_micros: self + .usd_micros + .checked_add(other.usd_micros) + .ok_or(BudgetError::Overflow)?, + }; + if amount.tokens > MAX_LEDGER_INTEGER || amount.usd_micros > MAX_LEDGER_INTEGER { + return Err(BudgetError::Overflow); + } + Ok(amount) + } + + pub fn subtract(self, other: Self) -> Result { + Ok(Self { + tokens: self + .tokens + .checked_sub(other.tokens) + .ok_or(BudgetError::Corrupt)?, + usd_micros: self + .usd_micros + .checked_sub(other.usd_micros) + .ok_or(BudgetError::Corrupt)?, + }) + } + + pub fn within(self, upper: Self) -> bool { + self.tokens <= upper.tokens && self.usd_micros <= upper.usd_micros + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct TaskAuthority { + pub task: ResourceIdentity, + pub parent_uid: Option, + pub root_task_uid: String, + pub authorization_digest: String, + /// Full effective snapshot supplied by the controller's authorization helper. + pub effective_authorization: serde_json::Value, + pub limits: Limits, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ExecutionIdentity { + pub task_uid: String, + pub authorization_digest: String, + pub sandbox: ResourceIdentity, + pub runtime_namespace_uid: String, + pub pod_name: String, + pub pod_uid: String, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AccountReference { + pub namespace: String, + pub name: String, + pub uid: String, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct TaskBudgetBinding { + pub scope: BudgetScope, + pub account: AccountReference, + pub root: RootIdentity, + pub task_uid: String, + pub parent_task_uid: Option, + pub root_task_uid: String, + pub authorization_digest: String, +} + +/// Controller-authored router configuration. Pod identity is supplied through +/// the downward API, not an agent request header. This object contains no token. +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct RouterBinding { + pub task: TaskBudgetBinding, + pub sandbox: ResourceIdentity, + pub runtime_namespace: String, + pub runtime_namespace_uid: String, + pub privacy_epoch: Option, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct BrokerRequest { + pub root: RootIdentity, + pub payload: T, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct SessionRequest { + pub account_uid: String, + pub identity: ExecutionIdentity, +} + +impl ExecutionIdentity { + pub fn validate(&self) -> Result<(), BudgetError> { + self.sandbox.validate()?; + if !valid_uid(&self.task_uid) + || !valid_uid(&self.runtime_namespace_uid) + || !valid_name(&self.pod_name) + || !valid_uid(&self.pod_uid) + || !valid_digest(&self.authorization_digest) + { + return Err(BudgetError::Identity); + } + Ok(()) + } +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum AccountPhase { + Active, + Closing, + Closed, + Frozen, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum AttemptPhase { + Reserved, + InFlight, + Settled, + Uncertain, + Expired, +} + +impl AttemptPhase { + pub fn terminal(self) -> bool { + matches!(self, Self::Settled | Self::Uncertain | Self::Expired) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AttemptKey { + pub pod_uid: String, + pub sequence: u64, +} + +impl AttemptKey { + pub fn storage_key(&self) -> String { + format!("{}:{}", self.pod_uid, self.sequence) + } +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct ReserveRequest { + pub account_uid: String, + pub identity: ExecutionIdentity, + pub sequence: u64, + pub wire_digest: String, + pub quote: super::tariffs::Quote, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct AttemptCommand { + pub account_uid: String, + pub key: AttemptKey, + pub identity: ExecutionIdentity, + pub wire_digest: String, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Usage { + pub input_tokens: u64, + pub output_tokens: u64, + pub cached_input_tokens: u64, + pub cache_creation_input_tokens: u64, + pub reasoning_output_tokens: u64, +} + +#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +pub struct Settlement { + pub attempt: AttemptCommand, + /// Absent/incomplete/malformed upstream usage commits the entire bound. + pub usage: Option, +} + +#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] +pub enum BudgetError { + #[error("Governed inference budget identity is invalid or changed")] + Identity, + #[error("Governed inference authorization is missing, stale, or amplified")] + Authorization, + #[error("Governed inference account is not accepting dispatch")] + Closed, + #[error("Governed inference budget capacity is exhausted")] + Capacity, + #[error("Governed inference budget ceiling would be exceeded")] + Exhausted, + #[error("Governed inference budget arithmetic overflow")] + Overflow, + #[error("Governed inference ledger is corrupt; it must not reset to zero")] + Corrupt, + #[error("Governed inference sequence is out of order or already retired")] + Sequence, + #[error("Governed inference attempt is already dispatched or finalized; no re-dispatch")] + AlreadyDispatched, + #[error("Governed inference reservation expired before dispatch")] + Expired, + #[error("Governed inference contract is missing, expired, invalid, or unsupported")] + Contract, + #[error("Governed inference usage exceeded its reserved contract bound")] + Breach, +} + +pub fn valid_uid(value: &str) -> bool { + !value.is_empty() + && value.len() <= 128 + && value + .bytes() + .all(|b| b.is_ascii_alphanumeric() || b == b'-') +} + +pub fn valid_label(value: &str) -> bool { + !value.is_empty() + && value.len() <= 63 + && value + .bytes() + .all(|b| b.is_ascii_lowercase() || b.is_ascii_digit() || b == b'-') + && value.as_bytes()[0].is_ascii_alphanumeric() + && value.as_bytes()[value.len() - 1].is_ascii_alphanumeric() +} + +pub fn valid_name(value: &str) -> bool { + !value.is_empty() && value.len() <= 253 && value.split('.').all(valid_label) +} + +pub fn valid_digest(value: &str) -> bool { + value + .strip_prefix("sha256:") + .is_some_and(|hex| hex.len() == 64 && hex.bytes().all(|b| b.is_ascii_hexdigit())) +} diff --git a/tests/e2e/inference-budget-api.mjs b/tests/e2e/inference-budget-api.mjs new file mode 100644 index 000000000..9a2111346 --- /dev/null +++ b/tests/e2e/inference-budget-api.mjs @@ -0,0 +1,176 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Disposable Kind API/CEL preflight. No controller/router image build is needed. +// Credentials stay in this process: never persist or print TokenRequest bodies. +import assert from "node:assert/strict"; +import { execFileSync } from "node:child_process"; +import { readFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { fileURLToPath } from "node:url"; + +const require = createRequire(new URL("../../cli/package.json", import.meta.url)); +const { parseAllDocuments } = require("yaml"); +const root = fileURLToPath(new URL("../../", import.meta.url)); +const context = "kind-kars-budget-api"; +const namespace = "budget-api-fixture"; +const controller = `system:serviceaccount:${namespace}:kars-controller`; +const principal = `system:serviceaccount:${namespace}:untrusted`; +const audience = "kars.azure.com/governed-inference-budget"; +const shared = JSON.parse(readFileSync(new URL("../../deploy/helm/kars/files/inference-budget-admission.json", import.meta.url), "utf8") + .replaceAll("__ACCOUNTING_NAMESPACE__", namespace)); + +function kubectl(args, input, publicSchema = false) { + try { + return execFileSync("kubectl", ["--context", context, "--request-timeout=20s", ...args], { + cwd: root, encoding: "utf8", input: input === undefined ? undefined : JSON.stringify(input), + stdio: ["pipe", "pipe", "pipe"], timeout: 30_000, + }); + } catch (error) { + if (publicSchema) { + // This opt-in is used ONLY for the four public CRDs and eight public VAPs + // below. Do not enable it for Secret/token/agent-response commands. + console.error(String(error.stderr ?? "").slice(0, 12_000)); + } + throw new Error("Disposable budget API assertion command failed", { cause: undefined }); + } +} + +function create(value, as, publicSchema = false) { + return JSON.parse(kubectl(["create", "-f", "-", "-o", "json", ...(as ? ["--as", as] : [])], value, publicSchema)); +} + +function denied(value, as) { + assert.throws(() => create(value, as), /Disposable budget API assertion command failed/); +} + +async function until(check, description) { + for (let attempt = 0; attempt < 60; attempt++) { + if (await check()) return; + await new Promise((resolve) => setTimeout(resolve, 1000)); + } + throw new Error(`Timed out: ${description}`); +} + +const clusterNodes = execFileSync("kind", ["get", "nodes", "--name", "kars-budget-api"], { + encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 20_000, +}).trim().split(/\s+/); +assert(clusterNodes.length > 0 && clusterNodes.every((name) => name.startsWith("kars-budget-api-"))); +const version = JSON.parse(kubectl(["get", "--raw", "/version"])); +assert.match(version.gitVersion, /^v1\.31\./, "Use the same pinned Kind v0.24/v1.31 apiserver as the supported harness"); +create({ apiVersion: "v1", kind: "Namespace", metadata: { name: namespace } }); + +const rendered = execFileSync("helm", ["template", "budget-api", "deploy/helm/kars", "--namespace", namespace], { + cwd: root, encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 30_000, +}); +const crdNames = [ + "karssandboxes.kars.azure.com", "karstasks.kars.azure.com", + "karsteams.kars.azure.com", "karsbudgetaccounts.kars.azure.com", +]; +const definitions = parseAllDocuments(rendered).map((document) => { + assert.equal(document.errors.length, 0); + return document.toJSON(); +}).filter((value) => value?.kind === "CustomResourceDefinition" && crdNames.includes(value.metadata.name)); +assert.equal(definitions.length, crdNames.length); +for (const definition of definitions) { + create(definition, undefined, true); + kubectl(["wait", "--for=condition=Established", `crd/${definition.metadata.name}`, "--timeout=60s"]); +} +for (const policy of shared.items) { + create({ apiVersion: "admissionregistration.k8s.io/v1", kind: "ValidatingAdmissionPolicy", + metadata: { name: policy.name }, spec: policy.spec }, undefined, true); + create({ apiVersion: "admissionregistration.k8s.io/v1", kind: "ValidatingAdmissionPolicyBinding", + metadata: { name: policy.name }, spec: { policyName: policy.name, validationActions: ["Deny", "Audit"] } }, undefined, true); + await until(() => { + const observed = JSON.parse(kubectl(["get", "validatingadmissionpolicy", policy.name, "-o", "json"])); + if (observed.status?.observedGeneration !== observed.metadata.generation) return false; + const warnings = observed.status?.typeChecking?.expressionWarnings ?? []; + assert.deepEqual(warnings, [], `Public CEL compilation warnings: ${JSON.stringify(warnings)}`); + return true; + }, `public policy compilation ${policy.name}`); +} + +for (const name of ["kars-controller", "untrusted"]) { + create({ apiVersion: "v1", kind: "ServiceAccount", metadata: { name, namespace } }); +} +// Deliberately over-grant only these fixture principals in the disposable +// namespace: negative tests must prove admission, not merely missing RBAC. +create({ apiVersion: "rbac.authorization.k8s.io/v1", kind: "Role", metadata: { name: "fixture", namespace }, + rules: [ + { apiGroups: ["kars.azure.com"], resources: ["karsbudgetaccounts", "karsbudgetaccounts/status", "karstasks/status"], verbs: ["get", "create", "patch", "update"] }, + { apiGroups: [""], resources: ["serviceaccounts/token", "configmaps"], verbs: ["create", "get", "patch"] }, + ] }); +create({ apiVersion: "rbac.authorization.k8s.io/v1", kind: "RoleBinding", metadata: { name: "fixture", namespace }, + roleRef: { apiGroup: "rbac.authorization.k8s.io", kind: "Role", name: "fixture" }, + subjects: ["kars-controller", "untrusted"].map((name) => ({ kind: "ServiceAccount", name, namespace })) }); + +const identity = { namespace, name: "root", uid: "immutable-root-uid" }; +const account = { apiVersion: "kars.azure.com/v1alpha1", kind: "KarsBudgetAccount", + metadata: { name: "inference-budget-root", namespace }, + spec: { scope: "GovernedInference", root: { kind: "KarsTask", resource: identity, + workspaceUid: "workspace-uid", clusterUid: "cluster-uid" }, limits: { tokens: 100, usdMicros: 10 } } }; +denied(account, principal); +const created = create(account, controller); +const changedRoot = structuredClone(created); +changedRoot.spec.root.resource.uid = "replacement-root-uid"; +assert.throws(() => kubectl(["replace", "-f", "-", "--as", controller], changedRoot)); +assert.equal(JSON.parse(kubectl(["get", "karsbudgetaccount", created.metadata.name, "-n", namespace, "-o", "json"])).metadata.uid, created.metadata.uid); + +const taskPlan = { + apiVersion:"kars.azure.com/v1alpha1", kind:"KarsTask", + metadata:{name:"finite-task", namespace}, + spec:{objective:"Public API budget fixture", envelope:{ + tier:3, authorityCeiling:3, delegationDepth:2, + budget:{scope:"GovernedInference", tokens:100, usdMicros:10}, + }, execution:{launch:true}}, +}; +const scopedTask = create(taskPlan); +assert.equal(scopedTask.spec.envelope.budget.scope, "GovernedInference"); +const legacyTask = structuredClone(taskPlan); +legacyTask.metadata.name = "legacy-finite"; +delete legacyTask.spec.envelope.budget.scope; +denied(legacyTask); +const unbounded = structuredClone(taskPlan); +unbounded.metadata.name = "existing-unbounded"; +unbounded.spec.envelope.budget = {tokens:0, usdMicros:0}; +const unboundedTask = create(unbounded); +unboundedTask.spec.envelope.budget = taskPlan.spec.envelope.budget; +assert.throws(() => kubectl(["replace", "-f", "-"], unboundedTask)); +const changedScope = structuredClone(scopedTask); +delete changedScope.spec.envelope.budget.scope; +changedScope.spec.execution.launch = false; +assert.throws(() => kubectl(["replace", "-f", "-"], changedScope)); + +const tokenRequest = { apiVersion: "authentication.k8s.io/v1", kind: "TokenRequest", + spec: { audiences: [audience], expirationSeconds: 600 } }; +assert.throws(() => kubectl(["create", "--raw", + `/api/v1/namespaces/${namespace}/serviceaccounts/untrusted/token`, "-f", "-", "--as", principal], tokenRequest)); + +const projection = { apiVersion: "v1", kind: "ConfigMap", metadata: { name: "kars-inference-budget-ca", namespace }, + data: { "ca.crt": "public-test-only" } }; +denied(projection, principal); +create(projection, controller); +const regular = { ...projection, metadata: { name: "unrelated-customer-data", namespace } }; +create(regular, principal); +assert.equal(JSON.parse(kubectl(["get", "configmap", regular.metadata.name, "-n", namespace, "-o", "json"])).data["ca.crt"], "public-test-only"); + +// Genuine kubelet issuance and TokenReview extra claims on the supported API. +// This fixture namespace is deliberately NOT a protected runtime namespace; +// production exec/attach into a finite runtime namespace is denied separately. +const pod = create({ apiVersion: "v1", kind: "Pod", metadata: { name: "identity", namespace }, + spec: { serviceAccountName: "untrusted", restartPolicy: "Never", + containers: [{ name: "router", image: "busybox:latest", command: ["sleep", "300"], + volumeMounts: [{ name: "audience", mountPath: "/private", readOnly: true }] }], + volumes: [{ name: "audience", projected: { sources: [{ serviceAccountToken: { + audience, expirationSeconds: 600, path: "token", + } }] } }] } }); +kubectl(["wait", "-n", namespace, "--for=condition=Ready", "pod/identity", "--timeout=90s"]); +const token = kubectl(["exec", "-n", namespace, "identity", "-c", "router", "--", "cat", "/private/token"]).trim(); +assert(token.length > 0); +const review = create({ apiVersion: "authentication.k8s.io/v1", kind: "TokenReview", + spec: { audiences: [audience], token } }); +assert.equal(review.status.authenticated, true); +assert.deepEqual(review.status.audiences, [audience]); +assert.deepEqual(review.status.user.extra["authentication.kubernetes.io/pod-uid"], [pod.metadata.uid]); +assert.deepEqual(review.status.user.extra["authentication.kubernetes.io/pod-name"], ["identity"]); +console.log("Budget public CRDs/CEL, controller-only accounting/CA, private audience issuance and UID claims passed on disposable Kubernetes v1.31"); diff --git a/tests/e2e/inference-budget-enforcement.mjs b/tests/e2e/inference-budget-enforcement.mjs new file mode 100644 index 000000000..92f8800a6 --- /dev/null +++ b/tests/e2e/inference-budget-enforcement.mjs @@ -0,0 +1,255 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Runs only after the existing E2E harness loads its real controller/router +// images into its disposable Kind cluster. No external provider/model is used. +import assert from "node:assert/strict"; +import { execFileSync, spawn } from "node:child_process"; +import { mkdirSync, readFileSync, rmSync } from "node:fs"; +import { fileURLToPath } from "node:url"; +import { join } from "node:path"; + +const root = fileURLToPath(new URL("../../", import.meta.url)); +const context = "kind-kars-e2e"; +const namespace = "kars-system"; +const source = "budget-provider-fixture"; +const endpoint = `http://provider.${source}.svc.cluster.local:8000`; +const scratch = join(root, `.budget-kind-${process.pid}`); +const forwards = []; + +function execute(binary, args, input) { + try { + return execFileSync(binary, args, { + cwd: root, encoding: "utf8", input, stdio: ["pipe", "pipe", "pipe"], timeout: 120_000, + }); + } catch { + // Do not persist/echo argv, API bodies, signing material or bearer tokens. + throw new Error("Disposable budget enforcement fixture command failed"); + } +} + +function k(args, value) { + return execute("kubectl", ["--context", context, "--request-timeout=30s", ...args], + value === undefined ? undefined : JSON.stringify(value)); +} +function create(value) { return JSON.parse(k(["create", "-f", "-", "-o", "json"], value)); } +function get(kind, name, ns = namespace) { return JSON.parse(k(["get", kind, name, "-n", ns, "-o", "json"])); } +const sleep = (milliseconds) => new Promise((resolve) => setTimeout(resolve, milliseconds)); + +async function until(check, message, seconds = 120) { + const deadline = Date.now() + seconds * 1000; + while (Date.now() < deadline) { + try { + const value = await check(); + if (value) return value; + } catch { /* The API resource may not have been materialized yet. */ } + await sleep(500); + } + throw new Error(`Budget fixture deadline: ${message}`); +} + +async function portForward(ns, target, port) { + const process = spawn("kubectl", ["--context", context, "port-forward", "--address", "127.0.0.1", + "-n", ns, target, `:${port}`], { cwd: root, stdio: ["ignore", "pipe", "pipe"] }); + forwards.push(process); + let output = ""; + process.stdout.on("data", (data) => { output = (output + data).slice(-2048); }); + process.stderr.on("data", () => {}); // Never persist transport/API diagnostics. + return until(() => { + if (process.exitCode !== null) throw new Error("Port-forward exited"); + const match = output.match(/Forwarding from 127\.0\.0\.1:(\d+) ->/); + return match && `http://127.0.0.1:${match[1]}`; + }, "port-forward ready", 30); +} + +async function request(url, path, body, timeout = 30_000) { + const response = await fetch(`${url}${path}`, { + method: body === undefined ? "GET" : "POST", + headers: body === undefined ? {} : { "content-type": "application/json" }, + body: body === undefined ? undefined : JSON.stringify(body), + signal: AbortSignal.timeout(timeout), + }); + const text = await response.text(); + const value = text && response.headers.get("content-type")?.includes("application/json") + ? JSON.parse(text) : text; + return { status: response.status, value }; +} + +const message = { messages: [{ role: "user", content: "fixture" }] }; +function blueprint() { + return { isolation: "standard", model: { provider: "ollama", deployment: "fixture" }, instructions: "Fixture only" }; +} + +function task(name, tokens, usdMicros, parent, launch) { + return create({ + apiVersion: "kars.azure.com/v1alpha1", kind: "KarsTask", metadata: { name, namespace }, + spec: { + objective: "Disposable governed inference accounting fixture", + envelope: { tier: parent ? 2 : 3, authorityCeiling: parent ? 2 : 3, delegationDepth: parent ? 1 : 2, + budget: { scope: "GovernedInference", tokens, usdMicros } }, + ...(parent ? { parentRef: { name: parent } } : {}), + execution: { launch }, blueprint: blueprint(), + }, + }); +} + +async function router(name) { + await until(() => { + const current = get("karstask", name); + return current.status?.phase === "Ready" && current.status?.inferenceBudget + && current.status?.observedGeneration === current.metadata.generation; + }, `Task ${name} account bound`); + const runtime = `kars-${name}`; + const pod = await until(() => { + const pods = JSON.parse(k(["get", "pods", "-n", runtime, "-l", `kars.azure.com/sandbox=${name}`, "-o", "json"])); + return pods.items.find((pod) => !pod.metadata.deletionTimestamp + && pod.status?.containerStatuses?.some((container) => container.name === "inference-router" && container.state?.running)); + }, `router container ${name}`); + const url = await portForward(runtime, `pod/${pod.metadata.name}`, 8443); + await until(async () => (await request(url, "/readyz")).status === 200, `private budget readiness ${name}`); + return { url, pod, runtime }; +} + +function accountFor(name) { + const reference = get("karstask", name).status.inferenceBudget.account; + const account = get("karsbudgetaccount", reference.name, reference.namespace); + assert.equal(account.metadata.uid, reference.uid, "Account UID must not reset"); + return account; +} + +async function scenario() { + const nodes = execute("kind", ["get", "nodes", "--name", "kars-e2e"]).trim().split(/\s+/); + assert(nodes.length > 0 && nodes.every((node) => node.startsWith("kars-e2e-"))); + const worker = nodes.find((node) => node.endsWith("-worker")) ?? nodes[0]; + const images = execute("docker", ["exec", worker, "ctr", "-n", "k8s.io", "images", "list"]); + const routerLine = images.split("\n").find((line) => line.split(/\s+/)[0] === "docker.io/library/kars-inference-router:e2e"); + assert(routerLine, "Existing harness must load the real router image first"); + const digest = routerLine.split(/\s+/).find((field) => /^sha256:[a-f0-9]{64}$/.test(field)); + assert(digest, "A manifest digest, not an image-config ID, is required"); + + mkdirSync(scratch, { mode: 0o700 }); + execute("openssl", ["req", "-x509", "-newkey", "rsa:2048", "-nodes", "-days", "1", + "-keyout", join(scratch, "tls.key"), "-out", join(scratch, "tls.crt"), + "-subj", "/CN=kars-inference-budget.kars-system.svc", + "-addext", "subjectAltName=DNS:kars-inference-budget.kars-system.svc"]); + const certificate = readFileSync(join(scratch, "tls.crt")); + create({ apiVersion: "v1", kind: "Secret", + metadata: { name: "budget-fixture-tls", namespace, annotations: { "kars.azure.com/inference-budget-tls": "v1" } }, + type: "kubernetes.io/tls", data: { "tls.crt": certificate.toString("base64"), + "tls.key": readFileSync(join(scratch, "tls.key")).toString("base64") } }); + rmSync(scratch, { recursive: true }); + + create({ apiVersion: "v1", kind: "Namespace", metadata: { name: source } }); + const providerCode = ` + const http = require('node:http'); + let count = 0, hold = false, waiting = []; + http.createServer((req, res) => { + res.setHeader('content-type', 'application/json'); + if (req.url === '/count') return res.end(JSON.stringify({count})); + if (req.url === '/hold') { hold = true; return res.end('{}'); } + if (req.url === '/release') { hold = false; for (const send of waiting.splice(0)) send(); return res.end('{}'); } + if (req.url !== '/v1/chat/completions' || req.method !== 'POST') { res.statusCode = 404; return res.end('{}'); } + let body = ''; + req.on('data', chunk => { body += chunk; }); + req.on('end', () => { + const request = JSON.parse(body); + if (request.model !== 'fixture' || request.max_tokens !== 20 + || request.messages.length !== 1 || request.messages[0].content !== 'fixture') { + res.statusCode = 400; return res.end('{}'); + } + count++; + const send = () => res.end(JSON.stringify({id:'fixture',object:'chat.completion',choices:[ + {index:0,message:{role:'assistant',content:'fixture'},finish_reason:'stop'}], + usage:{prompt_tokens:3,completion_tokens:5,total_tokens:8}})); + if (hold) waiting.push(send); else send(); + }); + }).listen(8000, '0.0.0.0'); + `; + create({ apiVersion: "apps/v1", kind: "Deployment", metadata: { name: "provider", namespace: source }, + spec: { replicas: 1, selector: { matchLabels: { app: "budget-provider" } }, + template: { metadata: { labels: { app: "budget-provider" } }, spec: { + containers: [{ name: "provider", image: "node:latest", command: ["node", "-e", providerCode], + ports: [{ containerPort: 8000 }] }], + } } } }); + create({ apiVersion: "v1", kind: "Service", metadata: { name: "provider", namespace: source }, + spec: { selector: { app: "budget-provider" }, ports: [{ port: 8000, targetPort: 8000 }] } }); + k(["rollout", "status", "-n", source, "deployment/provider", "--timeout=90s"]); + const provider = await portForward(source, "deployment/provider", 8000); + + const values = JSON.parse(execute("helm", ["get", "values", "kars", "--kube-context", context, "-n", namespace, "--all", "-o", "json"])); + values.inferenceBudget = { + enabled: true, routerImageDigest: digest, catalogVersion: "fixture-v1", + tlsSecretName: "budget-fixture-tls", caBundle: certificate.toString("utf8"), + nonInferenceEgressHosts: [], + contracts: [{ id: "fixture-chat", version: "v1", validUntil: "2030-01-01T00:00:00Z", + providerId: "ollama", endpoint, model: "fixture", operation: "ChatCompletions", outputField: "MaxTokens", + maximumInputTokens: 10, maximumOutputTokens: 20, maximumWireBytes: 4096, + outputBoundIncludesReasoning: true, maximumPrice: { kind: "perRequest", maximumMicros: 5 } }], + }; + values.controller.extraEnv = (values.controller.extraEnv ?? []).filter((entry) => entry.name !== "OLLAMA_ENDPOINT"); + values.controller.extraEnv.push({ name: "OLLAMA_ENDPOINT", value: endpoint }); + values.localInference = { namespaces: [source], targets: [{ namespace: source, matchLabels: { app: "budget-provider" }, ports: [8000] }] }; + execute("helm", ["upgrade", "kars", "deploy/helm/kars", "--kube-context", context, "-n", namespace, + "--reuse-values", "-f", "-", "--wait", "--timeout", "90s"], JSON.stringify(values)); + + task("budget-money-root", 1000, 12, undefined, false); + task("budget-money-left", 1000, 12, "budget-money-root", true); + task("budget-money-right", 1000, 12, "budget-money-root", true); + const left = await router("budget-money-left"); + const right = await router("budget-money-right"); + const results = await Promise.all([left, right, left].map((router) => + request(router.url, "/v1/chat/completions", message))); + assert.deepEqual(results.map((result) => result.status).sort(), [200, 200, 429]); + const money = accountFor("budget-money-left"); + assert.equal(money.metadata.uid, accountFor("budget-money-right").metadata.uid); + assert.equal(money.status.ledger.meters.settled.usdMicros, 10); + assert.equal(money.status.ledger.meters.settled.tokens, 16); + + task("budget-token-root", 50, 0, undefined, false); + task("budget-token-left", 50, 0, "budget-token-root", true); + task("budget-token-right", 50, 0, "budget-token-root", true); + const tokenLeft = await router("budget-token-left"); + const tokenRight = await router("budget-token-right"); + const before = (await request(provider, "/count")).value.count; + await request(provider, "/hold"); + const first = request(tokenLeft.url, "/v1/chat/completions", message, 90_000); + await until(async () => (await request(provider, "/count")).value.count === before + 1, "actual first dispatch"); + assert.equal((await request(tokenRight.url, "/v1/chat/completions", message)).status, 429); + assert.equal(accountFor("budget-token-left").status.ledger.meters.reserved.tokens, 30); + await request(provider, "/release"); + assert.equal((await first).status, 200); + assert.equal(accountFor("budget-token-left").status.ledger.meters.settled.tokens, 8); + + const count = (await request(provider, "/count")).value.count; + assert.equal((await request(tokenRight.url, "/v1/embeddings", { input: "fixture" })).status, 503); + assert.equal((await request(tokenRight.url, "/agents", {})).status, 403); + assert.equal((await request(provider, "/count")).value.count, count); + const binding = get("karstask", "budget-token-left").status.inferenceBudget; + await request(provider, "/hold"); + const accepted = request(tokenLeft.url, "/v1/chat/completions", message, 90_000).catch(() => ({ status: 0 })); + await until(async () => (await request(provider, "/count")).value.count === count + 1, "accepted work before cancellation"); + const current = get("karstask", "budget-token-left"); + k(["patch", "karstask", current.metadata.name, "-n", namespace, "--type=merge", "--patch-file", "-"], + { metadata: { uid: current.metadata.uid, resourceVersion: current.metadata.resourceVersion }, spec: { execution: { launch: false } } }); + await until(() => { + const account = get("karsbudgetaccount", binding.account.name, binding.account.namespace); + assert.equal(account.metadata.uid, binding.account.uid); + return account.status.ledger.meters.uncertain.tokens === 30; + }, "cancellation conservatively accounts for accepted work"); + await request(provider, "/release"); + await accepted; + const final = get("karsbudgetaccount", binding.account.name, binding.account.namespace); + assert.equal(final.status.ledger.meters.settled.tokens, 8); + assert.equal(final.status.ledger.meters.uncertain.tokens, 30); + assert.equal(final.status.ledger.meters.reserved.tokens, 0); + console.log("Real broker/Pod authentication, sibling token and maximum-price CAS, unsupported-route closure and accepted-work cancellation passed"); +} + +try { + await scenario(); +} finally { + for (const process of forwards) { + process.kill("SIGTERM"); + } + rmSync(scratch, { recursive: true, force: true }); +} diff --git a/tests/e2e/run.sh b/tests/e2e/run.sh index d57b47b47..3b60c0414 100755 --- a/tests/e2e/run.sh +++ b/tests/e2e/run.sh @@ -3124,6 +3124,11 @@ main() { test_sre_namespace_ownership || fail "SRE namespace lifecycle gate failed" test_credential_sources || fail "Credential-source lifecycle gate failed" + if node "$SCRIPT_DIR/inference-budget-enforcement.mjs"; then + pass "Durable governed inference: real broker, sibling caps, route closure and cancellation" + else + fail "Durable governed-inference enforcement gate failed" + fi echo "" echo "═══════════════════════════════════════════════════════" From dca0d23dde83b101bf6f4b329be228457eab9107 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 22:26:38 +0200 Subject: [PATCH 02/16] refactor(authority): share verified Task UID lineage with budget consumers Capture stable same-workspace UID/RV ancestry, canonical readiness/attenuation/full authorization, and live Team ownership without coupling identity resolution to budget state. Budget callers verify their persisted pins separately. Add focused API/race/owner tests; Rust execution remains pending the coordinated Cargo lease. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/inference_budget/auth.rs | 58 ++- controller/src/inference_budget/binding.rs | 121 ++----- controller/src/inference_budget/store.rs | 10 + controller/src/main.rs | 1 + controller/src/task_identity.rs | 334 ++++++++++++++++++ controller/src/task_identity_tests.rs | 263 ++++++++++++++ .../governed-inference-budgets-2026-09-08.md | 12 +- docs/governed-inference-budgets.md | 8 + 8 files changed, 704 insertions(+), 103 deletions(-) create mode 100644 controller/src/task_identity.rs create mode 100644 controller/src/task_identity_tests.rs diff --git a/controller/src/inference_budget/auth.rs b/controller/src/inference_budget/auth.rs index 0d11e3f37..a695530d9 100644 --- a/controller/src/inference_budget/auth.rs +++ b/controller/src/inference_budget/auth.rs @@ -29,7 +29,6 @@ use crate::{ BudgetError, ExecutionIdentity, RootKind, RouterBinding, ledger::Ledger, }, kars_task::KarsTask, - kars_team::KarsTeam, }; fn api_error(stage: &'static str, error: kube::Error) -> StoreError { @@ -520,17 +519,30 @@ async fn verify_task_binding( }; } let tasks: Api = Api::namespaced(client.clone(), &ledger.root.resource.namespace); - for uid in ledger.ancestors(&identity.task_uid)? { - let node = &ledger.nodes[&uid]; - let task = tasks - .get(&node.authority.task.name) - .await - .map_err(|e| api_error("verify current task authority", e))?; + let leaf = tasks + .get(&node.authority.task.name) + .await + .map_err(|error| api_error("verify current task authority", error))?; + let lineage = crate::task_identity::resolve( + client, + &leaf, + crate::task_identity::LeafReadiness::RequireReady, + ) + .await?; + if lineage.workspace_uid != ledger.root.workspace_uid + || lineage.nodes.len() != ledger.ancestors(&identity.task_uid)?.len() + { + return Err(denied()); + } + let mut pins = Vec::new(); + for current in &lineage.nodes { + let uid = ¤t.pin.task.uid; + let node = ledger.nodes.get(uid).ok_or_else(denied)?; + let task = ¤t.task; let status = task.status.as_ref().ok_or_else(denied)?; if task.metadata.uid.as_deref() != Some(uid.as_str()) - || !crate::kars_task_reconciler::task_is_ready(&task) - || task.envelope_digest() != node.authority.authorization_digest - || (uid == identity.task_uid + || current.authorization_digest != node.authority.authorization_digest + || (uid == &identity.task_uid && !task .spec .execution @@ -544,21 +556,28 @@ async fn verify_task_binding( || binding.account.name != account.name_any() || Some(binding.account.namespace.as_str()) != account.metadata.namespace.as_deref() || binding.root != ledger.root - || binding.task_uid != uid + || binding.task_uid != *uid || binding.parent_task_uid != node.authority.parent_uid || binding.root_task_uid != node.authority.root_task_uid || binding.authorization_digest != node.authority.authorization_digest { return Err(denied()); } + pins.push(crate::task_identity::TaskLineagePin { + task: crate::task_identity::ObjectUidRef { + namespace: node.authority.task.namespace.clone(), + name: node.authority.task.name.clone(), + uid: node.authority.task.uid.clone(), + }, + parent_task_uid: node.authority.parent_uid.clone(), + root_task_uid: node.authority.root_task_uid.clone(), + }); } + lineage.verify_pins(&pins)?; if ledger.root.kind == RootKind::KarsTeam { - let teams: Api = Api::namespaced(client.clone(), &ledger.root.resource.namespace); - let team = teams - .get(&ledger.root.resource.name) - .await - .map_err(|e| api_error("verify team lifetime budget", e))?; + let team = lineage.team.as_ref().ok_or_else(denied)?; if team.metadata.uid.as_deref() != Some(ledger.root.resource.uid.as_str()) + || team.name_any() != ledger.root.resource.name || team.metadata.deletion_timestamp.is_some() || team.spec.paused || super::binding::limits(&team.spec.envelope)? != ledger.limits @@ -575,6 +594,13 @@ async fn verify_task_binding( { return Err(denied()); } + } else if lineage.team.is_some() + || lineage.nodes.first().is_none_or(|node| { + node.pin.task.uid != ledger.root.resource.uid + || node.pin.task.name != ledger.root.resource.name + }) + { + return Err(denied()); } Ok(()) } diff --git a/controller/src/inference_budget/binding.rs b/controller/src/inference_budget/binding.rs index b3b62c30d..80be93e61 100644 --- a/controller/src/inference_budget/binding.rs +++ b/controller/src/inference_budget/binding.rs @@ -7,7 +7,6 @@ use kube::{ api::{Patch, PatchParams}, }; use serde_json::json; -use std::collections::BTreeSet; use super::{ account::{KarsBudgetAccountSpec, name_for_root}, @@ -72,62 +71,53 @@ fn resource(task: &KarsTask) -> Result { Ok(resource) } -async fn chain(client: &Client, task: &KarsTask) -> Result, StoreError> { - let namespace = task.namespace().ok_or(BudgetError::Identity)?; - let api: Api = Api::namespaced(client.clone(), &namespace); - let mut nodes = Vec::new(); - let mut current = task.clone(); - let mut seen = BTreeSet::new(); - loop { - let id = resource(¤t)?; - if !seen.insert(id.uid) - || nodes.len() >= crate::inference_budget_contract::MAX_NODES - || current.metadata.deletion_timestamp.is_some() - { - return Err(BudgetError::Authorization.into()); - } - limits(¤t.spec.envelope)?; - let parent = current - .spec - .parent_ref - .as_ref() - .map(|reference| reference.name.clone()); - nodes.push(current.clone()); - let Some(name) = parent else { break }; - let parent = api - .get(&name) - .await - .map_err(|e| api_error("resolve budget parent UID", e))?; - if current +async fn chain( + client: &Client, + task: &KarsTask, +) -> Result { + let lineage = crate::task_identity::resolve( + client, + task, + crate::task_identity::LeafReadiness::AllowPending, + ) + .await?; + let mut pins = Vec::new(); + for node in &lineage.nodes { + limits(&node.task.spec.envelope)?; + if let Some(binding) = node + .task .status .as_ref() .and_then(|status| status.inference_budget.as_ref()) - .and_then(|binding| binding.parent_task_uid.as_deref()) - .is_some_and(|uid| parent.metadata.uid.as_deref() != Some(uid)) - || !crate::kars_task::spec_attenuation_violations(¤t.spec, &parent.spec) - .is_empty() { - return Err(BudgetError::Authorization.into()); + pins.push(crate::task_identity::TaskLineagePin { + task: crate::task_identity::ObjectUidRef { + namespace: node.pin.task.namespace.clone(), + name: node.pin.task.name.clone(), + uid: binding.task_uid.clone(), + }, + parent_task_uid: binding.parent_task_uid.clone(), + root_task_uid: binding.root_task_uid.clone(), + }); } - current = parent; } - nodes.reverse(); - Ok(nodes) + lineage.verify_pins(&pins)?; + Ok(lineage) } pub(super) async fn needs_account(client: &Client, task: &KarsTask) -> Result { - let nodes = chain(client, task).await?; - if nodes.iter().any(|task| { - has_finite(&task.spec.envelope) - || task + let lineage = chain(client, task).await?; + if lineage.nodes.iter().any(|node| { + has_finite(&node.task.spec.envelope) + || node + .task .status .as_ref() .is_some_and(|status| status.inference_budget.is_some()) }) { return Ok(true); } - let first = nodes.first().ok_or(BudgetError::Identity)?; - Ok(team_for_root(client, first).await?.is_some_and(|team| { + Ok(lineage.team.as_ref().is_some_and(|team| { has_finite(&team.spec.envelope) || team .status @@ -272,35 +262,6 @@ async fn pin_task( Ok(updated) } -async fn team_for_root(client: &Client, task: &KarsTask) -> Result, StoreError> { - let owners = task - .metadata - .owner_references - .as_deref() - .unwrap_or_default(); - let Some(owner) = owners.iter().find(|owner| { - owner.controller == Some(true) - && owner.kind == "KarsTeam" - && owner.api_version == "kars.azure.com/v1alpha1" - }) else { - return Ok(None); - }; - let api: Api = Api::namespaced( - client.clone(), - &task.namespace().ok_or(BudgetError::Identity)?, - ); - let team = api - .get(&owner.name) - .await - .map_err(|e| api_error("resolve lifetime Team UID", e))?; - if team.metadata.uid.as_deref() != Some(owner.uid.as_str()) - || team.metadata.deletion_timestamp.is_some() - { - return Err(BudgetError::Identity.into()); - } - Ok(Some(team)) -} - /// Bind a complete immutable UID ancestry and current full authorization before /// execution. A missing/corrupt pinned account is never recreated as a balance /// of zero. This routine does not relax any launch gate by itself. @@ -312,15 +273,12 @@ pub async fn ensure_task( settings .catalog(client, chrono::Utc::now().timestamp()) .await?; - let nodes = chain(client, task).await?; + let lineage = chain(client, task).await?; + let nodes: Vec<_> = lineage.nodes.iter().map(|node| node.task.clone()).collect(); let first = nodes.first().ok_or(BudgetError::Identity)?; let first_id = resource(first)?; - let team = team_for_root(client, first).await?; + let team = lineage.team; let namespaces: Api = Api::all(client.clone()); - let workspace = namespaces - .get(&first_id.namespace) - .await - .map_err(|e| api_error("resolve budget workspace UID", e))?; let cluster = namespaces .get("kube-system") .await @@ -339,7 +297,7 @@ pub async fn ensure_task( }, None => first_id.clone(), }, - workspace_uid: workspace.uid().ok_or(BudgetError::Identity)?, + workspace_uid: lineage.workspace_uid, cluster_uid: cluster.uid().ok_or(BudgetError::Identity)?, }; root.validate()?; @@ -442,17 +400,14 @@ pub async fn ensure_task( }) .await?; } - let default_model = crate::kars_task::blueprint::controller_default_model(); let mut output = None; for (index, node) in nodes.iter().enumerate() { let authority = TaskAuthority { task: resource(node)?, parent_uid: index.checked_sub(1).and_then(|index| nodes[index].uid()), root_task_uid: first_id.uid.clone(), - authorization_digest: node.spec.authorization_digest_with_model(&default_model), - effective_authorization: node - .spec - .authorization_configuration_with_model(&default_model), + authorization_digest: lineage.nodes[index].authorization_digest.clone(), + effective_authorization: lineage.nodes[index].effective_authorization.clone(), limits: limits(&node.spec.envelope)?, }; let current = store.read(&root, &account.uid).await?; diff --git a/controller/src/inference_budget/store.rs b/controller/src/inference_budget/store.rs index fb2879f12..5cf4c580f 100644 --- a/controller/src/inference_budget/store.rs +++ b/controller/src/inference_budget/store.rs @@ -33,6 +33,16 @@ pub enum StoreError { Contention, } +impl From for StoreError { + fn from(error: crate::task_identity::Error) -> Self { + match error { + crate::task_identity::Error::Api { stage, code } => Self::Api { stage, code }, + crate::task_identity::Error::Changed => Self::Contention, + _ => Self::Ledger(BudgetError::Authorization), + } + } +} + fn api_error(stage: &'static str, error: kube::Error) -> StoreError { StoreError::Api { stage, diff --git a/controller/src/main.rs b/controller/src/main.rs index d7563ab62..a1758413f 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -81,6 +81,7 @@ mod sre_authority; mod sre_privacy; mod sre_registration; mod status; +pub(crate) mod task_identity; mod task_models; mod team_commons; mod team_digest; diff --git a/controller/src/task_identity.rs b/controller/src/task_identity.rs new file mode 100644 index 000000000..0d8b542f0 --- /dev/null +++ b/controller/src/task_identity.rs @@ -0,0 +1,334 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Shared, read-only Task UID ancestry. This captures live identity, not a +//! durable ancestry registry; consumers must verify and persist their own pins. + +use crate::{kars_task::KarsTask, kars_team::KarsTeam}; +use k8s_openapi::{api::core::v1::Namespace, apimachinery::pkg::apis::meta::v1::OwnerReference}; +use kube::{Api, Client, ResourceExt}; +use std::collections::BTreeSet; + +const MAX_CHAIN: usize = 64; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum LeafReadiness { + RequireReady, + /// Identity/bootstrap only. This never declares the leaf governance-Ready. + AllowPending, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ObjectUidRef { + pub namespace: String, + pub name: String, + pub uid: String, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct TaskLineagePin { + pub task: ObjectUidRef, + pub parent_task_uid: Option, + pub root_task_uid: String, +} + +#[derive(Clone, Debug)] +pub struct VerifiedTaskNode { + pub task: KarsTask, + pub pin: TaskLineagePin, + pub generation: i64, + pub resource_version: String, + pub authorization_digest: String, + pub effective_authorization: serde_json::Value, +} + +#[derive(Clone, Debug)] +pub struct VerifiedTaskLineage { + pub workspace_uid: String, + /// Root first, requested leaf last. + pub nodes: Vec, + pub team: Option, +} + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("Task UID lineage identity is absent, replaced, terminating or inconsistent")] + Identity, + #[error("Task UID lineage changed while being read; retry from current authority")] + Changed, + #[error("Task UID lineage is not governance-Ready")] + NotReady, + #[error("Task UID lineage amplifies its parent's effective authority")] + Attenuation, + #[error("Task UID lineage is cyclic or exceeds 64 nodes")] + Depth, + #[error("Task UID lineage API {stage} failed (status {code:?})")] + Api { + stage: &'static str, + code: Option, + }, +} + +fn api(stage: &'static str, error: kube::Error) -> Error { + Error::Api { + stage, + code: match error { + kube::Error::Api(status) => Some(status.code), + _ => None, + }, + } +} + +fn name(value: &str) -> bool { + !value.is_empty() + && value.len() <= 253 + && value.split('.').all(|label| { + !label.is_empty() + && label.len() <= 63 + && label + .bytes() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-') + && label.as_bytes()[0].is_ascii_alphanumeric() + && label.as_bytes()[label.len() - 1].is_ascii_alphanumeric() + }) +} + +fn identity(task: &KarsTask, namespace: &str) -> Result { + let uid = task + .uid() + .filter(|uid| !uid.is_empty()) + .ok_or(Error::Identity)?; + if task.metadata.namespace.as_deref() != Some(namespace) + || !name(&task.name_any()) + || task.metadata.deletion_timestamp.is_some() + || task + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || task + .metadata + .generation + .is_none_or(|generation| generation <= 0) + { + return Err(Error::Identity); + } + Ok(ObjectUidRef { + namespace: namespace.into(), + name: task.name_any(), + uid, + }) +} + +fn owner(task: &KarsTask) -> Result, Error> { + let mut owners = task + .owner_references() + .iter() + .filter(|owner| owner.controller == Some(true)); + let first = owners.next(); + if owners.next().is_some() { + return Err(Error::Identity); + } + Ok(first) +} + +impl VerifiedTaskLineage { + /// Empty input is only a live snapshot check, NOT immutable continuity. + /// Supplied pins must come from the consumer's authoritative persistence. + pub fn verify_pins(&self, pins: &[TaskLineagePin]) -> Result<(), Error> { + let mut seen = BTreeSet::new(); + for pin in pins { + if !seen.insert(&pin.task.uid) + || self + .nodes + .iter() + .find(|node| node.pin.task.uid == pin.task.uid) + .is_none_or(|node| node.pin != *pin) + { + return Err(Error::Identity); + } + } + Ok(()) + } +} + +/// Capture one stable same-workspace chain, using the canonical readiness, +/// attenuation and full-authorization helpers. A second UID/RV inventory +/// rejects changes across the collected snapshot rather than mixing epochs. +pub async fn resolve( + client: &Client, + leaf: &KarsTask, + readiness: LeafReadiness, +) -> Result { + let namespace = leaf + .namespace() + .filter(|namespace| name(namespace) && namespace.len() <= 63) + .ok_or(Error::Identity)?; + identity(leaf, &namespace)?; + let namespaces: Api = Api::all(client.clone()); + let workspace = namespaces + .get(&namespace) + .await + .map_err(|error| api("read workspace", error))?; + let workspace_uid = workspace + .uid() + .filter(|uid| !uid.is_empty()) + .ok_or(Error::Identity)?; + if workspace.metadata.deletion_timestamp.is_some() { + return Err(Error::Identity); + } + let tasks: Api = Api::namespaced(client.clone(), &namespace); + let mut current = tasks + .get(&leaf.name_any()) + .await + .map_err(|error| api("read leaf", error))?; + if current.metadata.uid != leaf.metadata.uid + || current.metadata.generation != leaf.metadata.generation + { + return Err(Error::Changed); + } + let mut path = Vec::new(); + let mut seen = BTreeSet::new(); + loop { + let id = identity(¤t, &namespace)?; + if path.len() >= MAX_CHAIN || !seen.insert(id.uid) { + return Err(Error::Depth); + } + if (readiness == LeafReadiness::RequireReady || !path.is_empty()) + && !crate::kars_task_reconciler::task_is_ready(¤t) + { + return Err(Error::NotReady); + } + let parent = current + .spec + .parent_ref + .as_ref() + .map(|reference| reference.name.clone()); + path.push(current.clone()); + let Some(parent) = parent else { break }; + if !name(&parent) { + return Err(Error::Identity); + } + let parent = tasks + .get(&parent) + .await + .map_err(|error| api("read parent", error))?; + identity(&parent, &namespace)?; + if !crate::kars_task::spec_attenuation_violations(¤t.spec, &parent.spec).is_empty() { + return Err(Error::Attenuation); + } + current = parent; + } + path.reverse(); + let root = path.first().ok_or(Error::Identity)?; + let root_uid = root.uid().ok_or(Error::Identity)?; + let team_owner = owner(root)? + .filter(|owner| owner.kind == "KarsTeam" && owner.api_version == "kars.azure.com/v1alpha1"); + let teams: Api = Api::namespaced(client.clone(), &namespace); + let team = if let Some(owner) = team_owner { + if !name(&owner.name) { + return Err(Error::Identity); + } + let team = teams + .get(&owner.name) + .await + .map_err(|error| api("read Team owner", error))?; + if team.metadata.namespace.as_deref() != Some(namespace.as_str()) + || team.metadata.uid.as_deref() != Some(owner.uid.as_str()) + || team.metadata.deletion_timestamp.is_some() + || team + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + { + return Err(Error::Identity); + } + Some(team) + } else { + None + }; + for (index, task) in path.iter().enumerate() { + if let Some(owner) = owner(task)? + && owner.api_version == "kars.azure.com/v1alpha1" + { + if owner.kind == "KarsTeam" + && team.as_ref().is_none_or(|team| { + team.metadata.uid.as_deref() != Some(owner.uid.as_str()) + || team.name_any() != owner.name + }) + { + return Err(Error::Identity); + } + if owner.kind == "KarsTask" + && index.checked_sub(1).is_none_or(|index| { + path[index].metadata.uid.as_deref() != Some(owner.uid.as_str()) + || path[index].name_any() != owner.name + }) + { + return Err(Error::Identity); + } + } + } + let model = crate::kars_task::blueprint::controller_default_model(); + let nodes: Vec<_> = path + .iter() + .enumerate() + .map(|(index, task)| { + Ok(VerifiedTaskNode { + task: task.clone(), + pin: TaskLineagePin { + task: identity(task, &namespace)?, + parent_task_uid: index.checked_sub(1).and_then(|index| path[index].uid()), + root_task_uid: root_uid.clone(), + }, + generation: task.metadata.generation.ok_or(Error::Identity)?, + resource_version: task.resource_version().ok_or(Error::Identity)?, + authorization_digest: task.spec.authorization_digest_with_model(&model), + effective_authorization: task.spec.authorization_configuration_with_model(&model), + }) + }) + .collect::>()?; + for node in &nodes { + let fresh = tasks + .get(&node.pin.task.name) + .await + .map_err(|error| api("recheck Task snapshot", error))?; + if fresh.metadata.uid.as_deref() != Some(node.pin.task.uid.as_str()) + || fresh.metadata.resource_version.as_deref() != Some(node.resource_version.as_str()) + || fresh.metadata.deletion_timestamp.is_some() + { + return Err(Error::Changed); + } + } + if let Some(team) = &team { + let fresh = teams + .get(&team.name_any()) + .await + .map_err(|error| api("recheck Team snapshot", error))?; + if fresh.metadata.uid != team.metadata.uid + || fresh.metadata.resource_version != team.metadata.resource_version + || fresh.metadata.deletion_timestamp.is_some() + { + return Err(Error::Changed); + } + } + let fresh = namespaces + .get(&namespace) + .await + .map_err(|error| api("recheck workspace", error))?; + if fresh.metadata.uid.as_deref() != Some(workspace_uid.as_str()) + || fresh.metadata.deletion_timestamp.is_some() + { + return Err(Error::Changed); + } + Ok(VerifiedTaskLineage { + workspace_uid, + nodes, + team, + }) +} + +#[cfg(test)] +#[path = "task_identity_tests.rs"] +mod tests; diff --git a/controller/src/task_identity_tests.rs b/controller/src/task_identity_tests.rs new file mode 100644 index 000000000..2f9dffb08 --- /dev/null +++ b/controller/src/task_identity_tests.rs @@ -0,0 +1,263 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use serde_json::json; +use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, +}; +use wiremock::{ + Mock, MockServer, Request, Respond, ResponseTemplate, + matchers::{method, path}, +}; + +fn task(name: &str, parent: Option<&str>, ready: bool) -> KarsTask { + let mut task: KarsTask = serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1", "kind":"KarsTask", + "metadata":{"name":name,"namespace":"workspace","uid":format!("{name}-uid"),"resourceVersion":"1","generation":1}, + "spec":{ + "objective":"Fixture", "envelope":{"tier":if parent.is_some(){2}else{3}, + "authorityCeiling":if parent.is_some(){2}else{3},"delegationDepth":if parent.is_some(){1}else{2}}, + "blueprint":{"model":{"provider":"azure-openai","deployment":"fixture"}}, + "parentRef":parent.map(|name| json!({"name":name})) + } + })).unwrap(); + task.status = Some( + serde_json::from_value(json!({ + "phase":if ready{"Ready"}else{"Pending"}, "observedGeneration":1, + "envelopeDigest":task.envelope_digest(), + "conditions":[{"type":"Ready","status":if ready{"True"}else{"False"}, + "reason":"Fixture","message":"Fixture","lastTransitionTime":"2026-09-08T00:00:00Z"}] + })) + .unwrap(), + ); + task +} + +async fn setup() -> (MockServer, Client) { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/v1/namespaces/workspace")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "apiVersion":"v1","kind":"Namespace", + "metadata":{"name":"workspace","uid":"workspace-uid","resourceVersion":"1"} + }))) + .mount(&server) + .await; + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + (server, client) +} + +async fn serve_task(server: &MockServer, task: &KarsTask) { + Mock::given(method("GET")) + .and(path(format!( + "/apis/kars.azure.com/v1alpha1/namespaces/workspace/karstasks/{}", + task.name_any() + ))) + .respond_with(ResponseTemplate::new(200).set_body_json(task)) + .mount(server) + .await; +} + +#[tokio::test] +async fn returns_uid_order_and_exact_canonical_authorization_without_budget_coupling() { + let (server, client) = setup().await; + let root = task("root", None, true); + let leaf = task("leaf", Some("root"), true); + serve_task(&server, &root).await; + serve_task(&server, &leaf).await; + let lineage = resolve(&client, &leaf, LeafReadiness::RequireReady) + .await + .unwrap(); + assert_eq!(lineage.workspace_uid, "workspace-uid"); + assert!(lineage.team.is_none()); + assert_eq!( + lineage + .nodes + .iter() + .map(|node| node.pin.task.uid.as_str()) + .collect::>(), + ["root-uid", "leaf-uid"] + ); + assert_eq!( + lineage.nodes[1].pin.parent_task_uid.as_deref(), + Some("root-uid") + ); + assert_eq!(lineage.nodes[1].pin.root_task_uid, "root-uid"); + assert_eq!( + lineage.nodes[1].authorization_digest, + leaf.envelope_digest() + ); + assert_eq!(lineage.nodes[1].generation, 1); + lineage + .verify_pins(&[lineage.nodes[1].pin.clone()]) + .unwrap(); + let mut stale = lineage.nodes[1].pin.clone(); + stale.parent_task_uid = Some("recreated-root".into()); + assert!(lineage.verify_pins(&[stale]).is_err()); + let pin = lineage.nodes[1].pin.clone(); + assert!(lineage.verify_pins(&[pin.clone(), pin]).is_err()); +} + +#[tokio::test] +async fn pending_leaf_is_explicit_but_pending_ancestors_never_grant_authority() { + let (server, client) = setup().await; + let root = task("root", None, true); + let leaf = task("leaf", Some("root"), false); + serve_task(&server, &root).await; + serve_task(&server, &leaf).await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::NotReady) + )); + assert!( + resolve(&client, &leaf, LeafReadiness::AllowPending) + .await + .is_ok() + ); + + let (server, client) = setup().await; + serve_task(&server, &task("root", None, false)).await; + serve_task(&server, &leaf).await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::AllowPending).await, + Err(Error::NotReady) + )); +} + +#[derive(Clone)] +struct ReplacedOnRecheck { + calls: Arc, + task: KarsTask, +} + +impl Respond for ReplacedOnRecheck { + fn respond(&self, _: &Request) -> ResponseTemplate { + let mut task = self.task.clone(); + if self.calls.fetch_add(1, Ordering::SeqCst) > 0 { + task.metadata.uid = Some("recreated-root-uid".into()); + } + ResponseTemplate::new(200).set_body_json(task) + } +} + +#[tokio::test] +async fn double_inventory_rejects_recreated_uid_instead_of_mixing_ancestry() { + let (server, client) = setup().await; + let leaf = task("leaf", Some("root"), true); + serve_task(&server, &leaf).await; + Mock::given(method("GET")) + .and(path( + "/apis/kars.azure.com/v1alpha1/namespaces/workspace/karstasks/root", + )) + .respond_with(ReplacedOnRecheck { + calls: Arc::new(AtomicUsize::new(0)), + task: task("root", None, true), + }) + .mount(&server) + .await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::Changed) + )); +} + +#[tokio::test] +async fn stale_leaf_generation_and_foreign_namespace_are_not_adopted() { + let (server, client) = setup().await; + let leaf = task("leaf", None, true); + let mut replaced = leaf.clone(); + replaced.metadata.generation = Some(2); + serve_task(&server, &replaced).await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::Changed) + )); + let mut invalid = leaf; + invalid.metadata.namespace = Some("../foreign".into()); + assert!(matches!( + resolve(&client, &invalid, LeafReadiness::RequireReady).await, + Err(Error::Identity) + )); +} + +#[tokio::test] +async fn non404_parent_failure_propagates_without_authority_or_writes() { + let (server, client) = setup().await; + let leaf = task("leaf", Some("root"), true); + serve_task(&server, &leaf).await; + Mock::given(method("GET")) + .and(path( + "/apis/kars.azure.com/v1alpha1/namespaces/workspace/karstasks/root", + )) + .respond_with(ResponseTemplate::new(503).set_body_json(json!({ + "kind":"Status","apiVersion":"v1","status":"Failure","code":503, + "reason":"ServiceUnavailable","message":"fixture" + }))) + .mount(&server) + .await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::Api { + code: Some(503), + .. + }) + )); + assert!( + server + .received_requests() + .await + .unwrap() + .iter() + .all(|request| request.method == "GET") + ); +} + +#[tokio::test] +async fn conflicting_team_owner_uid_cannot_select_another_lifetime_root() { + let (server, client) = setup().await; + let root = task("root", None, true); + let mut leaf = task("leaf", Some("root"), true); + leaf.metadata.owner_references = Some( + serde_json::from_value(json!([{ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTeam", + "name":"foreign-team","uid":"foreign-team-uid","controller":true + }])) + .unwrap(), + ); + serve_task(&server, &root).await; + serve_task(&server, &leaf).await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::Identity) + )); +} + +#[tokio::test] +async fn team_owner_is_live_uid_bound_and_not_inferred_from_display_names() { + let (server, client) = setup().await; + let mut root = task("root", None, true); + let mut leaf = task("leaf", Some("root"), true); + let owner = serde_json::from_value(json!([{ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTeam", + "name":"team","uid":"team-uid","controller":true + }])) + .unwrap(); + root.metadata.owner_references = Some(owner); + leaf.metadata.owner_references = root.metadata.owner_references.clone(); + serve_task(&server, &root).await; + serve_task(&server, &leaf).await; + Mock::given(method("GET")) + .and(path("/apis/kars.azure.com/v1alpha1/namespaces/workspace/karsteams/team")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTeam", + "metadata":{"name":"team","namespace":"workspace","uid":"team-uid","resourceVersion":"1","generation":1}, + "spec":{"charter":"Fixture","envelope":{"tier":3,"authorityCeiling":3,"delegationDepth":2}} + }))).mount(&server).await; + let lineage = resolve(&client, &leaf, LeafReadiness::RequireReady) + .await + .unwrap(); + assert_eq!(lineage.team.unwrap().uid().as_deref(), Some("team-uid")); + assert_eq!(lineage.nodes[1].pin.root_task_uid, "root-uid"); +} diff --git a/docs/audits/governed-inference-budgets-2026-09-08.md b/docs/audits/governed-inference-budgets-2026-09-08.md index aa2aadf9d..98a4401c9 100644 --- a/docs/audits/governed-inference-budgets-2026-09-08.md +++ b/docs/audits/governed-inference-budgets-2026-09-08.md @@ -27,13 +27,17 @@ all-in task spend, invoice accuracy, taxes, or exchange rates. | Router code identity | Finite mode requires an operator-qualified immutable router manifest digest | Helm configuration tested; runtime qualification pending | No Cargo command, dependency installation, local Docker test, deployment, customer -mutation, H100 operation, main-branch change, image publication, or budget commit/ -push was performed for this evidence. Existing authorized cached CLI dependencies -were used after the local runner was found missing. +mutation, H100 operation, main-branch change, image publication, or public push +was performed for this evidence. Authorized local checkpoint `eb26efd9` and +forward merge `0701baed` preserve the candidate and exact privacy parent +`7dc72810a2e3c87aa751cfa95d9152f8dcd10194`. Existing authorized cached CLI +dependencies were used after the local runner was found missing. ## Remaining release decisions/gates -1. Integrate the real privacy issuer prerequisite; no fallback implementation. +1. Qualify against the real, now-forwarded privacy issuer prerequisite; no + fallback implementation. Parent re-review and full SRE Kind remain separate + gates, not implied by this merge. 2. Qualify Team lifecycle/selective launch integration, first-opt-in transition constraints, and source/route/cancellation regressions, including the actual broker-in-Kind scenario and separate schema/identity preflight. diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index 4fe9ee08b..297134392 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -60,6 +60,14 @@ Workspace producers create Tasks/Teams; they do not create accounts, write balances, manufacture status, or inject this Sandbox field. Root UID recreation is a new grant identity, not continuation of an old account. +The shared controller `task_identity::resolve` helper captures a stable live +same-workspace UID chain, with canonical readiness, attenuation and effective +authorization, and a verified optional Team owner. It rechecks UID/resourceVersion +after collection. `VerifiedTaskLineage::verify_pins` compares a consumer's +authoritative persisted references; traversal alone does not create immutable +continuity. Budget enrollment owns persistence and financial state. Credential +consumers can reuse identity checks without depending on the budget ledger. + ## Reserve, authorize one send, settle All ancestor meters and the root meter are updated in **one Kubernetes object From 46b99da10767f21cae77c7dfd611236009aae59e Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 22:36:08 +0200 Subject: [PATCH 03/16] docs(budget): record static qualification and pending human audit gate Record actual post-068ae160 static and cached-test results without claiming budget Rust, Kind or human signoff completion. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .../2026-09-08-governed-inference-budgets.md | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index 0e22ddfa2..7e0196cbf 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -60,7 +60,10 @@ all-in task spend, invoice accuracy, taxes, or exchange rates. | CLI static validation | Existing TypeScript typecheck and targeted oxlint | **Passed locally** | | Public API/CEL | Independent pinned Kind v0.24 / Kubernetes v1.31 preflight added, no Rust image dependency | Not executed | | Complete broker Kind integration | Real loaded router digest, TLS broker, sibling token/price caps, route closure and cancellation scenario added to the existing E2E runner | Not executed | -| Affected-crate strict Clippy/static guards | Required, no waivers | Pending | +| Shared Task UID identity | Seven readiness, ancestry, UID/RV, Team-owner and API failure tests authored | Rust execution pending | +| Static gates | Existing LOC, no-custom-crypto and no-stubs scripts against `068ae160` | **Passed locally** | +| Source checks | Rustfmt parsing, JavaScript/shell/YAML/TOML syntax and diff checks | **Passed locally** | +| Affected-crate strict Clippy | Required, no waivers | Pending | | Launch/Team cadence gates | Explicit scope syntax, first-opt-in transition constraints, and mandatory asynchronous broker/account checks | Source implemented; Rust/API qualification pending | | Router code identity | Finite mode requires an operator-qualified immutable router manifest digest | Helm configuration tested; runtime qualification pending | @@ -76,6 +79,9 @@ ancestry `068ae16041ecf7bd2b8321dfeb22e381ebbd587b`, including `9d0f8e23` epoch transition repairs and the shipped-schema Kubernetes compatibility repair. Parent-reported prerequisite tests/Clippy and schema API successes do not qualify this budget implementation. Full SRE Kind remains a separate open gate. +The existing `security-audit-required` script now discovers this correctly placed +record and fails specifically for **0 of 2 required genuine signer emails**. +That failure is intentional until human review, not a waived or fabricated pass. ## Remaining release decisions/gates From aa34c3fe7ba3f73d26015ab2afe38404c071bb15 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 23:32:51 +0200 Subject: [PATCH 04/16] fix(budget): preserve funded work and require final stream usage Repair the five bounded source-review findings: module scopes, supported Team launches, non-destructive admission waits, final Anthropic usage evidence, and legacy planning selection. Preserve Task UIDs and existing funded executions while denying new admissions; retain explicit pause/policy/UID revocation. Add full Team/Task API interleaving and actual stream-settlement regressions. Rust tests are authored but unrun pending the parent's Cargo grant. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/inference_budget/binding.rs | 29 +- controller/src/inference_budget/launch.rs | 128 +++++ controller/src/inference_budget/mod.rs | 1 + controller/src/inference_budget/pod.rs | 25 +- controller/src/kars_task_budget_tests.rs | 269 +++++++++ controller/src/kars_task_execution.rs | 24 + controller/src/kars_task_reconciler.rs | 37 +- controller/src/kars_team_reconciler.rs | 5 +- .../budget_interleaving_tests.rs | 512 ++++++++++++++++++ .../kars_team_reconciler/persistence_tests.rs | 5 +- controller/src/kars_team_reconciler/tasks.rs | 2 +- docs/governed-inference-budgets.md | 11 + .../2026-09-08-governed-inference-budgets.md | 21 + .../src/inference_budget/anthropic_cases.rs | 65 +++ .../src/inference_budget/client_tests.rs | 82 ++- inference-router/src/inference_budget/mod.rs | 2 + .../src/inference_budget/usage.rs | 137 ++++- shared/inference_budget/ledger_tests.rs | 9 + shared/inference_budget/ledger_validation.rs | 21 +- 19 files changed, 1323 insertions(+), 62 deletions(-) create mode 100644 controller/src/inference_budget/launch.rs create mode 100644 controller/src/kars_task_budget_tests.rs create mode 100644 controller/src/kars_team_reconciler/budget_interleaving_tests.rs create mode 100644 inference-router/src/inference_budget/anthropic_cases.rs diff --git a/controller/src/inference_budget/binding.rs b/controller/src/inference_budget/binding.rs index 80be93e61..2c7dc06c6 100644 --- a/controller/src/inference_budget/binding.rs +++ b/controller/src/inference_budget/binding.rs @@ -61,6 +61,13 @@ pub fn has_finite(envelope: &TaskEnvelope) -> bool { }) } +fn explicitly_governed(envelope: &TaskEnvelope) -> bool { + envelope + .budget + .as_ref() + .is_some_and(|budget| budget.scope == Some(BudgetScope::GovernedInference)) +} + fn resource(task: &KarsTask) -> Result { let resource = ResourceIdentity { namespace: task.namespace().ok_or(BudgetError::Identity)?, @@ -106,9 +113,17 @@ async fn chain( } pub(super) async fn needs_account(client: &Client, task: &KarsTask) -> Result { - let lineage = chain(client, task).await?; + // Selection must not interpret a legacy planning declaration as an + // enforcement request. Only validate budget semantics after opt-in or an + // existing lineage/account pin establishes governed continuity. + let lineage = crate::task_identity::resolve( + client, + task, + crate::task_identity::LeafReadiness::AllowPending, + ) + .await?; if lineage.nodes.iter().any(|node| { - has_finite(&node.task.spec.envelope) + explicitly_governed(&node.task.spec.envelope) || node .task .status @@ -118,7 +133,7 @@ pub(super) async fn needs_account(client: &Client, task: &KarsTask) -> Result Result<(), StoreErr } pub async fn prepare_task(client: &Client, task: &KarsTask) -> Result { - if !has_finite(&task.spec.envelope) + if !explicitly_governed(&task.spec.envelope) && task .status .as_ref() .is_none_or(|status| status.inference_budget.is_none()) + && task.spec.parent_ref.is_none() + && !task.owner_references().iter().any(|owner| { + owner.controller == Some(true) + && owner.kind == "KarsTeam" + && owner.api_version == "kars.azure.com/v1alpha1" + }) { return Ok(task.clone()); } diff --git a/controller/src/inference_budget/launch.rs b/controller/src/inference_budget/launch.rs new file mode 100644 index 000000000..24b23f123 --- /dev/null +++ b/controller/src/inference_budget/launch.rs @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Admission unavailability is not revocation of an already funded execution. + +use super::store::{Store, StoreError}; +use crate::{ + inference_budget_contract::{AccountPhase, BudgetError, MAX_SESSIONS, ledger::Ledger}, + kars_task::{KarsTask, KarsTaskStatus}, + status::conditions, +}; +use kube::{Client, ResourceExt}; + +const CONDITION: &str = "GovernedInferenceReady"; +const PENDING: &str = "AdmissionUnavailable"; + +pub fn capacity(ledger: &Ledger, task_uid: &str) -> Result<(), BudgetError> { + ledger.validate()?; + if ledger.phase != AccountPhase::Active { + return Err(BudgetError::Closed); + } + let exhausted = |limits: crate::inference_budget_contract::Limits, + meters: &crate::inference_budget_contract::ledger::Meters| + -> Result { + let used = meters.total()?; + Ok(limits + .tokens + .filter(|limit| *limit > 0) + .is_some_and(|limit| used.tokens >= limit) + || limits + .usd_micros + .filter(|limit| *limit > 0) + .is_some_and(|limit| used.usd_micros >= limit)) + }; + if exhausted(ledger.limits, &ledger.meters)? { + return Err(BudgetError::Exhausted); + } + for uid in ledger.ancestors(task_uid)? { + let node = &ledger.nodes[&uid]; + if !node.active { + return Err(BudgetError::Closed); + } + if exhausted(node.authority.limits, &node.meters)? { + return Err(BudgetError::Exhausted); + } + } + if ledger.sessions.len() >= MAX_SESSIONS { + return Err(BudgetError::Capacity); + } + Ok(()) +} + +pub async fn admit_new(client: &Client, task: &KarsTask) -> Result<(), StoreError> { + let Some(binding) = task + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + else { + return Ok(()); + }; + let account = Store::new(client.clone(), &binding.account.namespace) + .read(&binding.root, &binding.account.uid) + .await?; + let ledger = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + .ok_or(StoreError::Missing)?; + capacity(ledger, &binding.task_uid)?; + Ok(()) +} + +pub fn mark_pending(status: &mut KarsTaskStatus, task: &KarsTask, detail: &str) { + let conditions = status.conditions.get_or_insert_with(Vec::new); + let prior = task + .status + .as_ref() + .and_then(|status| status.conditions.as_ref()) + .and_then(|conditions| conditions::find(conditions, CONDITION)); + conditions::set( + conditions, + conditions::preserve_transition_time( + prior, + CONDITION, + "False", + PENDING, + detail, + task.metadata.generation, + ), + ); +} + +/// Only the explicit budget-admission wait marker can retain a child's +/// execution. A changed UID/spec or failed attenuation remains revocation. +pub fn pending_parent(parent: &KarsTask, child: &KarsTask) -> bool { + let Some(status) = &parent.status else { + return false; + }; + let Some(binding) = &status.inference_budget else { + return false; + }; + parent.metadata.deletion_timestamp.is_none() + && status.observed_generation == parent.metadata.generation + && parent.uid().as_deref() == Some(binding.task_uid.as_str()) + && parent.envelope_digest() == binding.authorization_digest + && crate::kars_task::spec_attenuation_violations(&child.spec, &parent.spec).is_empty() + && child + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + .is_none_or(|binding| binding.parent_task_uid.as_ref() == parent.metadata.uid.as_ref()) + && status.conditions.iter().flatten().any(|condition| { + condition.type_ == CONDITION + && condition.status == "False" + && condition.reason == PENDING + }) +} + +pub fn retain_execution(task: &KarsTask, status: &mut KarsTaskStatus) { + let prior = task.status.as_ref(); + status.sandbox_ref = prior.and_then(|status| status.sandbox_ref.clone()); + status.execution_phase = prior + .and_then(|status| status.execution_phase.clone()) + .or_else(|| Some("Pending".into())); + status.execution_detail = Some( + "New budget admission is unavailable; existing owned execution and funded work are retained, not re-authorized".into() + ); +} diff --git a/controller/src/inference_budget/mod.rs b/controller/src/inference_budget/mod.rs index 7c7b26df5..0cc17d21c 100644 --- a/controller/src/inference_budget/mod.rs +++ b/controller/src/inference_budget/mod.rs @@ -7,6 +7,7 @@ pub mod auth; pub mod binding; pub mod claim; pub mod config; +pub mod launch; pub mod pod; pub mod recovery; pub mod scope; diff --git a/controller/src/inference_budget/pod.rs b/controller/src/inference_budget/pod.rs index 389a25660..40aebaa2f 100644 --- a/controller/src/inference_budget/pod.rs +++ b/controller/src/inference_budget/pod.rs @@ -260,6 +260,19 @@ async fn mirror_public_ca( Ok(()) } +pub async fn decorate( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, + pod: &mut Value, +) -> Result, StoreError> { + let mut annotations = std::collections::BTreeMap::new(); + if let Some(plan) = prepare(client, sandbox, namespace).await? { + plan.apply(pod, &mut annotations)?; + } + Ok(annotations) +} + impl Plan { pub fn apply( &self, @@ -279,18 +292,6 @@ impl Plan { return Err(BudgetError::Corrupt.into()); } - pub async fn decorate( - client: &Client, - sandbox: &KarsSandbox, - namespace: &Namespace, - pod: &mut Value, - ) -> Result, StoreError> { - let mut annotations = std::collections::BTreeMap::new(); - if let Some(plan) = prepare(client, sandbox, namespace).await? { - plan.apply(pod, &mut annotations)?; - } - Ok(annotations) - } volumes.push(json!({ "name": TOKEN_VOLUME, "projected": {"sources": [{"serviceAccountToken": {"audience": AUDIENCE, "expirationSeconds": 600, "path": "token"}}]} diff --git a/controller/src/kars_task_budget_tests.rs b/controller/src/kars_task_budget_tests.rs new file mode 100644 index 000000000..522deee2b --- /dev/null +++ b/controller/src/kars_task_budget_tests.rs @@ -0,0 +1,269 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{ + inference_budget_contract::{ + AccountReference, BudgetScope, ResourceIdentity, RootIdentity, RootKind, TaskBudgetBinding, + }, + kars_task::{KarsTaskSpec, TaskBlueprint, TaskBudget, TaskEnvelope, TaskExecution, TaskModel}, + mcp_server::LocalObjectRef, +}; +use serde_json::Value; +use std::sync::Mutex; +use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; + +struct Objects { + task: KarsTask, + parent: Option, + sandbox: bool, + deletes: usize, +} +#[derive(Clone)] +struct ApiServer(Arc>); +impl Respond for ApiServer { + fn respond(&self, request: &Request) -> ResponseTemplate { + let mut objects = self.0.lock().unwrap(); + let path = request.url.path(); + if request.method == "GET" && path == "/api/v1/namespaces/workspace" { + return ResponseTemplate::new(200).set_body_json(json!({ + "apiVersion":"v1","kind":"Namespace","metadata":{"name":"workspace","uid":"namespace-uid","resourceVersion":"1"} + })); + } + if request.method == "GET" && path.ends_with("/karstasks/task") { + return ResponseTemplate::new(200).set_body_json(&objects.task); + } + if request.method == "GET" && path.ends_with("/karstasks/parent") { + if let Some(parent) = &objects.parent { + return ResponseTemplate::new(200).set_body_json(parent); + } + } + if request.method == "PATCH" && path.ends_with("/karstasks/task/status") { + let patch: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(patch["metadata"]["uid"], objects.task.uid().unwrap()); + assert_eq!( + patch["metadata"]["resourceVersion"], + objects.task.resource_version().unwrap() + ); + objects.task.status = Some(serde_json::from_value(patch["status"].clone()).unwrap()); + objects.task.metadata.resource_version = Some("2".into()); + return ResponseTemplate::new(200).set_body_json(&objects.task); + } + if path.ends_with("/karssandboxes/task") && objects.sandbox { + let sandbox = json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsSandbox", + "metadata":{"name":"task","namespace":"workspace","uid":"sandbox-uid","resourceVersion":"1", + "ownerReferences":[{"apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTask", + "name":"task","uid":objects.task.uid(),"controller":true}]}, + "spec":{} + }); + if request.method == "GET" { + return ResponseTemplate::new(200).set_body_json(sandbox); + } + if request.method == "DELETE" { + let body: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(body["preconditions"]["uid"], "sandbox-uid"); + assert_eq!(body["preconditions"]["resourceVersion"], "1"); + objects.deletes += 1; + objects.sandbox = false; + return ResponseTemplate::new(200).set_body_json(sandbox); + } + } + ResponseTemplate::new(404).set_body_json(json!({ + "apiVersion":"v1","kind":"Status","status":"Failure", + "code":404,"reason":"NotFound","message":"fixture" + })) + } +} + +fn fixture(name: &str, parent: Option<&str>, governed: bool, launched: bool) -> KarsTask { + let mut task = KarsTask::new( + name, + KarsTaskSpec { + objective: "fixture".into(), + envelope: TaskEnvelope { + tier: if parent.is_some() { 2 } else { 3 }, + authority_ceiling: if parent.is_some() { 2 } else { 3 }, + delegation_depth: if parent.is_some() { 1 } else { 2 }, + budget: Some(TaskBudget { + scope: governed.then_some(BudgetScope::GovernedInference), + tokens: Some(30), + usd_micros: Some(50), + }), + ..Default::default() + }, + parent_ref: parent.map(|name| LocalObjectRef { name: name.into() }), + execution: Some(TaskExecution { + launch: launched, + runtime: None, + }), + blueprint: Some(TaskBlueprint { + model: Some(TaskModel { + provider: "azure-openai".into(), + deployment: "fixture".into(), + }), + ..Default::default() + }), + ..Default::default() + }, + ); + task.metadata.namespace = Some("workspace".into()); + task.metadata.uid = Some(format!("{name}-uid")); + task.metadata.generation = Some(1); + task.metadata.resource_version = Some("1".into()); + task.metadata.finalizers = Some(vec![FINALIZER.into()]); + task.status = Some(ready_status(None, Some(1), task.envelope_digest(), vec![])); + if governed { + let root_name = parent.unwrap_or(name); + let root_uid = format!("{root_name}-uid"); + task.status.as_mut().unwrap().inference_budget = Some(TaskBudgetBinding { + scope: BudgetScope::GovernedInference, + account: AccountReference { + namespace: "accounting".into(), + name: "inference-budget-root".into(), + uid: "account-uid".into(), + }, + root: RootIdentity { + kind: RootKind::KarsTask, + resource: ResourceIdentity { + namespace: "workspace".into(), + name: root_name.into(), + uid: root_uid.clone(), + }, + workspace_uid: "namespace-uid".into(), + cluster_uid: "cluster-uid".into(), + }, + task_uid: task.uid().unwrap(), + parent_task_uid: parent.map(|name| format!("{name}-uid")), + root_task_uid: root_uid, + authorization_digest: task.envelope_digest(), + }); + } + if launched { + let status = task.status.as_mut().unwrap(); + status.execution_phase = Some("Running".into()); + status.sandbox_ref = Some(LocalObjectRef { name: name.into() }); + } + task +} + +async fn run(task: KarsTask, parent: Option, sandbox: bool) -> Arc> { + let server = MockServer::start().await; + let objects = Arc::new(Mutex::new(Objects { + task: task.clone(), + parent, + sandbox, + deletes: 0, + })); + Mock::given(wiremock::matchers::any()) + .respond_with(ApiServer(objects.clone())) + .mount(&server) + .await; + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + reconcile( + Arc::new(task), + Arc::new(Ctx { + client, + signer: crate::providers::signing::ReceiptSigner::from_bytes(&[7; 32]), + }), + ) + .await + .unwrap(); + objects +} + +#[tokio::test] +async fn legacy_positive_plans_and_their_parent_readiness_are_unchanged_without_opt_in() { + let root = fixture("task", None, false, false); + let objects = run(root, None, false).await; + assert!(task_is_ready(&objects.lock().unwrap().task)); + let child = fixture("task", Some("parent"), false, false); + let parent = fixture("parent", None, false, false); + let objects = run(child, Some(parent), false).await; + let objects = objects.lock().unwrap(); + assert!(task_is_ready(&objects.task)); + assert!( + objects + .task + .status + .as_ref() + .unwrap() + .inference_budget + .is_none() + ); + assert_eq!(objects.deletes, 0); +} + +#[tokio::test] +async fn unavailable_budget_preparation_keeps_owned_execution_but_denies_new_authority() { + let objects = run(fixture("task", None, true, true), None, true).await; + let objects = objects.lock().unwrap(); + assert!(!task_is_ready(&objects.task)); + assert!(objects.sandbox); + assert_eq!(objects.deletes, 0); + assert_eq!(objects.task.uid().as_deref(), Some("task-uid")); + assert_eq!( + objects + .task + .status + .as_ref() + .unwrap() + .execution_phase + .as_deref(), + Some("Running") + ); + assert!( + objects + .task + .status + .as_ref() + .unwrap() + .inference_budget + .is_some() + ); +} + +#[tokio::test] +async fn budget_pending_parent_preserves_funded_child_but_pause_and_uid_revocation_still_stop() { + let mut parent = fixture("parent", None, true, false); + let mut status = parent.status.clone().unwrap(); + status.phase = Some("Degraded".into()); + crate::inference_budget::launch::mark_pending( + &mut status, + &parent, + "fixture budget unavailable", + ); + parent.status = Some(status); + let child = fixture("task", Some("parent"), true, true); + let held = run(child.clone(), Some(parent.clone()), true).await; + assert!(held.lock().unwrap().sandbox); + assert_eq!(held.lock().unwrap().deletes, 0); + let mut paused = child.clone(); + paused.spec.execution.as_mut().unwrap().launch = false; + let stopped = run(paused, Some(parent.clone()), true).await; + assert!(!stopped.lock().unwrap().sandbox); + assert_eq!(stopped.lock().unwrap().deletes, 1); + parent.metadata.uid = Some("recreated-parent".into()); + let revoked = run(child, Some(parent), true).await; + assert!(!revoked.lock().unwrap().sandbox); + assert_eq!(revoked.lock().unwrap().deletes, 1); +} + +#[tokio::test] +async fn a_pinned_account_cannot_escape_by_removing_governed_scope() { + let mut task = fixture("task", None, true, false); + task.spec.envelope.budget.as_mut().unwrap().scope = None; + let objects = run(task, None, false).await; + let objects = objects.lock().unwrap(); + assert!(!task_is_ready(&objects.task)); + assert!( + objects + .task + .status + .as_ref() + .unwrap() + .inference_budget + .is_some() + ); +} diff --git a/controller/src/kars_task_execution.rs b/controller/src/kars_task_execution.rs index ab55ff1c8..43dc9dd77 100644 --- a/controller/src/kars_task_execution.rs +++ b/controller/src/kars_task_execution.rs @@ -140,6 +140,30 @@ pub async fn materialize( let prepared = crate::inference_budget::binding::prepare_task(client, task) .await .map_err(|error| contract_error(error.to_string()))?; + if let Some(binding) = prepared + .status + .as_ref() + .and_then(|status| status.inference_budget.as_ref()) + { + let api: Api = + Api::namespaced_with(client.clone(), namespace, &sandbox_api_resource()); + let existing = api.get_opt(&task_name).await?; + if existing.as_ref().is_some_and(|sandbox| { + !owned_by_task(sandbox, task) || sandbox.metadata.deletion_timestamp.is_some() + }) { + return Err(contract_error( + "existing execution is foreign or still terminating".into(), + )); + } + let continues = existing.as_ref().is_some_and(|sandbox| { + sandbox.data.pointer("/spec/inferenceBudgetRef") == Some(&json!(binding)) + }); + if !continues { + crate::inference_budget::launch::admit_new(client, &prepared) + .await + .map_err(|error| contract_error(error.to_string()))?; + } + } // 1. InferencePolicy scoped to this sandbox. Model: blueprint wins, else // the controller default (required — without it the sandbox degrades). diff --git a/controller/src/kars_task_reconciler.rs b/controller/src/kars_task_reconciler.rs index 95be75c9e..133c4a22f 100644 --- a/controller/src/kars_task_reconciler.rs +++ b/controller/src/kars_task_reconciler.rs @@ -42,6 +42,10 @@ const REQUEUE_OK: Duration = Duration::from_secs(300); /// promptly once the parent reconciles, rather than waiting a full cycle. const REQUEUE_PENDING: Duration = Duration::from_secs(10); +#[cfg(test)] +#[path = "kars_task_budget_tests.rs"] +mod budget_tests; + #[derive(Debug, thiserror::Error)] enum ReconcileError { #[error(transparent)] @@ -172,6 +176,7 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result degraded_status( @@ -193,7 +198,11 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result { + Delegation::ParentNotReady { + parent, + admission_pending, + } => { + budget_pending = admission_pending; tracing::info!(karstask = %name, ns = %ns, %parent, "KarsTask parent not yet ready — waiting"); pending_status( prior_ready, @@ -241,6 +250,7 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result task = Arc::new(prepared), Err(error) => { + budget_pending = true; new_status = degraded_status( prior_ready, generation, @@ -263,12 +273,29 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result = tokio::sync::Mutex::const_new(()); +struct BudgetEnvironment(Vec<(&'static str, Option)>); +impl BudgetEnvironment { + fn enabled() -> Self { + let mut old = Vec::new(); + for (key, value) in [ + ("KARS_INFERENCE_BUDGET_ENABLED", "true".into()), + ("KARS_INFERENCE_BUDGET_TLS_SECRET", "fixture-tls".into()), + ( + "KARS_INFERENCE_BUDGET_ROUTER_DIGEST", + format!("sha256:{}", "a".repeat(64)), + ), + ] { + old.push((key, std::env::var(key).ok())); + // This module serializes its configuration mutations and restores + // them on unwind; run its targeted selector with --test-threads=1. + unsafe { std::env::set_var(key, value) }; + } + Self(old) + } +} +impl Drop for BudgetEnvironment { + fn drop(&mut self) { + for (key, value) in &self.0 { + unsafe { + match value { + Some(value) => std::env::set_var(key, value), + None => std::env::remove_var(key), + } + } + } + } +} + +fn contract() -> ModelContract { + ModelContract { + id: "fixture".into(), + version: "v1".into(), + valid_until: "2030-01-01T00:00:00Z".into(), + provider_id: "azure-openai".into(), + endpoint: "https://fixture.example".into(), + model: "reviewed-model".into(), + operation: Operation::ChatCompletions, + output_field: OutputField::MaxTokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 5 }), + } +} + +#[derive(Clone)] +struct BudgetApis { + account: Arc>>, + catalog_failed: Arc, + store_failed: Arc, + namespace: String, +} +impl Respond for BudgetApis { + fn respond(&self, request: &Request) -> ResponseTemplate { + if request.method != "GET" { + return failure(405); + } + let path = request.url.path(); + if path.contains("/admissionregistration.k8s.io/") { + let name = path.rsplit('/').next().unwrap(); + let bundle: Value = serde_json::from_str( + &include_str!("../../../deploy/helm/kars/files/inference-budget-admission.json") + .replace("__ACCOUNTING_NAMESPACE__", &self.namespace), + ) + .unwrap(); + let policy = bundle["items"] + .as_array() + .unwrap() + .iter() + .find(|policy| policy["name"] == name) + .unwrap(); + let (kind, spec) = if path.contains("/validatingadmissionpolicybindings/") { + ( + "ValidatingAdmissionPolicyBinding", + json!({"policyName":name,"validationActions":["Deny","Audit"]}), + ) + } else { + ("ValidatingAdmissionPolicy", policy["spec"].clone()) + }; + return response( + 200, + json!({ + "apiVersion":"admissionregistration.k8s.io/v1","kind":kind, + "metadata":{"name":name,"uid":format!("uid-{name}"),"generation":1,"resourceVersion":"1"}, + "spec":spec,"status":{"observedGeneration":1,"typeChecking":{"expressionWarnings":[]}} + }), + ); + } + if path.ends_with("/configmaps/kars-inference-budget-contracts") { + if self.catalog_failed.load(Ordering::SeqCst) { + return failure(503); + } + return response( + 200, + json!({ + "apiVersion":"v1","kind":"ConfigMap", + "metadata":{"name":"kars-inference-budget-contracts","namespace":self.namespace, + "uid":"catalog-uid","resourceVersion":"1", + "annotations":{"kars.azure.com/inference-budget-contracts":"v1"}}, + "data":{"contracts.json":serde_json::to_string(&json!({ + "version":"v1","contracts":[contract()],"nonInferenceEgressHosts":[] + })).unwrap()} + }), + ); + } + if self.store_failed.load(Ordering::SeqCst) { + return failure(503); + } + match self.account.lock().unwrap().as_ref() { + Some(account) => response(200, serde_json::to_value(account).unwrap()), + None => failure(404), + } + } +} + +async fn budget_apis(server: &MockServer) -> BudgetApis { + let settings = Settings::from_env().unwrap().unwrap(); + let apis = BudgetApis { + account: Arc::new(Mutex::new(None)), + catalog_failed: Arc::new(AtomicBool::new(false)), + store_failed: Arc::new(AtomicBool::new(false)), + namespace: settings.accounting_namespace, + }; + for prefix in [ + "/apis/admissionregistration.k8s.io/v1/".to_string(), + format!( + "/api/v1/namespaces/{}/configmaps/kars-inference-budget-contracts", + apis.namespace + ), + format!( + "/apis/kars.azure.com/v1alpha1/namespaces/{}/karsbudgetaccounts/", + apis.namespace + ), + ] { + Mock::given(wiremock::matchers::path_regex(format!( + "^{}", + regex::escape(&prefix) + ))) + .respond_with(apis.clone()) + .with_priority(1) + .mount(server) + .await; + } + apis +} + +fn latest(store: &Arc>) -> KarsTeam { + serde_json::from_value(store.lock().unwrap().team.clone()).unwrap() +} +fn ids(store: &Arc>) -> BTreeMap { + store + .lock() + .unwrap() + .tasks + .iter() + .map(|(name, task)| { + ( + name.clone(), + task["metadata"]["uid"].as_str().unwrap().to_string(), + ) + }) + .collect() +} +async fn cycle(client: &Client, store: &Arc>) { + reconcile( + Arc::new(latest(store)), + Arc::new(Ctx { + client: client.clone(), + }), + ) + .await + .unwrap(); +} + +fn enroll(store: &Arc>, apis: &BudgetApis) { + let mut state = store.lock().unwrap(); + let team: KarsTeam = serde_json::from_value(state.team.clone()).unwrap(); + let principal = specs::principal_name(&team); + let root_uid = state.tasks[&principal]["metadata"]["uid"] + .as_str() + .unwrap() + .to_string(); + let root = RootIdentity { + kind: RootKind::KarsTeam, + resource: ResourceIdentity { + namespace: "tenant-a".into(), + name: team.name_any(), + uid: team.uid().unwrap(), + }, + workspace_uid: "workspace-uid".into(), + cluster_uid: "cluster-uid".into(), + }; + let account_ref = AccountReference { + namespace: apis.namespace.clone(), + name: name_for_root(&root), + uid: "account-uid".into(), + }; + let limits = crate::inference_budget::binding::limits(&team.spec.envelope).unwrap(); + let prior = apis.account.lock().unwrap().clone(); + let mut ledger = prior + .and_then(|account| account.status.and_then(|status| status.ledger)) + .unwrap_or_else(|| Ledger::new(account_ref.uid.clone(), root.clone(), limits).unwrap()); + let mut names: Vec<_> = state.tasks.keys().cloned().collect(); + names.sort_by_key(|name| name != &principal); + for name in names { + let mut task: KarsTask = serde_json::from_value(state.tasks[&name].clone()).unwrap(); + let parent_uid = (name != principal).then(|| root_uid.clone()); + let model = crate::kars_task::blueprint::controller_default_model(); + let authority = TaskAuthority { + task: ResourceIdentity { + namespace: "tenant-a".into(), + name: name.clone(), + uid: task.uid().unwrap(), + }, + parent_uid: parent_uid.clone(), + root_task_uid: root_uid.clone(), + authorization_digest: task.spec.authorization_digest_with_model(&model), + effective_authorization: task.spec.authorization_configuration_with_model(&model), + limits: crate::inference_budget::binding::limits(&task.spec.envelope).unwrap(), + }; + ledger = ledger.register_task(authority.clone()).unwrap().next; + task.status = Some(KarsTaskStatus { + phase: Some("Ready".into()), + observed_generation: task.metadata.generation, + envelope_digest: Some(authority.authorization_digest.clone()), + conditions: Some(vec![ + serde_json::from_value(json!({ + "type":"Ready","status":"True","reason":"Fixture","message":"Fixture", + "lastTransitionTime":"2026-09-08T00:00:00Z" + })) + .unwrap(), + ]), + inference_budget: Some(TaskBudgetBinding { + scope: BudgetScope::GovernedInference, + account: account_ref.clone(), + root: root.clone(), + task_uid: task.uid().unwrap(), + parent_task_uid: parent_uid, + root_task_uid: root_uid.clone(), + authorization_digest: authority.authorization_digest, + }), + ..Default::default() + }); + state + .tasks + .insert(name, serde_json::to_value(task).unwrap()); + } + state.team["status"]["inferenceBudgetAccount"] = json!(account_ref); + let mut account = KarsBudgetAccount::new( + &account_ref.name, + KarsBudgetAccountSpec { + scope: BudgetScope::GovernedInference, + root, + limits, + }, + ); + account.metadata.namespace = Some(apis.namespace.clone()); + account.metadata.uid = Some(account_ref.uid); + account.metadata.resource_version = Some("1".into()); + account.metadata.labels = Some([(MANAGED_BY.into(), OWNER.into())].into()); + account.metadata.annotations = Some([(BOOTSTRAP.into(), "sealed".into())].into()); + account.status = Some(KarsBudgetAccountStatus { + ledger: Some(ledger), + }); + *apis.account.lock().unwrap() = Some(account); +} + +#[tokio::test] +async fn governed_team_admission_interleavings_preserve_uids_funding_and_launch_intent() { + let _lock = ENV_LOCK.lock().await; + let _environment = BudgetEnvironment::enabled(); + let mut team = team(); + team.spec.envelope.budget.as_mut().unwrap().scope = Some(BudgetScope::GovernedInference); + team.spec.envelope.budget.as_mut().unwrap().tokens = Some(30); + team.spec.blueprint = Some(TaskBlueprint { + model: Some(TaskModel { + provider: "azure-openai".into(), + deployment: "reviewed-model".into(), + }), + ..Default::default() + }); + team.spec.roster = vec![TeamRole { + name: "worker".into(), + ..Default::default() + }]; + team.spec.cadence = Some(TeamCadence { + every_minutes: Some(1), + ..Default::default() + }); + let (server, client, store) = setup(&team).await; + let apis = budget_apis(&server).await; + + cycle(&client, &store).await; // Principal created but not enrolled yet. + let initial = ids(&store); + assert_eq!(initial.len(), 2); + cycle(&client, &store).await; + assert_eq!( + ids(&store), + initial, + "Pending enrollment must not delete/recreate seats" + ); + assert!(apis.account.lock().unwrap().is_none()); + enroll(&store, &apis); + for task in store.lock().unwrap().tasks.values_mut() { + task["spec"]["execution"] = json!({"launch":true}); + } + cycle(&client, &store).await; + { + let state = store.lock().unwrap(); + assert_eq!( + state.tasks.len(), + 3, + "Qualified cadence creates exactly one run" + ); + assert!( + state + .tasks + .values() + .all(|task| task["spec"]["execution"]["launch"] == true) + ); + } + enroll(&store, &apis); + let stable_ids = ids(&store); + let run: KarsTask = { + let state = store.lock().unwrap(); + serde_json::from_value( + state + .tasks + .values() + .find(|task| task["metadata"]["annotations"][ANNOT_TEAM_ROLE] == "taskforce") + .unwrap() + .clone(), + ) + .unwrap() + }; + let identity = ExecutionIdentity { + task_uid: run.uid().unwrap(), + authorization_digest: run.envelope_digest(), + sandbox: ResourceIdentity { + namespace: "tenant-a".into(), + name: run.name_any(), + uid: "sandbox-uid".into(), + }, + runtime_namespace_uid: "runtime-uid".into(), + pod_name: "pod".into(), + pod_uid: "pod-uid".into(), + }; + let (_, quote) = contract() + .normalize( + br#"{"model":"reviewed-model","messages":[{"role":"user","content":"fixture"}]}"#, + 100, + true, + ) + .unwrap(); + let request = ReserveRequest { + account_uid: "account-uid".into(), + identity: identity.clone(), + sequence: 1, + wire_digest: format!("sha256:{}", "a".repeat(64)), + quote, + }; + let command = AttemptCommand { + account_uid: "account-uid".into(), + identity, + key: AttemptKey { + pod_uid: "pod-uid".into(), + sequence: 1, + }, + wire_digest: request.wire_digest.clone(), + }; + { + let mut account = apis.account.lock().unwrap(); + let ledger = account + .as_mut() + .unwrap() + .status + .as_mut() + .unwrap() + .ledger + .as_mut() + .unwrap(); + *ledger = ledger + .register_session(request.identity.clone()) + .unwrap() + .next; + *ledger = ledger.reserve(&request, 100).unwrap().next; + *ledger = ledger.begin_dispatch(&command, 101).unwrap().next; + } + store.lock().unwrap().team["status"]["lastRunAt"] = + json!((Utc::now() - chrono::Duration::minutes(2)).to_rfc3339()); + for (catalog_failure, store_failure) in + [(false, false), (true, false), (false, true), (false, false)] + { + apis.catalog_failed.store(catalog_failure, Ordering::SeqCst); + apis.store_failed.store(store_failure, Ordering::SeqCst); + cycle(&client, &store).await; + assert_eq!(ids(&store), stable_ids); + assert!( + store + .lock() + .unwrap() + .tasks + .values() + .all(|task| task["spec"]["execution"]["launch"] == true) + ); + if !store_failure { + assert!(matches!( + crate::inference_budget::launch::admit_new(&client, &run).await, + Err(crate::inference_budget::store::StoreError::Ledger( + crate::inference_budget_contract::BudgetError::Exhausted + )) + )); + } + let account = apis.account.lock().unwrap(); + let ledger = account + .as_ref() + .unwrap() + .status + .as_ref() + .unwrap() + .ledger + .as_ref() + .unwrap(); + assert_eq!(ledger.nodes.len(), 3); + assert_eq!(ledger.meters.reserved.tokens, 30); + assert_eq!(ledger.meters.uncertain.tokens, 0); + assert!(matches!( + crate::inference_budget::launch::capacity(ledger, &run.uid().unwrap()), + Err(crate::inference_budget_contract::BudgetError::Exhausted) + )); + } + { + let mut account = apis.account.lock().unwrap(); + let ledger = account + .as_mut() + .unwrap() + .status + .as_mut() + .unwrap() + .ledger + .as_mut() + .unwrap(); + *ledger = ledger + .settle(&Settlement { + attempt: command, + usage: Some(Usage { + input_tokens: 3, + output_tokens: 5, + cached_input_tokens: 0, + cache_creation_input_tokens: 0, + reasoning_output_tokens: 0, + }), + }) + .unwrap() + .next; + assert_eq!(ledger.meters.settled.tokens, 8); + crate::inference_budget::launch::capacity(ledger, &run.uid().unwrap()).unwrap(); + } + cycle(&client, &store).await; // Admission resumes, not a re-created principal/run. + assert_eq!(store.lock().unwrap().tasks.len(), 4); + for (name, uid) in stable_ids { + assert_eq!(ids(&store)[&name], uid); + } + store.lock().unwrap().team["spec"]["paused"] = json!(true); + let pause_ids = ids(&store); + cycle(&client, &store).await; + assert_eq!(ids(&store), pause_ids); + assert!( + store + .lock() + .unwrap() + .tasks + .values() + .all(|task| task["spec"]["execution"]["launch"] != true) + ); + store.lock().unwrap().team["spec"]["roster"] = json!([]); + cycle(&client, &store).await; + assert!( + !store.lock().unwrap().tasks.contains_key("eng-worker"), + "Actual removed-seat authority must still retire" + ); +} diff --git a/controller/src/kars_team_reconciler/persistence_tests.rs b/controller/src/kars_team_reconciler/persistence_tests.rs index 43e62ff3c..9fe007073 100644 --- a/controller/src/kars_team_reconciler/persistence_tests.rs +++ b/controller/src/kars_team_reconciler/persistence_tests.rs @@ -10,6 +10,9 @@ use serde_json::Value; use std::{collections::BTreeMap, sync::Mutex}; use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; +#[path = "budget_interleaving_tests.rs"] +mod budget_interleavings; + const TASKS_PATH: &str = "/apis/kars.azure.com/v1alpha1/namespaces/tenant-a/karstasks"; const CMS_PATH: &str = "/api/v1/namespaces/tenant-a/configmaps"; const TEAM_STATUS_PATH: &str = @@ -139,7 +142,7 @@ impl Respond for KubeServer { return failure(409); } store.version += 1; - body["metadata"]["uid"] = json!(format!("uid-{name}")); + body["metadata"]["uid"] = json!(format!("uid-{name}-{}", store.version)); body["metadata"]["resourceVersion"] = json!(store.version.to_string()); body["metadata"]["generation"] = json!(1); body["metadata"]["creationTimestamp"] = json!(Utc::now().to_rfc3339()); diff --git a/controller/src/kars_team_reconciler/tasks.rs b/controller/src/kars_team_reconciler/tasks.rs index 4415b63a9..33937e32b 100644 --- a/controller/src/kars_team_reconciler/tasks.rs +++ b/controller/src/kars_team_reconciler/tasks.rs @@ -238,7 +238,7 @@ pub(super) async fn apply_task( spec.execution = Some(execution); } } - if (team.spec.paused || specs::has_positive_budget(&spec.envelope)) + if (team.spec.paused || specs::unsupported_budget(&spec.envelope)) && let Some(execution) = &mut spec.execution { execution.launch = false; diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index d23e96e05..027dfc3bd 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -92,9 +92,20 @@ Only provably **undispatched Reserved** attempts expire/refund (30-second maximu An InFlight attempt never refunds merely because its TTL, router, Pod, or Task disappears. Recovery commits stale InFlight maxima after the transport deadline and a margin. Cancellation closes sessions/subtrees without erasing spend. +Quota held by active reservations, pending enrollment, and catalog/store outages +block **new admissions**, not revoke already funded executions. Team seats/runs +retain their Task UIDs and launch intent during such waits; the Task controller +does not tear down an existing execution solely for a budget-admission failure. +Explicit pause, removed authority, invalid policy and UID replacement retain +their normal fail-closed revocation paths. Every new send still needs the broker. Normal terminal rows compact behind durable Pod sequence high-water marks. Uncertain attempts retain their contract as tombstones: late over-bound usage freezes the account without granting another refund. +Anthropic streams require a valid final usage-bearing `message_delta` with a +recognized stop reason, followed by `message_stop` and clean transport completion. +Start-of-message usage or a terminal event alone never proves final output usage. +Missing, malformed, reset, decreasing or inconsistent final evidence commits the +complete token and maximum-price reservation. Arithmetic is checked, uses signed-Kubernetes-compatible integer ranges, and rounds input/output tariff categories upward separately. An observed provider diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index 7e0196cbf..e07ddd03b 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -98,6 +98,27 @@ That failure is intentional until human review, not a waived or fabricated pass. accessibility, and conservative uncertainty/capacity behavior. 5. Obtain genuine required human audit signoffs before protected publication. +## Bounded reviewer repair candidate + +The source-only independent review identified five blockers. This repair: + +1. Moves integer validation and Pod decoration to their intended module scopes. +2. Restricts legacy launch suppression to unsupported budget scopes. +3. Separates new-work admission from actual revocation; pending/exhausted/API + failures retain Task UIDs and already funded work. Explicit pause, removed + policy/owner authority and UID changes still revoke normally. +4. Requires authoritative final Anthropic stream usage before any refund. +5. Preserves unscoped legacy planning readiness without allowing a pinned account + to escape enforcement. + +New Rust regressions cover full Team reconcile interleavings with real ledger +reserve/begin/settle transitions, UID stability, supported launch intent, +pause/revocation, Task-controller budget waits and legacy parent readiness, plus +16 malformed/incomplete Anthropic protocol variants through the actual stream +and settlement path. **These Rust tests have not run**: the shared Cargo lease +remains with the credential integration owner until the parent directly grants it. +Formatting/source checks do not establish compilation or technical closure. + ## Signoffs - Implementation author: changes under active development; not a signoff. diff --git a/inference-router/src/inference_budget/anthropic_cases.rs b/inference-router/src/inference_budget/anthropic_cases.rs new file mode 100644 index 000000000..aaeed73a3 --- /dev/null +++ b/inference-router/src/inference_budget/anthropic_cases.rs @@ -0,0 +1,65 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +pub fn incomplete() -> Vec<(&'static str, String)> { + let start = r#"data: {"type":"message_start","message":{"usage":{"input_tokens":2,"output_tokens":0}}} + +"#; + let content = r#"data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"generated"}} + +"#; + let stop = "data: {\"type\":\"message_stop\"}\n\n"; + let final_usage = r#"data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":7}} + +"#; + [ + ("missing-final", format!("{start}{content}{stop}")), + ("no-content-missing-final", format!("{start}{stop}")), + ("empty-usage", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{}}}}\n\n{stop}")), + ("missing-usage", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}}}}\n\n{stop}")), + ("malformed-output", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"output_tokens\":\"7\"}}}}\n\n{stop}")), + ("negative-output", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"output_tokens\":-1}}}}\n\n{stop}")), + ("intermediate-only", format!("{start}{content}data: {{\"type\":\"message_delta\",\"usage\":{{\"output_tokens\":7}}}}\n\n{stop}")), + ("changed-input", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"input_tokens\":3,\"output_tokens\":7}}}}\n\n{stop}")), + ("zero-output-with-content", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\"}},\"usage\":{{\"output_tokens\":0}}}}\n\n{stop}")), + ("decreasing-output", format!("{start}{content}data: {{\"type\":\"message_delta\",\"usage\":{{\"output_tokens\":8}}}}\n\n{final_usage}{stop}")), + ("duplicate-final", format!("{start}{content}{final_usage}{final_usage}{stop}")), + ("content-after-final", format!("{start}{final_usage}{content}{stop}")), + ("start-reset", format!("{start}{content}{start}{final_usage}{stop}")), + ("no-terminal", format!("{start}{content}{final_usage}")), + ("final-after-stop", format!("{start}{content}{stop}{final_usage}")), + ("invalid-stop-reason", format!("{start}{content}data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":42}},\"usage\":{{\"output_tokens\":7}}}}\n\n{stop}")), + ].into_iter().collect() +} + +pub fn complete() -> String { + concat!( + "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":0}}}\n\n", + "data: {\"type\":\"content_block_delta\",\"delta\":{\"type\":\"text_delta\",\"text\":\"generated\"}}\n\n", + "data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":7}}\n\n", + "data: {\"type\":\"message_stop\"}\n\n", + ).into() +} + +#[test] +fn refunds_require_a_consistent_final_delta_and_complete_terminal_at_every_chunk_boundary() { + use super::usage::StreamUsage; + use crate::inference_budget_contract::tariffs::Operation; + for (name, wire) in incomplete() { + for size in [1, 2, 7, wire.len()] { + let mut parser = StreamUsage::new(Operation::AnthropicMessages); + for chunk in wire.as_bytes().chunks(size) { + parser.push(chunk); + } + assert!(parser.finish().is_none(), "{name}, chunk size {size}"); + } + } + for size in 1..complete().len() { + let mut parser = StreamUsage::new(Operation::AnthropicMessages); + for chunk in complete().as_bytes().chunks(size) { + parser.push(chunk); + } + let usage = parser.finish().unwrap(); + assert_eq!((usage.input_tokens, usage.output_tokens), (2, 7)); + } +} diff --git a/inference-router/src/inference_budget/client_tests.rs b/inference-router/src/inference_budget/client_tests.rs index 85866b5d7..781675bb8 100644 --- a/inference-router/src/inference_budget/client_tests.rs +++ b/inference-router/src/inference_budget/client_tests.rs @@ -17,6 +17,72 @@ use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; static NEXT: AtomicUsize = AtomicUsize::new(0); +#[tokio::test] +async fn native_stream_settlement_never_refunds_missing_or_inconsistent_final_usage() { + for (name, wire) in super::super::anthropic_cases::incomplete() { + let fixture = Fixture::for_operation(100, false, Operation::AnthropicMessages).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(wire), + ) + .mount(&fixture.provider) + .await; + let (_, _, stream) = crate::proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::new(), + fixture.upstream(), + "v1/messages", + HeaderMap::new(), + bytes::Bytes::from_static( + br#"{"stream":true,"messages":[{"role":"user","content":"text"}]}"#, + ), + ) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + let ledger = fixture.ledger.lock().unwrap(); + assert_eq!(ledger.meters.uncertain.tokens, 30, "{name}"); + assert_eq!(ledger.meters.uncertain.usd_micros, 30, "{name}"); + assert_eq!(ledger.meters.settled.tokens, 0, "{name}"); + assert_eq!(ledger.meters.reserved.tokens, 0, "{name}"); + } +} + +#[tokio::test] +async fn native_stream_with_final_usage_settles_exact_evidence_once() { + let fixture = Fixture::for_operation(100, false, Operation::AnthropicMessages).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/event-stream") + .set_body_string(super::super::anthropic_cases::complete()), + ) + .mount(&fixture.provider) + .await; + let (_, _, stream) = crate::proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::new(), + fixture.upstream(), + "v1/messages", + HeaderMap::new(), + bytes::Bytes::from_static( + br#"{"stream":true,"messages":[{"role":"user","content":"text"}]}"#, + ), + ) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + let ledger = fixture.ledger.lock().unwrap(); + assert_eq!(ledger.meters.settled.tokens, 9); + assert_eq!(ledger.meters.settled.usd_micros, 9); + assert_eq!(ledger.meters.uncertain.tokens, 0); + assert_eq!(ledger.meters.reserved.tokens, 0); +} + #[tokio::test] async fn finite_fence_cannot_be_bypassed_by_disabling_blocklist_or_enabling_learning() { let fixture = Fixture::new(100, false).await; @@ -126,6 +192,10 @@ impl Drop for Fixture { impl Fixture { async fn new(limit: u64, lose_begin_ack: bool) -> Self { + Self::for_operation(limit, lose_begin_ack, Operation::ChatCompletions).await + } + + async fn for_operation(limit: u64, lose_begin_ack: bool, operation: Operation) -> Self { let provider = MockServer::start().await; let broker = MockServer::start().await; let task = ResourceIdentity { @@ -164,13 +234,21 @@ impl Fixture { provider_id: "ollama".into(), endpoint: provider.uri(), model: "model".into(), - operation: Operation::ChatCompletions, + operation, output_field: OutputField::MaxTokens, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, output_bound_includes_reasoning: true, - maximum_price: Some(MaximumPrice::PerRequest { maximum_micros: 5 }), + maximum_price: Some(if operation == Operation::AnthropicMessages { + MaximumPrice::TokenRates { + input_micros_per_million: 1_000_000, + output_micros_per_million: 1_000_000, + fixed_micros: 0, + } + } else { + MaximumPrice::PerRequest { maximum_micros: 5 } + }), }; Mock::given(wiremock::matchers::method("POST")) .respond_with(BrokerFixture { diff --git a/inference-router/src/inference_budget/mod.rs b/inference-router/src/inference_budget/mod.rs index c1cbb0126..03797f662 100644 --- a/inference-router/src/inference_budget/mod.rs +++ b/inference-router/src/inference_budget/mod.rs @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +#[cfg(test)] +mod anthropic_cases; pub mod client; pub mod dispatch; pub mod egress; diff --git a/inference-router/src/inference_budget/usage.rs b/inference-router/src/inference_budget/usage.rs index a1d170da0..9592afa7d 100644 --- a/inference-router/src/inference_budget/usage.rs +++ b/inference-router/src/inference_budget/usage.rs @@ -95,6 +95,8 @@ pub struct StreamUsage { usage: Option, terminal: bool, invalid: bool, + native_final_usage: bool, + native_content: bool, } impl StreamUsage { @@ -105,6 +107,8 @@ impl StreamUsage { usage: None, terminal: false, invalid: false, + native_final_usage: false, + native_content: false, } } @@ -180,39 +184,122 @@ impl StreamUsage { Some("response.failed" | "error") => self.invalid = true, _ => {} }, - Operation::AnthropicMessages => match event.get("type").and_then(Value::as_str) { - Some("message_start") => { - self.usage = event - .pointer("/message/usage") - .and_then(|usage| parse(usage, self.operation)); - if self.usage.is_none() { - self.invalid = true; - } + Operation::AnthropicMessages => self.native_event(event), + } + } + + fn native_event(&mut self, event: &Value) { + match event.get("type").and_then(Value::as_str) { + Some("message_start") => { + if self.usage.is_some() { + self.invalid = true; + return; + } + self.usage = event + .pointer("/message/usage") + .and_then(|usage| parse(usage, self.operation)); + if self.usage.is_none() { + self.invalid = true; } - Some("message_delta") => { - match ( - self.usage.as_mut(), - event - .get("usage") - .and_then(|usage| number(usage, "output_tokens")), - ) { - (Some(usage), Some(output)) if output >= usage.output_tokens => { - usage.output_tokens = output - } - _ => self.invalid = true, + } + Some("content_block_start" | "content_block_delta" | "content_block_stop") => { + if self.usage.is_none() || self.native_final_usage { + self.invalid = true; + return; + } + for object in ["delta", "content_block"] { + if let Some(value) = event.get(object) { + self.native_content |= ["text", "partial_json", "thinking", "data"] + .iter() + .any(|key| { + value + .get(key) + .and_then(Value::as_str) + .is_some_and(|text| !text.is_empty()) + }) + || value.get("type").and_then(Value::as_str) == Some("tool_use"); } } - Some("message_stop") => self.terminal = true, - Some("error") => self.invalid = true, - _ => {} - }, + } + Some("message_delta") => { + let Some(usage) = self.usage.as_mut() else { + self.invalid = true; + return; + }; + let Some(final_usage) = event.get("usage") else { + self.invalid = true; + return; + }; + let Some(output) = number(final_usage, "output_tokens") else { + self.invalid = true; + return; + }; + let fresh = usage + .input_tokens + .checked_sub(usage.cached_input_tokens) + .and_then(|value| value.checked_sub(usage.cache_creation_input_tokens)); + let inconsistent = [ + ("input_tokens", fresh), + ("cache_read_input_tokens", Some(usage.cached_input_tokens)), + ( + "cache_creation_input_tokens", + Some(usage.cache_creation_input_tokens), + ), + ] + .iter() + .any(|(key, expected)| { + final_usage + .get(*key) + .is_some_and(|value| value.as_u64() != *expected) + }); + if self.native_final_usage || output < usage.output_tokens || inconsistent { + self.invalid = true; + return; + } + usage.output_tokens = output; + self.native_final_usage = matches!( + event.pointer("/delta/stop_reason").and_then(Value::as_str), + Some( + "end_turn" + | "max_tokens" + | "stop_sequence" + | "tool_use" + | "pause_turn" + | "refusal" + | "model_context_window_exceeded" + ) + ); + if event + .pointer("/delta/stop_reason") + .is_some_and(|reason| !reason.is_null() && !self.native_final_usage) + { + self.invalid = true; + } + } + Some("message_stop") => { + self.terminal = true; + if !self.native_final_usage + || self + .usage + .as_ref() + .is_none_or(|usage| self.native_content && usage.output_tokens == 0) + { + self.invalid = true; + } + } + Some("ping") => {} + _ => self.invalid = true, } } /// Only a cleanly completed transport with a terminal provider event is /// eligible for a refund. Client disconnect/drop uses None instead. pub fn finish(self) -> Option { - if self.invalid || !self.terminal || !self.buffer.iter().all(u8::is_ascii_whitespace) { + if self.invalid + || !self.terminal + || !self.buffer.iter().all(u8::is_ascii_whitespace) + || (self.operation == Operation::AnthropicMessages && !self.native_final_usage) + { return None; } self.usage @@ -261,7 +348,7 @@ mod tests { fn native_stream_requires_start_usage_and_terminal_stop() { let mut stream = StreamUsage::new(Operation::AnthropicMessages); stream.push(b"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":2,\"output_tokens\":0}}}\n\n"); - stream.push(b"data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":7}}\n\n"); + stream.push(b"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":7}}\n\n"); stream.push(b"data: {\"type\":\"message_stop\"}\n\n"); assert_eq!(stream.finish().unwrap().output_tokens, 7); } diff --git a/shared/inference_budget/ledger_tests.rs b/shared/inference_budget/ledger_tests.rs index 000d97ebf..fda49f730 100644 --- a/shared/inference_budget/ledger_tests.rs +++ b/shared/inference_budget/ledger_tests.rs @@ -7,6 +7,15 @@ use crate::inference_budget_contract::tariffs::{ }; use serde_json::json; +#[test] +fn integer_range_validation_applies_even_when_attempt_map_is_empty() { + let mut ledger = ledger(100); + assert!(ledger.attempts.is_empty()); + ledger.validate().unwrap(); + ledger.limits.tokens = Some(MAX_LEDGER_INTEGER + 1); + assert!(matches!(ledger.validate(), Err(BudgetError::Overflow))); +} + fn resource(name: &str, uid: &str) -> ResourceIdentity { ResourceIdentity { namespace: "workspace".into(), diff --git a/shared/inference_budget/ledger_validation.rs b/shared/inference_budget/ledger_validation.rs index 7ee7db829..c1bba4cbe 100644 --- a/shared/inference_budget/ledger_validation.rs +++ b/shared/inference_budget/ledger_validation.rs @@ -3,6 +3,17 @@ use super::*; +fn signed_integer_range(value: &serde_json::Value) -> bool { + match value { + serde_json::Value::Number(number) => number + .as_u64() + .is_none_or(|number| number <= MAX_LEDGER_INTEGER), + serde_json::Value::Array(values) => values.iter().all(signed_integer_range), + serde_json::Value::Object(values) => values.values().all(signed_integer_range), + _ => true, + } +} + impl Ledger { pub fn validate(&self) -> Result<(), BudgetError> { self.root.validate()?; @@ -117,16 +128,6 @@ impl Ledger { return Err(BudgetError::Corrupt); } - fn signed_integer_range(value: &serde_json::Value) -> bool { - match value { - serde_json::Value::Number(number) => number - .as_u64() - .is_none_or(|number| number <= MAX_LEDGER_INTEGER), - serde_json::Value::Array(values) => values.iter().all(signed_integer_range), - serde_json::Value::Object(values) => values.values().all(signed_integer_range), - _ => true, - } - } // Expired contracts remain accounting evidence; validate their // shape and arithmetic without making old charges disappear. attempt From dd888762c79490d5aff11b4744f98bbea7c5e0b1 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Wed, 9 Sep 2026 13:08:14 +0200 Subject: [PATCH 05/16] fix(budget): qualify compiled contracts and align Helm schemas Local-only qualification checkpoint on aa34c3fe7ba3f73d26015ab2afe38404c071bb15. Retain all previously uncommitted parent repairs: existing workspace hex and schemars dependency edges, kube 3 typed replace_status, ledger borrow/move, concurrent test root lifetime, required legacy inferenceRef fixture, and real text/event-stream MIME for all three wiremock SSE fixtures. The original HEAD was not compile/test qualified; these repairs are part of this checkpoint. Resolve strict Clippy without new allow/expect attributes, feature flags, or new dependencies. Rename Amounts.add to checked_add with identical arithmetic. Keep OutputField operator wire names via explicit serde renames and add exact serde/JsonSchema compatibility coverage. Move unchanged request normalization, shape validation and final route/egress classification into a shared router-only dispatch module, included by the controller only in tests. Keep the public VerifiedTaskNode.generation field for credential composition and explicitly verify that generation in the second live inventory, with regression coverage. Actual Helm/Rust drift tests revealed three pre-existing candidate failures: Task, Team and Profile budget scope nullability/descriptions, Task's untyped status budget binding, Team account reference schema, and CEL rule ordering. Align these templates to the existing generated Rust contract; retain the same CEL rule strings, lifetime opt-in/retention and legacy launch prohibition. No fallback prices, raw-token authority, all-compute scope, or feature expansion. All Cargo commands below used the existing guard, cwd this worktree, and only /Users/pallakatos/Private/Repos/kars/target. The guard sets CARGO_INCREMENTAL=0, CARGO_BUILD_JOBS=2, CARGO_NET_OFFLINE=true and the existing CARGO_TARGET_DIR; it enforces an 8.5 GiB free-space floor. Both packages kept default features. Exact invocation prefix: python3 /Users/pallakatos/.copilot/session-state/a7f32227-c055-4df4-81ee-ce27d7f2d74d/files/run-cargo-guard.py --cwd /Users/pallakatos/.copilot/session-state/a7f32227-c055-4df4-81ee-ce27d7f2d74d/kars-pr11-budgets -- Passing guarded commands: 1. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --lib --bins -- budget task_identity 75 controller + 57 router = 132 passed, zero failures/ignored. Includes all 16 incomplete Anthropic cases, positive final SSE settlement, complete funded-Team interleavings, new generation and wire-schema tests. 2. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --lib --bins -- helm_drift::tests::helm_ kars_task::tests:: kars_task::authorization_tests:: kars_team_reconciler::tests:: kars_team_reconciler::persistence_tests:: routes::model_routing routes::chat_completions routes::anthropic_messages routes::inference routes::spawn_policy forward_proxy:: failover:: blocklist:: 68 controller + 83 router = 151 passed, including all 15 Helm drift tests. Counts describe executions; some selectors overlap command 1. 3. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --test proxy_fake_upstream --test foundry_route_guard --test failover_walk --test multi_provider_guardrails --test anthropic_buffered_guardrail --test chat_output_guardrail_nonjson 24 real loopback HTTP tests passed: proxy 3, Foundry 9, failover 3, multi-provider 5, Anthropic 2, chat non-JSON 2; no external services. 4. cargo clippy --quiet --offline --locked -p kars-controller -p kars-inference-router --all-targets -- -D warnings Passed final all-target check, zero warnings and no new waivers. 5. cargo fmt --package kars-controller --package kars-inference-router 6. cargo fmt --package kars-controller --package kars-inference-router --check Passed. Only affected files required formatting. The compatibility/HTTP test invocations set TMPDIR=$PWD/.qualification-fixtures inside this worktree; that empty scratch directory has been removed. Minimum observed free space across qualification: 10.03 GiB; floor never hit. Passing existing Node/static commands, cwd this worktree unless noted: node cli/node_modules/vitest/vitest.mjs run --root cli --configLoader runner --no-cache src/commands/budget.test.ts 7 passed. Existing symlinked cache only; no dependency install. node cli/node_modules/vitest/vitest.mjs run --root deploy/helm/kars/tests --globals --no-cache src/inference-budget.test.ts src/local-inference.test.ts 18 passed (12 budget + 6 local-inference), rerun after template repair. cd cli && npm run typecheck Passed, no emit. git diff --check git diff --cached --check bash ci/check-copyright-headers.sh 736 source files passed with the new shared module staged. BASE_REF=068ae160 bash ci/check-loc.sh Passed against the reviewed parent. Additional source comparison confirmed normalize, shape validation, Quote accounting, mediated egress and operation classification unchanged modulo whitespace/relocation; Cargo.lock adds only hex and schemars 1.2.1 package edges. Limits: this is source/compile/strict-lint/test qualification, not production acceptance. Real budget/SRE Kind, human review of these repairs, combined credential/budget authority composition, and an independently fully-qualified router image digest remain pending. Privacy epoch remains an additional fence, not Pod UID/custom-audience TokenReview authentication. No public push, merge, main/customer change, Azure, Docker, image build/push or private Bridge actions. cli/node_modules remains untracked and is not part of this checkpoint. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- Cargo.lock | 2 + controller/src/inference_budget/pod_tests.rs | 3 +- controller/src/inference_budget/store.rs | 3 +- .../src/inference_budget/store_tests.rs | 7 +- controller/src/kars_task_budget_tests.rs | 9 +- .../budget_interleaving_tests.rs | 2 +- controller/src/main.rs | 3 + controller/src/task_identity.rs | 1 + controller/src/task_identity_tests.rs | 25 ++ .../helm/kars/templates/crd-karsprofile.yaml | 7 +- deploy/helm/kars/templates/crd-karstask.yaml | 74 +++- deploy/helm/kars/templates/crd-karsteam.yaml | 39 +- inference-router/Cargo.toml | 2 + .../src/inference_budget/client_tests.rs | 28 +- .../src/inference_budget/dispatch.rs | 4 +- inference-router/src/lib.rs | 2 + shared/inference_budget/catalog.rs | 40 +- shared/inference_budget/dispatch.rs | 352 ++++++++++++++++++ shared/inference_budget/ledger.rs | 25 +- shared/inference_budget/ledger_tests.rs | 2 +- shared/inference_budget/ledger_validation.rs | 14 +- shared/inference_budget/tariff_tests.rs | 32 +- shared/inference_budget/tariffs.rs | 326 +--------------- shared/inference_budget/types.rs | 2 +- 24 files changed, 574 insertions(+), 430 deletions(-) create mode 100644 shared/inference_budget/dispatch.rs diff --git a/Cargo.lock b/Cargo.lock index bb46557f5..92620f81e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2548,6 +2548,7 @@ dependencies = [ "ed25519-dalek", "flate2", "futures", + "hex", "hkdf", "jsonwebtoken", "k8s-openapi", @@ -2560,6 +2561,7 @@ dependencies = [ "reqwest 0.12.28", "rustls", "rustls-pemfile", + "schemars 1.2.1", "serde", "serde_json", "serde_yaml", diff --git a/controller/src/inference_budget/pod_tests.rs b/controller/src/inference_budget/pod_tests.rs index b9cab598e..3477cc2c9 100644 --- a/controller/src/inference_budget/pod_tests.rs +++ b/controller/src/inference_budget/pod_tests.rs @@ -100,7 +100,8 @@ async fn no_task_owner_no_reference_is_a_byte_safe_legacy_noop_without_api_acces let client = Client::try_from(config).unwrap(); let sandbox: KarsSandbox = serde_json::from_value(json!({ "apiVersion":"kars.azure.com/v1alpha1", "kind":"KarsSandbox", - "metadata":{"name":"legacy","namespace":"workspace","uid":"legacy-uid"}, "spec":{} + "metadata":{"name":"legacy","namespace":"workspace","uid":"legacy-uid"}, + "spec":{"inferenceRef":{"name":"legacy-inference"}} })) .unwrap(); let namespace: Namespace = serde_json::from_value(json!({ diff --git a/controller/src/inference_budget/store.rs b/controller/src/inference_budget/store.rs index 5cf4c580f..9bff905e6 100644 --- a/controller/src/inference_budget/store.rs +++ b/controller/src/inference_budget/store.rs @@ -286,11 +286,10 @@ impl Store { next.status = Some(KarsBudgetAccountStatus { ledger: Some(mutation.next), }); - let bytes = serde_json::to_vec(&next).map_err(|_| BudgetError::Corrupt)?; let committed = tokio::time::timeout_at( deadline, self.accounts - .replace_status(&name_for_root(root), &PostParams::default(), bytes), + .replace_status(&name_for_root(root), &PostParams::default(), &next), ) .await .map_err(|_| StoreError::Contention)?; diff --git a/controller/src/inference_budget/store_tests.rs b/controller/src/inference_budget/store_tests.rs index cfac368b0..1a9c2c91b 100644 --- a/controller/src/inference_budget/store_tests.rs +++ b/controller/src/inference_budget/store_tests.rs @@ -99,7 +99,7 @@ fn reserve(pod: &str) -> ReserveRequest { endpoint: "https://configured.example".into(), model: "model".into(), operation: Operation::ChatCompletions, - output_field: OutputField::MaxTokens, + output_field: OutputField::Tokens, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, @@ -269,9 +269,10 @@ async fn concurrent_siblings_cannot_both_reserve_past_a_shared_ceiling() { let (_server, store, state) = setup(Some(account()), Fault::None).await; let a = reserve("a"); let b = reserve("b"); + let root = root(); let (first, second) = tokio::join!( - store.transact(&root(), "account-uid", |ledger| ledger.reserve(&a, 100)), - store.transact(&root(), "account-uid", |ledger| ledger.reserve(&b, 100)), + store.transact(&root, "account-uid", |ledger| ledger.reserve(&a, 100)), + store.transact(&root, "account-uid", |ledger| ledger.reserve(&b, 100)), ); assert_eq!(usize::from(first.is_ok()) + usize::from(second.is_ok()), 1); assert_eq!( diff --git a/controller/src/kars_task_budget_tests.rs b/controller/src/kars_task_budget_tests.rs index 522deee2b..ab69b0582 100644 --- a/controller/src/kars_task_budget_tests.rs +++ b/controller/src/kars_task_budget_tests.rs @@ -33,10 +33,11 @@ impl Respond for ApiServer { if request.method == "GET" && path.ends_with("/karstasks/task") { return ResponseTemplate::new(200).set_body_json(&objects.task); } - if request.method == "GET" && path.ends_with("/karstasks/parent") { - if let Some(parent) = &objects.parent { - return ResponseTemplate::new(200).set_body_json(parent); - } + if request.method == "GET" + && path.ends_with("/karstasks/parent") + && let Some(parent) = &objects.parent + { + return ResponseTemplate::new(200).set_body_json(parent); } if request.method == "PATCH" && path.ends_with("/karstasks/task/status") { let patch: Value = serde_json::from_slice(&request.body).unwrap(); diff --git a/controller/src/kars_team_reconciler/budget_interleaving_tests.rs b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs index 6fa7b05d0..0b6f68987 100644 --- a/controller/src/kars_team_reconciler/budget_interleaving_tests.rs +++ b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs @@ -64,7 +64,7 @@ fn contract() -> ModelContract { endpoint: "https://fixture.example".into(), model: "reviewed-model".into(), operation: Operation::ChatCompletions, - output_field: OutputField::MaxTokens, + output_field: OutputField::Tokens, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, diff --git a/controller/src/main.rs b/controller/src/main.rs index a1758413f..425541064 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -39,6 +39,9 @@ mod helm_drift; mod inference_budget; #[path = "../../shared/inference_budget/mod.rs"] mod inference_budget_contract; +#[cfg(test)] +#[path = "../../shared/inference_budget/dispatch.rs"] +mod inference_budget_dispatch; mod inference_policy; mod inference_policy_compile; mod inference_policy_reconciler; diff --git a/controller/src/task_identity.rs b/controller/src/task_identity.rs index 0d8b542f0..30337ed6b 100644 --- a/controller/src/task_identity.rs +++ b/controller/src/task_identity.rs @@ -295,6 +295,7 @@ pub async fn resolve( .await .map_err(|error| api("recheck Task snapshot", error))?; if fresh.metadata.uid.as_deref() != Some(node.pin.task.uid.as_str()) + || fresh.metadata.generation != Some(node.generation) || fresh.metadata.resource_version.as_deref() != Some(node.resource_version.as_str()) || fresh.metadata.deletion_timestamp.is_some() { diff --git a/controller/src/task_identity_tests.rs b/controller/src/task_identity_tests.rs index 2f9dffb08..0b2bde1f0 100644 --- a/controller/src/task_identity_tests.rs +++ b/controller/src/task_identity_tests.rs @@ -163,6 +163,31 @@ async fn double_inventory_rejects_recreated_uid_instead_of_mixing_ancestry() { )); } +#[tokio::test] +async fn double_inventory_rechecks_the_verified_generation() { + let (server, client) = setup().await; + let leaf = task("leaf", None, true); + let snapshot = leaf.clone(); + let calls = AtomicUsize::new(0); + Mock::given(method("GET")) + .and(path( + "/apis/kars.azure.com/v1alpha1/namespaces/workspace/karstasks/leaf", + )) + .respond_with(move |_: &Request| { + let mut task = snapshot.clone(); + if calls.fetch_add(1, Ordering::SeqCst) > 0 { + task.metadata.generation = Some(2); + } + ResponseTemplate::new(200).set_body_json(task) + }) + .mount(&server) + .await; + assert!(matches!( + resolve(&client, &leaf, LeafReadiness::RequireReady).await, + Err(Error::Changed) + )); +} + #[tokio::test] async fn stale_leaf_generation_and_foreign_namespace_are_not_adopted() { let (server, client) = setup().await; diff --git a/deploy/helm/kars/templates/crd-karsprofile.yaml b/deploy/helm/kars/templates/crd-karsprofile.yaml index bd2022b78..08e51e599 100644 --- a/deploy/helm/kars/templates/crd-karsprofile.yaml +++ b/deploy/helm/kars/templates/crd-karsprofile.yaml @@ -59,8 +59,11 @@ spec: properties: scope: type: string - enum: [GovernedInference] - description: "Governed inference only; no all-in task cost claim" + enum: [GovernedInference, null] + nullable: true + description: |- + Explicit opt-in to governed-inference tokens/configured maximum prices. + An absent scope retains the foundation's planning-only interpretation. tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means diff --git a/deploy/helm/kars/templates/crd-karstask.yaml b/deploy/helm/kars/templates/crd-karstask.yaml index ba8276bfb..d9554b1a7 100644 --- a/deploy/helm/kars/templates/crd-karstask.yaml +++ b/deploy/helm/kars/templates/crd-karstask.yaml @@ -183,8 +183,11 @@ spec: properties: scope: type: string - enum: [GovernedInference] - description: "Explicit governed-inference token/configured maximum-price scope; other costs are excluded" + enum: [GovernedInference, null] + nullable: true + description: |- + Explicit opt-in to governed-inference tokens/configured maximum prices. + An absent scope retains the foundation's planning-only interpretation. tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -335,14 +338,6 @@ spec: reason: FieldValueForbidden rule: >- !has(self.execution) || !self.execution.launch || !has(self.envelope.budget) || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') - - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted - reason: FieldValueForbidden - rule: >- - !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) - - message: GovernedInference scope cannot be removed or changed for an existing UID - reason: FieldValueForbidden - rule: >- - !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') - message: spec.blueprint.isolation must be one of standard, enhanced, confidential reason: FieldValueInvalid rule: '!has(self.blueprint) || !has(self.blueprint.isolation) || self.blueprint.isolation in [''standard'',''enhanced'',''confidential'']' @@ -358,14 +353,69 @@ spec: - message: spec.blueprint.egress may list at most 32 destinations reason: FieldValueInvalid rule: '!has(self.blueprint) || !has(self.blueprint.egress) || size(self.blueprint.egress) <= 32' + - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted + reason: FieldValueForbidden + rule: >- + !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) + - message: GovernedInference scope cannot be removed or changed for an existing UID + reason: FieldValueForbidden + rule: >- + !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') status: description: '`KarsTask.status`.' nullable: true properties: inferenceBudget: type: object - x-kubernetes-preserve-unknown-fields: true - description: "Controller-owned immutable budget account and Task UID ancestry binding" + nullable: true + description: |- + Controller-owned immutable account/UID ancestry binding. Not a task-name + metering label and never supplied by an agent header. + properties: + account: + type: object + properties: + name: + type: string + namespace: + type: string + uid: + type: string + required: [name, namespace, uid] + authorizationDigest: + type: string + parentTaskUid: + type: string + nullable: true + root: + type: object + properties: + clusterUid: + type: string + kind: + type: string + enum: [KarsTask, KarsTeam] + resource: + type: object + properties: + name: + type: string + namespace: + type: string + uid: + type: string + required: [name, namespace, uid] + workspaceUid: + type: string + required: [clusterUid, kind, resource, workspaceUid] + rootTaskUid: + type: string + scope: + type: string + enum: [GovernedInference] + taskUid: + type: string + required: [account, authorizationDigest, root, rootTaskUid, scope, taskUid] conditions: description: |- Standard K8s conditions. `Ready` is set `True` once the envelope has diff --git a/deploy/helm/kars/templates/crd-karsteam.yaml b/deploy/helm/kars/templates/crd-karsteam.yaml index 070c88c11..867553928 100644 --- a/deploy/helm/kars/templates/crd-karsteam.yaml +++ b/deploy/helm/kars/templates/crd-karsteam.yaml @@ -220,8 +220,11 @@ spec: properties: scope: type: string - enum: [GovernedInference] - description: "Governed inference over the lifetime of this Team UID; other costs are excluded" + enum: [GovernedInference, null] + nullable: true + description: |- + Explicit opt-in to governed-inference tokens/configured maximum prices. + An absent scope retains the foundation's planning-only interpretation. tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -466,7 +469,11 @@ spec: properties: scope: type: string - enum: [GovernedInference] + enum: [GovernedInference, null] + nullable: true + description: |- + Explicit opt-in to governed-inference tokens/configured maximum prices. + An absent scope retains the foundation's planning-only interpretation. tokens: description: |- Maximum total tokens the task subtree may consume. `0`/absent means @@ -553,14 +560,6 @@ spec: - envelope type: object x-kubernetes-validations: - - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted - reason: FieldValueForbidden - rule: >- - !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) - - message: GovernedInference scope cannot be removed or changed for an existing UID - reason: FieldValueForbidden - rule: >- - !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') - message: spec.charter must be 1-8192 characters (or empty when spec.profileRef is set, to inherit the profile's charter) reason: FieldValueInvalid rule: (has(self.profileRef) && size(self.charter) == 0) || (size(self.charter) > 0 && size(self.charter) <= 8192) @@ -579,17 +578,27 @@ spec: - message: spec.cadence.everyMinutes, when set, must be >= 1 reason: FieldValueInvalid rule: '!has(self.cadence) || !has(self.cadence.everyMinutes) || self.cadence.everyMinutes >= 1' + - message: First finite GovernedInference opt-in requires a new Task/Team UID; an existing unbounded runtime cannot be silently converted + reason: FieldValueForbidden + rule: >- + !has(self.envelope.budget) || !has(self.envelope.budget.scope) || self.envelope.budget.scope != 'GovernedInference' || ((!has(self.envelope.budget.tokens) || self.envelope.budget.tokens == 0) && (!has(self.envelope.budget.usdMicros) || self.envelope.budget.usdMicros == 0)) || (has(oldSelf.envelope.budget) && has(oldSelf.envelope.budget.scope) && oldSelf.envelope.budget.scope == 'GovernedInference' && ((has(oldSelf.envelope.budget.tokens) && oldSelf.envelope.budget.tokens > 0) || (has(oldSelf.envelope.budget.usdMicros) && oldSelf.envelope.budget.usdMicros > 0))) + - message: GovernedInference scope cannot be removed or changed for an existing UID + reason: FieldValueForbidden + rule: >- + !has(oldSelf.envelope.budget) || !has(oldSelf.envelope.budget.scope) || oldSelf.envelope.budget.scope != 'GovernedInference' || (has(self.envelope.budget) && has(self.envelope.budget.scope) && self.envelope.budget.scope == 'GovernedInference') status: description: '`KarsTeam.status` — the controller is the sole writer.' nullable: true properties: inferenceBudgetAccount: type: object - required: [namespace, name, uid] + nullable: true + description: Lifetime Team-UID inference account; cadence runs never reset its balance. + required: [name, namespace, uid] properties: - namespace: {type: string, minLength: 1, maxLength: 63} - name: {type: string, minLength: 1, maxLength: 253} - uid: {type: string, minLength: 1, maxLength: 128} + namespace: {type: string} + name: {type: string} + uid: {type: string} commonsEntryCount: description: Number of entries in the team's knowledge commons (shared memory size). format: int64 diff --git a/inference-router/Cargo.toml b/inference-router/Cargo.toml index 726c9fb87..4bb1ee148 100644 --- a/inference-router/Cargo.toml +++ b/inference-router/Cargo.toml @@ -25,6 +25,7 @@ futures.workspace = true # Serialization serde.workspace = true serde_json.workspace = true +schemars.workspace = true # UTC calendar math for token budget tracker (daily / monthly reset # on UTC boundary). Used by `budget.rs` only. @@ -68,6 +69,7 @@ kars-a2a-core = { path = "../kars-a2a-core", version = "0.1.0" } aes-gcm = "0.10" hkdf = "0.12" sha2 = "0.10" +hex.workspace = true rand = "0.9" base64 = "0.22" flate2 = "1" diff --git a/inference-router/src/inference_budget/client_tests.rs b/inference-router/src/inference_budget/client_tests.rs index 781675bb8..ac6a5238e 100644 --- a/inference-router/src/inference_budget/client_tests.rs +++ b/inference-router/src/inference_budget/client_tests.rs @@ -22,11 +22,7 @@ async fn native_stream_settlement_never_refunds_missing_or_inconsistent_final_us for (name, wire) in super::super::anthropic_cases::incomplete() { let fixture = Fixture::for_operation(100, false, Operation::AnthropicMessages).await; Mock::given(wiremock::matchers::method("POST")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(wire), - ) + .respond_with(ResponseTemplate::new(200).set_body_raw(wire, "text/event-stream")) .mount(&fixture.provider) .await; let (_, _, stream) = crate::proxy::forward_stream( @@ -55,11 +51,10 @@ async fn native_stream_settlement_never_refunds_missing_or_inconsistent_final_us async fn native_stream_with_final_usage_settles_exact_evidence_once() { let fixture = Fixture::for_operation(100, false, Operation::AnthropicMessages).await; Mock::given(wiremock::matchers::method("POST")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("content-type", "text/event-stream") - .set_body_string(super::super::anthropic_cases::complete()), - ) + .respond_with(ResponseTemplate::new(200).set_body_raw( + super::super::anthropic_cases::complete(), + "text/event-stream", + )) .mount(&fixture.provider) .await; let (_, _, stream) = crate::proxy::forward_stream( @@ -235,7 +230,7 @@ impl Fixture { endpoint: provider.uri(), model: "model".into(), operation, - output_field: OutputField::MaxTokens, + output_field: OutputField::Tokens, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, @@ -387,10 +382,13 @@ async fn lost_begin_ack_funds_uncertain_work_but_never_sends_or_regrants_that_at #[tokio::test] async fn actual_stream_eof_requires_terminal_usage_and_charges_output_once() { let fixture = Fixture::new(100, false).await; - Mock::given(wiremock::matchers::method("POST")).respond_with( - ResponseTemplate::new(200).insert_header("content-type", "text/event-stream") - .set_body_string("data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5}}\n\ndata: [DONE]\n\n"), - ).mount(&fixture.provider).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + "data: {\"usage\":{\"prompt_tokens\":3,\"completion_tokens\":5}}\n\ndata: [DONE]\n\n", + "text/event-stream", + )) + .mount(&fixture.provider) + .await; let (_, _, stream) = crate::proxy::forward_stream( Arc::new(WorkloadIdentityAuth::new()), None, diff --git a/inference-router/src/inference_budget/dispatch.rs b/inference-router/src/inference_budget/dispatch.rs index 3bb42242e..014364ee2 100644 --- a/inference-router/src/inference_budget/dispatch.rs +++ b/inference-router/src/inference_budget/dispatch.rs @@ -5,7 +5,7 @@ use super::{ client::{AttemptGuard, Error}, usage, }; -use crate::{inference_budget_contract::catalog, proxy::UpstreamConfig}; +use crate::{inference_budget_dispatch, proxy::UpstreamConfig}; use axum::http::{Method, StatusCode}; use bytes::Bytes; use futures::{StreamExt, stream::BoxStream}; @@ -28,7 +28,7 @@ pub async fn begin( if method != Method::POST { return Err(denied().into()); } - let operation = catalog::operation(path).ok_or_else(denied)?; + let operation = inference_budget_dispatch::operation(path).ok_or_else(denied)?; let (wire, guard) = client .begin( upstream.telemetry_provider(), diff --git a/inference-router/src/lib.rs b/inference-router/src/lib.rs index 437e11781..e01917590 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -39,6 +39,8 @@ pub mod handoff; pub mod inference_budget; #[path = "../../shared/inference_budget/mod.rs"] pub mod inference_budget_contract; +#[path = "../../shared/inference_budget/dispatch.rs"] +mod inference_budget_dispatch; pub mod inference_policy_loader; pub mod mcp; pub mod memory_binding_loader; diff --git a/shared/inference_budget/catalog.rs b/shared/inference_budget/catalog.rs index 84aed096d..4662d0a69 100644 --- a/shared/inference_budget/catalog.rs +++ b/shared/inference_budget/catalog.rs @@ -103,24 +103,6 @@ impl Catalog { } quote.validate(now, money_required) } - - pub fn allows_mediated_non_inference( - &self, - host: &str, - configured_model_hosts: &[String], - ) -> bool { - let host = host.trim_end_matches('.').to_ascii_lowercase(); - self.non_inference_egress_hosts - .iter() - .any(|allowed| allowed == &host) - && !configured_model_hosts - .iter() - .any(|model| model.eq_ignore_ascii_case(&host)) - && !self - .contracts - .iter() - .any(|contract| endpoint_host(&contract.endpoint) == Some(host.as_str())) - } } pub fn endpoint_host(endpoint: &str) -> Option<&str> { @@ -129,30 +111,14 @@ pub fn endpoint_host(endpoint: &str) -> Option<&str> { if authority.is_empty() || authority.contains('@') || authority.starts_with('[') { return None; } - Some(authority.split(':').next()?) -} - -/// Exact, closed final-dispatch path classification. No substring such as -/// "completion" grants access to an unimplemented provider operation. -pub fn operation(path: &str) -> Option { - if path.contains(['?', '#']) { - return None; - } - let path = path.trim_matches('/'); - match path { - "chat/completions" | "v1/chat/completions" | "openai/v1/chat/completions" => { - Some(Operation::ChatCompletions) - } - "messages" | "v1/messages" | "anthropic/v1/messages" => Some(Operation::AnthropicMessages), - "responses" | "v1/responses" | "openai/v1/responses" => Some(Operation::Responses), - _ => None, - } + authority.split(':').next() } #[cfg(test)] mod tests { use super::*; use crate::inference_budget_contract::tariffs::{MaximumPrice, OutputField}; + use crate::inference_budget_dispatch::operation; fn catalog() -> Catalog { Catalog { @@ -165,7 +131,7 @@ mod tests { endpoint: "https://models.example".into(), model: "model".into(), operation: Operation::ChatCompletions, - output_field: OutputField::MaxTokens, + output_field: OutputField::Tokens, maximum_input_tokens: 100, maximum_output_tokens: 50, maximum_wire_bytes: 4096, diff --git a/shared/inference_budget/dispatch.rs b/shared/inference_budget/dispatch.rs new file mode 100644 index 000000000..5d13e161d --- /dev/null +++ b/shared/inference_budget/dispatch.rs @@ -0,0 +1,352 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Router-only request normalization and final-dispatch classification. +//! The controller compiles this module only for shared contract tests. + +use crate::inference_budget_contract::{ + catalog::{Catalog, endpoint_host}, + tariffs::{ModelContract, Operation, OutputField, Quote}, + types::*, +}; +use serde_json::Value; + +impl OutputField { + fn key(self) -> &'static str { + match self { + Self::Tokens => "max_tokens", + Self::Completion => "max_completion_tokens", + Self::Output => "max_output_tokens", + } + } +} + +impl ModelContract { + pub fn normalize( + &self, + body: &[u8], + now: i64, + money_required: bool, + ) -> Result<(Vec, Quote), BudgetError> { + self.validate(now)?; + if body.len() as u64 > self.maximum_wire_bytes + || (money_required && self.maximum_price.is_none()) + { + return Err(BudgetError::Contract); + } + let mut value: Value = serde_json::from_slice(body).map_err(|_| BudgetError::Contract)?; + validate_shape(&value, self.operation)?; + let map = value.as_object_mut().ok_or(BudgetError::Contract)?; + if map + .get("model") + .is_some_and(|model| model.as_str() != Some(self.model.as_str())) + { + return Err(BudgetError::Contract); + } + map.insert("model".into(), self.model.clone().into()); + let key = self.output_field.key(); + for other in ["max_tokens", "max_completion_tokens", "max_output_tokens"] { + if other != key && map.contains_key(other) { + return Err(BudgetError::Contract); + } + } + let output = match map.get(key) { + None => self.maximum_output_tokens, + Some(value) => value + .as_u64() + .filter(|value| *value > 0 && *value <= self.maximum_output_tokens) + .ok_or(BudgetError::Contract)?, + }; + map.insert(key.into(), output.into()); + let bytes = serde_json::to_vec(&value).map_err(|_| BudgetError::Contract)?; + if bytes.len() as u64 > self.maximum_wire_bytes { + return Err(BudgetError::Contract); + } + let maximum = Amounts { + tokens: self + .maximum_input_tokens + .checked_add(output) + .ok_or(BudgetError::Overflow)?, + usd_micros: self + .maximum_price + .as_ref() + .map(|price| price.price(self.maximum_input_tokens, output)) + .transpose()? + .unwrap_or(0), + }; + Ok(( + bytes, + Quote { + contract: self.clone(), + output_tokens: output, + maximum, + price_covered: self.maximum_price.is_some(), + }, + )) + } +} + +fn only_keys(value: &Value, keys: &[&str]) -> bool { + value + .as_object() + .is_some_and(|object| object.keys().all(|key| keys.contains(&key.as_str()))) +} + +fn cache_control(value: &Value) -> bool { + value.get("cache_control").is_none_or(|cache| { + only_keys(cache, &["type", "ttl"]) + && cache.get("type").and_then(Value::as_str) == Some("ephemeral") + && cache + .get("ttl") + .is_none_or(|ttl| matches!(ttl.as_str(), Some("5m" | "1h"))) + }) +} + +fn text_content(value: &Value, anthropic: bool) -> bool { + if value.is_null() || value.is_string() { + return true; + } + value.as_array().is_some_and(|blocks| { + blocks.len() <= 256 + && blocks + .iter() + .all(|block| match block.get("type").and_then(Value::as_str) { + Some("text" | "input_text" | "output_text") => { + block.get("text").is_some_and(Value::is_string) + && only_keys(block, &["type", "text", "cache_control"]) + && cache_control(block) + } + Some("tool_use") if anthropic => { + only_keys(block, &["type", "id", "name", "input", "cache_control"]) + && cache_control(block) + && block.get("name").is_some_and(Value::is_string) + && block.get("id").is_some_and(Value::is_string) + && block.get("input").is_some_and(Value::is_object) + } + Some("tool_result") if anthropic => { + only_keys( + block, + &[ + "type", + "tool_use_id", + "content", + "is_error", + "cache_control", + ], + ) && cache_control(block) + && block.get("tool_use_id").is_some_and(Value::is_string) + && block + .get("content") + .is_some_and(|value| text_content(value, false)) + } + _ => false, + }) + }) +} + +fn messages(value: &Value, anthropic: bool) -> bool { + value.as_array().is_some_and(|items| { + !items.is_empty() + && items.len() <= 256 + && items.iter().all(|item| { + let keys: &[&str] = if anthropic { + &["role", "content"] + } else { + &[ + "role", + "content", + "name", + "tool_calls", + "tool_call_id", + "refusal", + ] + }; + let role = item.get("role").and_then(Value::as_str); + only_keys(item, keys) + && matches!( + role, + Some("system" | "developer" | "user" | "assistant" | "tool") + ) + && item + .get("content") + .is_some_and(|content| text_content(content, anthropic)) + && item.get("tool_calls").is_none_or(|tools| { + tools.as_array().is_some_and(|items| { + items.iter().all(|tool| { + tool.get("type").and_then(Value::as_str) == Some("function") + }) + }) + }) + }) + }) +} + +fn validate_shape(value: &Value, operation: Operation) -> Result<(), BudgetError> { + let map = value.as_object().ok_or(BudgetError::Contract)?; + let allowed: &[&str] = match operation { + Operation::ChatCompletions => &[ + "model", + "messages", + "max_tokens", + "max_completion_tokens", + "stream", + "stream_options", + "temperature", + "top_p", + "frequency_penalty", + "presence_penalty", + "stop", + "seed", + "tools", + "tool_choice", + "parallel_tool_calls", + "response_format", + "user", + "metadata", + "n", + "reasoning_effort", + "store", + ], + Operation::AnthropicMessages => &[ + "model", + "messages", + "system", + "max_tokens", + "stream", + "temperature", + "top_p", + "top_k", + "stop_sequences", + "tools", + "tool_choice", + "thinking", + "metadata", + ], + Operation::Responses => &[ + "model", + "input", + "instructions", + "max_output_tokens", + "stream", + "tools", + "tool_choice", + "parallel_tool_calls", + "text", + "reasoning", + "store", + "background", + "metadata", + "user", + ], + }; + if map.keys().any(|key| !allowed.contains(&key.as_str())) + || map.get("n").is_some_and(|value| value.as_u64() != Some(1)) + || map + .get("background") + .is_some_and(|value| value.as_bool() != Some(false)) + || map + .get("store") + .is_some_and(|value| value.as_bool() != Some(false)) + || map.get("stream").is_some_and(|value| !value.is_boolean()) + { + return Err(BudgetError::Contract); + } + let valid_input = match operation { + Operation::ChatCompletions => map + .get("messages") + .is_some_and(|value| messages(value, false)), + Operation::AnthropicMessages => { + map.get("messages") + .is_some_and(|value| messages(value, true)) + && map + .get("system") + .is_none_or(|value| text_content(value, false)) + } + Operation::Responses => map.get("input").is_some_and(|input| { + input.is_string() + || input.as_array().is_some_and(|items| { + items.len() <= 256 + && items + .iter() + .all(|item| match item.get("type").and_then(Value::as_str) { + Some("function_call") => { + only_keys( + item, + &["type", "id", "call_id", "name", "arguments", "status"], + ) && item.get("name").is_some_and(Value::is_string) + && item.get("arguments").is_some_and(Value::is_string) + } + Some("function_call_output") => { + only_keys(item, &["type", "id", "call_id", "output", "status"]) + && item.get("call_id").is_some_and(Value::is_string) + && item.get("output").is_some_and(Value::is_string) + } + None | Some("message") => { + only_keys(item, &["type", "role", "content", "id", "status"]) + && item + .get("content") + .is_some_and(|content| text_content(content, false)) + } + _ => false, + }) + }) + }), + }; + if !valid_input { + return Err(BudgetError::Contract); + } + if let Some(tools) = map.get("tools") { + let tools = tools.as_array().ok_or(BudgetError::Contract)?; + if tools.len() > 128 + || tools.iter().any(|tool| { + if operation == Operation::AnthropicMessages { + tool.get("type").is_some() + || !tool.get("name").is_some_and(Value::is_string) + || !tool.get("input_schema").is_some_and(Value::is_object) + } else { + tool.get("type").and_then(Value::as_str) != Some("function") + } + }) + { + return Err(BudgetError::Contract); + } + } + Ok(()) +} + +impl Catalog { + pub fn allows_mediated_non_inference( + &self, + host: &str, + configured_model_hosts: &[String], + ) -> bool { + let host = host.trim_end_matches('.').to_ascii_lowercase(); + self.non_inference_egress_hosts + .iter() + .any(|allowed| allowed == &host) + && !configured_model_hosts + .iter() + .any(|model| model.eq_ignore_ascii_case(&host)) + && !self + .contracts + .iter() + .any(|contract| endpoint_host(&contract.endpoint) == Some(host.as_str())) + } +} + +/// Exact, closed final-dispatch path classification. No substring such as +/// "completion" grants access to an unimplemented provider operation. +pub(crate) fn operation(path: &str) -> Option { + if path.contains(['?', '#']) { + return None; + } + let path = path.trim_matches('/'); + match path { + "chat/completions" | "v1/chat/completions" | "openai/v1/chat/completions" => { + Some(Operation::ChatCompletions) + } + "messages" | "v1/messages" | "anthropic/v1/messages" => Some(Operation::AnthropicMessages), + "responses" | "v1/responses" | "openai/v1/responses" => Some(Operation::Responses), + _ => None, + } +} diff --git a/shared/inference_budget/ledger.rs b/shared/inference_budget/ledger.rs index bb3f69695..04a517b4d 100644 --- a/shared/inference_budget/ledger.rs +++ b/shared/inference_budget/ledger.rs @@ -25,7 +25,9 @@ pub struct Meters { impl Meters { pub fn total(&self) -> Result { - self.reserved.add(self.settled)?.add(self.uncertain) + self.reserved + .checked_add(self.settled)? + .checked_add(self.uncertain) } } @@ -287,10 +289,11 @@ impl Ledger { return Err(BudgetError::Contract); } let mut next = self.clone(); - next.nodes + let node = next + .nodes .get_mut(&authority.task.uid) - .ok_or(BudgetError::Corrupt)? - .authority = authority; + .ok_or(BudgetError::Corrupt)?; + node.authority = authority; self.mutation(next, ()) } @@ -380,12 +383,14 @@ impl Ledger { .quote .validate(now, self.requires_price(&request.identity.task_uid)?)?; let maximum = request.quote.maximum; - if !self.limits.allows(self.meters.total()?.add(maximum)?) + if !self + .limits + .allows(self.meters.total()?.checked_add(maximum)?) || path.iter().any(|uid| { let node = &self.nodes[uid]; node.meters .total() - .and_then(|total| total.add(maximum)) + .and_then(|total| total.checked_add(maximum)) .map_or(true, |total| !node.authority.limits.allows(total)) }) { @@ -400,10 +405,10 @@ impl Ledger { .timestamp(), ); let mut next = self.clone(); - next.meters.reserved = next.meters.reserved.add(maximum)?; + next.meters.reserved = next.meters.reserved.checked_add(maximum)?; for uid in path { let node = next.nodes.get_mut(&uid).ok_or(BudgetError::Corrupt)?; - node.meters.reserved = node.meters.reserved.add(maximum)?; + node.meters.reserved = node.meters.reserved.checked_add(maximum)?; } next.sessions .get_mut(&key.pod_uid) @@ -510,9 +515,9 @@ impl Ledger { let apply = |meters: &mut Meters| -> Result<(), BudgetError> { meters.reserved = meters.reserved.subtract(attempt.quote.maximum)?; if phase == AttemptPhase::Uncertain { - meters.uncertain = meters.uncertain.add(charged)?; + meters.uncertain = meters.uncertain.checked_add(charged)?; } else { - meters.settled = meters.settled.add(charged)?; + meters.settled = meters.settled.checked_add(charged)?; } if phase != AttemptPhase::Expired && !attempt.quote.price_covered { meters.unpriced_attempts = meters diff --git a/shared/inference_budget/ledger_tests.rs b/shared/inference_budget/ledger_tests.rs index fda49f730..39577d526 100644 --- a/shared/inference_budget/ledger_tests.rs +++ b/shared/inference_budget/ledger_tests.rs @@ -84,7 +84,7 @@ fn request(task: &str, pod: &str, sequence: u64) -> ReserveRequest { endpoint: "https://provider.example".into(), model: "model".into(), operation: Operation::ChatCompletions, - output_field: OutputField::MaxTokens, + output_field: OutputField::Tokens, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, diff --git a/shared/inference_budget/ledger_validation.rs b/shared/inference_budget/ledger_validation.rs index c1bba4cbe..674b4393e 100644 --- a/shared/inference_budget/ledger_validation.rs +++ b/shared/inference_budget/ledger_validation.rs @@ -66,24 +66,24 @@ impl Ledger { if *root_uid != node.authority.root_task_uid { return Err(BudgetError::Corrupt); } - let used = node.meters.settled.add(node.meters.uncertain)?; + let used = node.meters.settled.checked_add(node.meters.uncertain)?; if let Some(parent) = &node.authority.parent_uid { let prior = children_used.get(parent).copied().unwrap_or_default(); - children_used.insert(parent.clone(), prior.add(used)?); + children_used.insert(parent.clone(), prior.checked_add(used)?); } else { if self.root.kind == RootKind::KarsTask && node.authority.task != self.root.resource { return Err(BudgetError::Corrupt); } - root_nodes_used = root_nodes_used.add(used)?; + root_nodes_used = root_nodes_used.checked_add(used)?; } } - if !root_nodes_used.within(self.meters.settled.add(self.meters.uncertain)?) { + if !root_nodes_used.within(self.meters.settled.checked_add(self.meters.uncertain)?) { return Err(BudgetError::Corrupt); } for (uid, used) in children_used { let node = self.nodes.get(&uid).ok_or(BudgetError::Corrupt)?; - if !used.within(node.meters.settled.add(node.meters.uncertain)?) { + if !used.within(node.meters.settled.checked_add(node.meters.uncertain)?) { return Err(BudgetError::Corrupt); } } @@ -137,10 +137,10 @@ impl Ledger { if attempt.charged != Amounts::default() { return Err(BudgetError::Corrupt); } - root_reserved = root_reserved.add(attempt.quote.maximum)?; + root_reserved = root_reserved.checked_add(attempt.quote.maximum)?; for uid in self.ancestors(&attempt.identity.task_uid)? { let previous = node_reserved.get(&uid).copied().unwrap_or_default(); - node_reserved.insert(uid, previous.add(attempt.quote.maximum)?); + node_reserved.insert(uid, previous.checked_add(attempt.quote.maximum)?); } } } diff --git a/shared/inference_budget/tariff_tests.rs b/shared/inference_budget/tariff_tests.rs index b0aefdcdb..4de4748fa 100644 --- a/shared/inference_budget/tariff_tests.rs +++ b/shared/inference_budget/tariff_tests.rs @@ -2,7 +2,7 @@ // Licensed under the MIT License. use super::*; -use serde_json::json; +use serde_json::{Value, json}; pub(super) fn contract() -> ModelContract { ModelContract { @@ -13,7 +13,7 @@ pub(super) fn contract() -> ModelContract { endpoint: "https://operator.example/inference".into(), model: "model-revision-1".into(), operation: Operation::ChatCompletions, - output_field: OutputField::MaxCompletionTokens, + output_field: OutputField::Completion, maximum_input_tokens: 10, maximum_output_tokens: 20, maximum_wire_bytes: 4096, @@ -33,6 +33,30 @@ fn request() -> Vec { .unwrap() } +#[test] +fn output_field_wire_names_and_schema_remain_operator_compatible() { + let variants = [ + (OutputField::Tokens, "MaxTokens"), + (OutputField::Completion, "MaxCompletionTokens"), + (OutputField::Output, "MaxOutputTokens"), + ]; + for (variant, wire) in variants { + assert_eq!(serde_json::to_value(variant).unwrap(), json!(wire)); + assert_eq!( + serde_json::from_value::(json!(wire)).unwrap(), + variant + ); + } + let schema = serde_json::to_value(schemars::schema_for!(OutputField)).unwrap(); + assert_eq!( + schema["enum"], + json!(["MaxTokens", "MaxCompletionTokens", "MaxOutputTokens"]) + ); + for internal_name in ["Tokens", "Completion", "Output"] { + assert!(serde_json::from_value::(json!(internal_name)).is_err()); + } +} + #[test] fn missing_maximum_is_injected_and_input_uses_operator_bound_not_character_estimate() { let (wire, quote) = contract().normalize(&request(), 100, true).unwrap(); @@ -133,7 +157,7 @@ fn client_function_tools_remain_available_without_enabling_hosted_generation() { fn native_and_responses_shapes_have_explicit_distinct_output_contracts() { let mut contract = contract(); contract.operation = Operation::AnthropicMessages; - contract.output_field = OutputField::MaxTokens; + contract.output_field = OutputField::Tokens; let body = json!({"model": contract.model, "messages": [{"role": "user", "content": "hello"}], "tools": [{"name": "read_file", "input_schema": {"type": "object"}}]}); let (wire, _) = contract @@ -144,7 +168,7 @@ fn native_and_responses_shapes_have_explicit_distinct_output_contracts() { 20 ); contract.operation = Operation::Responses; - contract.output_field = OutputField::MaxOutputTokens; + contract.output_field = OutputField::Output; let body = json!({"model": contract.model, "input": "hello", "store": false, "background": false}); let (wire, _) = contract diff --git a/shared/inference_budget/tariffs.rs b/shared/inference_budget/tariffs.rs index 7c06303cf..7ec71fe18 100644 --- a/shared/inference_budget/tariffs.rs +++ b/shared/inference_budget/tariffs.rs @@ -6,7 +6,6 @@ use super::types::*; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; -use serde_json::Value; #[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] pub enum Operation { @@ -17,19 +16,12 @@ pub enum Operation { #[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] pub enum OutputField { - MaxTokens, - MaxCompletionTokens, - MaxOutputTokens, -} - -impl OutputField { - pub fn key(self) -> &'static str { - match self { - Self::MaxTokens => "max_tokens", - Self::MaxCompletionTokens => "max_completion_tokens", - Self::MaxOutputTokens => "max_output_tokens", - } - } + #[serde(rename = "MaxTokens")] + Tokens, + #[serde(rename = "MaxCompletionTokens")] + Completion, + #[serde(rename = "MaxOutputTokens")] + Output, } #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] @@ -141,9 +133,9 @@ impl ModelContract { (self.operation, self.output_field), ( Operation::ChatCompletions, - OutputField::MaxTokens | OutputField::MaxCompletionTokens - ) | (Operation::AnthropicMessages, OutputField::MaxTokens) - | (Operation::Responses, OutputField::MaxOutputTokens) + OutputField::Tokens | OutputField::Completion + ) | (Operation::AnthropicMessages, OutputField::Tokens) + | (Operation::Responses, OutputField::Output) ); if !valid_name(&self.id) || !valid_uid(&self.version) @@ -175,77 +167,13 @@ impl ModelContract { if tokens > MAX_LEDGER_INTEGER { return Err(BudgetError::Overflow); } - if let Some(price) = &self.maximum_price { - if price.price(self.maximum_input_tokens, self.maximum_output_tokens)? + if let Some(price) = &self.maximum_price + && price.price(self.maximum_input_tokens, self.maximum_output_tokens)? > MAX_LEDGER_INTEGER - { - return Err(BudgetError::Overflow); - } - } - Ok(()) - } - - pub fn normalize( - &self, - body: &[u8], - now: i64, - money_required: bool, - ) -> Result<(Vec, Quote), BudgetError> { - self.validate(now)?; - if body.len() as u64 > self.maximum_wire_bytes - || (money_required && self.maximum_price.is_none()) - { - return Err(BudgetError::Contract); - } - let mut value: Value = serde_json::from_slice(body).map_err(|_| BudgetError::Contract)?; - validate_shape(&value, self.operation)?; - let map = value.as_object_mut().ok_or(BudgetError::Contract)?; - if map - .get("model") - .is_some_and(|model| model.as_str() != Some(self.model.as_str())) { - return Err(BudgetError::Contract); - } - map.insert("model".into(), self.model.clone().into()); - let key = self.output_field.key(); - for other in ["max_tokens", "max_completion_tokens", "max_output_tokens"] { - if other != key && map.contains_key(other) { - return Err(BudgetError::Contract); - } - } - let output = match map.get(key) { - None => self.maximum_output_tokens, - Some(value) => value - .as_u64() - .filter(|value| *value > 0 && *value <= self.maximum_output_tokens) - .ok_or(BudgetError::Contract)?, - }; - map.insert(key.into(), output.into()); - let bytes = serde_json::to_vec(&value).map_err(|_| BudgetError::Contract)?; - if bytes.len() as u64 > self.maximum_wire_bytes { - return Err(BudgetError::Contract); + return Err(BudgetError::Overflow); } - let maximum = Amounts { - tokens: self - .maximum_input_tokens - .checked_add(output) - .ok_or(BudgetError::Overflow)?, - usd_micros: self - .maximum_price - .as_ref() - .map(|price| price.price(self.maximum_input_tokens, output)) - .transpose()? - .unwrap_or(0), - }; - Ok(( - bytes, - Quote { - contract: self.clone(), - output_tokens: output, - maximum, - price_covered: self.maximum_price.is_some(), - }, - )) + Ok(()) } } @@ -315,234 +243,6 @@ impl Quote { } } -fn only_keys(value: &Value, keys: &[&str]) -> bool { - value - .as_object() - .is_some_and(|object| object.keys().all(|key| keys.contains(&key.as_str()))) -} - -fn cache_control(value: &Value) -> bool { - value.get("cache_control").is_none_or(|cache| { - only_keys(cache, &["type", "ttl"]) - && cache.get("type").and_then(Value::as_str) == Some("ephemeral") - && cache - .get("ttl") - .is_none_or(|ttl| matches!(ttl.as_str(), Some("5m" | "1h"))) - }) -} - -fn text_content(value: &Value, anthropic: bool) -> bool { - if value.is_null() || value.is_string() { - return true; - } - value.as_array().is_some_and(|blocks| { - blocks.len() <= 256 - && blocks - .iter() - .all(|block| match block.get("type").and_then(Value::as_str) { - Some("text" | "input_text" | "output_text") => { - block.get("text").is_some_and(Value::is_string) - && only_keys(block, &["type", "text", "cache_control"]) - && cache_control(block) - } - Some("tool_use") if anthropic => { - only_keys(block, &["type", "id", "name", "input", "cache_control"]) - && cache_control(block) - && block.get("name").is_some_and(Value::is_string) - && block.get("id").is_some_and(Value::is_string) - && block.get("input").is_some_and(Value::is_object) - } - Some("tool_result") if anthropic => { - only_keys( - block, - &[ - "type", - "tool_use_id", - "content", - "is_error", - "cache_control", - ], - ) && cache_control(block) - && block.get("tool_use_id").is_some_and(Value::is_string) - && block - .get("content") - .is_some_and(|value| text_content(value, false)) - } - _ => false, - }) - }) -} - -fn messages(value: &Value, anthropic: bool) -> bool { - value.as_array().is_some_and(|items| { - !items.is_empty() - && items.len() <= 256 - && items.iter().all(|item| { - let keys: &[&str] = if anthropic { - &["role", "content"] - } else { - &[ - "role", - "content", - "name", - "tool_calls", - "tool_call_id", - "refusal", - ] - }; - let role = item.get("role").and_then(Value::as_str); - only_keys(item, keys) - && matches!( - role, - Some("system" | "developer" | "user" | "assistant" | "tool") - ) - && item - .get("content") - .is_some_and(|content| text_content(content, anthropic)) - && item.get("tool_calls").is_none_or(|tools| { - tools.as_array().is_some_and(|items| { - items.iter().all(|tool| { - tool.get("type").and_then(Value::as_str) == Some("function") - }) - }) - }) - }) - }) -} - -fn validate_shape(value: &Value, operation: Operation) -> Result<(), BudgetError> { - let map = value.as_object().ok_or(BudgetError::Contract)?; - let allowed: &[&str] = match operation { - Operation::ChatCompletions => &[ - "model", - "messages", - "max_tokens", - "max_completion_tokens", - "stream", - "stream_options", - "temperature", - "top_p", - "frequency_penalty", - "presence_penalty", - "stop", - "seed", - "tools", - "tool_choice", - "parallel_tool_calls", - "response_format", - "user", - "metadata", - "n", - "reasoning_effort", - "store", - ], - Operation::AnthropicMessages => &[ - "model", - "messages", - "system", - "max_tokens", - "stream", - "temperature", - "top_p", - "top_k", - "stop_sequences", - "tools", - "tool_choice", - "thinking", - "metadata", - ], - Operation::Responses => &[ - "model", - "input", - "instructions", - "max_output_tokens", - "stream", - "tools", - "tool_choice", - "parallel_tool_calls", - "text", - "reasoning", - "store", - "background", - "metadata", - "user", - ], - }; - if map.keys().any(|key| !allowed.contains(&key.as_str())) - || map.get("n").is_some_and(|value| value.as_u64() != Some(1)) - || map - .get("background") - .is_some_and(|value| value.as_bool() != Some(false)) - || map - .get("store") - .is_some_and(|value| value.as_bool() != Some(false)) - || map.get("stream").is_some_and(|value| !value.is_boolean()) - { - return Err(BudgetError::Contract); - } - let valid_input = match operation { - Operation::ChatCompletions => map - .get("messages") - .is_some_and(|value| messages(value, false)), - Operation::AnthropicMessages => { - map.get("messages") - .is_some_and(|value| messages(value, true)) - && map - .get("system") - .is_none_or(|value| text_content(value, false)) - } - Operation::Responses => map.get("input").is_some_and(|input| { - input.is_string() - || input.as_array().is_some_and(|items| { - items.len() <= 256 - && items - .iter() - .all(|item| match item.get("type").and_then(Value::as_str) { - Some("function_call") => { - only_keys( - item, - &["type", "id", "call_id", "name", "arguments", "status"], - ) && item.get("name").is_some_and(Value::is_string) - && item.get("arguments").is_some_and(Value::is_string) - } - Some("function_call_output") => { - only_keys(item, &["type", "id", "call_id", "output", "status"]) - && item.get("call_id").is_some_and(Value::is_string) - && item.get("output").is_some_and(Value::is_string) - } - None | Some("message") => { - only_keys(item, &["type", "role", "content", "id", "status"]) - && item - .get("content") - .is_some_and(|content| text_content(content, false)) - } - _ => false, - }) - }) - }), - }; - if !valid_input { - return Err(BudgetError::Contract); - } - if let Some(tools) = map.get("tools") { - let tools = tools.as_array().ok_or(BudgetError::Contract)?; - if tools.len() > 128 - || tools.iter().any(|tool| { - if operation == Operation::AnthropicMessages { - tool.get("type").is_some() - || !tool.get("name").is_some_and(Value::is_string) - || !tool.get("input_schema").is_some_and(Value::is_object) - } else { - tool.get("type").and_then(Value::as_str) != Some("function") - } - }) - { - return Err(BudgetError::Contract); - } - } - Ok(()) -} - #[cfg(test)] #[path = "tariff_tests.rs"] mod tests; diff --git a/shared/inference_budget/types.rs b/shared/inference_budget/types.rs index efe7b139f..4e1ca8746 100644 --- a/shared/inference_budget/types.rs +++ b/shared/inference_budget/types.rs @@ -112,7 +112,7 @@ pub struct Amounts { } impl Amounts { - pub fn add(self, other: Self) -> Result { + pub fn checked_add(self, other: Self) -> Result { let amount = Self { tokens: self .tokens From 45b9df8d3c8557290fa76a2c41475a1ecc6eff0e Mon Sep 17 00:00:00 2001 From: pallakatos Date: Wed, 9 Sep 2026 15:23:45 +0200 Subject: [PATCH 06/16] test(authority): initialize TLS explicitly for isolated identity cases Hosted nextest runs each case independently; the pending-leaf fixture constructed a Kubernetes client before selecting the existing AWS-LC provider. Initialize it in the shared test setup so the real authority assertions execute without depending on another test process. No production authority or TLS policy changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/task_identity_tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/controller/src/task_identity_tests.rs b/controller/src/task_identity_tests.rs index 0b2bde1f0..848f45a01 100644 --- a/controller/src/task_identity_tests.rs +++ b/controller/src/task_identity_tests.rs @@ -36,6 +36,7 @@ fn task(name: &str, parent: Option<&str>, ready: bool) -> KarsTask { } async fn setup() -> (MockServer, Client) { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/api/v1/namespaces/workspace")) From abbfc7d19213eb00d09bee9b8f84841fdf82a3d0 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Wed, 9 Sep 2026 17:01:07 +0200 Subject: [PATCH 07/16] fix(budget): report durable account state and standard conditions Bounded local source checkpoint on 2b2e0b8018b5059b05467ac76d3aa6115575be78. Addresses the actual #553 hosted CNCF failures c3_conditions_array_present and c12_state_printer_column without changing conformance checks or scope. Add observational status.phase, observedGeneration, and standard Kubernetes Ready/LedgerValid conditions to both Rust and Helm. The Phase printer points to .status.phase; the existing .status.ledger.phase and ledger wire remain unchanged. Conditions use the existing project helper and standard Condition and jiff-backed Time types, list-map keys by type, and stable transition times. Populate status after UID-anchor creation, during ledger initialization, after sealing, with every changed ledger transaction, and in periodic recovery/backfill. Derive bootstrap, active, temporarily reserved, exhausted, revoked, capacity-blocked, frozen, closing/retired, invalid-ledger and API- unavailable observations from actual account state. No blanket Ready signal, router-health assertion, fallback pricing, private errors, or spend authority. Reporting reloads a fresh account and uses bounded UID/RV whole-status PUT. It preserves the exact ledger, including corrupt or missing data; it cannot initialize/reset money. Initialization also uses complete status PUT. Lost acknowledgements still fail; unchanged ledger mutations retain their original no-write/idempotent return. A failed reporting read/write is not successful publication; an unreachable API necessarily retains its last observation. No funded Task deletion, authority reset, or changed account limit/schema. Add 11 Rust regressions covering real store transitions and fault injection: - reserved headroom versus durable exhaustion, preserving funded dispatch; - expiry/headroom recovery and retained replay fences; - revoked/retired authority and accepted-work liability; - provider breach plus forged observational Ready not authorizing spending; - live recovery API failure, stable Unknown, and subsequent root retirement; - corrupt/missing ledger reporting without repair or reinitialization; - failed reads/writes, UID replacement, conflicts and lost reporting acks; - legacy observation backfill and stable timestamp across generations; - Helm/generated-Rust reporting schema and printer-column equality; - report CAS conflict with concurrent reservation, preserving new money; - lost bootstrap ledger-status acknowledgement without reset before sealing. Extend the existing Helm tests for new reporting fields in enabled/default shapes. Update the directly related operator documentation. FAST GATES ACTUALLY RUN AND PASSED, cwd this worktree: rustfmt --check --edition 2024 controller/src/inference_budget/store.rs Includes nested store/status test modules; syntax/format only, not compilation. rustfmt --check --edition 2024 --config skip_children=true controller/src/inference_budget/account.rs controller/src/inference_budget/status.rs controller/src/inference_budget/status_tests.rs controller/src/inference_budget/store.rs controller/src/inference_budget/store_tests.rs controller/src/inference_budget/recovery.rs controller/src/inference_budget/mod.rs controller/src/kars_team_reconciler/budget_interleaving_tests.rs node cli/node_modules/vitest/vitest.mjs run --root deploy/helm/kars/tests --globals --no-cache src/inference-budget.test.ts src/local-inference.test.ts 20 passed: 14 budget and 6 local-inference. node cli/node_modules/vitest/vitest.mjs run --root cli --configLoader runner --no-cache src/commands/budget.test.ts 7 passed. cd cli && npm run typecheck Passed, no emit. git diff --check git diff --cached --check bash ci/check-copyright-headers.sh 739 tracked source files passed with new files staged. BASE_REF=2b2e0b8018b5059b05467ac76d3aa6115575be78 bash ci/check-loc.sh Passed; every touched Rust module is <=800 lines; no new whitelist/waiver. Parsed-YAML comparison proved the prior financial spec, immutable-grant CEL and status.ledger schema exactly unchanged. No Cargo manifest or lock edits. NOT YET RUST/CNCF QUALIFIED: no Cargo grant was held or used for this checkpoint. The root target remains exclusively leased to the credential RPC owner. After explicit parent grant, run the existing guarded offline/locked commands with existing root target, CARGO_INCREMENTAL=0 and CARGO_BUILD_JOBS=2: cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --lib --bins -- budget task_identity helm_drift cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router -p kars-cncf-conformance --test criteria The latter is the actual 17-case CNCF criteria/report runner, not a substitute. cargo clippy --quiet --offline --locked -p kars-controller -p kars-inference-router --all-targets -- -D warnings cargo fmt --package kars-controller --package kars-inference-router --check No public push, parent review claim, merge, customer/main/private Bridge edit, Azure/cloud action, Docker/image action, install or new target. Pending: actual Rust compile/tests, generated schema drift, CNCF17, strict Clippy, parent review, and independent existing SRE Kind/cross-credential production qualification. cli/node_modules remains an unstaged existing symlink. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/inference_budget/account.rs | 35 +- controller/src/inference_budget/mod.rs | 1 + controller/src/inference_budget/recovery.rs | 27 +- controller/src/inference_budget/status.rs | 172 ++++++ .../src/inference_budget/status_tests.rs | 514 ++++++++++++++++++ controller/src/inference_budget/store.rs | 122 ++++- .../src/inference_budget/store_tests.rs | 37 ++ .../budget_interleaving_tests.rs | 1 + .../kars/templates/crd-karsbudgetaccount.yaml | 29 +- .../kars/tests/src/inference-budget.test.ts | 24 + docs/governed-inference-budgets.md | 30 + 11 files changed, 960 insertions(+), 32 deletions(-) create mode 100644 controller/src/inference_budget/status.rs create mode 100644 controller/src/inference_budget/status_tests.rs diff --git a/controller/src/inference_budget/account.rs b/controller/src/inference_budget/account.rs index 462578755..9ca9974de 100644 --- a/controller/src/inference_budget/account.rs +++ b/controller/src/inference_budget/account.rs @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +use k8s_openapi::apimachinery::pkg::apis::meta::v1::Condition; use kube::CustomResource; use schemars::JsonSchema; use serde::{Deserialize, Serialize}; @@ -16,7 +17,7 @@ use crate::inference_budget_contract::{BudgetScope, Limits, RootIdentity, ledger status = "KarsBudgetAccountStatus", shortname = "kbudget", printcolumn = r#"{"name":"Scope","type":"string","jsonPath":".spec.scope"}"#, - printcolumn = r#"{"name":"Phase","type":"string","jsonPath":".status.ledger.phase"}"# + printcolumn = r#"{"name":"Phase","type":"string","jsonPath":".status.phase"}"# )] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct KarsBudgetAccountSpec { @@ -26,13 +27,43 @@ pub struct KarsBudgetAccountSpec { pub limits: Limits, } -#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema)] +#[derive(Clone, Copy, Debug, Deserialize, Serialize, JsonSchema, PartialEq, Eq)] +pub enum AccountStatusPhase { + Bootstrap, + Active, + Blocked, + Closing, + Closed, + Frozen, + Corrupt, + Unknown, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize, JsonSchema, PartialEq)] #[serde(rename_all = "camelCase", deny_unknown_fields)] pub struct KarsBudgetAccountStatus { + /// Observational ledger admission state, never spend authority or router readiness. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub phase: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub observed_generation: Option, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[schemars(schema_with = "conditions_schema")] + pub conditions: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] pub ledger: Option, } +fn conditions_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { + let mut schema = Vec::::json_schema(generator); + schema.insert("x-kubernetes-list-type".into(), serde_json::json!("map")); + schema.insert( + "x-kubernetes-list-map-keys".into(), + serde_json::json!(["type"]), + ); + schema +} + pub const BOOTSTRAP: &str = "kars.azure.com/inference-budget-bootstrap"; pub const MANAGED_BY: &str = "app.kubernetes.io/managed-by"; pub const OWNER: &str = "kars-inference-budget"; diff --git a/controller/src/inference_budget/mod.rs b/controller/src/inference_budget/mod.rs index 0cc17d21c..be7c55094 100644 --- a/controller/src/inference_budget/mod.rs +++ b/controller/src/inference_budget/mod.rs @@ -12,6 +12,7 @@ pub mod pod; pub mod recovery; pub mod scope; pub mod service; +mod status; pub mod store; pub mod team; pub mod transport; diff --git a/controller/src/inference_budget/recovery.rs b/controller/src/inference_budget/recovery.rs index 445933085..a69b029f0 100644 --- a/controller/src/inference_budget/recovery.rs +++ b/controller/src/inference_budget/recovery.rs @@ -2,7 +2,7 @@ // Licensed under the MIT License. use super::{ - account::{KarsBudgetAccount, MANAGED_BY, OWNER}, + account::{BOOTSTRAP, KarsBudgetAccount, MANAGED_BY, OWNER}, config::Settings, store::{Store, StoreError}, }; @@ -47,10 +47,7 @@ async fn scan(client: &Client, settings: &Settings) -> Result<(), StoreError> { loop { let page = api.list(¶ms).await.map_err(api_error)?; for account in page.items { - if account.status.is_none() { - continue; - } - if recover(client, &store, &account).await.is_err() { + if reconcile_account(client, &store, &account).await.is_err() { tracing::error!(account = %account.name_any(), "Governed inference account recovery failed closed; no balances reset"); } @@ -63,6 +60,26 @@ async fn scan(client: &Client, settings: &Settings) -> Result<(), StoreError> { Ok(()) } +pub(super) async fn reconcile_account( + client: &Client, + store: &Store, + account: &KarsBudgetAccount, +) -> Result<(), StoreError> { + let result = if account.annotations().get(BOOTSTRAP).map(String::as_str) == Some("pending") { + Ok(()) + } else { + recover(client, store, account).await + }; + store + .refresh_status( + &account.spec.root, + account.metadata.uid.as_deref().ok_or(StoreError::Missing)?, + result.as_ref().err(), + ) + .await?; + result +} + async fn recover( client: &Client, store: &Store, diff --git a/controller/src/inference_budget/status.rs b/controller/src/inference_budget/status.rs new file mode 100644 index 000000000..ba8bafc7a --- /dev/null +++ b/controller/src/inference_budget/status.rs @@ -0,0 +1,172 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Observations only. Admission continues to validate the sealed UID-bound +//! ledger; editing a phase or Ready condition cannot change a grant or balance. + +use super::{ + account::{AccountStatusPhase as Phase, BOOTSTRAP, KarsBudgetAccount, KarsBudgetAccountStatus}, + store::{Store, StoreError}, +}; +use crate::{ + inference_budget_contract::{AccountPhase, Amounts, Limits, MAX_ATTEMPTS}, + status::conditions, +}; +use kube::ResourceExt; + +struct Observation { + phase: Phase, + ready: &'static str, + valid: &'static str, + reason: &'static str, + message: &'static str, +} + +fn blocked(phase: Phase, reason: &'static str, message: &'static str) -> Observation { + Observation { + phase, + ready: "False", + valid: "True", + reason, + message, + } +} + +fn at_limit(limits: Limits, amount: Amounts) -> bool { + let limits = limits.normalized(); + limits.tokens.is_some_and(|cap| amount.tokens >= cap) + || limits + .usd_micros + .is_some_and(|cap| amount.usd_micros >= cap) +} + +fn observe(account: &KarsBudgetAccount, error: Option<&StoreError>) -> Observation { + if account.annotations().get(BOOTSTRAP).map(String::as_str) == Some("pending") + && Store::bootstrap_ledger(account).is_ok() + { + return Observation { + phase: Phase::Bootstrap, + ready: "False", + valid: "Unknown", + reason: "BootstrapPending", + message: "The account is not sealed; no governed-inference dispatch is authorized.", + }; + } + let Ok(ledger) = Store::ledger(account) else { + return Observation { + phase: Phase::Corrupt, + ready: "False", + valid: "False", + reason: "LedgerInvalid", + message: "The sealed account identity or ledger is invalid; balances are retained and admission fails closed.", + }; + }; + match ledger.phase { + AccountPhase::Frozen => { + return blocked( + Phase::Frozen, + "ContractBreach", + "Provider usage breached the governed-inference bound; the account is frozen.", + ); + } + AccountPhase::Closing => { + return blocked( + Phase::Closing, + "AccountClosing", + "Account authority is closing; outstanding governed-inference liabilities remain funded.", + ); + } + AccountPhase::Closed => { + return blocked( + Phase::Closed, + "AccountRetired", + "Account authority is retired; historical governed-inference charges are retained.", + ); + } + AccountPhase::Active => {} + } + if let Some(error) = error { + return match error { + StoreError::Api { .. } | StoreError::Contention => Observation { + phase: Phase::Unknown, + ready: "Unknown", + valid: "True", + reason: "ReconciliationUnavailable", + message: "Live authority reconciliation could not complete; recorded liabilities are retained, not re-authorized.", + }, + _ => blocked( + Phase::Blocked, + "ReconciliationBlocked", + "Live authority reconciliation failed closed; recorded liabilities are retained.", + ), + }; + } + let durable = ledger.meters.settled.checked_add(ledger.meters.uncertain); + if durable.is_ok_and(|amount| at_limit(ledger.limits, amount)) { + return blocked( + Phase::Blocked, + "BudgetExhausted", + "Settled or uncertain governed-inference charges have exhausted a declared account ceiling.", + ); + } + if ledger + .meters + .total() + .is_ok_and(|amount| at_limit(ledger.limits, amount)) + { + return blocked( + Phase::Blocked, + "BudgetReserved", + "Outstanding reservations occupy a declared ceiling; already funded work is retained.", + ); + } + if !ledger.nodes.is_empty() && ledger.nodes.values().all(|node| !node.active) { + return blocked( + Phase::Blocked, + "AuthorityRevoked", + "All enrolled task authorities are closed; reopening requires controller-verified authority.", + ); + } + if ledger.attempts.len() >= MAX_ATTEMPTS { + return blocked( + Phase::Blocked, + "AttemptCapacityReached", + "The bounded attempt ledger is full; replay fences and funded work are retained.", + ); + } + Observation { + phase: Phase::Active, + ready: "True", + valid: "True", + reason: "LedgerAvailable", + message: "The sealed governed-inference ledger has admission headroom; each request still requires live authority and a funded contract.", + } +} + +pub(super) fn project( + account: &KarsBudgetAccount, + error: Option<&StoreError>, +) -> KarsBudgetAccountStatus { + let observation = observe(account, error); + let mut status = account.status.clone().unwrap_or_default(); + let prior = &status.conditions; + status.conditions = [ + (conditions::TYPE_READY, observation.ready), + ("LedgerValid", observation.valid), + ] + .into_iter() + .map(|(kind, value)| { + conditions::preserve_transition_time( + conditions::find(prior, kind), + kind, + value, + observation.reason, + observation.message, + account.metadata.generation, + ) + }) + .collect(); + status.phase = Some(observation.phase); + status.observed_generation = account.metadata.generation; + status +} diff --git a/controller/src/inference_budget/status_tests.rs b/controller/src/inference_budget/status_tests.rs new file mode 100644 index 000000000..736238451 --- /dev/null +++ b/controller/src/inference_budget/status_tests.rs @@ -0,0 +1,514 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::inference_budget::account::AccountStatusPhase as Phase; +use crate::inference_budget_contract::{Settlement, Usage}; +use k8s_openapi::apimachinery::pkg::apis::meta::v1::Condition; +use kube::CustomResourceExt; + +fn snapshot(state: &Arc>) -> KarsBudgetAccount { + state.lock().unwrap().account.clone().unwrap() +} + +fn ready(account: &KarsBudgetAccount) -> &Condition { + account + .status + .as_ref() + .unwrap() + .conditions + .iter() + .find(|condition| condition.type_ == "Ready") + .unwrap() +} + +fn check(account: &KarsBudgetAccount, phase: Phase, value: &str, reason: &str) { + let status = account.status.as_ref().unwrap(); + assert_eq!(status.phase, Some(phase)); + assert_eq!(status.observed_generation, account.metadata.generation); + assert_eq!(status.conditions.len(), 2); + assert_eq!(ready(account).status, value); + assert_eq!(ready(account).reason, reason); + for condition in &status.conditions { + assert_eq!(condition.observed_generation, account.metadata.generation); + assert!(!condition.message.contains("private-request-body")); + assert!(!condition.reason.is_empty()); + } +} + +fn command(request: &ReserveRequest) -> AttemptCommand { + AttemptCommand { + account_uid: request.account_uid.clone(), + key: AttemptKey { + pod_uid: request.identity.pod_uid.clone(), + sequence: request.sequence, + }, + identity: request.identity.clone(), + wire_digest: request.wire_digest.clone(), + } +} + +#[tokio::test] +async fn reserved_headroom_is_not_revocation_and_settlement_reports_durable_exhaustion() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + store + .transact(&root(), "account-uid", |ledger| { + ledger.narrow_root_limits(Limits { + tokens: Some(30), + usd_micros: Some(10), + }) + }) + .await + .unwrap(); + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + let reserved = snapshot(&state); + check(&reserved, Phase::Blocked, "False", "BudgetReserved"); + let ledger = reserved.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert!(ledger.nodes["root"].active); + assert!(!ledger.sessions["a"].closed); + assert_eq!(ledger.meters.reserved.tokens, 30); + let writes = state.lock().unwrap().writes; + assert!(matches!( + store + .transact(&root(), "account-uid", |ledger| ledger + .reserve(&reserve("b"), 100)) + .await, + Err(StoreError::Ledger(BudgetError::Exhausted)) + )); + assert_eq!(state.lock().unwrap().writes, writes); + store + .transact(&root(), "account-uid", |ledger| { + ledger.begin_dispatch(&command(&request), 101) + }) + .await + .unwrap(); + store + .transact(&root(), "account-uid", |ledger| { + ledger.settle(&Settlement { + attempt: command(&request), + usage: Some(Usage { + input_tokens: 10, + output_tokens: 20, + cached_input_tokens: 0, + cache_creation_input_tokens: 0, + reasoning_output_tokens: 0, + }), + }) + }) + .await + .unwrap(); + let settled = snapshot(&state); + check(&settled, Phase::Blocked, "False", "BudgetExhausted"); + assert_eq!( + ready(&settled).last_transition_time, + ready(&reserved).last_transition_time + ); + let ledger = settled.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert_eq!(ledger.meters.settled.tokens, 30); + assert_eq!(ledger.meters.reserved.tokens, 0); + assert!(ledger.attempts.is_empty()); + assert!(ledger.nodes["root"].active); +} + +#[tokio::test] +async fn expiry_restores_headroom_without_resetting_uid_or_replay_fences() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + store + .transact(&root(), "account-uid", |ledger| { + ledger.narrow_root_limits(Limits { + tokens: Some(30), + usd_micros: Some(10), + }) + }) + .await + .unwrap(); + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + check(&snapshot(&state), Phase::Blocked, "False", "BudgetReserved"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.expire_undispatched(200) + }) + .await + .unwrap(); + let current = snapshot(&state); + check(¤t, Phase::Active, "True", "LedgerAvailable"); + let ledger = current.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert_eq!(ledger.account_uid, "account-uid"); + assert_eq!(ledger.sessions["a"].closed_through, 1); + assert!( + store + .transact(&root(), "account-uid", |ledger| ledger + .reserve(&request, 201)) + .await + .is_err() + ); +} + +#[tokio::test] +async fn revoked_and_retired_authority_retains_inflight_liability() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + store + .transact(&root(), "account-uid", |ledger| { + ledger.begin_dispatch(&command(&request), 101) + }) + .await + .unwrap(); + store + .transact(&root(), "account-uid", |ledger| { + ledger.close_subtree("root") + }) + .await + .unwrap(); + let revoked = snapshot(&state); + check(&revoked, Phase::Blocked, "False", "AuthorityRevoked"); + let ledger = revoked.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert_eq!(ledger.meters.uncertain.tokens, 30); + assert!(ledger.sessions.values().all(|session| session.closed)); + store + .transact(&root(), "account-uid", |ledger| ledger.close_account()) + .await + .unwrap(); + let closed = snapshot(&state); + check(&closed, Phase::Closed, "False", "AccountRetired"); + assert_eq!( + closed + .status + .as_ref() + .unwrap() + .ledger + .as_ref() + .unwrap() + .meters, + ledger.meters + ); +} + +#[tokio::test] +async fn provider_breach_freezes_even_if_observational_ready_is_forged() { + let (_server, store, state) = setup(Some(account()), Fault::None).await; + let request = reserve("a"); + store + .transact(&root(), "account-uid", |ledger| { + ledger.reserve(&request, 100) + }) + .await + .unwrap(); + store + .transact(&root(), "account-uid", |ledger| { + ledger.begin_dispatch(&command(&request), 101) + }) + .await + .unwrap(); + store + .transact(&root(), "account-uid", |ledger| { + ledger.settle(&Settlement { + attempt: command(&request), + usage: Some(Usage { + input_tokens: 11, + output_tokens: 0, + cached_input_tokens: 0, + cache_creation_input_tokens: 0, + reasoning_output_tokens: 0, + }), + }) + }) + .await + .unwrap(); + check(&snapshot(&state), Phase::Frozen, "False", "ContractBreach"); + { + let mut state = state.lock().unwrap(); + let status = state.account.as_mut().unwrap().status.as_mut().unwrap(); + status.phase = Some(Phase::Active); + status + .conditions + .iter_mut() + .find(|condition| condition.type_ == "Ready") + .unwrap() + .status = "True".into(); + } + assert!(matches!( + store + .transact(&root(), "account-uid", |ledger| ledger + .reserve(&reserve("b"), 102)) + .await, + Err(StoreError::Ledger(BudgetError::Closed)) + )); + let reported = store + .refresh_status(&root(), "account-uid", None) + .await + .unwrap(); + check(&reported, Phase::Frozen, "False", "ContractBreach"); + assert_eq!( + reported + .status + .unwrap() + .ledger + .unwrap() + .meters + .uncertain + .tokens, + 30 + ); +} + +#[tokio::test] +async fn recovery_reports_unknown_on_live_api_failure_without_churn_or_private_errors() { + let original = account(); + let (server, store, state) = setup(Some(original.clone()), Fault::None).await; + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + assert!( + super::super::super::recovery::reconcile_account(&client, &store, &original) + .await + .is_err() + ); + let unknown = snapshot(&state); + check( + &unknown, + Phase::Unknown, + "Unknown", + "ReconciliationUnavailable", + ); + assert_eq!( + unknown.status.as_ref().unwrap().ledger, + original.status.as_ref().unwrap().ledger + ); + let writes = state.lock().unwrap().writes; + assert!( + super::super::super::recovery::reconcile_account(&client, &store, &unknown) + .await + .is_err() + ); + assert_eq!(state.lock().unwrap().writes, writes); + assert_eq!(snapshot(&state).status, unknown.status); + Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path("/api/v1/namespaces/workspace")) + .respond_with(failure(404)) + .with_priority(1) + .mount(&server) + .await; + super::super::super::recovery::reconcile_account(&client, &store, &unknown) + .await + .unwrap(); + check(&snapshot(&state), Phase::Closed, "False", "AccountRetired"); +} + +#[tokio::test] +async fn reporting_corruption_and_missing_ledger_never_initializes_or_repairs_money() { + for missing in [false, true] { + let mut original = account(); + if missing { + original.status.as_mut().unwrap().ledger = None; + } else { + original + .status + .as_mut() + .unwrap() + .ledger + .as_mut() + .unwrap() + .meters + .reserved + .tokens = 1; + } + let (_server, store, state) = setup(Some(original.clone()), Fault::None).await; + let reported = store + .refresh_status(&root(), "account-uid", None) + .await + .unwrap(); + check(&reported, Phase::Corrupt, "False", "LedgerInvalid"); + assert_eq!( + reported.status.as_ref().unwrap().ledger, + original.status.as_ref().unwrap().ledger + ); + assert!(store.read(&root(), "account-uid").await.is_err()); + assert!(store.initialize(&root(), "account-uid").await.is_err()); + assert_eq!(snapshot(&state).status, reported.status); + } +} + +#[tokio::test] +async fn observation_write_failures_and_replacements_never_fake_success() { + for fault in [ + Fault::FailRead, + Fault::FailWrite, + Fault::RecreateOnWrite, + Fault::CommitThenFail, + Fault::ConflictOnce, + ] { + let original = account(); + let (_server, store, state) = setup(Some(original.clone()), fault).await; + let error = StoreError::Api { + stage: "private-request-body", + code: Some(503), + }; + let result = store + .refresh_status(&root(), "account-uid", Some(&error)) + .await; + assert_eq!(result.is_ok(), matches!(fault, Fault::ConflictOnce)); + let stored = snapshot(&state); + assert_eq!( + stored.status.as_ref().unwrap().ledger, + original.status.as_ref().unwrap().ledger + ); + if matches!(fault, Fault::ConflictOnce | Fault::CommitThenFail) { + check( + &stored, + Phase::Unknown, + "Unknown", + "ReconciliationUnavailable", + ); + } else { + assert_eq!(stored.status, original.status); + } + } +} + +#[tokio::test] +async fn observation_backfill_and_generation_update_preserve_stable_transition_time() { + let mut original = account(); + let ledger = original.status.as_ref().unwrap().ledger.clone(); + original.status = Some(KarsBudgetAccountStatus { + ledger, + ..Default::default() + }); + let (_server, store, state) = setup(Some(original), Fault::None).await; + let first = store + .refresh_status(&root(), "account-uid", None) + .await + .unwrap(); + check(&first, Phase::Active, "True", "LedgerAvailable"); + state + .lock() + .unwrap() + .account + .as_mut() + .unwrap() + .metadata + .generation = Some(2); + let second = store + .refresh_status(&root(), "account-uid", None) + .await + .unwrap(); + check(&second, Phase::Active, "True", "LedgerAvailable"); + assert_eq!( + ready(&first).last_transition_time, + ready(&second).last_transition_time + ); + assert_eq!(ready(&second).observed_generation, Some(2)); + let writes = state.lock().unwrap().writes; + store + .refresh_status(&root(), "account-uid", None) + .await + .unwrap(); + assert_eq!(state.lock().unwrap().writes, writes); +} + +#[test] +fn helm_budget_account_reporting_matches_generated_schema() { + fn without_descriptions(mut value: serde_json::Value) -> serde_json::Value { + match &mut value { + serde_json::Value::Object(fields) => { + fields.remove("description"); + for field in fields.values_mut() { + *field = without_descriptions(field.take()); + } + + #[tokio::test] + async fn observation_conflict_reloads_concurrently_reserved_money() { + let (_server, store, state) = + setup(Some(account()), Fault::ReserveOnConflict).await; + let error = StoreError::Api { + stage: "recovery", + code: Some(503), + }; + let reported = store + .refresh_status(&root(), "account-uid", Some(&error)) + .await + .unwrap(); + check( + &reported, + Phase::Unknown, + "Unknown", + "ReconciliationUnavailable", + ); + let ledger = reported.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert_eq!(ledger.meters.reserved.tokens, 30); + assert_eq!(ledger.attempts.len(), 1); + assert_eq!(ledger.sessions["a"].next_sequence, 2); + assert_eq!(snapshot(&state).status, reported.status); + } + + #[tokio::test] + async fn lost_bootstrap_status_ack_reuses_the_same_ledger_before_sealing() { + let (_server, store, state) = setup(None, Fault::None).await; + store + .create_anchor( + account().spec, + &crate::providers::signing::ReceiptSigner::from_bytes(&[7; 32]), + ) + .await + .unwrap(); + state.lock().unwrap().fault = Fault::CommitThenFail; + assert!(store.initialize(&root(), "account-uid").await.is_err()); + let pending = snapshot(&state); + check(&pending, Phase::Bootstrap, "False", "BootstrapPending"); + assert!(store.read(&root(), "account-uid").await.is_err()); + let initialized = store.initialize(&root(), "account-uid").await.unwrap(); + check(&initialized, Phase::Active, "True", "LedgerAvailable"); + assert_eq!( + initialized.status.as_ref().unwrap().ledger, + pending.status.as_ref().unwrap().ledger + ); + assert_eq!(initialized.metadata.uid, pending.metadata.uid); + } + } + serde_json::Value::Array(items) => { + for item in items { + *item = without_descriptions(item.take()); + } + } + _ => {} + } + value + } + let helm: serde_json::Value = serde_yaml::from_str(include_str!( + "../../../deploy/helm/kars/templates/crd-karsbudgetaccount.yaml" + )) + .unwrap(); + let generated = serde_json::to_value(KarsBudgetAccount::crd()).unwrap(); + let helm = &helm["spec"]["versions"][0]; + let generated = &generated["spec"]["versions"][0]; + assert_eq!( + helm["additionalPrinterColumns"], + generated["additionalPrinterColumns"] + ); + for field in ["phase", "observedGeneration", "conditions"] { + let path = |schema: &serde_json::Value| { + schema["schema"]["openAPIV3Schema"]["properties"]["status"]["properties"][field].clone() + }; + assert_eq!( + without_descriptions(path(helm)), + without_descriptions(path(generated)), + "{field}" + ); + } +} diff --git a/controller/src/inference_budget/store.rs b/controller/src/inference_budget/store.rs index 9bff905e6..3b8ae678e 100644 --- a/controller/src/inference_budget/store.rs +++ b/controller/src/inference_budget/store.rs @@ -93,7 +93,7 @@ impl Store { Ok(()) } - fn ledger(account: &KarsBudgetAccount) -> Result<&Ledger, StoreError> { + pub(super) fn ledger(account: &KarsBudgetAccount) -> Result<&Ledger, StoreError> { let ledger = account .status .as_ref() @@ -120,6 +120,29 @@ impl Store { Ok(ledger) } + pub(super) fn bootstrap_ledger(account: &KarsBudgetAccount) -> Result<(), StoreError> { + if let Some(ledger) = account + .status + .as_ref() + .and_then(|status| status.ledger.as_ref()) + { + ledger.validate()?; + if account.metadata.uid.as_deref() != Some(ledger.account_uid.as_str()) + || ledger.root != account.spec.root + || ledger.scope != account.spec.scope + || ledger.limits.normalized() != account.spec.limits.normalized() + || !ledger.nodes.is_empty() + || !ledger.sessions.is_empty() + || !ledger.attempts.is_empty() + || ledger.meters != Default::default() + || ledger.phase != crate::inference_budget_contract::AccountPhase::Active + { + return Err(BudgetError::Corrupt.into()); + } + } + Ok(()) + } + /// Reserve an account anchor. The owning root's protected status must pin /// the returned UID BEFORE `initialize` is called. Existing callers with a /// pin use `read`, never this bootstrap path. @@ -140,7 +163,7 @@ impl Store { (BOOTSTRAP.into(), "pending".into()), (super::claim::ANNOTATION.into(), authority), ])); - match self.accounts.create(&PostParams::default(), &account).await { + let created = match self.accounts.create(&PostParams::default(), &account).await { Ok(created) => Ok(created), Err(kube::Error::Api(status)) if status.code == 409 => { let existing = self @@ -168,7 +191,13 @@ impl Store { Ok(existing) } Err(error) => Err(api_error("create account anchor", error)), - } + }?; + self.refresh_status( + &created.spec.root, + created.metadata.uid.as_deref().ok_or(StoreError::Missing)?, + None, + ) + .await } pub async fn initialize( @@ -190,43 +219,41 @@ impl Store { Self::validate_identity(&account, root, pinned_uid)?; if account.annotations().get(BOOTSTRAP).map(String::as_str) == Some("sealed") { Self::ledger(&account)?; - return Ok(account); + return self.refresh_status(root, pinned_uid, None).await; } if account.annotations().get(BOOTSTRAP).map(String::as_str) != Some("pending") { return Err(StoreError::Missing); } - if let Some(ledger) = account + Self::bootstrap_ledger(&account)?; + if account .status .as_ref() .and_then(|status| status.ledger.as_ref()) + .is_some() { // A crash after status initialization must validate the existing // ledger, not reset it. Pending accounts cannot dispatch. - ledger.validate()?; - if ledger.account_uid != pinned_uid - || ledger.root != *root - || ledger.limits.normalized() != account.spec.limits.normalized() - || !ledger.nodes.is_empty() - || !ledger.sessions.is_empty() - || !ledger.attempts.is_empty() - || ledger.meters != Default::default() - { - return Err(BudgetError::Corrupt.into()); - } match self.accounts.patch(&name, &PatchParams::default(), &Patch::Merge(json!({ "metadata": {"uid": pinned_uid, "resourceVersion": account.resource_version(), "annotations": {BOOTSTRAP: "sealed"}} }))).await { - Ok(sealed) => { Self::ledger(&sealed)?; return Ok(sealed); } + Ok(sealed) => { + Self::ledger(&sealed)?; + return self.refresh_status(root, pinned_uid, None).await; + } Err(kube::Error::Api(status)) if status.code == 409 => continue, Err(error) => return Err(api_error("seal account bootstrap", error)), } } let ledger = Ledger::new(pinned_uid.into(), root.clone(), account.spec.limits)?; - match self.accounts.patch_status(&name, &PatchParams::default(), &Patch::Merge(json!({ - "metadata": {"uid": pinned_uid, "resourceVersion": account.resource_version()}, - "status": {"ledger": ledger} - }))).await { + let mut next = account.clone(); + next.status.get_or_insert_with(Default::default).ledger = Some(ledger); + next.status = Some(super::status::project(&next, None)); + match self + .accounts + .replace_status(&name, &PostParams::default(), &next) + .await + { Ok(_) => {} Err(kube::Error::Api(status)) if status.code == 409 => continue, Err(error) => return Err(api_error("initialize account ledger", error)), @@ -283,9 +310,8 @@ impl Store { // Merge-patching maps would retain compacted attempt rows. Replace // the status as one value using PUT with UID/RV preconditions. let mut next = account.clone(); - next.status = Some(KarsBudgetAccountStatus { - ledger: Some(mutation.next), - }); + next.status.get_or_insert_with(Default::default).ledger = Some(mutation.next); + next.status = Some(super::status::project(&next, None)); let committed = tokio::time::timeout_at( deadline, self.accounts @@ -312,6 +338,54 @@ impl Store { } Err(StoreError::Contention) } + + /// Report only from a fresh UID/RV snapshot. Even corrupt or unavailable + /// accounts keep their exact ledger; a report can never initialize funding. + pub(super) async fn refresh_status( + &self, + root: &RootIdentity, + account_uid: &str, + error: Option<&StoreError>, + ) -> Result { + root.validate()?; + if !crate::inference_budget_contract::valid_uid(account_uid) { + return Err(BudgetError::Identity.into()); + } + let deadline = + tokio::time::Instant::now() + std::time::Duration::from_secs(MAX_OPERATION_SECONDS); + for _ in 0..RETRIES { + let account = + tokio::time::timeout_at(deadline, self.accounts.get(&name_for_root(root))) + .await + .map_err(|_| StoreError::Contention)? + .map_err(|error| api_error("read account observation", error))?; + Self::validate_identity(&account, root, account_uid)?; + let mut next = account.clone(); + next.status = Some(super::status::project(&account, error)); + if next.status == account.status { + return Ok(account); + } + match tokio::time::timeout_at( + deadline, + self.accounts + .replace_status(&name_for_root(root), &PostParams::default(), &next), + ) + .await + .map_err(|_| StoreError::Contention)? + { + Ok(stored) => { + Self::validate_identity(&stored, root, account_uid)?; + if stored.status != next.status { + return Err(BudgetError::Corrupt.into()); + } + return Ok(stored); + } + Err(kube::Error::Api(status)) if status.code == 409 => continue, + Err(error) => return Err(api_error("commit account observation", error)), + } + } + Err(StoreError::Contention) + } } #[cfg(test)] diff --git a/controller/src/inference_budget/store_tests.rs b/controller/src/inference_budget/store_tests.rs index 1a9c2c91b..20242f654 100644 --- a/controller/src/inference_budget/store_tests.rs +++ b/controller/src/inference_budget/store_tests.rs @@ -76,6 +76,7 @@ fn account() -> KarsBudgetAccount { account.metadata.namespace = Some("kars-system".into()); account.metadata.uid = Some("account-uid".into()); account.metadata.resource_version = Some("1".into()); + account.metadata.generation = Some(1); account.labels_mut().insert(MANAGED_BY.into(), OWNER.into()); account .annotations_mut() @@ -86,7 +87,9 @@ fn account() -> KarsBudgetAccount { let ledger = ledger.register_session(identity("b")).unwrap().next; account.status = Some(KarsBudgetAccountStatus { ledger: Some(ledger), + ..Default::default() }); + account.status = Some(super::super::status::project(&account, None)); account } @@ -128,6 +131,9 @@ enum Fault { ConflictOnce, CommitThenFail, RecreateOnWrite, + FailWrite, + FailRead, + ReserveOnConflict, } struct State { @@ -156,6 +162,7 @@ impl Respond for Server { fn respond(&self, request: &Request) -> ResponseTemplate { let mut state = self.0.lock().unwrap(); match (request.method.as_str(), request.url.path()) { + ("GET", OBJECT) if matches!(state.fault, Fault::FailRead) => failure(503), ("GET", OBJECT) => state .account .as_ref() @@ -170,6 +177,7 @@ impl Respond for Server { account.metadata.uid = Some("account-uid".into()); account.metadata.namespace = Some("kars-system".into()); account.metadata.resource_version = Some(state.version.to_string()); + account.metadata.generation = Some(1); let output = response(&account); state.account = Some(account); state.writes += 1; @@ -179,6 +187,24 @@ impl Respond for Server { let mut incoming: KarsBudgetAccount = serde_json::from_slice(&request.body).unwrap(); match state.fault { + Fault::FailWrite => return failure(503), + Fault::ReserveOnConflict => { + state.fault = Fault::None; + let status = state.account.as_mut().unwrap().status.as_mut().unwrap(); + status.ledger = Some( + status + .ledger + .as_ref() + .unwrap() + .reserve(&reserve("a"), 100) + .unwrap() + .next, + ); + state.version += 1; + let version = state.version.to_string(); + state.account.as_mut().unwrap().metadata.resource_version = Some(version); + return failure(409); + } Fault::ConflictOnce => { state.fault = Fault::None; state.version += 1; @@ -440,6 +466,10 @@ async fn bootstrap_is_metadata_first_and_sealed_only_after_uid_bound_status() { .await .unwrap(); assert!(store.read(&root(), "account-uid").await.is_err()); + let pending = anchor.status.as_ref().unwrap(); + assert_eq!(pending.phase, Some(AccountStatusPhase::Bootstrap)); + assert!(pending.ledger.is_none()); + assert_eq!(pending.conditions[0].status, "False"); let initialized = store .initialize(&root(), anchor.metadata.uid.as_deref().unwrap()) .await @@ -448,6 +478,10 @@ async fn bootstrap_is_metadata_first_and_sealed_only_after_uid_bound_status() { initialized.annotations().get(BOOTSTRAP).map(String::as_str), Some("sealed") ); + assert_eq!( + initialized.status.as_ref().unwrap().phase, + Some(AccountStatusPhase::Active) + ); assert!( initialized .status @@ -474,3 +508,6 @@ async fn replacement_between_read_and_commit_cannot_receive_a_grant() { ); assert_eq!(state.lock().unwrap().writes, 0); } + +#[path = "status_tests.rs"] +mod status_tests; diff --git a/controller/src/kars_team_reconciler/budget_interleaving_tests.rs b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs index 0b6f68987..e25ff58a3 100644 --- a/controller/src/kars_team_reconciler/budget_interleaving_tests.rs +++ b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs @@ -291,6 +291,7 @@ fn enroll(store: &Arc>, apis: &BudgetApis) { account.metadata.annotations = Some([(BOOTSTRAP.into(), "sealed".into())].into()); account.status = Some(KarsBudgetAccountStatus { ledger: Some(ledger), + ..Default::default() }); *apis.account.lock().unwrap() = Some(account); } diff --git a/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml b/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml index 7fd263daf..7d3655307 100644 --- a/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml +++ b/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml @@ -29,7 +29,7 @@ spec: jsonPath: .spec.scope - name: Phase type: string - jsonPath: .status.ledger.phase + jsonPath: .status.phase schema: openAPIV3Schema: type: object @@ -75,6 +75,33 @@ spec: status: type: object properties: + phase: + type: string + nullable: true + enum: [Bootstrap, Active, Blocked, Closing, Closed, Frozen, Corrupt, Unknown, null] + description: "Observational ledger admission state, never spend authority or router readiness." + observedGeneration: + type: integer + format: int64 + nullable: true + conditions: + type: array + x-kubernetes-list-type: map + x-kubernetes-list-map-keys: [type] + items: + type: object + required: [lastTransitionTime, message, reason, status, type] + properties: + lastTransitionTime: + type: string + format: date-time + message: {type: string} + observedGeneration: + type: integer + format: int64 + reason: {type: string} + status: {type: string} + type: {type: string} ledger: type: object x-kubernetes-preserve-unknown-fields: true diff --git a/deploy/helm/kars/tests/src/inference-budget.test.ts b/deploy/helm/kars/tests/src/inference-budget.test.ts index f2ee1f9c2..cd82a913a 100644 --- a/deploy/helm/kars/tests/src/inference-budget.test.ts +++ b/deploy/helm/kars/tests/src/inference-budget.test.ts @@ -55,6 +55,30 @@ const configured = { }; describe("governed inference budget Helm contract", () => { + it.each([undefined, configured])( + "publishes observational account phase and keyed standard conditions with %s", (value) => { + const definition = render(value).find((doc: { kind: string; spec?: { names?: { kind?: string } } }) => + doc.kind === "CustomResourceDefinition" && doc.spec?.names?.kind === "KarsBudgetAccount"); + const version = definition.spec.versions[0]; + expect(version.additionalPrinterColumns).toContainEqual({ + name: "Phase", type: "string", jsonPath: ".status.phase", + }); + const status = version.schema.openAPIV3Schema.properties.status.properties; + expect(status.phase.enum).toEqual([ + "Bootstrap", "Active", "Blocked", "Closing", "Closed", "Frozen", "Corrupt", "Unknown", null, + ]); + expect(status.conditions.type).toBe("array"); + expect(status.conditions["x-kubernetes-list-type"]).toBe("map"); + expect(status.conditions["x-kubernetes-list-map-keys"]).toEqual(["type"]); + expect(status.conditions.items.required).toEqual([ + "lastTransitionTime", "message", "reason", "status", "type", + ]); + expect(status.conditions.items.properties.lastTransitionTime.format).toBe("date-time"); + expect(status.conditions.items.properties.observedGeneration.format).toBe("int64"); + expect(status.ledger.properties.phase.enum).toEqual(["Active", "Closing", "Closed", "Frozen"]); + }, + ); + it("keeps generated Task/Team CEL in sync with the Rust scope contract", () => { const source = readFileSync(join(chart, "../../../controller/src/inference_budget/scope.rs"), "utf8"); const rule = (name: string): string => { diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index 027dfc3bd..5183c524a 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -118,6 +118,36 @@ smaller than the model's hard context/output maximum can refuse even a short prompt; concurrency needs room for all simultaneous maxima. This version does not substitute a character heuristic for exact pre-dispatch token evidence. +## Account observations + +`KarsBudgetAccount.status.phase` (the `kubectl` Phase column) and the standard +`Ready`/`LedgerValid` conditions describe the last observed account state. +They include `observedGeneration`; condition transition timestamps change only +when the corresponding True/False/Unknown value changes. + +- **Bootstrap:** the UID anchor is not sealed; it cannot dispatch. +- **Active / LedgerAvailable:** the validated, sealed ledger has headroom. + This is not router health, provider availability, or permission for a send. +- **Blocked / BudgetReserved:** reservations occupy a declared ceiling. + Already funded executions and their Task UIDs remain intact. +- **Blocked / BudgetExhausted:** settled or uncertain charges occupy a ceiling. +- **Blocked / AuthorityRevoked** or **AttemptCapacityReached:** enrolled + authorities are closed, or bounded attempt storage has no room. +- **Frozen / ContractBreach**, **Closing**, or **Closed / AccountRetired:** + the ledger's enforcement phase is retained, including historical liabilities. +- **Corrupt / LedgerInvalid:** identity, sealing or ledger validation failed. + Reporting does not repair, reinitialize, or zero the ledger. +- **Unknown / ReconciliationUnavailable:** a live API read or reconciliation + could not complete. If the account API itself is unavailable, an error cannot + be persisted: the last stored observation remains, and the operation fails. + +Bootstrap, ledger transitions, and periodic recovery populate these observations. +Reporting writes use a fresh UID/resourceVersion and replace the complete status +without changing its ledger. The original `status.ledger.phase` wire contract +remains unchanged. Neither a Phase value nor `Ready=True` is spend authority: +every send still requires live Pod/Task authority, the sealed validated ledger, +an operator contract, and an atomic reservation/begin transition. + ## Operator contracts and unavoidable configuration Enable the optional Helm `inferenceBudget` section only after supplying: From a6a8d75f0d33a3b20989045523b72db90f8f2ec0 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Wed, 9 Sep 2026 17:09:46 +0200 Subject: [PATCH 08/16] fix(budget): qualify account observations and CNCF status schema Actual guarded Rust qualification of abbfc7d19213eb00d09bee9b8f84841fdf82a3d0 found and closed four bounded issues rather than treating source checks as compilation evidence: - Whole-status acknowledgement comparison now compares exact serialized wire values: standard Kubernetes Time serializes at second precision, so comparing unpublished subsecond clock values falsely reported committed status corrupt. Ledger counters and all other fields remain exact in this comparison. - KarsBudgetAccount::crd() hit kube's structural-schema panic on the existing tagged MaximumPrice union. Its generated ledger schema now uses the same bounded structural envelope already deployed by Helm, with existing contract version and capacity constants. The actual typed Ledger and all runtime financial/authority validation are unchanged. The schema regression now checks the ledger envelope too; only documentation, unordered required sets and equivalent JSON-Schema numeric minima are canonicalized for comparison. - Move the concurrent-reservation and bootstrap-lost-ack tests to module scope. Both are now genuinely registered/executed, not nested inactive functions. - Explicitly install the controller's AWS-LC Rustls provider in the legacy Pod fixture before creating its Client. Verify this fixture in isolation too. PASSING ACTUAL CARGO COMMANDS, cwd kars-pr11-budgets: All invoked through: python3 /Users/pallakatos/.copilot/session-state/a7f32227-c055-4df4-81ee-ce27d7f2d74d/files/run-cargo-guard.py --cwd /Users/pallakatos/.copilot/session-state/a7f32227-c055-4df4-81ee-ce27d7f2d74d/kars-pr11-budgets -- The guard used only /Users/pallakatos/Private/Repos/kars/target, default core package features, CARGO_INCREMENTAL=0, CARGO_BUILD_JOBS=2, CARGO_NET_OFFLINE=true, and the existing 8.5 GiB free-space floor. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --lib --bins -- budget task_identity helm_drift 116 controller + 57 router = 173 passed; zero failed/ignored. Includes all 11 new status regressions, bootstrap/active/reserved/exhausted/revoked/frozen/ corrupt/unavailable transitions, report CAS races and lost acks, timestamp stability, exact generated/Helm reporting+ledger schema checks, and existing budget/interleaving/identity/Helm gates. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router -p kars-cncf-conformance --test criteria 17/17 passed, including c3_conditions_array_present, c12_state_printer_column, full_report_is_all_pass and report stability. cargo test --quiet --offline --locked -p kars-controller -p kars-inference-router --lib --bins -- no_task_owner_no_reference_is_a_byte_safe_legacy_noop_without_api_access 1 controller fixture passed in isolation, without another test installing TLS. cargo clippy --quiet --offline --locked -p kars-controller -p kars-inference-router --all-targets -- -D warnings Passed with no new lint waivers. cargo fmt --package kars-controller --package kars-inference-router --check Passed. FAST GATES: node cli/node_modules/vitest/vitest.mjs run --root deploy/helm/kars/tests --globals --no-cache src/inference-budget.test.ts src/local-inference.test.ts Re-run: 20 passed (14 budget + 6 local-inference). CLI budget 7 and CLI typecheck passed at abbfc7d1; no CLI source changed since. git diff --check; git diff --cached --check bash ci/check-copyright-headers.sh 739 source files passed. BASE_REF=2b2e0b8018b5059b05467ac76d3aa6115575be78 bash ci/check-loc.sh Passed; all changed Rust modules remain <=800 lines; no whitelist changes. No Cargo manifest or lock changes, no dependency installs or alternate targets. Minimum observed free space during this lease: 9.84 GiB; guard floor not hit. This qualifies the local budget-account status/CNCF repair, not production rollout. Parent review, hosted CI, real SRE/Kind and cross-credential composition remain independent gates; no claim of a qualified production router digest. No public push/merge, main/customer/private Bridge edits, Azure, Docker or image operations. Existing cli/node_modules symlink remains untracked. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/inference_budget/account.rs | 43 +++++++ controller/src/inference_budget/pod_tests.rs | 1 + .../src/inference_budget/status_tests.rs | 113 +++++++++--------- controller/src/inference_budget/store.rs | 5 +- 4 files changed, 107 insertions(+), 55 deletions(-) diff --git a/controller/src/inference_budget/account.rs b/controller/src/inference_budget/account.rs index 9ca9974de..abed5a674 100644 --- a/controller/src/inference_budget/account.rs +++ b/controller/src/inference_budget/account.rs @@ -51,9 +51,52 @@ pub struct KarsBudgetAccountStatus { #[schemars(schema_with = "conditions_schema")] pub conditions: Vec, #[serde(default, skip_serializing_if = "Option::is_none")] + #[schemars(schema_with = "ledger_schema")] pub ledger: Option, } +// Keep the bounded Helm ledger envelope. Runtime Ledger::validate checks its +// full tagged-union contents; kube's structural union rewrite cannot represent +// the distinct MaximumPrice discriminator schemas without changing that wire. +fn ledger_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema { + use crate::inference_budget_contract::{ + CONTRACT_VERSION, MAX_ATTEMPTS, MAX_NODES, MAX_SESSIONS, + }; + schemars::json_schema!({ + "type": "object", + "x-kubernetes-preserve-unknown-fields": true, + "required": ["version", "scope", "accountUid", "root", "limits", "phase", + "meters", "nodes", "sessions", "attempts"], + "properties": { + "version": {"type": "string", "enum": [CONTRACT_VERSION]}, + "scope": {"type": "string", "enum": ["GovernedInference"]}, + "accountUid": {"type": "string", "minLength": 1, "maxLength": 128}, + "phase": {"type": "string", "enum": ["Active", "Closing", "Closed", "Frozen"]}, + "root": {"type": "object", "x-kubernetes-preserve-unknown-fields": true}, + "limits": { + "type": "object", + "properties": { + "tokens": {"type": "integer", "format": "int64", "minimum": 0}, + "usdMicros": {"type": "integer", "format": "int64", "minimum": 0} + } + }, + "meters": {"type": "object", "x-kubernetes-preserve-unknown-fields": true}, + "nodes": { + "type": "object", "maxProperties": MAX_NODES, + "additionalProperties": {"type": "object", "x-kubernetes-preserve-unknown-fields": true} + }, + "sessions": { + "type": "object", "maxProperties": MAX_SESSIONS, + "additionalProperties": {"type": "object", "x-kubernetes-preserve-unknown-fields": true} + }, + "attempts": { + "type": "object", "maxProperties": MAX_ATTEMPTS, + "additionalProperties": {"type": "object", "x-kubernetes-preserve-unknown-fields": true} + } + } + }) +} + fn conditions_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema { let mut schema = Vec::::json_schema(generator); schema.insert("x-kubernetes-list-type".into(), serde_json::json!("map")); diff --git a/controller/src/inference_budget/pod_tests.rs b/controller/src/inference_budget/pod_tests.rs index 3477cc2c9..27e318b8a 100644 --- a/controller/src/inference_budget/pod_tests.rs +++ b/controller/src/inference_budget/pod_tests.rs @@ -97,6 +97,7 @@ fn every_runtime_receives_only_a_router_private_token_mount() { async fn no_task_owner_no_reference_is_a_byte_safe_legacy_noop_without_api_access() { let server = wiremock::MockServer::start().await; let config = kube::Config::new(server.uri().parse().unwrap()); + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); let client = Client::try_from(config).unwrap(); let sandbox: KarsSandbox = serde_json::from_value(json!({ "apiVersion":"kars.azure.com/v1alpha1", "kind":"KarsSandbox", diff --git a/controller/src/inference_budget/status_tests.rs b/controller/src/inference_budget/status_tests.rs index 736238451..f1470bbd7 100644 --- a/controller/src/inference_budget/status_tests.rs +++ b/controller/src/inference_budget/status_tests.rs @@ -424,66 +424,23 @@ async fn observation_backfill_and_generation_update_preserve_stable_transition_t #[test] fn helm_budget_account_reporting_matches_generated_schema() { - fn without_descriptions(mut value: serde_json::Value) -> serde_json::Value { + fn canonical_schema(mut value: serde_json::Value) -> serde_json::Value { match &mut value { serde_json::Value::Object(fields) => { fields.remove("description"); - for field in fields.values_mut() { - *field = without_descriptions(field.take()); + if let Some(serde_json::Value::Array(required)) = fields.get_mut("required") { + required.sort_by(|a, b| a.as_str().cmp(&b.as_str())); } - - #[tokio::test] - async fn observation_conflict_reloads_concurrently_reserved_money() { - let (_server, store, state) = - setup(Some(account()), Fault::ReserveOnConflict).await; - let error = StoreError::Api { - stage: "recovery", - code: Some(503), - }; - let reported = store - .refresh_status(&root(), "account-uid", Some(&error)) - .await - .unwrap(); - check( - &reported, - Phase::Unknown, - "Unknown", - "ReconciliationUnavailable", - ); - let ledger = reported.status.as_ref().unwrap().ledger.as_ref().unwrap(); - assert_eq!(ledger.meters.reserved.tokens, 30); - assert_eq!(ledger.attempts.len(), 1); - assert_eq!(ledger.sessions["a"].next_sequence, 2); - assert_eq!(snapshot(&state).status, reported.status); + if let Some(minimum) = fields.get_mut("minimum") { + *minimum = json!(minimum.as_f64().unwrap()); } - - #[tokio::test] - async fn lost_bootstrap_status_ack_reuses_the_same_ledger_before_sealing() { - let (_server, store, state) = setup(None, Fault::None).await; - store - .create_anchor( - account().spec, - &crate::providers::signing::ReceiptSigner::from_bytes(&[7; 32]), - ) - .await - .unwrap(); - state.lock().unwrap().fault = Fault::CommitThenFail; - assert!(store.initialize(&root(), "account-uid").await.is_err()); - let pending = snapshot(&state); - check(&pending, Phase::Bootstrap, "False", "BootstrapPending"); - assert!(store.read(&root(), "account-uid").await.is_err()); - let initialized = store.initialize(&root(), "account-uid").await.unwrap(); - check(&initialized, Phase::Active, "True", "LedgerAvailable"); - assert_eq!( - initialized.status.as_ref().unwrap().ledger, - pending.status.as_ref().unwrap().ledger - ); - assert_eq!(initialized.metadata.uid, pending.metadata.uid); + for field in fields.values_mut() { + *field = canonical_schema(field.take()); } } serde_json::Value::Array(items) => { for item in items { - *item = without_descriptions(item.take()); + *item = canonical_schema(item.take()); } } _ => {} @@ -501,14 +458,62 @@ fn helm_budget_account_reporting_matches_generated_schema() { helm["additionalPrinterColumns"], generated["additionalPrinterColumns"] ); - for field in ["phase", "observedGeneration", "conditions"] { + for field in ["phase", "observedGeneration", "conditions", "ledger"] { let path = |schema: &serde_json::Value| { schema["schema"]["openAPIV3Schema"]["properties"]["status"]["properties"][field].clone() }; assert_eq!( - without_descriptions(path(helm)), - without_descriptions(path(generated)), + canonical_schema(path(helm)), + canonical_schema(path(generated)), "{field}" ); } } + +#[tokio::test] +async fn observation_conflict_reloads_concurrently_reserved_money() { + let (_server, store, state) = setup(Some(account()), Fault::ReserveOnConflict).await; + let error = StoreError::Api { + stage: "recovery", + code: Some(503), + }; + let reported = store + .refresh_status(&root(), "account-uid", Some(&error)) + .await + .unwrap(); + check( + &reported, + Phase::Unknown, + "Unknown", + "ReconciliationUnavailable", + ); + let ledger = reported.status.as_ref().unwrap().ledger.as_ref().unwrap(); + assert_eq!(ledger.meters.reserved.tokens, 30); + assert_eq!(ledger.attempts.len(), 1); + assert_eq!(ledger.sessions["a"].next_sequence, 2); + assert_eq!(snapshot(&state).status, reported.status); +} + +#[tokio::test] +async fn lost_bootstrap_status_ack_reuses_the_same_ledger_before_sealing() { + let (_server, store, state) = setup(None, Fault::None).await; + store + .create_anchor( + account().spec, + &crate::providers::signing::ReceiptSigner::from_bytes(&[7; 32]), + ) + .await + .unwrap(); + state.lock().unwrap().fault = Fault::CommitThenFail; + assert!(store.initialize(&root(), "account-uid").await.is_err()); + let pending = snapshot(&state); + check(&pending, Phase::Bootstrap, "False", "BootstrapPending"); + assert!(store.read(&root(), "account-uid").await.is_err()); + let initialized = store.initialize(&root(), "account-uid").await.unwrap(); + check(&initialized, Phase::Active, "True", "LedgerAvailable"); + assert_eq!( + initialized.status.as_ref().unwrap().ledger, + pending.status.as_ref().unwrap().ledger + ); + assert_eq!(initialized.metadata.uid, pending.metadata.uid); +} diff --git a/controller/src/inference_budget/store.rs b/controller/src/inference_budget/store.rs index 3b8ae678e..6fd959993 100644 --- a/controller/src/inference_budget/store.rs +++ b/controller/src/inference_budget/store.rs @@ -375,7 +375,10 @@ impl Store { { Ok(stored) => { Self::validate_identity(&stored, root, account_uid)?; - if stored.status != next.status { + // Kubernetes Time drops subsecond precision on the wire. + if serde_json::to_value(&stored.status).map_err(|_| BudgetError::Corrupt)? + != serde_json::to_value(&next.status).map_err(|_| BudgetError::Corrupt)? + { return Err(BudgetError::Corrupt.into()); } return Ok(stored); From c4291b0f3808747d7fe221af99f47b37e17d16a9 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 05:24:40 +0200 Subject: [PATCH 09/16] fix(budget): admit primary controller workloads with omitted subresource Forward the exact shared admission repair and native probes from 63c9a06403e6ae9264fefd3c3074bd7c55338d99. Preserve the exact controller allowlist, core-only Deployments, namespace fence, and ephemeral-container denial. Helm and Rust keep consuming the same JSON policy. Native composed run 34428691641/job 102719417174 proves all 19 actor and actual controller-chain cases. Original-branch targeted CLI/probe tests (12), TypeScript, Helm and syntax checks pass. Record honest limits: composed standalone remains 23 passed/1 router image-pull failure; budget spending/cancellation and fresh exact-head hosted qualification remain outstanding. No local Cargo, customer deployment, main change or signature waiver. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- cli/src/testing/budget-workload-api.test.ts | 153 +++++++++++ .../files/inference-budget-admission.json | 2 +- .../2026-09-08-governed-inference-budgets.md | 38 +++ tests/e2e/budget-workload-cases.mjs | 250 ++++++++++++++++++ tests/e2e/inference-budget-api.mjs | 5 + 5 files changed, 447 insertions(+), 1 deletion(-) create mode 100644 cli/src/testing/budget-workload-api.test.ts create mode 100644 tests/e2e/budget-workload-cases.mjs diff --git a/cli/src/testing/budget-workload-api.test.ts b/cli/src/testing/budget-workload-api.test.ts new file mode 100644 index 000000000..e2e14dd55 --- /dev/null +++ b/cli/src/testing/budget-workload-api.test.ts @@ -0,0 +1,153 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { readFileSync } from "node:fs"; +import { describe, expect, it } from "vitest"; + +const { workloadCases, policyDenial } = await import( + new URL("../../../tests/e2e/budget-workload-cases.mjs", import.meta.url).href +); +const policy = JSON.parse(readFileSync( + new URL("../../../deploy/helm/kars/files/inference-budget-admission.json", import.meta.url), "utf8", +)).items.find((item: { name: string }) => item.name === "kars-inference-budget-workloads"); +const namespace = "budget-api-fixture"; +const controller = `system:serviceaccount:${namespace}:kars-controller`; +const principal = `system:serviceaccount:${namespace}:untrusted`; + +function rejection(name: string, message = policy.spec.validations[0].message) { + return { status: 422, body: { kind: "Status", status: "Failure", reason: "Invalid", + details: { name, causes: [{ message: + `ValidatingAdmissionPolicy '${policy.name}' with binding '${policy.name}' denied request: ${message}` }] } } }; +} + +class ApiFixture { + objects = new Map(); + calls: Array<{ method: string; path: string; body: any; actor?: string }> = []; + serial = 0; + allowTenant = false; + conflict = false; + + remove(path: string) { + const uid = this.objects.get(path)?.metadata.uid; + this.objects.delete(path); + for (const [child, body] of this.objects) { + if (body.metadata.ownerReferences?.some((owner: any) => owner.uid === uid)) this.remove(child); + } + } + + request = async (method: string, input: string, body?: any, actor?: string) => { + const url = new URL(input, "http://fixture"); + const path = url.pathname; + this.calls.push({ method, path: input, body: structuredClone(body), actor }); + if (method === "GET" && path.includes("/serviceaccounts/")) { + return { status: 200, body: { metadata: { uid: "real-service-account" } } }; + } + if (method === "GET" && url.searchParams.has("labelSelector")) { + return { status: 200, body: { items: [...this.objects.entries()] + .filter(([key, value]) => key.startsWith(path + "/") && value.metadata.labels?.["budget-api-case"] === "actual-chain") + .map(([, value]) => structuredClone(value)) } }; + } + if (method === "GET") { + return this.objects.has(path) ? { status: 200, body: structuredClone(this.objects.get(path)) } + : { status: 404, body: {} }; + } + if (path.endsWith("/subjectaccessreviews")) return { status: 201, body: { status: { allowed: true } } }; + if (method === "DELETE") { + const existing = this.objects.get(path); + expect(body.preconditions).toEqual({ uid: existing.metadata.uid, resourceVersion: existing.metadata.resourceVersion }); + if (this.conflict) { + this.conflict = false; + existing.metadata.resourceVersion = "2"; + return { status: 409, body: {} }; + } + this.remove(path); + return { status: 200, body: {} }; + } + if (url.searchParams.get("dryRun") === "All") { + if (this.allowTenant && actor === principal) return { status: 201, body }; + return rejection(body.metadata.name); + } + if (method === "PATCH") { + const existing = this.objects.get(path); + expect(body.metadata.uid).toBe(existing.metadata.uid); + existing.metadata.labels = body.metadata.labels; + return { status: 200, body: structuredClone(existing) }; + } + expect(method).toBe("POST"); + const created = structuredClone(body); + created.metadata = { ...body.metadata, uid: `native-${++this.serial}`, resourceVersion: "1" }; + this.objects.set(path + "/" + body.metadata.name, created); + if (body.kind === "Deployment") { + const ns = body.metadata.namespace; + const labels = { "budget-api-case": "actual-chain" }; + this.objects.set(`/apis/apps/v1/namespaces/${ns}/replicasets/actual-rs`, { + kind: "ReplicaSet", metadata: { name: "actual-rs", uid: "controller-rs", labels, + ownerReferences: [{ controller: true, uid: created.metadata.uid }] }, + }); + this.objects.set(`/api/v1/namespaces/${ns}/pods/actual-pod`, { + kind: "Pod", metadata: { name: "actual-pod", uid: "controller-pod", labels, + ownerReferences: [{ controller: true, uid: "controller-rs" }] }, + spec: { automountServiceAccountToken: false }, + }); + } + return { status: 201, body: structuredClone(created) }; + }; +} + +const wait = async (check: () => Promise) => expect(await check()).toBeTruthy(); +const options = { namespace, controller, principal, policy }; + +describe("native budget workload proof orchestration (not native execution evidence)", () => { + it("keeps the exact shared primary/ephemeral distinction and controller allowlist", () => { + const spec = policy.spec; + expect(spec.validations).toHaveLength(1); + expect(spec.validations[0].expression).toContain("request.?subResource.orValue('') != 'ephemeralcontainers'"); + expect(spec.validations[0].expression).toContain("request.resource.resource != 'deployments'"); + expect(spec.validations[0].expression).toContain( + "['system:kube-controller-manager', 'system:serviceaccount:kube-system:deployment-controller', 'system:serviceaccount:kube-system:replicaset-controller']"); + expect(spec.namespaceSelector).toBeUndefined(); + expect(spec.matchConstraints.namespaceSelector).toEqual({ matchLabels: { "kars.azure.com/inference-budget": "v1" } }); + expect(spec.failurePolicy).toBe("Fail"); + }); + + it("requires exact native policy denial rather than RBAC, missing-field or other errors", () => { + expect(policyDenial(rejection("fixture"), policy, "fixture")).toBe(true); + expect(policyDenial({ status: 403, body: { kind: "Status", reason: "Forbidden" } }, policy, "fixture")).toBe(false); + expect(policyDenial(rejection("other"), policy, "fixture")).toBe(false); + expect(policyDenial(rejection("fixture", "evaluation failed: no such key: subResource"), policy, "fixture")).toBe(false); + }); + + it("exercises the native-controller chain, primary actor matrix and exact forbidden paths", async () => { + const api = new ApiFixture(); + const results: any[] = []; + await workloadCases(api.request, wait, options, (result: any) => results.push(result)); + expect(results).toHaveLength(19); + expect(results.every(result => result.matched)).toBe(true); + expect(results.find(result => result.case === "actual-controller-chain-created").readinessClaimed).toBe(false); + expect(results.filter(result => result.case.endsWith("-primary"))).toHaveLength(8); + expect(results.filter(result => result.case.endsWith("-ephemeral-denied"))).toHaveLength(4); + expect(api.calls.some(call => call.path.endsWith("/status"))).toBe(false); + expect(api.calls.every(call => !call.body?.status || call.body.kind === "SubjectAccessReview")).toBe(true); + for (const call of api.calls.filter(call => call.body?.kind === "Pod" && call.method === "POST")) { + expect(call.body.spec.automountServiceAccountToken).toBe(false); + expect(call.body.spec.schedulerName).toBe("budget-api-never-schedule"); + } + expect(api.objects.size).toBe(0); + }); + + it("fails if a forbidden tenant path is admitted, without leaving owned fixtures", async () => { + const api = new ApiFixture(); + api.allowTenant = true; + await expect(workloadCases(api.request, wait, options, () => {})).rejects.toThrow("Wrong native denial"); + expect(api.objects.size).toBe(0); + }); + + it("retries cleanup conflicts with fresh UID/RV rather than discarding the fence", async () => { + const api = new ApiFixture(); + api.conflict = true; + await workloadCases(api.request, wait, options, () => {}); + expect(api.objects.size).toBe(0); + expect(api.calls.filter(call => call.method === "DELETE").some( + call => call.body.preconditions.resourceVersion === "2")).toBe(true); + }); +}); diff --git a/deploy/helm/kars/files/inference-budget-admission.json b/deploy/helm/kars/files/inference-budget-admission.json index b2a01b42e..dad745d8e 100644 --- a/deploy/helm/kars/files/inference-budget-admission.json +++ b/deploy/helm/kars/files/inference-budget-admission.json @@ -112,7 +112,7 @@ }] }, "validations": [{ - "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller' || (request.resource.resource != 'deployments' && request.subResource != 'ephemeralcontainers' && request.userInfo.username in ['system:kube-controller-manager', 'system:serviceaccount:kube-system:deployment-controller', 'system:serviceaccount:kube-system:replicaset-controller'])", + "expression": "request.userInfo.username == 'system:serviceaccount:__ACCOUNTING_NAMESPACE__:kars-controller' || (request.resource.resource != 'deployments' && request.?subResource.orValue('') != 'ephemeralcontainers' && request.userInfo.username in ['system:kube-controller-manager', 'system:serviceaccount:kube-system:deployment-controller', 'system:serviceaccount:kube-system:replicaset-controller'])", "message": "Budget runtime workloads may be produced only by the core controller and Kubernetes Deployment/ReplicaSet controllers" }] } diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index e07ddd03b..87f5eb1c2 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -49,6 +49,44 @@ all-in task spend, invoice accuracy, taxes, or exchange rates. ## Verification +### Native primary-workload repair (2026-09-10) + +Standalone qualification found a real ReplicaSet admission failure: +`AdmissionRequest.subResource` is absent on primary operations. The shared +`kars-inference-budget-workloads` expression now uses +`request.?subResource.orValue('')` rather than a raw field access. The exact +controller usernames, core-only Deployment restriction, namespace selector, +ephemeral-container restriction, failure policy and all other predicates remain +unchanged. Helm and Rust consume the same JSON; no second policy copy was added. + +The isolated repair and native regressions originate at +`63c9a06403e6ae9264fefd3c3074bd7c55338d99`. Actual Kubernetes qualification at +composed head `c8f7a143a3242484ea782d5c3170ae00295074f6` passed all **19 workload +records**, including the real Deployment/ReplicaSet/Pod owner-UID chain, allowed +primary controller creation, and exact intended tenant/Deployment/ephemeral +denials. The probe proves each principal has the necessary RBAC rather than +mistaking an unrelated authorization error for admission enforcement. The +tokenless, deliberately unscheduled fixtures do not claim runtime readiness. + +Evidence: [native API job 102719417174](https://github.com/Azure/kars/actions/runs/34428691641/job/102719417174). +The same composed run passed hosted Rust, CLI, schema and benchmark jobs. +Its separate [standalone job](https://github.com/Azure/kars/actions/runs/34428691641/job/102725472891) +finished **23 passed / 1 failed**: budget Pods are created, but the router reports +`ErrImagePull` / `ImagePullBackOff`. Spending and cancellation assertions were +not reached. No image-resolution cause, full SRE or CNI enforcement is inferred. + +Forwarding the exact repair to this budget branch passed 12 focused CLI/probe +tests, TypeScript checking, native JavaScript syntax, Helm lint and diff checks. +No local Cargo build was used for this forward. This is bounded admission +closure, not complete budget enforcement or new-head hosted qualification. +Required human signatures, exact composed-runtime qualification and publication +approval remain outstanding. + +### Initial authoring snapshot (historical) + +The table and initial local-build statements below record the earlier authoring +state. They are not the current result of the native workload repair above. + | Area | Evidence | Status | |---|---|---| | Shared arithmetic/state engine | UID ancestry, ancestor reservation, BeginDispatch, settlement, cancellation, replay and breach tests written | Rust execution pending | diff --git a/tests/e2e/budget-workload-cases.mjs b/tests/e2e/budget-workload-cases.mjs new file mode 100644 index 000000000..e3405330c --- /dev/null +++ b/tests/e2e/budget-workload-cases.mjs @@ -0,0 +1,250 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Native admission/creation proof only. No account state or Pod readiness is fabricated. +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; + +const kubeControllers = [ + ["manager", "system:kube-controller-manager"], + ["deployment", "system:serviceaccount:kube-system:deployment-controller"], + ["replicaset", "system:serviceaccount:kube-system:replicaset-controller"], +]; +const api = "/api/v1"; +const apps = "/apis/apps/v1"; +const rbac = "/apis/rbac.authorization.k8s.io/v1"; + +export function policyDenial(response, policy, name) { + const validation = policy.spec.validations[0]; + const reason = validation.reason ?? "Invalid"; + const expected = { Invalid: 422, Forbidden: 403 }[reason]; + const body = response.body; + const message = `ValidatingAdmissionPolicy '${policy.name}' with binding '${policy.name}' denied request: ${validation.message}`; + return response.status === expected && body?.kind === "Status" && body.status === "Failure" + && body.reason === reason && body.details?.name === name + && body.details?.causes?.some(cause => cause.message === message) === true; +} + +export async function workloadCases(request, wait, { namespace, controller, principal, policy }, report) { + assert(policy?.name === "kars-inference-budget-workloads" && policy.spec.validations.length === 1, + "Expected the exact shared budget workload policy"); + const protectedNamespace = `${namespace}-runtime`; + const owned = []; + let primaryFailure = false; + const emit = (name, response, expected) => report({ + case: name, httpStatus: response.status, expectedStatus: expected, matched: response.status === expected, + }); + const create = async (path, body, actor) => { + const response = await request("POST", path, body, actor); + assert(response.status === 201 && response.body?.kind === body.kind + && response.body.metadata?.name === body.metadata.name + && response.body.metadata.namespace === body.metadata.namespace + && typeof response.body.metadata.uid === "string" + && typeof response.body.metadata.resourceVersion === "string", "Native workload fixture CREATE failed"); + owned.push([`${path}/${body.metadata.name}`, response.body.metadata.uid]); + return response.body; + }; + const podSpec = { + automountServiceAccountToken: false, schedulerName: "budget-api-never-schedule", + containers: [{ name: "probe", image: "registry.invalid/budget-api:never", imagePullPolicy: "Never" }], + }; + const workload = (kind, name, label = name) => ({ + apiVersion: kind === "Pod" ? "v1" : "apps/v1", kind, + metadata: { name, namespace: protectedNamespace, labels: { "budget-api-case": label } }, + spec: kind === "Pod" ? structuredClone(podSpec) : { + replicas: kind === "Deployment" ? 1 : 0, + selector: { matchLabels: { "budget-api-case": label } }, + template: { metadata: { labels: { "budget-api-case": label } }, spec: structuredClone(podSpec) }, + }, + }); + const collection = kind => kind === "Pod" + ? `${api}/namespaces/${protectedNamespace}/pods` + : `${apps}/namespaces/${protectedNamespace}/${kind === "Deployment" ? "deployments" : "replicasets"}`; + const expectDenied = async (caseName, method, path, body, actor) => { + const response = await request(method, path + "?dryRun=All", body, actor); + assert(policyDenial(response, policy, body.metadata.name), `Wrong native denial for ${caseName}`); + emit(caseName, response, policy.spec.validations[0].reason === "Forbidden" ? 403 : 422); + }; + try { + for (const actor of [controller, principal, ...kubeControllers.map(([, actor]) => actor)]) { + if (actor.startsWith("system:serviceaccount:")) { + const [, , ns, name] = actor.split(":"); + const actual = await request("GET", `${api}/namespaces/${ns}/serviceaccounts/${name}`); + assert(actual.status === 200 && actual.body?.metadata?.uid, "Native controller ServiceAccount missing"); + } + } + const ns = await create(`${api}/namespaces`, { + apiVersion: "v1", kind: "Namespace", metadata: { name: protectedNamespace }, + }); + const labelRole = `${protectedNamespace}-label`; + await create(`${rbac}/clusterroles`, { + apiVersion: "rbac.authorization.k8s.io/v1", kind: "ClusterRole", metadata: { name: labelRole }, + rules: [{ apiGroups: [""], resources: ["namespaces"], resourceNames: [protectedNamespace], verbs: ["get", "patch"] }], + }); + await create(`${rbac}/clusterrolebindings`, { + apiVersion: "rbac.authorization.k8s.io/v1", kind: "ClusterRoleBinding", metadata: { name: labelRole }, + roleRef: { apiGroup: "rbac.authorization.k8s.io", kind: "ClusterRole", name: labelRole }, + subjects: [{ kind: "User", apiGroup: "rbac.authorization.k8s.io", name: controller }], + }); + await create(`${rbac}/namespaces/${protectedNamespace}/roles`, { + apiVersion: "rbac.authorization.k8s.io/v1", kind: "Role", + metadata: { name: "admission-proof", namespace: protectedNamespace }, + rules: [ + { apiGroups: [""], resources: ["pods", "pods/ephemeralcontainers"], verbs: ["get", "create", "update", "patch"] }, + { apiGroups: ["apps"], resources: ["deployments", "replicasets"], verbs: ["get", "create"] }, + ], + }); + const actors = [controller, principal, ...kubeControllers.map(([, actor]) => actor)]; + await create(`${rbac}/namespaces/${protectedNamespace}/rolebindings`, { + apiVersion: "rbac.authorization.k8s.io/v1", kind: "RoleBinding", + metadata: { name: "admission-proof", namespace: protectedNamespace }, + roleRef: { apiGroup: "rbac.authorization.k8s.io", kind: "Role", name: "admission-proof" }, + subjects: actors.map(name => ({ kind: "User", apiGroup: "rbac.authorization.k8s.io", name })), + }); + const permission = async (actor, group, resource, verb, subresource, ns = protectedNamespace, name) => { + const response = await request("POST", "/apis/authorization.k8s.io/v1/subjectaccessreviews", { + apiVersion: "authorization.k8s.io/v1", kind: "SubjectAccessReview", + spec: { user: actor, resourceAttributes: { namespace: ns, group, resource, verb, + ...(subresource ? { subresource } : {}), ...(name ? { name } : {}) } }, + }); + return response.status === 201 && response.body?.status?.allowed === true + && !response.body.status.evaluationError; + }; + await wait(() => permission(controller, "", "namespaces", "patch", undefined, "", protectedNamespace), + "core fixture namespace-label authority"); + const labeled = await request("PATCH", `${api}/namespaces/${protectedNamespace}`, { + metadata: { uid: ns.metadata.uid, resourceVersion: ns.metadata.resourceVersion, + labels: { "kars.azure.com/inference-budget": "v1" } }, + }, controller); + assert(labeled.status === 200 && labeled.body?.metadata?.uid === ns.metadata.uid + && labeled.body.metadata.labels?.["kars.azure.com/inference-budget"] === "v1", + "Native budget namespace fence was not installed"); + for (const actor of actors) { + for (const [group, resource, verb, subresource] of [ + ["", "pods", "create"], ["apps", "deployments", "create"], ["apps", "replicasets", "create"], + ["", "pods", "update", "ephemeralcontainers"], + ]) { + await wait(() => permission(actor, group, resource, verb, subresource), "workload fixture RBAC"); + } + } + await wait(async () => { + const response = await request("GET", `${api}/namespaces/${protectedNamespace}/serviceaccounts/default`); + return response.status === 200 && response.body?.metadata?.uid; + }, "native default ServiceAccount"); + // This is the real built-in controller chain, not an impersonated/fabricated child. + const deployment = await create(collection("Deployment"), workload("Deployment", "actual-chain"), controller); + emit("core-deployment-primary", { status: 201 }, 201); + let replicaSet, pod; + await wait(async () => { + const response = await request("GET", collection("ReplicaSet") + "?labelSelector=budget-api-case%3Dactual-chain"); + replicaSet = response.body?.items?.find(item => item.metadata?.ownerReferences?.some( + owner => owner.controller === true && owner.uid === deployment.metadata.uid)); + return response.status === 200 && replicaSet?.metadata?.uid; + }, "native Deployment controller creates budget ReplicaSet"); + await wait(async () => { + const response = await request("GET", collection("Pod") + "?labelSelector=budget-api-case%3Dactual-chain"); + pod = response.body?.items?.find(item => item.metadata?.ownerReferences?.some( + owner => owner.controller === true && owner.uid === replicaSet.metadata.uid)); + return response.status === 200 && pod?.metadata?.uid; + }, "native ReplicaSet controller creates budget Pod"); + assert(pod.spec.automountServiceAccountToken === false && !pod.metadata.deletionTimestamp, + "Native chain must retain its non-executing tokenless fixture"); + report({ case: "actual-controller-chain-created", matched: true, readinessClaimed: false }); + + const target = await create(collection("Pod"), workload("Pod", "ephemeral-target"), controller); + emit("core-pod-primary", { status: 201 }, 201); + for (const [label, actor] of kubeControllers) { + for (const kind of ["ReplicaSet", "Pod"]) { + const created = await create(collection(kind), workload(kind, `${label}-${kind.toLowerCase()}`), actor); + emit(`${label}-${kind.toLowerCase()}-primary`, { status: 201 }, 201); + assert(created.metadata.uid, "Native primary creation omitted UID"); + } + await expectDenied(`${label}-deployment-denied`, "POST", collection("Deployment"), + workload("Deployment", `denied-${label}-deployment`), actor); + } + for (const kind of ["Deployment", "ReplicaSet", "Pod"]) { + await expectDenied(`tenant-${kind.toLowerCase()}-denied`, "POST", collection(kind), + workload(kind, `denied-tenant-${kind.toLowerCase()}`), principal); + } + for (const [label, actor] of [...kubeControllers, ["tenant", principal]]) { + const response = await request("GET", collection("Pod") + "/" + target.metadata.name); + assert(response.status === 200 && response.body?.metadata?.uid === target.metadata.uid, + "Ephemeral-container target changed"); + const update = structuredClone(response.body); + update.spec.ephemeralContainers = [{ + name: "denied-ephemeral", image: "registry.invalid/budget-api:never", + imagePullPolicy: "Never", targetContainerName: "probe", + }]; + await expectDenied(`${label}-ephemeral-denied`, "PUT", + collection("Pod") + "/" + target.metadata.name + "/ephemeralcontainers", update, actor); + } + } catch (error) { + primaryFailure = true; + throw error; + } finally { + let cleanupFailed = false; + for (const [path, uid] of owned.reverse()) { + try { + let removed = false; + for (let attempt = 0; attempt < 3; attempt++) { + const current = await request("GET", path); + if (current.status === 404) { removed = true; break; } + assert(current.status === 200 && current.body?.metadata?.uid === uid, "Workload fixture cleanup UID changed"); + const result = await request("DELETE", path, { + apiVersion: "v1", kind: "DeleteOptions", + preconditions: { uid, resourceVersion: current.body.metadata.resourceVersion }, + propagationPolicy: "Background", + }); + if (result.status === 409) continue; + assert([200, 202].includes(result.status), "Workload fixture cleanup failed"); + removed = true; + break; + } + assert(removed, "Workload fixture cleanup contention exceeded its bound"); + } catch { + cleanupFailed = true; + } + } + if (cleanupFailed && !primaryFailure) throw new Error("Native workload fixture cleanup incomplete"); + } +} + +export async function runWorkloadProof(options) { + const { root, context, kubectl, until } = options; + const config = JSON.parse(kubectl(["config", "view", "--minify", "-o", "json"])); + assert(config.contexts?.length === 1 && config.contexts[0].name === context + && config.clusters?.length === 1 + && ["localhost", "127.0.0.1", "[::1]"].includes(new URL(config.clusters[0].cluster.server).hostname), + "Budget workload proof requires the exact loopback Kind context"); + const proxy = spawn("kubectl", ["--context", context, "--request-timeout=20s", "proxy", + "--address=127.0.0.1", "--port=0"], { cwd: root, stdio: ["ignore", "pipe", "pipe"] }); + let output = "", port; + proxy.stdout.on("data", data => { output = (output + data).slice(-2048); }); + proxy.stderr.on("data", () => {}); + try { + await until(() => { + assert(proxy.exitCode === null, "Budget API proxy exited"); + port = output.match(/127\.0\.0\.1:(\d+)/)?.[1]; + return Boolean(port); + }, "budget workload API proxy"); + const request = async (method, path, body, actor) => { + const response = await fetch(`http://127.0.0.1:${port}${path}`, { + method, signal: AbortSignal.timeout(15_000), + headers: { "Content-Type": method === "PATCH" ? "application/merge-patch+json" : "application/json", + Accept: "application/json", ...(actor ? { "Impersonate-User": actor, "Impersonate-Group": "system:authenticated" } : {}) }, + body: body === undefined ? undefined : JSON.stringify(body), + }); + return { status: response.status, body: await response.json().catch(() => null) }; + }; + await workloadCases(request, until, options, + result => console.log("BUDGET-WORKLOAD " + JSON.stringify(result))); + } finally { + if (proxy.exitCode === null) { + proxy.kill("SIGTERM"); + await new Promise(resolve => { + const timer = setTimeout(() => { proxy.kill("SIGKILL"); resolve(); }, 5000); + proxy.once("exit", () => { clearTimeout(timer); resolve(); }); + }); + } + } +} diff --git a/tests/e2e/inference-budget-api.mjs b/tests/e2e/inference-budget-api.mjs index 9a2111346..55a60cc97 100644 --- a/tests/e2e/inference-budget-api.mjs +++ b/tests/e2e/inference-budget-api.mjs @@ -8,6 +8,7 @@ import { execFileSync } from "node:child_process"; import { readFileSync } from "node:fs"; import { createRequire } from "node:module"; import { fileURLToPath } from "node:url"; +import { runWorkloadProof } from "./budget-workload-cases.mjs"; const require = createRequire(new URL("../../cli/package.json", import.meta.url)); const { parseAllDocuments } = require("yaml"); @@ -173,4 +174,8 @@ assert.equal(review.status.authenticated, true); assert.deepEqual(review.status.audiences, [audience]); assert.deepEqual(review.status.user.extra["authentication.kubernetes.io/pod-uid"], [pod.metadata.uid]); assert.deepEqual(review.status.user.extra["authentication.kubernetes.io/pod-name"], ["identity"]); +await runWorkloadProof({ + root, context, kubectl, until, namespace, controller, principal, + policy: shared.items.find(policy => policy.name === "kars-inference-budget-workloads"), +}); console.log("Budget public CRDs/CEL, controller-only accounting/CA, private audience issuance and UID claims passed on disposable Kubernetes v1.31"); From e89040dfa76ae3e3b015f2ec7bc39b7c3b24ffc3 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 16:52:27 +0200 Subject: [PATCH 10/16] fix(budget): fence recovery revocation against refreshed authority Recheck the full captured TaskAuthority inside Store.transact before stale-source recovery closes a subtree, including each UID/resourceVersion CAS retry. If a newer enrollment is current, make no ledger mutation and defer its explicit live Task check to the next scan. Unchanged invalid authority still revokes, and accepted old work stays fully funded. Add deterministic wiremock interleavings for B enrollment/new Pod/InFlight work before the revocation transaction and during a 409 retry. Assert B survives, the next scan rechecks its live source, subsequent invalid B revokes, unchanged invalid A still revokes, and old/new maximum liabilities, limits and account UID remain intact. Reuse the existing Store test server. Validation completed without Cargo: rustfmt check for recovery source and regression tests, diff checks, and byte-exact preservation of all seven native fixture files from 5ae7b99a. Rust test execution is pending the observer-held sole Cargo lease; requested batch: cargo test --locked --offline --package kars-controller inference_budget. No new target, cleanup, installs, push or CI dispatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/inference_budget/recovery.rs | 16 +- .../src/inference_budget/recovery_tests.rs | 362 ++++++++++++++++++ .../src/inference_budget/store_tests.rs | 3 + docs/governed-inference-budgets.md | 6 + 4 files changed, 386 insertions(+), 1 deletion(-) create mode 100644 controller/src/inference_budget/recovery_tests.rs diff --git a/controller/src/inference_budget/recovery.rs b/controller/src/inference_budget/recovery.rs index a69b029f0..47d2465ca 100644 --- a/controller/src/inference_budget/recovery.rs +++ b/controller/src/inference_budget/recovery.rs @@ -7,7 +7,7 @@ use super::{ store::{Store, StoreError}, }; use crate::{ - inference_budget_contract::{BudgetError, RootKind}, + inference_budget_contract::{BudgetError, RootKind, ledger::Mutation}, kars_task::KarsTask, kars_team::KarsTeam, }; @@ -178,6 +178,20 @@ async fn recover( }) { store .transact(root, &uid, |ledger| { + let current = ledger + .nodes + .get(&node.authority.task.uid) + .ok_or(BudgetError::Identity)?; + if current.authority != node.authority { + // A newer enrollment may already have funded sessions. + // Defer its live-source check to the next recovery scan, + // including when the authority changed during CAS retry. + return Ok(Mutation { + next: ledger.clone(), + value: (), + changed: false, + }); + } ledger.close_subtree(&node.authority.task.uid) }) .await?; diff --git a/controller/src/inference_budget/recovery_tests.rs b/controller/src/inference_budget/recovery_tests.rs new file mode 100644 index 000000000..6d9da8f1a --- /dev/null +++ b/controller/src/inference_budget/recovery_tests.rs @@ -0,0 +1,362 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{ + inference_budget::{recovery::reconcile_account, status::project}, + inference_budget_contract::AttemptPhase, + kars_task::{ + KarsTask, KarsTaskSpec, TaskBlueprint, TaskBudget, TaskEnvelope, TaskExecution, TaskModel, + }, +}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; + +fn task_authority(revision: &str) -> (KarsTask, TaskAuthority) { + let model = TaskModel { + provider: "azure-openai".into(), + deployment: "fixture".into(), + }; + let mut task = KarsTask::new( + "root", + KarsTaskSpec { + objective: "recovery interleaving".into(), + envelope: TaskEnvelope { + tier: 3, + authority_ceiling: 3, + delegation_depth: 2, + budget: Some(TaskBudget { + scope: Some(BudgetScope::GovernedInference), + tokens: Some(100), + usd_micros: Some(20), + }), + ..Default::default() + }, + execution: Some(TaskExecution { + launch: true, + runtime: None, + }), + blueprint: Some(TaskBlueprint { + model: Some(model.clone()), + instructions: Some(revision.into()), + ..Default::default() + }), + ..Default::default() + }, + ); + task.metadata.namespace = Some("workspace".into()); + task.metadata.uid = Some("root".into()); + task.metadata.resource_version = Some(revision.into()); + let authority = TaskAuthority { + task: root().resource, + parent_uid: None, + root_task_uid: "root".into(), + authorization_digest: task.envelope_digest(), + effective_authorization: task.spec.authorization_configuration_with_model(&model), + limits: Limits { + tokens: Some(100), + usd_micros: Some(20), + }, + }; + (task, authority) +} + +fn inflight(ledger: Ledger, authority: &TaskAuthority, pod: &str, now: i64) -> Ledger { + let mut request = reserve(pod); + request.identity.authorization_digest = authority.authorization_digest.clone(); + request.identity.sandbox.name = authority.task.name.clone(); + let command = AttemptCommand { + account_uid: request.account_uid.clone(), + key: AttemptKey { + pod_uid: pod.into(), + sequence: request.sequence, + }, + identity: request.identity.clone(), + wire_digest: request.wire_digest.clone(), + }; + ledger + .register_session(request.identity.clone()) + .unwrap() + .next + .reserve(&request, now) + .unwrap() + .next + .begin_dispatch(&command, now) + .unwrap() + .next +} + +fn original_account(now: i64) -> KarsBudgetAccount { + let (_, authority) = task_authority("A"); + let mut original = account(); + original.spec.limits = authority.limits; + let ledger = Ledger::new("account-uid".into(), root(), authority.limits) + .unwrap() + .register_task(authority.clone()) + .unwrap() + .next; + original.status.as_mut().unwrap().ledger = Some(inflight(ledger, &authority, "old", now)); + original.status = Some(project(&original, None)); + original +} + +fn enroll_new_authority(ledger: &Ledger, authority: &TaskAuthority, now: i64) -> Ledger { + // These are ensure_task's close/update/resume transitions, followed by a + // genuinely new Pod session and accepted work under the refreshed authority. + let next = ledger + .close_subtree(&authority.task.uid) + .unwrap() + .next + .update_authority(authority.clone()) + .unwrap() + .next + .resume_task(&authority.task.uid, &authority.authorization_digest) + .unwrap() + .next; + inflight(next, authority, "new", now) +} + +#[derive(Clone)] +struct Enrollment { + state: Arc>, + authority: TaskAuthority, + now: i64, + occurred: Arc, +} + +impl Enrollment { + fn install(&self) { + assert!(!self.occurred.swap(true, Ordering::SeqCst)); + let mut state = self.state.lock().unwrap(); + let next = enroll_new_authority( + Store::ledger(state.account.as_ref().unwrap()).unwrap(), + &self.authority, + self.now, + ); + state.version += 1; + state.writes += 1; + let version = state.version.to_string(); + let account = state.account.as_mut().unwrap(); + account.metadata.resource_version = Some(version); + account.status.as_mut().unwrap().ledger = Some(next); + account.status = Some(project(account, None)); + } +} + +struct LiveTask { + task: Arc>, + reads: Arc, + enrollment: Option, +} + +impl Respond for LiveTask { + fn respond(&self, _: &Request) -> ResponseTemplate { + // Root-lifetime GET precedes recovery's ledger snapshot; the per-node + // authority GET follows it. Interleave deterministically at the latter. + if self.reads.fetch_add(1, Ordering::SeqCst) == 1 + && let Some(enrollment) = &self.enrollment + { + enrollment.install(); + } + ResponseTemplate::new(200).set_body_json(&*self.task.lock().unwrap()) + } +} + +struct Conflict(Enrollment); + +impl Respond for Conflict { + fn respond(&self, request: &Request) -> ResponseTemplate { + let incoming: KarsBudgetAccount = serde_json::from_slice(&request.body).unwrap(); + if !self.0.occurred.load(Ordering::SeqCst) + && !Store::ledger(&incoming).unwrap().nodes["root"].active + { + { + let state = self.0.state.lock().unwrap(); + let current = state.account.as_ref().unwrap(); + assert_eq!(incoming.metadata.uid, current.metadata.uid); + assert_eq!( + incoming.metadata.resource_version, + current.metadata.resource_version + ); + } + self.0.install(); + return failure(409); + } + Server(self.0.state.clone()).respond(request) + } +} + +async fn live_sources( + server: &MockServer, + task: Arc>, + enrollment: Option, +) -> Arc { + let reads = Arc::new(AtomicUsize::new(0)); + for (path, object) in [ + ( + "/api/v1/namespaces/workspace", + json!({"apiVersion":"v1", "kind":"Namespace", + "metadata":{"name":"workspace", "uid":"workspace-uid"}}), + ), + ( + "/api/v1/namespaces/kars-root/pods/pod-new", + json!({"apiVersion":"v1", "kind":"Pod", + "metadata":{"name":"pod-new", "namespace":"kars-root", "uid":"new"}}), + ), + ] { + Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path(path)) + .respond_with(ResponseTemplate::new(200).set_body_json(object)) + .with_priority(1) + .mount(server) + .await; + } + Mock::given(wiremock::matchers::method("GET")) + .and(wiremock::matchers::path( + "/apis/kars.azure.com/v1alpha1/namespaces/workspace/karstasks/root", + )) + .respond_with(LiveTask { + task, + reads: reads.clone(), + enrollment, + }) + .with_priority(1) + .mount(server) + .await; + reads +} + +async fn changed_authority_survives(conflict: bool) { + let now = chrono::Utc::now().timestamp(); + let original = original_account(now); + let (task, authority) = task_authority("B"); + let expected = enroll_new_authority(Store::ledger(&original).unwrap(), &authority, now); + let (server, store, state) = setup(Some(original.clone()), Fault::None).await; + let enrollment = Enrollment { + state: state.clone(), + authority, + now, + occurred: Arc::new(AtomicBool::new(false)), + }; + let live = Arc::new(Mutex::new(task)); + let reads = live_sources( + &server, + live.clone(), + (!conflict).then_some(enrollment.clone()), + ) + .await; + if conflict { + Mock::given(wiremock::matchers::method("PUT")) + .and(wiremock::matchers::path(format!("{OBJECT}/status"))) + .respond_with(Conflict(enrollment.clone())) + .with_priority(1) + .mount(&server) + .await; + } + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + reconcile_account(&client, &store, &original).await.unwrap(); + assert!(enrollment.occurred.load(Ordering::SeqCst)); + assert_eq!(reads.load(Ordering::SeqCst), 2); + let current = store.read(&root(), "account-uid").await.unwrap(); + assert_eq!(Store::ledger(¤t).unwrap(), &expected); + assert!(expected.nodes["root"].active); + assert!(expected.sessions["old"].closed); + assert!(!expected.sessions["new"].closed); + assert_eq!( + expected.meters.uncertain, + Amounts { + tokens: 30, + usd_micros: 5 + } + ); + assert_eq!( + expected.meters.reserved, + Amounts { + tokens: 30, + usd_micros: 5 + } + ); + for (pod, phase) in [ + ("old", AttemptPhase::Uncertain), + ("new", AttemptPhase::InFlight), + ] { + let key = AttemptKey { + pod_uid: pod.into(), + sequence: 1, + } + .storage_key(); + assert_eq!(expected.attempts[&key].phase, phase); + } + + // Deferral must actually re-read B's live source on the next scan, and must + // still revoke B if that source becomes invalid without another enrollment. + reconcile_account(&client, &store, ¤t).await.unwrap(); + assert_eq!(reads.load(Ordering::SeqCst), 4); + let current = store.read(&root(), "account-uid").await.unwrap(); + assert_eq!(Store::ledger(¤t).unwrap(), &expected); + *live.lock().unwrap() = task_authority("C").0; + reconcile_account(&client, &store, ¤t).await.unwrap(); + assert_eq!(reads.load(Ordering::SeqCst), 6); + let closed = store.read(&root(), "account-uid").await.unwrap(); + let ledger = Store::ledger(&closed).unwrap(); + assert!(!ledger.nodes["root"].active); + assert!(ledger.sessions.values().all(|session| session.closed)); + assert_eq!( + ledger.meters.total().unwrap(), + expected.meters.total().unwrap() + ); + assert_eq!(ledger.meters.reserved, Amounts::default()); + assert_eq!( + ledger.meters.uncertain, + Amounts { + tokens: 60, + usd_micros: 10 + } + ); + assert_eq!(ledger.limits, expected.limits); + assert_eq!(closed.metadata.uid, original.metadata.uid); + assert_ne!( + closed.metadata.resource_version, + original.metadata.resource_version + ); +} + +#[tokio::test] +async fn recovery_authority_change_before_transaction_preserves_new_session_and_inflight_work() { + changed_authority_survives(false).await; +} + +#[tokio::test] +async fn recovery_authority_change_on_cas_retry_preserves_new_session_and_inflight_work() { + changed_authority_survives(true).await; +} + +#[tokio::test] +async fn recovery_unchanged_authority_with_invalid_live_source_still_revokes_and_funds_old_work() { + let now = chrono::Utc::now().timestamp(); + let original = original_account(now); + let (server, store, state) = setup(Some(original.clone()), Fault::None).await; + live_sources(&server, Arc::new(Mutex::new(task_authority("B").0)), None).await; + let client = Client::try_from(kube::Config::new(server.uri().parse().unwrap())).unwrap(); + reconcile_account(&client, &store, &original).await.unwrap(); + let current = state.lock().unwrap().account.clone().unwrap(); + let ledger = Store::ledger(¤t).unwrap(); + assert_eq!( + ledger, + &Store::ledger(&original) + .unwrap() + .close_subtree("root") + .unwrap() + .next + ); + assert!(!ledger.nodes["root"].active); + assert!(ledger.sessions["old"].closed); + assert_eq!( + ledger.meters.uncertain, + Amounts { + tokens: 30, + usd_micros: 5 + } + ); + assert_eq!(ledger.meters.reserved, Amounts::default()); + assert_eq!(current.metadata.uid, original.metadata.uid); +} diff --git a/controller/src/inference_budget/store_tests.rs b/controller/src/inference_budget/store_tests.rs index 20242f654..69cf62687 100644 --- a/controller/src/inference_budget/store_tests.rs +++ b/controller/src/inference_budget/store_tests.rs @@ -511,3 +511,6 @@ async fn replacement_between_read_and_commit_cannot_receive_a_grant() { #[path = "status_tests.rs"] mod status_tests; + +#[path = "recovery_tests.rs"] +mod recovery_tests; diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index 4565f5644..1309ef30b 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -148,6 +148,12 @@ remains unchanged. Neither a Phase value nor `Ready=True` is spend authority: every send still requires live Pod/Task authority, the sealed validated ledger, an operator contract, and an atomic reservation/begin transition. +Recovery revokes a stale Task authority only while that exact captured authority +is still current inside the account CAS, including after resourceVersion retries. +If concurrent enrollment installed a newer authority, recovery leaves its sessions +and funded work intact and checks its live Task source on the next scan. Old +accepted work remains conservatively charged; deferral never refunds liabilities. + ## Operator contracts and unavoidable configuration Enable the optional Helm `inferenceBudget` section only after supplying: From 2a95d02a5baf22c4c3381c4c0d5f9c49f0501550 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 17:07:06 +0200 Subject: [PATCH 11/16] docs(budgets): record delegated source review and recovery closure Record the maintainer-authorized, explicitly AI-attributed audit after focused review and deterministic recovery CAS regression qualification. Preserve failed historical native evidence and require fresh complete budget/SRE acceptance; no technical check or audit rule is waived. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .../2026-09-08-governed-inference-budgets.md | 86 ++++++++++++++++--- 1 file changed, 76 insertions(+), 10 deletions(-) diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index 87f5eb1c2..fe1f331fe 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -1,7 +1,8 @@ # Security Audit — Governed inference budgets (v1) Date: **2026-09-08 UTC** -Status: **Implementation/integration candidate — publication not approved** +Status: **Source audit approved under explicit maintainer delegation**; +exact-head combined native qualification remains mandatory before publication. Scope: `shared/inference_budget/`, `controller/src/inference_budget/`, `controller/src/task_identity.rs`, `inference-router/src/inference_budget/`, @@ -10,6 +11,67 @@ controller Task/Team/runtime wiring, router dispatch/route wiring, CLI and Helm. Gated paths: controller CRDs, router providers/routes, CLI commands, `deploy/helm/kars/files/`, `shared/inference_budget/`. +## Current delegated approval (2026-09-10) + +The maintainer explicitly authorized publication sign-offs after additional +focused review rounds, recorded in +[comment 5615522306](https://github.com/Azure/kars/pull/551#issuecomment-5615522306). +The AI attestation below is delegated review, not a claim that a second human +personally reviewed or signed this capability. Earlier, separately limited +signature waivers are not extended, and no technical check is waived. + +A fresh focused review of `c4291b0f3808747d7fe221af99f47b37e17d16a9` +identified a recovery race: after capturing authority A, recovery could observe +that A was no longer live, then close replacement authority B and its new Pod +session inside a transaction against the latest ledger. The repair at +`e89040dfa76ae3e3b015f2ec7bc39b7c3b24ffc3` compares the captured authority +inside every store CAS attempt. Changed authority is deferred for a fresh +live check; genuinely invalid unchanged authority is still revoked. + +Deterministic regressions cover replacement before the transaction, replacement +during a 409 CAS retry, and unchanged-A revocation. They check the replacement +Pod session and InFlight work, plus conservative maximum liabilities for old +accepted work. Independent re-review of the exact repair against `5ae7b99a` +reported no significant issues. This is bounded source-review closure, not +an independent native execution or financial certification. + +Actual guarded qualification at `e89040df` passed: + +- Controller binary: 82 budget cases, including the five selected recovery + cases and all three new interleaving/revocation regressions. These selectors + overlap and are not added together. +- Router library: 40 inference-budget cases. +- Strict paired controller/router all-target Clippy and workspace formatting. + +The existing offline/locked shared target was used, with two jobs, incremental +compilation disabled and an 8.5 GiB stop floor; minimum observed free space was +9.38 GiB. No dependency installation or target cleanup was needed. + +The current source `66c4ce3a674317d9329ba1d7dda7b8763d0834ae` additionally +merges services `920f9f0e` and its actual landed SRE ancestry. Every tracked +byte outside three CLI test-harness files is identical to qualified `e89040df`. +The forwarded 39 CLI cases and typecheck passed using the existing compatible +cache (Vitest 4.1.10 versus lockfile 4.1.8). The services head's actual hosted +CLI also passed; that result does not qualify this budget composition. + +The same candidate retains the reviewed native fixture repairs: verified +same-image digest aliases, consistent named-provider routing, a proper +CA:false/serverAuth/SAN server leaf and negative TLS checks, and bounded +UID/generation/image/budget-binding checks when reconnecting a port-forward +after legitimate Pod replacement. A genuine 503 or identity/authorization +failure cannot become success. Spending/cancellation assertions and the +native readiness deadline remain unchanged. + +**Exact-head hosted Rust/CLI/security and complete native budget/SRE lifecycle +qualification remain required.** Earlier image, TLS-profile and stale +port-forward failures remain failed historical results, not proof of current +spending or cancellation enforcement. Shared credential/MCP composition is a +separate gate. No main promotion, customer/H100 deployment, public image or +private Bridge publication is authorized by this source sign-off. + +Signed-off-by: pallakatos (maintainer delegation recorded above) +Signed-off-by: GitHub Copilot (delegated AI audit, not an independent human) <223556219+Copilot@users.noreply.github.com> + ## Summary Durable token ceilings and operator-configured maximum-price caps for governed @@ -49,7 +111,7 @@ all-in task spend, invoice accuracy, taxes, or exchange rates. ## Verification -### Native primary-workload repair (2026-09-10) +### Historical native primary-workload repair (2026-09-10) Standalone qualification found a real ReplicaSet admission failure: `AdmissionRequest.subResource` is absent on primary operations. The shared @@ -121,7 +183,11 @@ The existing `security-audit-required` script now discovers this correctly place record and fails specifically for **0 of 2 required genuine signer emails**. That failure is intentional until human review, not a waived or fabricated pass. -## Remaining release decisions/gates +## Original release checklist (historical) + +The following records the initial requirements, not the current approval or +Rust execution status. The current delegated approval and still-open native +gates are stated above. 1. Qualify against the real, now-forwarded privacy issuer prerequisite; no fallback implementation. Parent re-review and full SRE Kind remain separate @@ -136,7 +202,7 @@ That failure is intentional until human review, not a waived or fabricated pass. accessibility, and conservative uncertainty/capacity behavior. 5. Obtain genuine required human audit signoffs before protected publication. -## Bounded reviewer repair candidate +## Earlier bounded reviewer repair candidate (historical) The source-only independent review identified five blockers. This repair: @@ -157,7 +223,7 @@ and settlement path. **These Rust tests have not run**: the shared Cargo lease remains with the credential integration owner until the parent directly grants it. Formatting/source checks do not establish compilation or technical closure. -## Signoffs +## Original unsigned authoring state (historical) - Implementation author: changes under active development; not a signoff. - Independent technical reviewer: **pending**. @@ -170,8 +236,8 @@ fabricated. Prior feature waivers do not apply to this capability. ## Verdict -**Pending — not approved for publication or Ready.** Budget Rust/Clippy, actual -Kind enforcement, independent review and two genuine human signoffs remain -required. No `Signed-off-by` identity is supplied until those people actually -review and approve; the existing security-audit gate is expected to remain red -for missing signatures, without a waiver or altered rule. +Source approval is recorded above under the maintainer's explicit delegation. +Publication and Ready remain pending exact-head technical and native +qualification. Historical unsigned states and failed runs in this document +must not be represented as current blockers already repaired or as successful +runtime acceptance. The existing audit gate is unchanged. From 65d0ecaf9bb2aeaaff66021fa6c21c998f45c70c Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 18:35:16 +0200 Subject: [PATCH 12/16] fix(budget): reject unsupported operations before provider authentication Native existing-PR run 34493554446 failed only E2E Kind at inference-budget-enforcement.mjs:286: POST /v1/embeddings returned 502 instead of 503. Before this assertion, real monetary siblings passed [200,200,429] and settled 10 usdMicros/16 tokens; held-token siblings passed 429, reserved 30 tokens and settled 8 after release. Four owned router UIDs passed healthz and governed readyz. Agents closure, the following provider-count check and cancellation were not reached. Bounded artifacts contain no detailed proxy error; source places legacy-default credential-scope rejection before the existing unsupported budget-operation classifier. Extract the same canonical method/path classifier into a side-effect-free preflight before buffered/streaming provider credential acquisition. Unsupported finite operations retain the typed 503 inference_budget_unavailable contract. Supported requests still obtain credentials and finish wire normalization before the unchanged final-dispatch reserve/begin boundary; legacy errors and preceding policy denials are unchanged. Add actual Axum endpoint regression for named-primary/default-auth handoff, buffered/streaming/method negative cases, supported-auth-before-grant ordering, and broker-loss no-new-dispatch with accepted unknown work still funded. Reuse the existing broker/ledger fixture and make its ephemeral token directory private. Keep native 503 expectation and strengthen it with code/type assertions plus fixed safe status/category facts. Accounting, cancellation, UID/readiness, image, recovery and shared-ledger source parity verified. Validation: 71 targeted CLI/fixture tests, TypeScript typecheck, Rustfmt checks, JS syntax, diff checks and targeted lint (one pre-existing warning). Existing Node cache is Vitest 4.1.10 versus lock 4.1.8. Four new Rust regressions remain unexecuted pending a fresh immutable-SHA Cargo grant; request router LIB inference_budget::client::tests::unsupported_routes_tests followed by relevant budget regressions and strict checks under the existing guard. No Cargo, new target, cleanup, lease SQL updates, push or native dispatch. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- cli/src/testing/budget-fixture-route.test.ts | 22 ++ docs/governed-inference-budgets.md | 7 + .../src/inference_budget/client_tests.rs | 11 +- .../src/inference_budget/dispatch.rs | 37 ++- .../unsupported_routes_tests.rs | 311 ++++++++++++++++++ inference-router/src/proxy.rs | 2 + tests/e2e/budget-fixture-route.mjs | 10 + tests/e2e/inference-budget-enforcement.mjs | 9 +- 8 files changed, 397 insertions(+), 12 deletions(-) create mode 100644 inference-router/src/inference_budget/unsupported_routes_tests.rs diff --git a/cli/src/testing/budget-fixture-route.test.ts b/cli/src/testing/budget-fixture-route.test.ts index 1a2407e9f..71f996620 100644 --- a/cli/src/testing/budget-fixture-route.test.ts +++ b/cli/src/testing/budget-fixture-route.test.ts @@ -101,6 +101,28 @@ describe("budget named-provider fixture and private readiness diagnostics", () = }))).toEqual([{ stage: "router-contract-match", provider_matches: false, endpoint_matches: false, model_matches: true, bounds_valid: true, count: 1 }]); }); + it("requires the unsupported-operation 503 contract without exporting arbitrary error bodies", () => { + const secret = "DO-NOT-EXPORT-PRIVATE-ERROR-CODE-OR-BODY"; + const error = { code: "inference_budget_unavailable", type: "inference_budget_unavailable", message: secret }; + expect(route.unsupportedOperationFact({ status: 503, value: { error } })).toEqual({ + stage: "unsupported-inference-operation", httpStatus: 503, + category: "inference-budget-unavailable", matchesContract: true, + }); + for (const response of [ + { status: 502, value: { error } }, + { status: 503, value: { error: { code: secret, type: secret, message: secret } } }, + { status: secret, value: secret }, + { status: 503, value: { error: { code: error.code, type: secret } } }, + ]) { + const fact = route.unsupportedOperationFact(response); + expect(fact.matchesContract).toBe(false); + expect(JSON.stringify(fact)).not.toContain(secret); + } + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).toContain("assert.equal(unsupported.status, 503)"); + expect(harness).toContain("assert(denial.matchesContract"); + expect(harness).toContain('assert.equal((await request(provider, "/count")).value.count, count)'); + }); it("reports actual UID and live-template mismatches as booleans without workload inputs", () => { const router = { name: "inference-router", image: "router@sha256:fixture", env: [{ name: "KARS_INFERENCE_BUDGET_BINDING", value: "PRIVATE-BINDING" }], args: ["PRIVATE-ARGV"] }; diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index 1309ef30b..30d2feedc 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -191,6 +191,13 @@ unknown request options, and incompatible output-limit fields. Embeddings, legacy completions, image/audio/video generation, Foundry internal generation, fine-tuning, agent runs, evaluations, memory generation, and generic Foundry proxy operations are not authorized by these text contracts. +At the proxy's finite-budget gate, unsupported inference operations return HTTP +503 with the OpenAI error code/type `inference_budget_unavailable`. The operation check +runs before unrelated default-provider credential acquisition, so a missing +credential cannot mask this budget denial as a generic 502. Supported operations +still resolve credentials before the final normalized-wire reservation/begin +boundary; this early rejection does not grant authority, contact the broker, +refund work, or change legacy unbounded error handling or preceding policy denials. Standalone moderation API stages also lack a v1 bounded contract: when a policy requires one, the entire finite request is rejected before that API call. The stage is never silently disabled or asserted to have zero cost. Native diff --git a/inference-router/src/inference_budget/client_tests.rs b/inference-router/src/inference_budget/client_tests.rs index ac6a5238e..a2626acb3 100644 --- a/inference-router/src/inference_budget/client_tests.rs +++ b/inference-router/src/inference_budget/client_tests.rs @@ -17,6 +17,9 @@ use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate}; static NEXT: AtomicUsize = AtomicUsize::new(0); +#[path = "unsupported_routes_tests.rs"] +mod unsupported_routes_tests; + #[tokio::test] async fn native_stream_settlement_never_refunds_missing_or_inconsistent_final_usage() { for (name, wire) in super::super::anthropic_cases::incomplete() { @@ -274,7 +277,13 @@ impl Fixture { std::process::id(), NEXT.fetch_add(1, Ordering::Relaxed) )); - std::fs::create_dir(&directory).unwrap(); + let mut builder = std::fs::DirBuilder::new(); + #[cfg(unix)] + { + use std::os::unix::fs::DirBuilderExt; + builder.mode(0o700); + } + builder.create(&directory).unwrap(); let token_path = directory.join("token"); std::fs::write(&token_path, "private-test-token").unwrap(); let client = Arc::new(Client { diff --git a/inference-router/src/inference_budget/dispatch.rs b/inference-router/src/inference_budget/dispatch.rs index 014364ee2..338bf6e8c 100644 --- a/inference-router/src/inference_budget/dispatch.rs +++ b/inference-router/src/inference_budget/dispatch.rs @@ -5,11 +5,37 @@ use super::{ client::{AttemptGuard, Error}, usage, }; -use crate::{inference_budget_dispatch, proxy::UpstreamConfig}; +use crate::{ + inference_budget_contract::tariffs::Operation, inference_budget_dispatch, proxy::UpstreamConfig, +}; use axum::http::{Method, StatusCode}; use bytes::Bytes; use futures::{StreamExt, stream::BoxStream}; +fn operation(method: &Method, path: &str) -> Result { + let denied = || Error { + stage: "unsupported inference operation", + status: None, + }; + if method != Method::POST { + return Err(denied()); + } + inference_budget_dispatch::operation(path).ok_or_else(denied) +} + +/// Classify only: unsupported finite routes must not consult unrelated provider +/// credentials first. Supported sends still acquire their grant at final dispatch. +pub(crate) fn preflight( + upstream: &UpstreamConfig, + method: &Method, + path: &str, +) -> Result<(), Error> { + if upstream.inference_budget.is_some() { + operation(method, path)?; + } + Ok(()) +} + /// Runs after provider resolution and all wire transformations, immediately /// before the actual transport send. Each retry needs a separate grant. pub async fn begin( @@ -21,14 +47,7 @@ pub async fn begin( let Some(client) = &upstream.inference_budget else { return Ok((body, None)); }; - let denied = || Error { - stage: "unsupported inference operation", - status: None, - }; - if method != Method::POST { - return Err(denied().into()); - } - let operation = inference_budget_dispatch::operation(path).ok_or_else(denied)?; + let operation = operation(method, path)?; let (wire, guard) = client .begin( upstream.telemetry_provider(), diff --git a/inference-router/src/inference_budget/unsupported_routes_tests.rs b/inference-router/src/inference_budget/unsupported_routes_tests.rs new file mode 100644 index 000000000..2e1b50884 --- /dev/null +++ b/inference-router/src/inference_budget/unsupported_routes_tests.rs @@ -0,0 +1,311 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{ + config::{Config, ProviderEndpoint}, + inference_policy_loader::{LoadedInferencePolicy, ModelPreference, ModelRef}, + proxy::failure::{Acceptance, FailureCategory, ForwardFailure}, + routes::AppState, +}; +use axum::{ + Router, + body::{Body, to_bytes}, + http::{Request as HttpRequest, StatusCode}, +}; +use tower::ServiceExt; + +fn default_upstream(fixture: &Fixture, governed: bool) -> UpstreamConfig { + let mut upstream = UpstreamConfig::azure(fixture.provider.uri(), "model".into(), "task".into()); + upstream.inference_budget = governed.then(|| fixture.client.clone()); + upstream +} + +async fn assert_budget_denial(error: anyhow::Error, stage: &'static str) { + let typed = error.downcast_ref::().expect("typed budget denial"); + assert_eq!(typed.stage, stage); + assert!(!crate::proxy::failure::retryable_failure(&error)); + let response = crate::inference_budget::response::denial(&error).unwrap(); + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["error"]["code"], "inference_budget_unavailable"); +} + +async fn endpoint_router(fixture: &Fixture, governed: bool) -> Router { + let mut config = Config::from_env().unwrap(); + config.azure_openai_endpoint = Some(fixture.provider.uri()); + config.default_model = "model".into(); + config.providers = std::collections::HashMap::from([( + "budget-fixture".into(), + ProviderEndpoint { + tag: "budget-fixture".into(), + endpoint: fixture.provider.uri(), + api_key: None, + }, + )]); + let policy_status = Arc::new(crate::policy_status::PolicyStatusRegistry::new()); + let governance = Arc::new(crate::governance::Governance::new_with_status( + "task", + policy_status.clone(), + )); + let state = AppState { + services: Default::default(), + auth: Arc::new(WorkloadIdentityAuth::for_test(None, None)), + copilot: Arc::new(crate::copilot_auth::CopilotTokenCache::with_test_exchange( + "unused-default-seat", + format!("{}/unexpected-token-exchange", fixture.provider.uri()), + )), + client: reqwest::Client::builder().no_proxy().build().unwrap(), + config: Arc::new(config), + budget: crate::budget::TokenBudgetTracker::new(0, 0), + inference_budget: governed.then(|| fixture.client.clone()), + policy_provider: governance.clone(), + audit_sink: governance.clone(), + signing_provider: governance.clone(), + governance, + blocklist: crate::blocklist::Blocklist::disabled(), + blocked_egress: Arc::new(crate::egress_blocked::BlockedBuffer::with_defaults()), + sandbox_name: Arc::new("task".into()), + inbox: Arc::new(crate::mesh::MeshInbox::new()), + mesh_metrics: Arc::new(crate::mesh::MeshMetrics::new()), + model_override: Default::default(), + responses_only_models: Default::default(), + unavailable_models: Default::default(), + admin_token: None, + handoff_tokens: crate::handoff::HandoffTokenStore::new(), + handoff_session: crate::handoff::HandoffSession::new(), + drain_state: crate::handoff::DrainState::new(), + pending_handoff: crate::handoff::PendingHandoffStore::new(), + policy_status, + inference_policy: crate::inference_policy_loader::empty_handle(), + memory_binding: crate::memory_binding_loader::empty_handle(), + egress_allowlist: crate::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new(crate::deployment_health::DeploymentHealthRegistry::new()), + }; + // Match the native fixture: a named primary, not an explicit legacy provider. + // Embeddings still starts from the default Azure upstream, without a key. + *state.inference_policy.write().await = Some(LoadedInferencePolicy { + digest: "budget-fixture".into(), + source_path: "budget-fixture".into(), + per_request_tokens: None, + daily_tokens: None, + monthly_tokens: None, + content_safety: None, + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "budget-fixture".into(), + deployment: "model".into(), + }, + fallback: vec![], + }), + provider: None, + guardrails: vec![], + raw: json!({}), + }); + Router::new() + .merge(crate::routes::inference_routes()) + .with_state(state) +} + +#[tokio::test] +async fn native_embeddings_returns_budget_503_before_unrelated_default_authentication() { + for governed in [false, true] { + let fixture = Fixture::new(100, false).await; + let original = fixture.ledger.lock().unwrap().clone(); + let app = endpoint_router(&fixture, governed).await; + for path in ["/v1/embeddings", "/v1/completions"] { + let response = app + .clone() + .oneshot( + HttpRequest::builder() + .method("POST") + .uri(path) + .header("content-type", "application/json") + .header("x-kars-sandbox", "task") + .body(Body::from(r#"{"input":"fixture"}"#)) + .unwrap(), + ) + .await + .unwrap(); + if governed { + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = to_bytes(response.into_body(), 4096).await.unwrap(); + let body: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body["error"]["code"], "inference_budget_unavailable"); + assert_eq!(body["error"]["type"], "inference_budget_unavailable"); + } else { + assert_eq!(response.status(), StatusCode::BAD_GATEWAY); + } + } + assert!( + fixture + .provider + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert!( + fixture + ._broker + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert_eq!(*fixture.ledger.lock().unwrap(), original); + } +} + +#[tokio::test] +async fn unsupported_finite_buffered_and_streaming_operations_never_reach_auth_or_broker() { + let fixture = Fixture::new(100, false).await; + let original = fixture.ledger.lock().unwrap().clone(); + for path in [ + "embeddings", + "/v1/embeddings", + "completions", + "images/generations", + "responses?background=true", + ] { + let upstream = default_upstream(&fixture, true); + let auth = Arc::new(WorkloadIdentityAuth::for_test(None, None)); + let error = crate::proxy::forward( + &auth, + None, + &reqwest::Client::new(), + &upstream, + Method::POST, + path, + &HeaderMap::new(), + bytes::Bytes::from_static(br#"{"input":"fixture"}"#), + ) + .await + .err() + .unwrap(); + assert_budget_denial(error, "unsupported inference operation").await; + let error = crate::proxy::forward_stream( + auth, + None, + reqwest::Client::new(), + upstream, + path, + HeaderMap::new(), + bytes::Bytes::from_static(br#"{"input":"fixture","stream":true}"#), + ) + .await + .err() + .unwrap(); + assert_budget_denial(error, "unsupported inference operation").await; + } + let error = crate::proxy::forward( + &WorkloadIdentityAuth::for_test(None, None), + None, + &reqwest::Client::new(), + &default_upstream(&fixture, true), + Method::GET, + "chat/completions", + &HeaderMap::new(), + bytes::Bytes::new(), + ) + .await + .err() + .unwrap(); + assert_budget_denial(error, "unsupported inference operation").await; + assert!( + fixture + .provider + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert!( + fixture + ._broker + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert_eq!(*fixture.ledger.lock().unwrap(), original); +} + +#[tokio::test] +async fn supported_finite_operation_still_resolves_credentials_before_acquiring_any_grant() { + let fixture = Fixture::new(100, false).await; + let original = fixture.ledger.lock().unwrap().clone(); + let error = crate::proxy::forward( + &WorkloadIdentityAuth::for_test(None, None), + None, + &reqwest::Client::new(), + &default_upstream(&fixture, true), + Method::POST, + "chat/completions", + &HeaderMap::new(), + bytes::Bytes::from_static(br#"{"messages":[{"role":"user","content":"text"}]}"#), + ) + .await + .err() + .unwrap(); + let failure = error.downcast_ref::().unwrap(); + assert_eq!(failure.category, FailureCategory::Authentication); + assert_eq!(failure.acceptance, Acceptance::NotAccepted); + assert!(crate::inference_budget::response::denial(&error).is_none()); + assert!(!crate::proxy::failure::retryable_failure(&error)); + assert!( + fixture + .provider + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert!( + fixture + ._broker + .received_requests() + .await + .unwrap() + .is_empty() + ); + assert_eq!(*fixture.ledger.lock().unwrap(), original); +} + +#[tokio::test] +async fn supported_broker_loss_blocks_new_sends_without_refunding_accepted_unknown_work() { + let fixture = Fixture::new(100, false).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[]}))) + .mount(&fixture.provider) + .await; + fixture.buffered().await.unwrap(); + let funded = fixture.ledger.lock().unwrap().clone(); + assert_eq!(funded.meters.uncertain.tokens, 30); + assert_eq!(funded.meters.uncertain.usd_micros, 5); + assert_eq!(fixture.provider.received_requests().await.unwrap().len(), 1); + Mock::given(wiremock::matchers::any()) + .respond_with(ResponseTemplate::new(503)) + .with_priority(1) + .mount(&fixture._broker) + .await; + assert_budget_denial(fixture.buffered().await.err().unwrap(), "/v1/catalog").await; + let error = crate::proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::for_test(None, None)), + None, + reqwest::Client::new(), + fixture.upstream(), + "chat/completions", + HeaderMap::new(), + bytes::Bytes::from_static( + br#"{"messages":[{"role":"user","content":"text"}],"stream":true}"#, + ), + ) + .await + .err() + .unwrap(); + assert_budget_denial(error, "/v1/catalog").await; + assert_eq!(fixture.provider.received_requests().await.unwrap().len(), 1); + assert_eq!(*fixture.ledger.lock().unwrap(), funded); +} diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 2f48a90fa..fdf038525 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -264,6 +264,7 @@ pub async fn forward( request_headers: &HeaderMap, request_body: Bytes, ) -> Result<(StatusCode, HeaderMap, Bytes)> { + crate::inference_budget::dispatch::preflight(upstream, &method, path)?; let start = Instant::now(); let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { telemetry.begin( @@ -538,6 +539,7 @@ pub async fn forward_stream( HeaderMap, futures::stream::BoxStream<'static, Result>, )> { + crate::inference_budget::dispatch::preflight(&upstream, &Method::POST, path)?; let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { telemetry.begin( path, diff --git a/tests/e2e/budget-fixture-route.mjs b/tests/e2e/budget-fixture-route.mjs index db026adbe..34f4dcb2d 100644 --- a/tests/e2e/budget-fixture-route.mjs +++ b/tests/e2e/budget-fixture-route.mjs @@ -45,6 +45,16 @@ export function readinessFact(endpoint, response, error) { ? "budget-authority-provider-contract" : "unexpected-readiness-response" }; } +export function unsupportedOperationFact(response) { + const httpStatus = Number.isInteger(response?.status) && response.status >= 100 && response.status <= 599 + ? response.status : 0; + const coded = response?.value?.error?.code === "inference_budget_unavailable" + && response.value.error.type === "inference_budget_unavailable"; + return { stage: "unsupported-inference-operation", httpStatus, + category: coded ? "inference-budget-unavailable" : "unexpected-error-contract", + matchesContract: httpStatus === 503 && coded }; +} + const STAGES = { kars_inference_router: { readiness: ["router-binding", "router-guardrails", "router-provider", "router-candidate", diff --git a/tests/e2e/inference-budget-enforcement.mjs b/tests/e2e/inference-budget-enforcement.mjs index 8cb5bd7e4..68e96995c 100644 --- a/tests/e2e/inference-budget-enforcement.mjs +++ b/tests/e2e/inference-budget-enforcement.mjs @@ -10,7 +10,8 @@ import { fileURLToPath } from "node:url"; import { join } from "node:path"; import { prepareRouterImage } from "./kind-router-image.mjs"; import { PROVIDER, ENDPOINT, providerSource, - budgetStageFacts, routerTemplateFacts, TLS_SERVER_EXTENSIONS, verifyFixtureCertificate } from "./budget-fixture-route.mjs"; + budgetStageFacts, routerTemplateFacts, TLS_SERVER_EXTENSIONS, verifyFixtureCertificate, + unsupportedOperationFact } from "./budget-fixture-route.mjs"; import { ownedRouterResolver, startForward, waitForOwnedRouter } from "./budget-router-readiness.mjs"; const root = fileURLToPath(new URL("../../", import.meta.url)); @@ -283,7 +284,11 @@ async function scenario() { assert.equal(accountFor("budget-token-left").status.ledger.meters.settled.tokens, 8); const count = (await request(provider, "/count")).value.count; - assert.equal((await request(tokenRight.url, "/v1/embeddings", { input: "fixture" })).status, 503); + const unsupported = await request(tokenRight.url, "/v1/embeddings", { input: "fixture" }); + const denial = unsupportedOperationFact(unsupported); + console.log("BUDGET-FAILURE-CONTRACT " + JSON.stringify(denial)); + assert.equal(unsupported.status, 503); + assert(denial.matchesContract, "Unsupported inference must return the coded budget denial"); assert.equal((await request(tokenRight.url, "/agents", {})).status, 403); assert.equal((await request(provider, "/count")).value.count, count); const binding = get("karstask", "budget-token-left").status.inferenceBudget; From 058b79731b1e8a5aa28f9723dd3512ef381f0661 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 18:49:38 +0200 Subject: [PATCH 13/16] test(budget): initialize the concrete content safety floor The exact-65d0ecaf guarded router-library regression build found E0308 before executing tests: LoadedInferencePolicy.content_safety is ContentSafetyFloor, not Option. Initialize its default floor without changing production code, route checks, budget semantics or assertions. Guard remained above its 8.5 GiB floor (minimum observed 8.95 GiB). Formatting and diff checks pass; Rust qualification remains pending a fresh immutable-SHA grant. No Cargo retry, cleanup, new target, profile change, public push or lease SQL write. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .../src/inference_budget/unsupported_routes_tests.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/inference-router/src/inference_budget/unsupported_routes_tests.rs b/inference-router/src/inference_budget/unsupported_routes_tests.rs index 2e1b50884..d37755040 100644 --- a/inference-router/src/inference_budget/unsupported_routes_tests.rs +++ b/inference-router/src/inference_budget/unsupported_routes_tests.rs @@ -91,7 +91,7 @@ async fn endpoint_router(fixture: &Fixture, governed: bool) -> Router { per_request_tokens: None, daily_tokens: None, monthly_tokens: None, - content_safety: None, + content_safety: Default::default(), model_preference: Some(ModelPreference { primary: ModelRef { provider: "budget-fixture".into(), From 4c2c89134bd1ab660aa13c633757f00173d2673b Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 19:34:28 +0200 Subject: [PATCH 14/16] docs(budgets): record qualified unsupported-route failure repair Record the actual late native502/503 failure, reviewed side-effect-free classification fix, four new Rust regressions and61budget cases. Preserve all native accounting/failure expectations and require fresh current-base qualification; no result is relabeled or gate waived. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .../2026-09-08-governed-inference-budgets.md | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index fe1f331fe..cf5728ed7 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -72,6 +72,43 @@ private Bridge publication is authorized by this source sign-off. Signed-off-by: pallakatos (maintainer delegation recorded above) Signed-off-by: GitHub Copilot (delegated AI audit, not an independent human) <223556219+Copilot@users.noreply.github.com> +## Unsupported finite-route failure contract (2026-09-10) + +Published `2a95d02a` passed Rust, CLI, benchmarks and the other gates, but +native job +[102935839436](https://github.com/Azure/kars/actions/runs/34493554446/job/102935839436) +finished with 166 passes and one failure. Four actual owned routers passed +both health and governed readiness. The later `/v1/embeddings` assertion +received502 rather than the required503; that result remains failed. + +The handler attempted default-provider authentication before classifying an +unsupported finite-budget operation. Repair `65d0ecaf9bb2aeaaff66021fa6c21c998f45c70c` +moves only side-effect-free support classification ahead of that credential +lookup. Unsupported finite requests retain the explicit budget failure rather +than an unrelated provider-authentication failure. Request authentication, +governance and final-send funding are not bypassed. Supported and unbounded +paths retain their existing behavior. + +The native503 expectation is preserved and strengthened with error code/type +checks. Independent focused re-review of the exact production repair reported +no significant issues. The first Rust attempt stopped at a test initializer +type error before running tests; `058b79731b1e8a5aa28f9723dd3512ef381f0661` +corrects only that initializer, without changing production or assertions. + +Exact058 qualification passed all four new unsupported-route regressions and +61 router budget cases including those four, strict router all-target Clippy +and formatting. Prior JS/CLI qualification passed71 cases and type checking. +The guarded shared-target Rust batch observed at least8.91GiB free; no floor +or dependency workaround was used. Controller source was not changed or +unnecessarily rebuilt for this router-only repair. + +Current source `2225fab6f789b9e954221e2686f331b4df2a4448` adds actual +integration ancestor `af4deba7` with an entire tree identical to qualified058. +The bounded source approval above applies to this reviewed repair as well. +Fresh current-base native accounting, cancellation and failure-contract +acceptance remain mandatory; successful readiness is not substituted for +those assertions. + ## Summary Durable token ceilings and operator-configured maximum-price caps for governed From 1ee92de3b1df5ef5349abd90581d956bc07849fa Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 21:12:15 +0200 Subject: [PATCH 15/16] test(budget): use fenced API cancellation instead of a literal patch file Continue from joined existing-PR base 1af895bc98dad70f09d2467bef34a01b7ae6fb45 without changing landed MCP or budget production code. Old-base native run 34509062774 passed the coded unsupported-operation 503, agents 403 and no-dispatch count assertion, then observed accepted provider work before cancellation failed at the kubectl PATCH handoff. The original command error discarded stderr; retained artifacts do not establish an API 409 or broker outage. Inspect the exact pinned kubectl v1.30.5 patch implementation: it calls os.ReadFile(PatchFile), so --patch-file - names a literal file rather than stdin. Reproduce this local-file failure with installed kubectl v1.35.3 against a controlled API, proving zero API requests. Replace the bad handoff with the existing loopback kubectl-proxy API pattern, extracted once for the workload proof and cancellation. Cancel only the created Task UID in its pinned workspace, with unchanged generation/spec/budget binding and fresh resourceVersion. Retry at most three times inside 30 seconds only for actual HTTP 409 with matching Kubernetes Status Conflict. Any identity/intent change, 403/422/transport, malformed status, stale RV or failed result verification stays fatal. Emit only fixed API verb/resource/status/reason facts. Provider, spending, no-dispatch, accepted-work and final uncertain/settled/reserved assertions and the 19-case native workload proof remain byte-identical. Validation: 94 targeted JS/CLI tests including 23 cancellation cases and real kubectl proxy conflict/denial cleanup; TypeScript typecheck; JS syntax and diff checks; lint with only two pre-existing warnings. Existing Node cache uses Vitest 4.1.10 versus lock 4.1.8. Private test configs removed and owned proxy children reaped. Native cancellation/accounting success is still unproven until parent-approved hosted E2E. No Cargo grant used, production changes, lease SQL updates, new agents, push or native cycle. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- cli/src/testing/budget-cancellation.test.ts | 280 ++++++++++++++++++++ docs/governed-inference-budgets.md | 12 + tests/e2e/budget-api-client.mjs | 111 ++++++++ tests/e2e/budget-cancellation.mjs | 93 +++++++ tests/e2e/budget-workload-cases.mjs | 41 +-- tests/e2e/inference-budget-enforcement.mjs | 11 +- 6 files changed, 507 insertions(+), 41 deletions(-) create mode 100644 cli/src/testing/budget-cancellation.test.ts create mode 100644 tests/e2e/budget-api-client.mjs create mode 100644 tests/e2e/budget-cancellation.mjs diff --git a/cli/src/testing/budget-cancellation.test.ts b/cli/src/testing/budget-cancellation.test.ts new file mode 100644 index 000000000..0facb34bb --- /dev/null +++ b/cli/src/testing/budget-cancellation.test.ts @@ -0,0 +1,280 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { execFileSync, spawn } from "node:child_process"; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { createServer } from "node:http"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; +import { describe, expect, it } from "vitest"; + +const { apiRequest, withKindApi } = await import( + new URL("../../../tests/e2e/budget-api-client.mjs", import.meta.url).href); +const { cancelAcceptedTask, cancellationFact } = await import( + new URL("../../../tests/e2e/budget-cancellation.mjs", import.meta.url).href); +const root = fileURLToPath(new URL("../../../", import.meta.url)); +const context = "kind-kars-e2e"; +const namespace = "kars-system", name = "budget-token-left"; +const taskPath = `/apis/kars.azure.com/v1alpha1/namespaces/${namespace}/karstasks/${name}`; +const secret = "DO-NOT-EXPORT-API-MESSAGE-BODY-OR-DETAILS"; + +function status(code = 409, reason = "Conflict") { + return { apiVersion: "v1", kind: "Status", status: "Failure", code, reason, + message: secret, details: { name: secret, causes: [{ message: secret }] } }; +} + +async function fixture() { + const binding = { scope: "GovernedInference", taskUid: "task-uid", + account: { namespace, name: "account", uid: "account-uid" }, + root: { workspaceUid: "workspace-uid" }, authorizationDigest: "fixed-authority" }; + const created = { apiVersion: "kars.azure.com/v1alpha1", kind: "KarsTask", + metadata: { name, namespace, uid: "task-uid", resourceVersion: "1", generation: 1 }, + spec: { objective: "fixture", envelope: { budget: { scope: "GovernedInference", tokens: 50 } }, + blueprint: { instructions: "unchanged" }, execution: { launch: true } }, + status: { inferenceBudget: binding } }; + const state = { + task: structuredClone(created), namespaceUid: "workspace-uid", patches: [] as any[], + requests: [] as string[], failure: false, conflicts: 0, + reply: undefined as { code: number; body: any } | undefined, + afterConflict: () => {}, afterCommit: () => {}, loseResponse: false, + }; + const server = createServer((request, response) => { + void (async () => { + state.requests.push(`${request.method} ${request.url}`); + const send = (code: number, body: unknown) => { + response.writeHead(code, { "Content-Type": "application/json" }); + response.end(typeof body === "string" ? body : JSON.stringify(body)); + }; + if (request.method === "GET" && request.url === `/api/v1/namespaces/${namespace}`) + return send(200, { apiVersion: "v1", kind: "Namespace", + metadata: { name: namespace, uid: state.namespaceUid } }); + if (request.method === "GET" && request.url === taskPath) return send(200, state.task); + if (request.method !== "PATCH" || request.url !== taskPath) return send(404, status(404, "NotFound")); + expect(request.headers["content-type"]).toBe("application/merge-patch+json"); + expect(request.headers["impersonate-user"]).toBeUndefined(); + const chunks: Buffer[] = []; + for await (const chunk of request) chunks.push(Buffer.from(chunk)); + const patch = JSON.parse(Buffer.concat(chunks).toString()); + state.patches.push(patch); + expect(patch).toEqual({ metadata: { uid: created.metadata.uid, + resourceVersion: state.task.metadata.resourceVersion }, spec: { execution: { launch: false } } }); + if (state.reply) return send(state.reply.code, state.reply.body); + if (state.conflicts-- > 0) { + state.task.metadata.resourceVersion = String(Number(state.task.metadata.resourceVersion) + 1); + state.afterConflict(); + return send(409, status()); + } + state.task.spec.execution.launch = false; + state.task.metadata.generation++; + state.task.metadata.resourceVersion = String(Number(state.task.metadata.resourceVersion) + 1); + state.afterCommit(); + if (state.loseResponse) return request.socket.destroy(); + send(200, state.task); + })().catch(() => { state.failure = true; response.writeHead(500); response.end(); }); + }); + await new Promise(resolve => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + if (!address || typeof address === "string") throw new Error("Controlled API listener failed"); + const url = `http://127.0.0.1:${address.port}`; + const facts: any[] = []; + const options = { created, binding: structuredClone(binding), request: apiRequest(url), + report: (fact: unknown) => facts.push(fact), deadline: Date.now() + 3000 }; + return { state, options, facts, url, close: async () => { + server.closeAllConnections(); + await new Promise((resolve, reject) => server.close(error => error ? reject(error) : resolve())); + expect(state.failure).toBe(false); + } }; +} + +async function kubeconfig(url: string, run: (file: string, directory: string) => Promise) { + const directory = mkdtempSync(join(root, ".budget-cancel-test-")); + const file = join(directory, "config.json"); + writeFileSync(file, JSON.stringify({ + apiVersion: "v1", kind: "Config", "current-context": context, + clusters: [{ name: "fixture", cluster: { server: url } }], + contexts: [{ name: context, context: { cluster: "fixture", user: "fixture" } }], + users: [{ name: "fixture", user: {} }], + }), { mode: 0o600 }); + try { await run(file, directory); } + finally { rmSync(directory, { recursive: true, force: true }); } +} + +describe("native accepted-work cancellation API contract", () => { + it("reproduces literal --patch-file '-' failure before any API request with real kubectl", async () => { + const f = await fixture(); + try { + await kubeconfig(f.url, async (file, directory) => { + const child = spawn("kubectl", ["--kubeconfig", file, "--cache-dir", join(directory, "cache"), + "--context", context, "patch", "karstask", name, "-n", namespace, "--type=merge", "--patch-file", "-"], + { cwd: directory, stdio: ["pipe", "pipe", "pipe"] }); + let stderr = ""; + child.stderr.on("data", data => { stderr = (stderr + data).slice(-8192); }); + child.stdout.on("data", () => {}); + child.stdin.end(JSON.stringify({ metadata: { uid: "task-uid", resourceVersion: "1" }, + spec: { execution: { launch: false } } })); + const code = await new Promise((resolve, reject) => { + child.once("close", resolve); + child.once("error", reject); + }); + expect(code).not.toBe(0); + expect(/unable to read patch file:.*open -:/.test(stderr)).toBe(true); + expect(f.state.requests).toEqual([]); + }); + } finally { await f.close(); } + }); + + it.each(["conflict", "forbidden"])("uses a real kubectl proxy, exact %s status and owned cleanup", async mode => { + const f = await fixture(); + f.state.conflicts = 1; + if (mode === "forbidden") f.state.reply = { code: 403, body: status(403, "Forbidden") }; + const children: ReturnType[] = []; + try { + await kubeconfig(f.url, async (file, directory) => { + const extra = ["--kubeconfig", file, "--cache-dir", join(directory, "cache")]; + const pending = withKindApi({ + root: directory, context, deadline: f.options.deadline, + kubectl: (args: string[]) => execFileSync("kubectl", [...extra, "--context", context, ...args], + { encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 2000 }), + spawnProcess: (binary: string, args: string[], options: any) => { + const child = spawn(binary, [...extra, ...args], options); + children.push(child); + return child; + }, + }, (request: any) => cancelAcceptedTask({ ...f.options, request })); + if (mode === "forbidden") { + await expect(pending).rejects.toThrow("cancellation-patch"); + expect(f.state.patches).toHaveLength(1); + expect(f.facts.at(-1)).toMatchObject({ httpStatus: 403, reason: "Forbidden" }); + } else { + const result = await pending; + expect(result.metadata.uid).toBe("task-uid"); + expect(result.spec.execution.launch).toBe(false); + expect(f.state.patches.map(patch => patch.metadata.resourceVersion)).toEqual(["1", "2"]); + expect(f.facts.filter(fact => fact.verb === "PATCH").map(fact => [fact.httpStatus, fact.reason])) + .toEqual([[409, "Conflict"], [200, "Success"]]); + } + expect(JSON.stringify(f.facts)).not.toContain(secret); + }); + expect(children).toHaveLength(1); + expect(children[0].signalCode).toBe("SIGTERM"); + } finally { await f.close(); } + }); + + it.each([[403, "Forbidden"], [422, "Invalid"], [429, "TooManyRequests"], [503, "ServiceUnavailable"]])( + "retains actual HTTP %s and fails without retry", async (code, reason) => { + const f = await fixture(); + f.state.reply = { code: Number(code), body: status(Number(code), String(reason)) }; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("cancellation-patch"); + expect(f.state.patches).toHaveLength(1); + expect(f.state.task.spec.execution.launch).toBe(true); + expect(f.facts.at(-1)).toMatchObject({ verb: "PATCH", httpStatus: code, reason }); + expect(JSON.stringify(f.facts)).not.toContain(secret); + } finally { await f.close(); } + }, + ); + + it.each([ + { code: 403, body: status(409, "Conflict") }, + { code: 409, body: status(409, "Invalid") }, + { code: 409, body: status(403, "Conflict") }, + { code: 409, body: secret }, + ])("never manufactures a retry from mismatched or malformed Status", async reply => { + const f = await fixture(); + f.state.reply = reply; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow(); + expect(f.state.patches).toHaveLength(1); + expect(JSON.stringify(f.facts)).not.toContain(secret); + } finally { await f.close(); } + }); + + it.each(["task-uid", "namespace-uid", "generation", "spec", "binding", "already-stopped"])( + "rejects changed %s after a confirmed conflict instead of adopting it", async change => { + const f = await fixture(); + f.state.conflicts = 1; + f.state.afterConflict = () => { + if (change === "task-uid") f.state.task.metadata.uid = "replacement"; + if (change === "namespace-uid") f.state.namespaceUid = "replacement"; + if (change === "generation") f.state.task.metadata.generation++; + if (change === "spec") f.state.task.spec.blueprint.instructions = "changed"; + if (change === "binding") f.state.task.status.inferenceBudget.account.uid = "replacement"; + if (change === "already-stopped") f.state.task.spec.execution.launch = false; + }; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("identity-or-intent"); + expect(f.state.patches).toHaveLength(1); + } finally { await f.close(); } + }, + ); + + it("does not repeat a stale RV when the confirmed conflict yields no new version", async () => { + const f = await fixture(); + f.state.conflicts = 1; + f.state.afterConflict = () => { f.state.task.metadata.resourceVersion = "1"; }; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("without-fresh-version"); + expect(f.state.patches).toHaveLength(1); + } finally { await f.close(); } + }); + + it("limits real conflicts to three attempts without removing UID/RV preconditions", async () => { + const f = await fixture(); + f.state.conflicts = 10; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("conflict-limit"); + expect(f.state.patches.map(patch => patch.metadata)).toEqual([ + { uid: "task-uid", resourceVersion: "1" }, { uid: "task-uid", resourceVersion: "2" }, + { uid: "task-uid", resourceVersion: "3" }, + ]); + } finally { await f.close(); } + }); + + it("does not retry or claim success after a committed patch loses its response", async () => { + const f = await fixture(); + f.state.loseResponse = true; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("transport"); + expect(f.state.patches).toHaveLength(1); + expect(f.state.task.spec.execution.launch).toBe(false); + expect(f.facts.at(-1)).toMatchObject({ httpStatus: 0, reason: "Unclassified" }); + } finally { await f.close(); } + }); + + it("retains the original deadline and never claims a late successful response", async () => { + const f = await fixture(); + let time = 0; + f.state.afterCommit = () => { time = 100; }; + try { + await expect(cancelAcceptedTask({ ...f.options, now: () => time, deadline: 100 })) + .rejects.toThrow("deadline"); + expect(f.state.patches).toHaveLength(1); + await expect(cancelAcceptedTask({ ...f.options, deadline: Date.now() - 1 })).rejects.toThrow("deadline"); + expect(f.state.patches).toHaveLength(1); + } finally { await f.close(); } + }); + + it("does not treat a 200 response with wrong intent as successful cancellation", async () => { + const f = await fixture(); + f.state.afterCommit = () => { f.state.task.spec.execution.launch = true; }; + try { + await expect(cancelAcceptedTask(f.options)).rejects.toThrow("identity-or-intent"); + expect(f.state.patches).toHaveLength(1); + } finally { await f.close(); } + }); + + it("exports only allowlisted status/reason facts and leaves funding assertions in place", () => { + expect(cancellationFact("PATCH", "karstasks", 1, { status: 422, + body: { ...status(422, secret), message: secret } })).toEqual({ + stage: "accepted-work-cancellation", verb: "PATCH", resource: "karstasks", + attempt: 1, httpStatus: 422, reason: "Unclassified", + }); + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).not.toContain('"--patch-file", "-"'); + expect(harness).toContain('created: createdTasks.get("budget-token-left"), binding'); + expect(harness).toContain("const cancellationDeadline = Date.now() + 30_000"); + expect(harness).toContain("assert.equal(final.status.ledger.meters.settled.tokens, 8)"); + expect(harness).toContain("assert.equal(final.status.ledger.meters.uncertain.tokens, 30)"); + expect(harness).toContain("assert.equal(final.status.ledger.meters.reserved.tokens, 0)"); + }); +}); diff --git a/docs/governed-inference-budgets.md b/docs/governed-inference-budgets.md index 30d2feedc..8f3a4482a 100644 --- a/docs/governed-inference-budgets.md +++ b/docs/governed-inference-budgets.md @@ -292,3 +292,15 @@ does not authorize a new target, and replaced parent identities fail immediately This is test-harness lifecycle handling, not a production readiness bypass: provider-attempt, sibling denial, pending/settled spending, and conservatively funded cancellation assertions remain unchanged. + +Accepted-work cancellation uses the same loopback `kubectl proxy` API client as +the native workload proof. The pinned kubectl v1.30.5 reads `--patch-file -` as a +literal filename, not stdin; cancellation therefore sends a real merge-patch +request instead. Only an HTTP 409 with a matching Kubernetes `Status/Conflict` +permits another attempt (at most three, within 30 seconds). Every attempt rereads +the workspace and Task, retaining the created Task UID, generation, complete spec, +budget binding, workspace UID, and fresh resourceVersion. Changed intent or +identity, 403/422, malformed responses, and transport failures remain fatal. +`BUDGET-CANCELLATION` prints only fixed verb/resource/status/reason facts, not API +messages, headers, or bodies. These client checks do not replace the subsequent +real uncertain-liability and settled/reserved accounting assertions. diff --git a/tests/e2e/budget-api-client.mjs b/tests/e2e/budget-api-client.mjs new file mode 100644 index 000000000..74660c98c --- /dev/null +++ b/tests/e2e/budget-api-client.mjs @@ -0,0 +1,111 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import assert from "node:assert/strict"; +import { spawn } from "node:child_process"; + +const sleep = ms => new Promise(resolve => setTimeout(resolve, ms)); +const loopback = host => ["localhost", "127.0.0.1", "[::1]"].includes(host); + +export class BudgetApiError extends Error { + constructor(category, httpStatus = 0) { + super(`Budget fixture API ${category}`); + this.category = category; + this.httpStatus = httpStatus; + } +} + +export function apiRequest(base, deadline = Infinity) { + const origin = new URL(base); + assert(origin.protocol === "http:" && loopback(origin.hostname) && !origin.username && !origin.password + && origin.pathname === "/" && !origin.search && !origin.hash, "Exact loopback API proxy required"); + return async (method, path, body, actor) => { + assert(["GET", "POST", "PUT", "PATCH", "DELETE"].includes(method), "Unexpected fixture API verb"); + const url = new URL(path, origin); + assert(path.startsWith("/") && url.origin === origin.origin, "Fixture API path changed origin"); + const remaining = Math.min(15_000, deadline - Date.now()); + if (remaining <= 0) throw new BudgetApiError("deadline"); + let response; + try { + response = await fetch(url, { + method, signal: AbortSignal.timeout(Math.ceil(remaining)), redirect: "error", + headers: { + "Content-Type": method === "PATCH" ? "application/merge-patch+json" : "application/json", + Accept: "application/json", + ...(actor ? { "Impersonate-User": actor, "Impersonate-Group": "system:authenticated" } : {}), + }, + ...(body === undefined ? {} : { body: JSON.stringify(body) }), + }); + } catch { + throw new BudgetApiError("transport"); + } + const chunks = []; + let size = 0; + try { + for await (const chunk of response.body ?? []) { + size += chunk.length; + if (size > 1024 * 1024) throw new BudgetApiError("response-size", response.status); + chunks.push(chunk); + } + } catch (error) { + if (error instanceof BudgetApiError) throw error; + throw new BudgetApiError("response-transport", response.status); + } + let value; + try { value = JSON.parse(Buffer.concat(chunks).toString("utf8")); } + catch { throw new BudgetApiError("response-json", response.status); } + return { status: response.status, body: value }; + }; +} + +export async function withKindApi({ root, context, kubectl, deadline = Infinity, spawnProcess = spawn }, run) { + assert(["kind-kars-e2e", "kind-kars-budget-api"].includes(context), "Owned budget Kind context required"); + let config, host; + try { + config = JSON.parse(kubectl(["config", "view", "--minify", "-o", "json"])); + host = new URL(config.clusters?.[0]?.cluster?.server).hostname; + } catch { throw new BudgetApiError("proxy-config"); } + assert(config.contexts?.length === 1 && config.contexts[0].name === context + && config.clusters?.length === 1 && loopback(host), + "Budget API requires the exact loopback Kind context"); + if (Date.now() >= deadline) throw new BudgetApiError("deadline"); + const proxy = spawnProcess("kubectl", ["--context", context, "--request-timeout=20s", "proxy", + "--address=127.0.0.1", "--port=0"], { cwd: root, stdio: ["ignore", "pipe", "pipe"] }); + let output = "", failed = false; + proxy.stdout.on("data", data => { output = (output + data).slice(-2048); }); + proxy.stderr.on("data", () => {}); + proxy.on("error", () => { failed = true; }); + const exited = () => proxy.exitCode !== null || proxy.signalCode !== null; + const invoke = async () => { + const startup = Math.min(deadline, Date.now() + 30_000); + while (Date.now() < startup) { + if (failed || exited()) throw new BudgetApiError("proxy-exited"); + const port = output.match(/Starting to serve on 127\.0\.0\.1:(\d+)/)?.[1]; + if (port) return await run(apiRequest(`http://127.0.0.1:${port}`, deadline)); + await sleep(Math.min(50, startup - Date.now())); + } + throw new BudgetApiError("proxy-deadline"); + }; + const stop = async () => { + if (!exited() && proxy.pid) { + proxy.kill("SIGTERM"); + const grace = Date.now() + 2000; + while (!exited() && Date.now() < grace) await sleep(25); + if (!exited()) { + proxy.kill("SIGKILL"); + const killed = Date.now() + 2000; + while (!exited() && Date.now() < killed) await sleep(25); + if (!exited()) throw new BudgetApiError("proxy-cleanup"); + } + } + }; + let result, failure, rejected = false; + try { result = await invoke(); } + catch (error) { failure = error; rejected = true; } + try { await stop(); } + catch (error) { + throw new AggregateError(rejected ? [failure, error] : [error], "Budget fixture API proxy cleanup failed"); + } + if (rejected) throw failure; + return result; +} diff --git a/tests/e2e/budget-cancellation.mjs b/tests/e2e/budget-cancellation.mjs new file mode 100644 index 000000000..c4b0ee075 --- /dev/null +++ b/tests/e2e/budget-cancellation.mjs @@ -0,0 +1,93 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { isDeepStrictEqual as equal } from "node:util"; +import { BudgetApiError } from "./budget-api-client.mjs"; + +const reasons = new Set(["BadRequest", "Unauthorized", "Forbidden", "NotFound", "AlreadyExists", + "Conflict", "Invalid", "Timeout", "ServerTimeout", "TooManyRequests", "InternalError", "ServiceUnavailable"]); +const validName = value => typeof value === "string" && /^[a-z0-9](?:[-a-z0-9]*[a-z0-9])?$/.test(value) + && value.length <= 63; +const requireIdentity = condition => { + if (!condition) throw new BudgetApiError("cancellation-identity-or-intent"); +}; + +export function cancellationFact(verb, resource, attempt, response, error) { + requireIdentity(["GET", "PATCH"].includes(verb) && ["namespaces", "karstasks"].includes(resource) + && Number.isInteger(attempt) && attempt >= 1 && attempt <= 3); + const status = response?.status ?? (error instanceof BudgetApiError ? error.httpStatus : 0); + const httpStatus = Number.isInteger(status) && status >= 0 && status <= 599 ? status : 0; + const body = response?.body; + const reason = body?.kind === "Status" && reasons.has(body.reason) ? body.reason + : response?.status === 200 ? "Success" : "Unclassified"; + return { stage: "accepted-work-cancellation", verb, resource, attempt, httpStatus, reason }; +} + +export async function cancelAcceptedTask({ created, binding, request, report, deadline, + maxAttempts = 3, now = Date.now, pause = ms => new Promise(resolve => setTimeout(resolve, ms)) }) { + requireIdentity(Number.isInteger(maxAttempts) && maxAttempts > 0 && maxAttempts <= 3 + && Number.isFinite(deadline)); + const name = created?.metadata?.name, namespace = created?.metadata?.namespace; + requireIdentity(validName(name) && validName(namespace) && created.metadata.uid + && Number.isInteger(created.metadata.generation) && created.spec?.execution?.launch === true + && created.spec.envelope?.budget?.scope === "GovernedInference" + && binding?.scope === "GovernedInference" && binding.taskUid === created.metadata.uid + && binding.account?.uid && binding.root?.workspaceUid); + const expected = structuredClone(created); + const authority = structuredClone(binding); + const desired = structuredClone(expected.spec); + desired.execution.launch = false; + const path = `/apis/kars.azure.com/v1alpha1/namespaces/${namespace}/karstasks/${name}`; + let failedVersion; + const call = async (verb, resource, attempt, path, body) => { + if (now() >= deadline) throw new BudgetApiError("cancellation-deadline"); + let response; + try { response = await request(verb, path, body); } + catch (error) { + report(cancellationFact(verb, resource, attempt, undefined, error)); + throw error instanceof BudgetApiError ? error : new BudgetApiError("transport"); + } + report(cancellationFact(verb, resource, attempt, response)); + if (now() >= deadline) throw new BudgetApiError("cancellation-deadline"); + return response; + }; + for (let attempt = 1; attempt <= maxAttempts; attempt++) { + const workspace = await call("GET", "namespaces", attempt, `/api/v1/namespaces/${namespace}`); + if (workspace.status !== 200) throw new BudgetApiError("cancellation-read", workspace.status); + requireIdentity(workspace.body?.kind === "Namespace" && workspace.body.apiVersion === "v1" + && workspace.body.metadata?.uid === authority.root.workspaceUid + && workspace.body.metadata.name === namespace && !workspace.body.metadata.deletionTimestamp); + const live = await call("GET", "karstasks", attempt, path); + if (live.status !== 200) throw new BudgetApiError("cancellation-read", live.status); + const current = live.body; + requireIdentity(current?.kind === "KarsTask" && current.apiVersion === "kars.azure.com/v1alpha1" + && current.metadata?.uid === expected.metadata.uid && current.metadata.name === name + && current.metadata.namespace === namespace && !current.metadata.deletionTimestamp + && current.metadata.generation === expected.metadata.generation + && equal(current.spec, expected.spec) && equal(current.status?.inferenceBudget, authority) + && typeof current.metadata.resourceVersion === "string" && current.metadata.resourceVersion.length > 0); + if (failedVersion !== undefined && current.metadata.resourceVersion === failedVersion) + throw new BudgetApiError("cancellation-conflict-without-fresh-version", 409); + const result = await call("PATCH", "karstasks", attempt, path, { + metadata: { uid: expected.metadata.uid, resourceVersion: current.metadata.resourceVersion }, + spec: { execution: { launch: false } }, + }); + // Only a genuine API Status Conflict authorizes a fresh, fully fenced retry. + if (result.status === 409 && result.body?.apiVersion === "v1" && result.body.kind === "Status" + && result.body.status === "Failure" && result.body.reason === "Conflict" && result.body.code === 409) { + failedVersion = current.metadata.resourceVersion; + if (attempt < maxAttempts && now() < deadline) await pause(Math.min(100, deadline - now())); + continue; + } + if (result.status !== 200) throw new BudgetApiError("cancellation-patch", result.status); + requireIdentity(result.body?.kind === "KarsTask" && result.body.apiVersion === "kars.azure.com/v1alpha1" + && result.body.metadata?.uid === expected.metadata.uid && result.body.metadata.name === name + && result.body.metadata.namespace === namespace && !result.body.metadata.deletionTimestamp + && result.body.metadata.generation === expected.metadata.generation + 1 + && typeof result.body.metadata.resourceVersion === "string" + && result.body.metadata.resourceVersion !== current.metadata.resourceVersion + && equal(result.body.spec, desired)); + return result.body; + } + throw new BudgetApiError("cancellation-conflict-limit", 409); +} diff --git a/tests/e2e/budget-workload-cases.mjs b/tests/e2e/budget-workload-cases.mjs index e3405330c..c4dbaa233 100644 --- a/tests/e2e/budget-workload-cases.mjs +++ b/tests/e2e/budget-workload-cases.mjs @@ -3,7 +3,7 @@ // Native admission/creation proof only. No account state or Pod readiness is fabricated. import assert from "node:assert/strict"; -import { spawn } from "node:child_process"; +import { withKindApi } from "./budget-api-client.mjs"; const kubeControllers = [ ["manager", "system:kube-controller-manager"], @@ -210,41 +210,6 @@ export async function workloadCases(request, wait, { namespace, controller, prin } export async function runWorkloadProof(options) { - const { root, context, kubectl, until } = options; - const config = JSON.parse(kubectl(["config", "view", "--minify", "-o", "json"])); - assert(config.contexts?.length === 1 && config.contexts[0].name === context - && config.clusters?.length === 1 - && ["localhost", "127.0.0.1", "[::1]"].includes(new URL(config.clusters[0].cluster.server).hostname), - "Budget workload proof requires the exact loopback Kind context"); - const proxy = spawn("kubectl", ["--context", context, "--request-timeout=20s", "proxy", - "--address=127.0.0.1", "--port=0"], { cwd: root, stdio: ["ignore", "pipe", "pipe"] }); - let output = "", port; - proxy.stdout.on("data", data => { output = (output + data).slice(-2048); }); - proxy.stderr.on("data", () => {}); - try { - await until(() => { - assert(proxy.exitCode === null, "Budget API proxy exited"); - port = output.match(/127\.0\.0\.1:(\d+)/)?.[1]; - return Boolean(port); - }, "budget workload API proxy"); - const request = async (method, path, body, actor) => { - const response = await fetch(`http://127.0.0.1:${port}${path}`, { - method, signal: AbortSignal.timeout(15_000), - headers: { "Content-Type": method === "PATCH" ? "application/merge-patch+json" : "application/json", - Accept: "application/json", ...(actor ? { "Impersonate-User": actor, "Impersonate-Group": "system:authenticated" } : {}) }, - body: body === undefined ? undefined : JSON.stringify(body), - }); - return { status: response.status, body: await response.json().catch(() => null) }; - }; - await workloadCases(request, until, options, - result => console.log("BUDGET-WORKLOAD " + JSON.stringify(result))); - } finally { - if (proxy.exitCode === null) { - proxy.kill("SIGTERM"); - await new Promise(resolve => { - const timer = setTimeout(() => { proxy.kill("SIGKILL"); resolve(); }, 5000); - proxy.once("exit", () => { clearTimeout(timer); resolve(); }); - }); - } - } + await withKindApi(options, request => workloadCases(request, options.until, options, + result => console.log("BUDGET-WORKLOAD " + JSON.stringify(result)))); } diff --git a/tests/e2e/inference-budget-enforcement.mjs b/tests/e2e/inference-budget-enforcement.mjs index 68e96995c..d672c5e71 100644 --- a/tests/e2e/inference-budget-enforcement.mjs +++ b/tests/e2e/inference-budget-enforcement.mjs @@ -13,6 +13,8 @@ import { PROVIDER, ENDPOINT, providerSource, budgetStageFacts, routerTemplateFacts, TLS_SERVER_EXTENSIONS, verifyFixtureCertificate, unsupportedOperationFact } from "./budget-fixture-route.mjs"; import { ownedRouterResolver, startForward, waitForOwnedRouter } from "./budget-router-readiness.mjs"; +import { withKindApi } from "./budget-api-client.mjs"; +import { cancelAcceptedTask } from "./budget-cancellation.mjs"; const root = fileURLToPath(new URL("../../", import.meta.url)); const context = "kind-kars-e2e"; @@ -295,9 +297,12 @@ async function scenario() { await request(provider, "/hold"); const accepted = request(tokenLeft.url, "/v1/chat/completions", message, 90_000).catch(() => ({ status: 0 })); await until(async () => (await request(provider, "/count")).value.count === count + 1, "accepted work before cancellation"); - const current = get("karstask", "budget-token-left"); - k(["patch", "karstask", current.metadata.name, "-n", namespace, "--type=merge", "--patch-file", "-"], - { metadata: { uid: current.metadata.uid, resourceVersion: current.metadata.resourceVersion }, spec: { execution: { launch: false } } }); + const cancellationDeadline = Date.now() + 30_000; + await withKindApi({ root, context, deadline: cancellationDeadline, + kubectl: args => k(args, undefined, cancellationDeadline) }, request => cancelAcceptedTask({ + created: createdTasks.get("budget-token-left"), binding, request, deadline: cancellationDeadline, + report: fact => console.log("BUDGET-CANCELLATION " + JSON.stringify(fact)), + })); await until(() => { const account = get("karsbudgetaccount", binding.account.name, binding.account.namespace); assert.equal(account.metadata.uid, binding.account.uid); From a5a1aea4f86e634fbb4cf36b6671ac1805ee9ad9 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Thu, 10 Sep 2026 21:28:33 +0200 Subject: [PATCH 16/16] test(budgets): qualify real-client cancellation with bounded integration setup Record the real kubectl literal patch-file reproduction, fenced API replacement, unchanged accounting assertions and combined MCP/budget qualification. Bound only external-tool tests using the established integration-test pattern; native cancellation deadlines and pure HTTP tests stay unchanged. Parent validation passed80budget cases, types/lint/syntax. Full current-base native cancellation acceptance remains required. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- cli/src/testing/budget-cancellation.test.ts | 9 +++-- .../2026-09-08-governed-inference-budgets.md | 39 +++++++++++++++++++ 2 files changed, 44 insertions(+), 4 deletions(-) diff --git a/cli/src/testing/budget-cancellation.test.ts b/cli/src/testing/budget-cancellation.test.ts index 0facb34bb..79408d74f 100644 --- a/cli/src/testing/budget-cancellation.test.ts +++ b/cli/src/testing/budget-cancellation.test.ts @@ -106,7 +106,7 @@ describe("native accepted-work cancellation API contract", () => { await kubeconfig(f.url, async (file, directory) => { const child = spawn("kubectl", ["--kubeconfig", file, "--cache-dir", join(directory, "cache"), "--context", context, "patch", "karstask", name, "-n", namespace, "--type=merge", "--patch-file", "-"], - { cwd: directory, stdio: ["pipe", "pipe", "pipe"] }); + { cwd: directory, stdio: ["pipe", "pipe", "pipe"], timeout: 20_000 }); let stderr = ""; child.stderr.on("data", data => { stderr = (stderr + data).slice(-8192); }); child.stdout.on("data", () => {}); @@ -121,10 +121,11 @@ describe("native accepted-work cancellation API contract", () => { expect(f.state.requests).toEqual([]); }); } finally { await f.close(); } - }); + }, 30_000); it.each(["conflict", "forbidden"])("uses a real kubectl proxy, exact %s status and owned cleanup", async mode => { const f = await fixture(); + f.options.deadline = Date.now() + 20_000; f.state.conflicts = 1; if (mode === "forbidden") f.state.reply = { code: 403, body: status(403, "Forbidden") }; const children: ReturnType[] = []; @@ -134,7 +135,7 @@ describe("native accepted-work cancellation API contract", () => { const pending = withKindApi({ root: directory, context, deadline: f.options.deadline, kubectl: (args: string[]) => execFileSync("kubectl", [...extra, "--context", context, ...args], - { encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 2000 }), + { encoding: "utf8", stdio: ["ignore", "pipe", "pipe"], timeout: 10_000 }), spawnProcess: (binary: string, args: string[], options: any) => { const child = spawn(binary, [...extra, ...args], options); children.push(child); @@ -158,7 +159,7 @@ describe("native accepted-work cancellation API contract", () => { expect(children).toHaveLength(1); expect(children[0].signalCode).toBe("SIGTERM"); } finally { await f.close(); } - }); + }, 30_000); it.each([[403, "Forbidden"], [422, "Invalid"], [429, "TooManyRequests"], [503, "ServiceUnavailable"]])( "retains actual HTTP %s and fails without retry", async (code, reason) => { diff --git a/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md index cf5728ed7..a8398692b 100644 --- a/docs/security-audits/2026-09-08-governed-inference-budgets.md +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -109,6 +109,45 @@ Fresh current-base native accounting, cancellation and failure-contract acceptance remain mandatory; successful readiness is not substituted for those assertions. +## Current MCP composition and cancellation fixture repair + +The actual MCP integration landing `12d1f2a3` is composed with the budget +candidate at `1af895bc98dad70f09d2467bef34a01b7ae6fb45`. Independent focused +review found no significant issues in the join. Exact combined qualification +passed2,453 Rust cases: controller binary1,282, router library1,138, router +binary15, governed telemetry12 and managed MCP6. Both crates independently +passed all-target/all-feature strict Clippy and formatting. The same source +passed93 CLI cases, type checking, six Python fixtures, four Helm renders, +lint and source gates. Minimum guarded free space was8.77GiB. + +The preceding4c2 native run `34509062774` remained failed (166 passed, one +failed). It now confirmed the required unsupported-route503/code/type response, +then stopped at the accepted-work cancellation command. The fixture used +`kubectl --patch-file -`, which treats `-` as a literal filename rather than +stdin. A controlled reproduction with the real kubectl made zero API requests. +The original native wrapper discarded the cause; no fictitious API409 is +claimed from that old result. + +Fixture-only repair `1ee92de3b1df5ef5349abd90581d956bc07849fa` sends a real +merge PATCH through the existing loopback Kind API client. Every attempt pins +the original Task UID/generation/spec, workspace UID and full budget binding, +then uses a fresh resourceVersion. Only an actual matching Status/Conflict409 +permits a retry, at most three attempts within the unchanged30-second bound. +Changed identity/intent, malformed responses,403/422, transport loss and a +lost acknowledgement remain fatal. Diagnostics retain fixed verb/resource/ +status/reason only, never API bodies, messages, headers or tokens. + +The settled8/uncertain30/reserved0 accepted-work assertions are unchanged. +Parent review and validation passed80 budget-related CLI cases including23 +cancellation cases, with the actual kubectl/proxy, plus types, lint and syntax. +Only real-tool integration tests use bounded30-second test deadlines and +10/20-second subprocess limits; pure HTTP test and native acceptance deadlines +were not relaxed. The joined production code is unchanged from qualified1af. + +Fresh current-base full native accounting and cancellation acceptance is still +required. These source and controlled-client results do not relabel either +historical failed native run or establish a successful final cluster run. + ## Summary Durable token ceilings and operator-configured maximum-price caps for governed