diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 7157bc56c..5d399ab15 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -491,6 +491,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 @@ -636,7 +667,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/|shared/|runtimes/hermes/src/kars_runtime_hermes/plugin/sre|cli/src/(commands/sre|lib/sre|lib/namespace-ownership)|Cargo\.toml|Cargo\.lock|Makefile)' >/dev/null; then + | grep -E '^(controller/|inference-router/|a2a-gateway/|kars-a2a-core/|deploy/helm/|sandbox-images/|tests/e2e/|shared/|runtimes/hermes/src/kars_runtime_hermes/plugin/sre|cli/src/(commands/(sre|budget)|lib/sre|lib/namespace-ownership)|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/Cargo.lock b/Cargo.lock index c9a7577b7..92620f81e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2505,6 +2505,7 @@ dependencies = [ "regex", "reqwest 0.12.28", "rustls", + "rustls-pemfile", "schemars 1.2.1", "serde", "serde_json", @@ -2514,6 +2515,7 @@ dependencies = [ "thiserror 2.0.18", "time", "tokio", + "tokio-rustls", "tokio-tungstenite 0.28.0", "tracing", "tracing-subscriber", @@ -2546,6 +2548,7 @@ dependencies = [ "ed25519-dalek", "flate2", "futures", + "hex", "hkdf", "jsonwebtoken", "k8s-openapi", @@ -2558,6 +2561,7 @@ dependencies = [ "reqwest 0.12.28", "rustls", "rustls-pemfile", + "schemars 1.2.1", "serde", "serde_json", "serde_yaml", 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/cli/src/testing/budget-cancellation.test.ts b/cli/src/testing/budget-cancellation.test.ts new file mode 100644 index 000000000..79408d74f --- /dev/null +++ b/cli/src/testing/budget-cancellation.test.ts @@ -0,0 +1,281 @@ +// 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"], timeout: 20_000 }); + 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(); } + }, 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[] = []; + 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: 10_000 }), + 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(); } + }, 30_000); + + 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/cli/src/testing/budget-fixture-route.test.ts b/cli/src/testing/budget-fixture-route.test.ts new file mode 100644 index 000000000..71f996620 --- /dev/null +++ b/cli/src/testing/budget-fixture-route.test.ts @@ -0,0 +1,152 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { describe, expect, it } from "vitest"; +import { readFileSync, mkdirSync, rmSync } from "node:fs"; +import { execFileSync } from "node:child_process"; +import { randomUUID } from "node:crypto"; +import { fileURLToPath } from "node:url"; +const route = await import(new URL("../../../tests/e2e/budget-fixture-route.mjs", import.meta.url).href); + +describe("budget named-provider fixture and private readiness diagnostics", () => { + it("issues a real server leaf and rejects a CA used as the server certificate", () => { + const directory = fileURLToPath(new URL(`../../../.budget-tls-test-${randomUUID()}/`, import.meta.url)); + mkdirSync(directory, { mode: 0o700 }); + try { + const generate = (extensions: string[], host = "kars-inference-budget.kars-system.svc", san = true) => { + try { + execFileSync("openssl", ["req", "-x509", "-newkey", "rsa:2048", "-nodes", "-days", "1", + "-keyout", `${directory}/tls.key`, "-out", `${directory}/tls.crt`, + "-subj", `/CN=${host}`, ...(san ? ["-addext", `subjectAltName=DNS:${host}`] : []), ...extensions], + { stdio: "ignore", timeout: 30_000 }); + } catch { throw new Error("Budget fixture certificate generation failed"); } + return readFileSync(`${directory}/tls.crt`); + }; + expect(() => route.verifyFixtureCertificate(generate(route.TLS_SERVER_EXTENSIONS))).not.toThrow(); + expect(() => route.verifyFixtureCertificate(generate([ + "-addext", "basicConstraints=critical,CA:TRUE", + "-addext", "extendedKeyUsage=serverAuth", + ]))).toThrow("end entity, not a CA"); + expect(() => route.verifyFixtureCertificate(generate(route.TLS_SERVER_EXTENSIONS, "wrong.invalid"))) + .toThrow("SAN must match"); + expect(() => route.verifyFixtureCertificate(generate(route.TLS_SERVER_EXTENSIONS, undefined, false))) + .toThrow("SAN must match"); + expect(() => route.verifyFixtureCertificate(generate([ + "-addext", "basicConstraints=critical,CA:FALSE", "-addext", "extendedKeyUsage=clientAuth", + ]))).toThrow("server authentication"); + expect(() => route.verifyFixtureCertificate(generate([ + "-addext", "basicConstraints=critical,CA:FALSE", + ]))).toThrow("server authentication"); + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).toContain("...TLS_SERVER_EXTENSIONS"); + expect(harness.indexOf("verifyFixtureCertificate(certificate)")).toBeLessThan( + harness.indexOf('type: "kubernetes.io/tls"')); + } finally { + rmSync(directory, { recursive: true, force: true }); + } + }); + it("registers the exact anonymous local endpoint, without credentials or a native-provider fallback", () => { + const source = route.providerSource("kars-system"); + expect(source.metadata).toEqual({ name: "kars-inference-providers", namespace: "kars-system" }); + expect(source.stringData).toEqual({ KARS_PROVIDER_BUDGET_FIXTURE_ENDPOINT: route.ENDPOINT }); + expect(route.ENDPOINT).toBe("http://provider.budget-provider-fixture.svc.cluster.local:8000/v1"); + expect(route.PROVIDER).toBe("budget-fixture"); + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).toContain("providerId: PROVIDER, endpoint"); + expect(harness).toContain("create(providerSource(namespace))"); + expect(harness).not.toContain('provider: "ollama"'); + }); + it("requires actual generated policy identity to match the registered operator contract", () => { + expect(() => route.verifyFixturePolicy({ spec: { modelPreference: { primary: { + provider: route.PROVIDER, deployment: "fixture", + } } } })).not.toThrow(); + expect(() => route.verifyFixturePolicy({ spec: { modelPreference: { primary: { + provider: "ollama", deployment: "fixture", + } } } })).toThrow(); + }); + it("distinguishes HTTP readiness, transport timeout and the private budget gate without body/error output", () => { + const secret = "DO-NOT-EXPORT-PRIVATE-BODY"; + const unavailable = route.readinessFact("/readyz", { status: 503, value: secret }); + expect(unavailable.httpStatus).toBe(503); + expect(unavailable.ready).toBe(false); + expect(JSON.stringify(unavailable)).not.toContain(secret); + expect(route.readinessFact("/readyz", undefined, { name: "TimeoutError", message: secret }).category).toBe("timeout"); + expect(route.readinessFact("/readyz", { status: 200, value: "ok" }).ready).toBe(false); + expect(route.readinessFact("/readyz", { + status: 200, value: "governed inference authority and model contracts available", + }).ready).toBe(true); + expect(route.readinessFact("/healthz", { status: 200, value: "ok" }).ready).toBe(true); + }); + it("exports only fixed production diagnostic stages, numeric facts and counts", () => { + const secret = "DO-NOT-EXPORT-ENV-ARGV-BODY"; + const record = { target: "kars_controller::inference_budget::auth", level: "WARN", + fields: { budget_stage: "broker-authorization", source_line: 123, http_status: 403, + message: secret, env: secret, argv: [secret] } }; + const log = [secret, JSON.stringify(record), JSON.stringify(record), + JSON.stringify({ ...record, fields: { budget_stage: secret } }), + JSON.stringify({ ...record, target: "untrusted", fields: record.fields }), + JSON.stringify({ ...record, target: "__proto__::inference_budget::constructor" }), + JSON.stringify({ ...record, fields: { ...record.fields, source_line: secret, http_status: secret } }), + ].join("\n"); + expect(route.budgetStageFacts(log)).toEqual([ + { stage: "broker-authorization", sourceLine: 123, httpStatus: 403, count: 2 }, + { stage: "broker-authorization", count: 1 }, + ]); + expect(JSON.stringify(route.budgetStageFacts(log))).not.toContain(secret); + expect(route.budgetStageFacts(JSON.stringify({ + target: "kars_inference_router::inference_budget::client", level: "WARN", fields: { + budget_stage: "router-contract-match", provider_matches: false, endpoint_matches: false, + model_matches: true, bounds_valid: true, price_available: secret, provider_id: secret, + }, + }))).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"] }; + const pod = { metadata: { ownerReferences: [{ kind: "ReplicaSet", controller: true, uid: "rs" }] }, + spec: { containers: [router] } }; + const replica = { metadata: { uid: "rs", + ownerReferences: [{ kind: "Deployment", controller: true, uid: "deployment" }] } }; + const deployment = { metadata: { uid: "deployment" }, + spec: { template: { spec: { containers: [structuredClone(router)] } } } }; + expect(route.routerTemplateFacts(pod, replica, deployment).budgetBindingMatches).toBe(true); + deployment.spec.template.spec.containers[0].env[0].value = "PRIVATE-NEW-BINDING"; + const facts = route.routerTemplateFacts(pod, replica, deployment); + expect(facts.podOwnerUidMatches).toBe(true); + expect(facts.replicaOwnerUidMatches).toBe(true); + expect(facts.budgetBindingMatches).toBe(false); + expect(facts.environmentMatches).toBe(false); + expect(JSON.stringify(facts)).not.toContain("PRIVATE"); + }); + it("collects bounded filtered diagnostics before cleanup even when the readiness gate fails", () => { + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).toMatch(/^function diagnostics\(\)/m); + expect(harness).toContain('"--tail=256", "--limit-bytes=131072"'); + expect(harness).toContain("facts: budgetStageFacts(log)"); + expect(harness).toMatch(/finally \{\s+try \{\s+diagnostics\(\);\s+\} finally \{/); + expect(harness.indexOf("diagnostics();")).toBeLessThan(harness.indexOf("handle.stop()")); + }); +}); diff --git a/cli/src/testing/budget-router-readiness.test.ts b/cli/src/testing/budget-router-readiness.test.ts new file mode 100644 index 000000000..61138d3b9 --- /dev/null +++ b/cli/src/testing/budget-router-readiness.test.ts @@ -0,0 +1,394 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { EventEmitter } from "node:events"; +import { readFileSync } from "node:fs"; +import { spawn } from "node:child_process"; +import { describe, expect, it, vi } from "vitest"; + +const { ownedRouterResolver, waitForOwnedRouter, startForward, FixtureIdentityError, ForwardUnavailable } = + await import(new URL("../../../tests/e2e/budget-router-readiness.mjs", import.meta.url).href); + +const name = "budget-money-left"; +const workspace = "kars-system"; +const runtime = `kars-${name}`; +const image = `kars-inference-router:e2e@sha256:${"a".repeat(64)}`; +const ready = { status: 200, value: "governed inference authority and model contracts available" }; +const unavailable = { status: 503, value: "not ready — governed inference authority, provider or contracts unavailable" }; +const secret = "DO-NOT-EXPORT-PRIVATE-BINDING-BODY-OR-ARGV"; +const owner = (kind: string, name: string, uid: string, apiVersion = "apps/v1") => + ({ kind, apiVersion, name, uid, controller: true }); +const annotation = (suffix: string) => `kars.azure.com/${suffix}`; + +function fixture() { + const task = { + metadata: { name, namespace: workspace, uid: "task-uid", generation: 1 }, + spec: { execution: { launch: true } }, + status: { phase: "Ready", observedGeneration: 1, inferenceBudget: { + taskUid: "task-uid", account: { name: "account", namespace: workspace, uid: "account-uid" }, + authorizationDigest: secret, + } }, + }; + const sandbox = { + metadata: { name, namespace: workspace, uid: "sandbox-uid", + ownerReferences: [owner("KarsTask", name, task.metadata.uid, "kars.azure.com/v1alpha1")], + annotations: { [annotation("namespace-uid")]: "namespace-uid" } }, + spec: { inferenceBudgetRef: structuredClone(task.status.inferenceBudget) }, + }; + const namespace = { metadata: { name: runtime, uid: "namespace-uid", annotations: { + [annotation("namespace-claim-version")]: "v1", [annotation("sandbox-namespace")]: workspace, + [annotation("sandbox-name")]: name, [annotation("sandbox-uid")]: sandbox.metadata.uid, + }, labels: { [annotation("inference-budget")]: "v1" } } }; + const binding = { task: task.status.inferenceBudget, + sandbox: { name, namespace: workspace, uid: sandbox.metadata.uid }, + runtimeNamespace: runtime, runtimeNamespaceUid: namespace.metadata.uid, privacyEpoch: null }; + const labels = { [annotation("sandbox")]: name }; + const template = { metadata: { labels, annotations: { revision: "1" } }, + spec: { serviceAccountName: "sandbox", containers: [{ name: "inference-router", image, + env: [{ name: "KARS_INFERENCE_BUDGET_BINDING", value: JSON.stringify(binding) }], + args: [secret], volumeMounts: [] as object[] }] } }; + const deployment = { + metadata: { name, namespace: runtime, uid: "deployment-uid", generation: 1, + labels: { ...labels, [annotation("parent-namespace")]: workspace } }, + spec: { replicas: 1, selector: { matchLabels: labels }, template }, + status: { observedGeneration: 1 }, + }; + const policy = { spec: { modelPreference: { primary: { provider: "budget-fixture", deployment: "fixture" } } } }; + // Deliberately mutable API snapshots let each test model real UID/rollout races. + const objects = new Map([ + ["karstask", task], ["karssandbox", sandbox], ["namespace", namespace], ["deployment", deployment], + ["inferencepolicy", policy], + ]); + const state = { pods: [] as any[], time: 0 }; + const read = vi.fn(async (kind: string, resourceName: string) => + objects.get(kind === "replicaset" ? resourceName : kind) ?? null); + const pods = vi.fn(async () => state.pods); + function addPod(revision: number) { + const replicaName = `router-rs-${revision}`; + const replica = { metadata: { name: replicaName, namespace: runtime, uid: `rs-${revision}`, + ownerReferences: [owner("Deployment", name, deployment.metadata.uid)] }, + spec: { template: structuredClone(deployment.spec.template) } }; + const pod = { metadata: { name: `router-pod-${revision}`, namespace: runtime, uid: `pod-${revision}`, + labels, ownerReferences: [owner("ReplicaSet", replicaName, replica.metadata.uid)] }, + spec: structuredClone(replica.spec.template.spec), + status: { phase: "Running", containerStatuses: [{ name: "inference-router", state: { running: {} } }] } }; + pod.spec.containers[0].volumeMounts.push({ name: "kube-api-access-12345", readOnly: true, + mountPath: "/var/run/secrets/kubernetes.io/serviceaccount" }); + objects.set(replicaName, replica); + state.pods.push(pod); + return pod; + } + const pod = addPod(1); + const roll = () => { + const revision = ++deployment.metadata.generation; + deployment.status.observedGeneration = revision; + deployment.spec.template.metadata.annotations.revision = String(revision); + return addPod(revision); + }; + const resolve = ownedRouterResolver({ task: structuredClone(task), image, read, pods, deadline: 2000 }); + const forwards: { alive: ReturnType; stop: ReturnType; url: string }[] = []; + const openForward = vi.fn(async () => { + const forward = { alive: vi.fn(() => true), stop: vi.fn(async () => {}), + url: `http://127.0.0.1:${10000 + forwards.length}` }; + forwards.push(forward); + return forward; + }); + const report = vi.fn(); + const options = { resolve, openForward, report, deadline: 2000, now: () => state.time, + sleep: async (ms: number) => { state.time += ms; }, + probe: vi.fn(async (_url: string, path: string) => path === "/healthz" ? { status: 200, value: "ok" } : ready) }; + return { task, sandbox, namespace, deployment, pod, objects, state, read, pods, roll, + resolve, openForward, forwards, report, options }; +} + +describe("UID-fenced native budget router readiness", () => { + it("waits for late creation without swallowing API errors or adopting identities", async () => { + const f = fixture(); + const delayed = f.objects.get("deployment"); + f.objects.delete("deployment"); + f.options.sleep = async ms => { + f.state.time += ms; + f.objects.set("deployment", delayed); + }; + const result = await waitForOwnedRouter(f.options); + expect(result.pod.metadata.uid).toBe("pod-1"); + expect(f.openForward).toHaveBeenCalledTimes(1); + expect(f.state.time).toBe(500); + expect(f.options.probe.mock.calls.map(call => call[1])).toEqual(["/healthz", "/readyz"]); + expect(f.forwards[0].stop).not.toHaveBeenCalled(); + }); + + it("waits for late Pod creation and refuses an ambiguous initial Pod choice", async () => { + const f = fixture(); + f.state.pods = []; + expect(await f.resolve()).toBeNull(); + f.state.pods = [f.pod, { ...f.pod, metadata: { ...f.pod.metadata, name: "second", uid: "pod-2" } }]; + expect(await f.resolve()).toBeNull(); + expect((await f.resolve("pod-1"))?.pod.metadata.uid).toBe("pod-1"); + f.state.pods = [f.pod]; + expect((await waitForOwnedRouter(f.options)).pod.metadata.uid).toBe("pod-1"); + }); + + it("reselects a legitimate current-template rollout and stops only its old tunnel", async () => { + const f = fixture(); + f.options.probe = vi.fn(async (_url, path) => { + if (path === "/healthz") return { status: 200, value: "ok" }; + if (f.forwards.length === 1) { f.roll(); return unavailable; } + return ready; + }); + const result = await waitForOwnedRouter(f.options); + expect(result.pod.metadata.uid).toBe("pod-2"); + expect(f.openForward).toHaveBeenCalledTimes(2); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + expect(f.forwards[1].stop).not.toHaveBeenCalled(); + expect(JSON.stringify(f.report.mock.calls)).not.toContain(secret); + }); + + it("ignores terminating Pods and stale ReplicaSet templates, not current ownership", async () => { + const f = fixture(); + f.roll(); + expect((await f.resolve())?.pod.metadata.uid).toBe("pod-2"); + f.state.pods[1].metadata.deletionTimestamp = "2026-09-10T00:00:00Z"; + expect(await f.resolve()).toBeNull(); + expect(f.openForward).not.toHaveBeenCalled(); + }); + + it.each(["karstask", "karssandbox", "namespace", "deployment"])( + "rejects a same-name recreated %s, even after a transport failure", async kind => { + const f = fixture(); + f.options.probe = vi.fn(async () => { + f.objects.get(kind).metadata.uid = "foreign-uid"; + throw new Error(secret); + }); + await expect(waitForOwnedRouter(f.options)).rejects.toBeInstanceOf(FixtureIdentityError); + expect(f.openForward).toHaveBeenCalledTimes(1); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + expect(JSON.stringify(f.report.mock.calls)).not.toContain(secret); + }, + ); + + it.each(["sandbox-owner", "namespace-claim", "deployment-scope", "replica-owner", "pod-owner", + "image", "binding", "pod-environment", "mount", "account"])("fails closed for %s", async fault => { + const f = fixture(); + const replica = f.objects.get("router-rs-1"); + if (fault === "sandbox-owner") f.sandbox.metadata.ownerReferences[0].uid = "foreign"; + if (fault === "namespace-claim") f.namespace.metadata.annotations[annotation("sandbox-uid")] = "foreign"; + if (fault === "deployment-scope") f.deployment.metadata.labels[annotation("parent-namespace")] = "foreign"; + if (fault === "replica-owner") replica.metadata.ownerReferences[0].uid = "foreign"; + if (fault === "pod-owner") f.pod.metadata.ownerReferences[0].uid = "foreign"; + if (fault === "image") f.deployment.spec.template.spec.containers[0].image = "other:latest"; + if (fault === "binding") f.deployment.spec.template.spec.containers[0].env[0].value = secret; + if (fault === "pod-environment") f.pod.spec.containers[0].env.push({ name: "UNTRUSTED", value: secret }); + if (fault === "mount") f.pod.spec.containers[0].volumeMounts.push({ name: "foreign", mountPath: "/private" }); + if (fault === "account") { + await f.resolve(); + f.task.status.inferenceBudget.account.uid = "new-zero-account"; + } + let error: unknown; + try { await waitForOwnedRouter(f.options); } catch (caught) { error = caught; } + expect(error).toBeInstanceOf(FixtureIdentityError); + expect(String(error)).not.toContain(secret); + expect(f.openForward).not.toHaveBeenCalled(); + }); + + it("cannot return Ready when authority changes during successful HTTP responses", async () => { + const f = fixture(); + f.options.probe = vi.fn(async () => { + f.objects.get("karstask").metadata.uid = "replacement-task"; + return ready; + }); + await expect(waitForOwnedRouter(f.options)).rejects.toBeInstanceOf(FixtureIdentityError); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + }); + + it("checks ownership after tunnel startup and before issuing any HTTP probe", async () => { + const f = fixture(); + const open = f.options.openForward; + await expect(waitForOwnedRouter({ ...f.options, openForward: async () => { + const forward = await open(); + f.objects.get("namespace").metadata.uid = "replacement-namespace"; + return forward; + } })).rejects.toBeInstanceOf(FixtureIdentityError); + expect(f.options.probe).not.toHaveBeenCalled(); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + }); + + it("reconnects an expired owned forward to the same verified Pod", async () => { + const f = fixture(); + f.options.probe = vi.fn(async (_url, path) => { + if (f.forwards.length === 1) f.forwards[0].alive.mockReturnValue(false); + return path === "/healthz" ? { status: 200, value: "ok" } : ready; + }); + expect((await waitForOwnedRouter(f.options)).pod.metadata.uid).toBe("pod-1"); + expect(f.openForward).toHaveBeenCalledTimes(2); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + }); + + it.each(["503", "transport", "unqualified-200", "healthz-503"])( + "%s remains blocked until the original deadline, without reconnecting a live tunnel", async failure => { + const f = fixture(); + f.options.probe = vi.fn(async (_url, path) => { + if (failure === "transport") throw new Error(secret); + if (failure === "unqualified-200") return { status: 200, value: "ok" }; + if (failure === "healthz-503") return path === "/healthz" ? unavailable : ready; + return path === "/healthz" ? { status: 200, value: "ok" } : unavailable; + }); + await expect(waitForOwnedRouter(f.options)).rejects.toThrow("deadline"); + expect(f.state.time).toBe(2000); + expect(f.openForward).toHaveBeenCalledTimes(1); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + expect(JSON.stringify(f.report.mock.calls)).not.toContain(secret); + }, + ); + + it.each([401, 403])("fails immediately on real HTTP %s authorization rejection", async status => { + const f = fixture(); + f.options.probe = vi.fn(async () => ({ status, value: secret })); + await expect(waitForOwnedRouter(f.options)).rejects.toThrow("authorization rejected"); + expect(f.state.time).toBe(0); + expect(f.openForward).toHaveBeenCalledTimes(1); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + }); + + it("limits reconnections despite repeated legitimate rolls", async () => { + const f = fixture(); + f.options.probe = vi.fn(async (_url, path) => { + if (path === "/readyz") f.roll(); + return unavailable; + }); + await expect(waitForOwnedRouter({ ...f.options, deadline: 10_000 })).rejects.toThrow("reconnect bound"); + expect(f.openForward).toHaveBeenCalledTimes(4); + expect(f.forwards.every(forward => forward.stop.mock.calls.length === 1)).toBe(true); + }); + + it("bounds startup failures and retries only the exact owned target", async () => { + const f = fixture(); + f.openForward.mockRejectedValue(new ForwardUnavailable("Port-forward exited")); + await expect(waitForOwnedRouter({ ...f.options, deadline: 10_000 })).rejects.toThrow("reconnect bound"); + expect(f.openForward).toHaveBeenCalledTimes(4); + expect(f.options.probe).not.toHaveBeenCalled(); + expect(f.openForward.mock.calls.every(call => (call as any[])[0].pod.metadata.uid === "pod-1")).toBe(true); + }); + + it("cannot reset the deadline by receiving a late successful response", async () => { + const f = fixture(); + f.options.probe = vi.fn(async () => { + f.state.time = 2000; + return ready; + }); + await expect(waitForOwnedRouter(f.options)).rejects.toThrow("deadline"); + expect(f.openForward).toHaveBeenCalledTimes(1); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + }); + + it("does not open a forward when no owned Pod appears before the deadline", async () => { + const f = fixture(); + f.state.pods = []; + await expect(waitForOwnedRouter(f.options)).rejects.toThrow("deadline"); + expect(f.state.time).toBe(2000); + expect(f.openForward).not.toHaveBeenCalled(); + }); + + it("propagates non-NotFound API failures and closes the owned tunnel", async () => { + const f = fixture(); + f.options.probe = vi.fn(async () => { + f.read.mockRejectedValue(new Error("Sanitized API failure")); + return unavailable; + }); + await expect(waitForOwnedRouter(f.options)).rejects.toThrow("Sanitized API failure"); + expect(f.forwards[0].stop).toHaveBeenCalledTimes(1); + const harness = readFileSync(new URL("../../../tests/e2e/inference-budget-enforcement.mjs", import.meta.url), "utf8"); + expect(harness).toContain('"--ignore-not-found"'); + expect(harness.slice(harness.indexOf("async function until"), harness.indexOf("async function portForward"))) + .not.toContain("catch"); + expect(harness).toContain("task: createdTasks.get(name)"); + expect(harness).toContain("preconditions: { uid: secret.uid }"); + }); +}); + +describe("owned port-forward process lifecycle", () => { + it("serves a real local probe and reaps only the child that owns that connection", async () => { + const children: ReturnType[] = []; + const handles: { stop: () => Promise }[] = []; + try { + const forward = await startForward({ context: "kind-kars-e2e", cwd: ".", namespace: runtime, + target: "pod/owned", port: 8443, deadline: Date.now() + 3000, + register: (handle: { stop: () => Promise }) => handles.push(handle), + spawnProcess: () => { + const child = spawn(process.execPath, ["-e", ` + const server = require("node:http").createServer((_, response) => response.end("ok")); + server.listen(0, "127.0.0.1", () => + console.log("Forwarding from 127.0.0.1:" + server.address().port + " -> 8443")); + `], { stdio: ["ignore", "pipe", "pipe"] }); + children.push(child); + return child; + } }); + const response = await fetch(`${forward.url}/healthz`, { signal: AbortSignal.timeout(1000) }); + expect(response.status).toBe(200); + expect(await response.text()).toBe("ok"); + await forward.stop(); + expect(forward.alive()).toBe(false); + expect(children[0].signalCode).toBe("SIGTERM"); + } finally { + await Promise.all(handles.map(handle => handle.stop())); + } + }); + + function processFixture(output = true) { + let time = 0; + const child = Object.assign(new EventEmitter(), { + pid: 12345, exitCode: null as number | null, signalCode: null as string | null, + stdout: new EventEmitter(), stderr: new EventEmitter(), + kill: vi.fn((signal: string) => { child.signalCode = signal; return true; }), + }); + const register = vi.fn(); + const spawnProcess = vi.fn(() => child); + const options = { context: "kind-kars-e2e", cwd: ".", namespace: runtime, target: "pod/owned", port: 8443, + deadline: 1000, register, spawnProcess, now: () => time, + sleep: async (ms: number) => { + time += ms; + if (output) child.stdout.emit("data", Buffer.from("Forwarding from 127.0.0.1:12345 -> 8443\n")); + } }; + return { child, options, spawnProcess, register }; + } + + it("owns only its created process and cleanup is idempotent", async () => { + const f = processFixture(); + const unrelated = vi.fn(); + const forward = await startForward(f.options); + expect(forward.url).toBe("http://127.0.0.1:12345"); + expect(forward.alive()).toBe(true); + expect(f.register).toHaveBeenCalledWith(forward); + await Promise.all([forward.stop(), forward.stop()]); + expect(f.child.kill).toHaveBeenCalledExactlyOnceWith("SIGTERM"); + expect(unrelated).not.toHaveBeenCalled(); + expect(forward.alive()).toBe(false); + }); + + it("cleans up a startup deadline without exposing stderr or command output", async () => { + const f = processFixture(false); + const pending = startForward(f.options); + f.child.stderr.emit("data", Buffer.from(secret)); + await expect(pending).rejects.toThrow("startup deadline"); + expect(f.child.kill).toHaveBeenCalledExactlyOnceWith("SIGTERM"); + }); + + it("detects an exited forward immediately rather than swallowing it until timeout", async () => { + const f = processFixture(); + f.child.exitCode = 1; + await expect(startForward(f.options)).rejects.toBeInstanceOf(ForwardUnavailable); + expect(f.child.kill).not.toHaveBeenCalled(); + }); + + it("escalates only its unresponsive child and awaits its exit", async () => { + const f = processFixture(); + f.child.kill.mockImplementation(signal => { + if (signal === "SIGKILL") f.child.signalCode = signal; + return true; + }); + const forward = await startForward(f.options); + await forward.stop(); + expect(f.child.kill.mock.calls).toEqual([["SIGTERM"], ["SIGKILL"]]); + expect(forward.alive()).toBe(false); + }); +}); 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/cli/src/testing/kind-router-image.test.ts b/cli/src/testing/kind-router-image.test.ts new file mode 100644 index 000000000..ef04ba796 --- /dev/null +++ b/cli/src/testing/kind-router-image.test.ts @@ -0,0 +1,144 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { createHash } from "node:crypto"; +import { describe, expect, it } from "vitest"; + +const { TAG, FIXTURE, imageRow, checkedObject, platformConfig, qualifyNodes } = await import( + new URL("../../../tests/e2e/kind-router-image.mjs", import.meta.url).href +); +const manifestType = "application/vnd.oci.image.manifest.v1+json"; +const configType = "application/vnd.oci.image.config.v1+json"; +const indexType = "application/vnd.oci.image.index.v1+json"; +const PRIVATE = "DO-NOT-EXPORT-CONFIG-ENV-OR-ARGS"; +const nodes = ["kars-e2e-control-plane", "kars-e2e-worker"]; + +function blob(value: unknown) { + const bytes = Buffer.from(JSON.stringify(value)); + return { bytes, digest: `sha256:${createHash("sha256").update(bytes).digest("hex")}` }; +} + +function fixture(options: Record = {}) { + const config = blob({ os: "linux", architecture: "amd64", config: { Env: [PRIVATE], Cmd: [PRIVATE] }, + rootfs: { type: "layers", diff_ids: [`sha256:${"1".repeat(64)}`] } }); + const manifest = blob({ schemaVersion: 2, mediaType: manifestType, + config: { mediaType: configType, digest: config.digest, size: config.bytes.length }, + layers: [{ mediaType: "application/vnd.oci.image.layer.v1.tar+gzip", digest: `sha256:${"2".repeat(64)}`, size: 100 }] }); + const index = blob({ schemaVersion: 2, mediaType: indexType, manifests: [{ + mediaType: manifestType, digest: manifest.digest, size: manifest.bytes.length, + platform: { os: "linux", architecture: "amd64" }, + }] }); + const target = options.index ? index : manifest; + const mediaType = options.index ? indexType : manifestType; + const canonical = `docker.io/library/kars-inference-router@${target.digest}`; + const aliases = new Map(nodes.map(node => [node, options.existing === true])); + const calls: string[][] = []; + const blobs = new Map([config, manifest, index].map(value => [value.digest, value.bytes])); + const rows = (node: string) => [ + `REF TYPE DIGEST SIZE PLATFORMS LABELS`, + ...(options.missingNode && node.endsWith("worker") ? [] : [ + `${TAG} ${mediaType} ${options.configTarget ? config.digest : + options.differentNode && node.endsWith("worker") ? index.digest : target.digest} 100 B linux/amd64 -`]), + `import-date@${target.digest} ${mediaType} ${target.digest} 100 B linux/amd64 -`, + ...(aliases.get(node) ? [`${canonical} ${mediaType} ${options.foreignAlias ? config.digest : target.digest} 100 B linux/amd64 -`] : []), + ].join("\n"); + const run = (_binary: string, args: string[]) => { + calls.push(args); + const node = args[1]; + if (args[2] === "ctr") { + const action = args.slice(5); + if (action[0] === "images" && action[1] === "list") return Buffer.from(rows(node)); + if (action[0] === "content" && action[1] === "get") { + const bytes = blobs.get(action[2]); + if (!bytes) throw new Error("Missing metadata fixture"); + return options.tamper ? Buffer.concat([bytes, Buffer.from(" ")]) : bytes; + } + if (action[0] === "images" && action[1] === "check") { + const name = JSON.parse(action.at(-1)!.replace("name==", "")); + return Buffer.from(`REF TYPE DIGEST STATUS SIZE UNPACKED\n${name} ${mediaType} ${target.digest} ${options.incomplete ? "incomplete" : "complete"} (3/3) 100B/100B ${!options.notUnpacked}\n`); + } + if (action[0] === "images" && action[1] === "tag") { + expect(action).toEqual(["images", "tag", TAG, canonical]); + aliases.set(node, true); + return Buffer.from(canonical); + } + } + if (args[2] === "crictl") { + const ref = args.at(-1)!; + const visible = aliases.get(node) && !options.noCriEvent; + const status = { id: options.wrongConfig && args[3] === "inspecti" && ref !== TAG ? `sha256:${"4".repeat(64)}` : config.digest, + repoTags: [TAG], repoDigests: visible ? [canonical] : [] }; + if (args[3] === "images") return Buffer.from(JSON.stringify({ images: [status] })); + if (args[3] === "inspecti") return Buffer.from(JSON.stringify({ status, + info: { imageSpec: { config: { Env: [PRIVATE] } } } })); + } + throw new Error("Unexpected fixture command"); + }; + const optionsFor = (report: (value: unknown) => void) => ({ + nodes, expectedConfig: config.digest, run, report, pause: async () => {}, + platformFor: () => ({ os: "linux", architecture: options.wrongPlatform ? "arm64" : "amd64" }), + }); + return { config, manifest, index, target, canonical, calls, run, optionsFor }; +} + +describe("same-image Kind CRI preload fixture (unit orchestration, not native evidence)", () => { + it("proves missing canonical references and aliases the identical manifest on every node", async () => { + const f = fixture(); + const proofs: any[] = []; + const result = await qualifyNodes(f.optionsFor(value => proofs.push(value))); + expect(result.manifestDigest).toBe(f.target.digest); + expect(result.manifestDigest).not.toBe(f.config.digest); + expect(result.reference).toBe(`${FIXTURE}@${f.target.digest}`); + expect(proofs.filter(proof => proof.phase === "before").every( + proof => !proof.aliasPresent && !proof.criDigestPresent)).toBe(true); + expect(proofs.filter(proof => proof.phase === "verified")).toHaveLength(2); + expect(f.calls.filter(args => args.includes("tag"))).toHaveLength(2); + expect(f.calls.some(args => args.includes("--force") || args.includes("pull") || args.includes("delete"))).toBe(false); + expect(f.calls.filter(args => args.at(-1) === result.reference)).toHaveLength(2); + expect(JSON.stringify(proofs)).not.toContain(PRIVATE); + }); + + it("checks index platform content and preserves an already-correct alias without rewriting", async () => { + const f = fixture({ index: true, existing: true }); + const result = await qualifyNodes(f.optionsFor(() => {})); + expect(result.manifestDigest).toBe(f.index.digest); + expect(f.calls.some(args => args.includes("tag"))).toBe(false); + }); + + it.each(["configTarget", "tamper", "wrongPlatform", "incomplete", "notUnpacked"])( + "rejects %s before adding any reference", async fault => { + const f = fixture({ [fault]: true }); + await expect(qualifyNodes(f.optionsFor(() => {}))).rejects.toThrow(); + expect(f.calls.some(args => args.includes("tag"))).toBe(false); + }, + ); + + it("refuses a conflicting existing canonical alias rather than force-overwriting it", async () => { + const f = fixture({ existing: true, foreignAlias: true }); + await expect(qualifyNodes(f.optionsFor(() => {}))).rejects.toThrow("another image"); + expect(f.calls.some(args => args.includes("tag"))).toBe(false); + }); + + it.each(["noCriEvent", "wrongConfig"])("does not qualify %s after metadata tagging", async fault => { + const f = fixture({ [fault]: true }); + await expect(qualifyNodes(f.optionsFor(() => {}))).rejects.toThrow(); + }); + + it.each(["missingNode", "differentNode"])("requires the same complete image on every node: %s", async fault => { + const f = fixture({ [fault]: true }); + await expect(qualifyNodes(f.optionsFor(() => {}))).rejects.toThrow(); + }); + + it("rejects arbitrary image references, config descriptors and mismatched metadata hashes", () => { + expect(imageRow(`other ${manifestType} sha256:${"1".repeat(64)}`, TAG)).toBeNull(); + expect(() => imageRow(`${TAG} ${configType} sha256:${"1".repeat(64)}`, TAG)).toThrow(); + expect(() => checkedObject(Buffer.from("{}"), `sha256:${"1".repeat(64)}`)).toThrow(); + }); + + it("rejects ambiguous platform descriptors instead of choosing an arbitrary image", () => { + const f = fixture({ index: true }); + const value = JSON.parse(f.index.bytes.toString()); + value.manifests.push(structuredClone(value.manifests[0])); + expect(() => platformConfig(value, () => ({}), { os: "linux", architecture: "amd64" })).toThrow("ambiguous"); + }); +}); diff --git a/controller/Cargo.toml b/controller/Cargo.toml index 1216bee75..8e0a30bfd 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.workspace = true +rustls-pemfile.workspace = true # 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 faa028353..7299b34b2 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -608,8 +608,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() }, @@ -651,7 +651,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") } @@ -704,7 +706,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..abed5a674 --- /dev/null +++ b/controller/src/inference_budget/account.rs @@ -0,0 +1,116 @@ +// 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}; + +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.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, 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")] + #[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")); + 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"; + +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..a695530d9 --- /dev/null +++ b/controller/src/inference_budget/auth.rs @@ -0,0 +1,606 @@ +// 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, +}; + +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); + 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()) + || current.authorization_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()); + } + 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 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 + || 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()); + } + } 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/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..2c7dc06c6 --- /dev/null +++ b/controller/src/inference_budget/binding.rs @@ -0,0 +1,485 @@ +// 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 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 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)?, + name: task.name_any(), + uid: task.uid().ok_or(BudgetError::Identity)?, + }; + resource.validate()?; + Ok(resource) +} + +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()) + { + 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(), + }); + } + } + lineage.verify_pins(&pins)?; + Ok(lineage) +} + +pub(super) async fn needs_account(client: &Client, task: &KarsTask) -> Result { + // 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| { + explicitly_governed(&node.task.spec.envelope) + || node + .task + .status + .as_ref() + .is_some_and(|status| status.inference_budget.is_some()) + }) { + return Ok(true); + } + Ok(lineage.team.as_ref().is_some_and(|team| { + explicitly_governed(&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 !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()); + } + 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) +} + +/// 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 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 = lineage.team; + let namespaces: Api = Api::all(client.clone()); + 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: lineage.workspace_uid, + 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 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: 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?; + 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/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 new file mode 100644 index 000000000..be7c55094 --- /dev/null +++ b/controller/src/inference_budget/mod.rs @@ -0,0 +1,43 @@ +// 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 launch; +pub mod pod; +pub mod recovery; +pub mod scope; +pub mod service; +mod status; +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..40aebaa2f --- /dev/null +++ b/controller/src/inference_budget/pod.rs @@ -0,0 +1,349 @@ +// 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(()) +} + +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, + 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()); + } + + 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..27e318b8a --- /dev/null +++ b/controller/src/inference_budget/pod_tests.rs @@ -0,0 +1,130 @@ +// 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 _ = 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", + "metadata":{"name":"legacy","namespace":"workspace","uid":"legacy-uid"}, + "spec":{"inferenceRef":{"name":"legacy-inference"}} + })) + .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..47d2465ca --- /dev/null +++ b/controller/src/inference_budget/recovery.rs @@ -0,0 +1,229 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + account::{BOOTSTRAP, KarsBudgetAccount, MANAGED_BY, OWNER}, + config::Settings, + store::{Store, StoreError}, +}; +use crate::{ + inference_budget_contract::{BudgetError, RootKind, ledger::Mutation}, + 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 reconcile_account(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(()) +} + +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, + 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| { + 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?; + 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/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/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/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..f1470bbd7 --- /dev/null +++ b/controller/src/inference_budget/status_tests.rs @@ -0,0 +1,519 @@ +// 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 canonical_schema(mut value: serde_json::Value) -> serde_json::Value { + match &mut value { + serde_json::Value::Object(fields) => { + fields.remove("description"); + if let Some(serde_json::Value::Array(required)) = fields.get_mut("required") { + required.sort_by(|a, b| a.as_str().cmp(&b.as_str())); + } + if let Some(minimum) = fields.get_mut("minimum") { + *minimum = json!(minimum.as_f64().unwrap()); + } + for field in fields.values_mut() { + *field = canonical_schema(field.take()); + } + } + serde_json::Value::Array(items) => { + for item in items { + *item = canonical_schema(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", "ledger"] { + let path = |schema: &serde_json::Value| { + schema["schema"]["openAPIV3Schema"]["properties"]["status"]["properties"][field].clone() + }; + assert_eq!( + 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 new file mode 100644 index 000000000..6fd959993 --- /dev/null +++ b/controller/src/inference_budget/store.rs @@ -0,0 +1,396 @@ +// 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, +} + +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, + 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(()) + } + + pub(super) 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) + } + + 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. + 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), + ])); + 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 + .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)), + }?; + self.refresh_status( + &created.spec.root, + created.metadata.uid.as_deref().ok_or(StoreError::Missing)?, + None, + ) + .await + } + + 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 self.refresh_status(root, pinned_uid, None).await; + } + if account.annotations().get(BOOTSTRAP).map(String::as_str) != Some("pending") { + return Err(StoreError::Missing); + } + 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. + 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 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)?; + 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)), + } + } + 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.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 + .replace_status(&name_for_root(root), &PostParams::default(), &next), + ) + .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) + } + + /// 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)?; + // 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); + } + 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)] +#[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..69cf62687 --- /dev/null +++ b/controller/src/inference_budget/store_tests.rs @@ -0,0 +1,516 @@ +// 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.metadata.generation = Some(1); + 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), + ..Default::default() + }); + account.status = Some(super::super::status::project(&account, None)); + 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::Tokens, + 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, + FailWrite, + FailRead, + ReserveOnConflict, +} + +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) if matches!(state.fault, Fault::FailRead) => failure(503), + ("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()); + account.metadata.generation = Some(1); + 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::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; + 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 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)), + ); + 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 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 + .unwrap(); + assert_eq!( + initialized.annotations().get(BOOTSTRAP).map(String::as_str), + Some("sealed") + ); + assert_eq!( + initialized.status.as_ref().unwrap().phase, + Some(AccountStatusPhase::Active) + ); + 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); +} + +#[path = "status_tests.rs"] +mod status_tests; + +#[path = "recovery_tests.rs"] +mod recovery_tests; 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_budget_tests.rs b/controller/src/kars_task_budget_tests.rs new file mode 100644 index 000000000..ab69b0582 --- /dev/null +++ b/controller/src/kars_task_budget_tests.rs @@ -0,0 +1,270 @@ +// 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") + && 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 9e44e0f59..43dc9dd77 100644 --- a/controller/src/kars_task_execution.rs +++ b/controller/src/kars_task_execution.rs @@ -137,6 +137,33 @@ 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()))?; + 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). @@ -160,6 +187,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..133c4a22f 100644 --- a/controller/src/kars_task_reconciler.rs +++ b/controller/src/kars_task_reconciler.rs @@ -42,8 +42,14 @@ 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)] + InferenceBudget(#[from] crate::inference_budget::store::StoreError), #[error("Kubernetes API error: {0}")] Kube(#[from] kube::Error), #[error("JSON serialization error: {0}")] @@ -53,6 +59,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 +115,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 +123,7 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result degraded_status( @@ -188,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, @@ -232,11 +246,56 @@ async fn reconcile(task: Arc, ctx: Arc) -> Result task = Arc::new(prepared), + Err(error) => { + budget_pending = true; + 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()); + if budget_pending { + crate::inference_budget::launch::mark_pending( + &mut new_status, + &task, + "Budget admission or ancestor funding availability is pending", + ); + } + // 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 // reconcile — the governance status is already durable. - reconcile_execution(&ctx.client, &ns, &task, &mut new_status).await; + if budget_pending + && task + .spec + .execution + .as_ref() + .is_some_and(|execution| execution.launch) + { + crate::inference_budget::launch::retain_execution(&task, &mut new_status); + } else { + reconcile_execution(&ctx.client, &ns, &task, &mut new_status).await; + } let status_patch = json!({ "apiVersion": "kars.azure.com/v1alpha1", @@ -263,7 +322,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..0f41ecc59 100644 --- a/controller/src/kars_team_reconciler.rs +++ b/controller/src/kars_team_reconciler.rs @@ -212,7 +212,16 @@ 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 + }; + 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 +316,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 admissions paused: {error}. No new cadence or launch is admitted; existing Task UIDs and funded work are retained." + ) + } 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 +364,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/budget_interleaving_tests.rs b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs new file mode 100644 index 000000000..e25ff58a3 --- /dev/null +++ b/controller/src/kars_team_reconciler/budget_interleaving_tests.rs @@ -0,0 +1,513 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Full Team reconcile against a stateful API, with real ledger transitions. + +use super::*; +use crate::inference_budget::{ + account::{ + BOOTSTRAP, KarsBudgetAccount, KarsBudgetAccountSpec, KarsBudgetAccountStatus, MANAGED_BY, + OWNER, name_for_root, + }, + config::Settings, +}; +use crate::inference_budget_contract::{ + AccountReference, AttemptCommand, AttemptKey, BudgetScope, ExecutionIdentity, ReserveRequest, + ResourceIdentity, RootIdentity, RootKind, Settlement, TaskAuthority, TaskBudgetBinding, Usage, + ledger::Ledger, + tariffs::{MaximumPrice, ModelContract, Operation, OutputField}, +}; +use crate::kars_task::{KarsTaskStatus, TaskBlueprint, TaskModel}; +use crate::kars_team::TeamRole; +use std::sync::atomic::{AtomicBool, Ordering}; + +static ENV_LOCK: tokio::sync::Mutex<()> = 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::Tokens, + 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), + ..Default::default() + }); + *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/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..33937e32b 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(), )); @@ -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/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 01ff8dbdb..425541064 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -36,6 +36,12 @@ mod fedcred; mod fedcred_reaper; mod field_managers; 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; @@ -78,6 +84,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; @@ -137,6 +144,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 f719d93b9..28f0573ab 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)) +} + /// Preserve the existing 128-bit content identifier used by skills, profiles /// and commons. Signed receipt payloads continue using the full SHA-256 digest. pub fn content_digest(bytes: &[u8]) -> String { diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 417adad6e..a37e386c7 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -65,6 +65,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. @@ -3185,6 +3180,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/controller/src/task_identity.rs b/controller/src/task_identity.rs new file mode 100644 index 000000000..30337ed6b --- /dev/null +++ b/controller/src/task_identity.rs @@ -0,0 +1,335 @@ +// 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.generation != Some(node.generation) + || 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..848f45a01 --- /dev/null +++ b/controller/src/task_identity_tests.rs @@ -0,0 +1,289 @@ +// 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 _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + 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 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; + 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/deploy/helm/kars/files/inference-budget-admission.json b/deploy/helm/kars/files/inference-budget-admission.json new file mode 100644 index 000000000..dad745d8e --- /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.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" + }] + } + }, + { + "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 7de147b27..e5fd50ee2 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..7d3655307 --- /dev/null +++ b/deploy/helm/kars/templates/crd-karsbudgetaccount.yaml @@ -0,0 +1,142 @@ +# 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.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: + 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 + 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..08e51e599 100644 --- a/deploy/helm/kars/templates/crd-karsprofile.yaml +++ b/deploy/helm/kars/templates/crd-karsprofile.yaml @@ -57,6 +57,13 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + 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 78fd122f2..d9554b1a7 100644 --- a/deploy/helm/kars/templates/crd-karstask.yaml +++ b/deploy/helm/kars/templates/crd-karstask.yaml @@ -181,6 +181,13 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + 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 @@ -327,9 +334,10 @@ 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))' + 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: 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'']' @@ -345,10 +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 + 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 bbd68ceb7..867553928 100644 --- a/deploy/helm/kars/templates/crd-karsteam.yaml +++ b/deploy/helm/kars/templates/crd-karsteam.yaml @@ -218,6 +218,13 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + 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 @@ -460,6 +467,13 @@ spec: description: Optional resource budget for the whole task subtree. nullable: true properties: + scope: + type: string + 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 @@ -564,10 +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 + nullable: true + description: Lifetime Team-UID inference account; cadence runs never reset its balance. + required: [name, namespace, uid] + properties: + 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/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 585b0f867..4480574fc 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..cd82a913a --- /dev/null +++ b/deploy/helm/kars/tests/src/inference-budget.test.ts @@ -0,0 +1,148 @@ +// 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.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 => { + 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 c7e686b6f..25dd91f99 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/governed-inference-budgets.md b/docs/governed-inference-budgets.md new file mode 100644 index 000000000..8f3a4482a --- /dev/null +++ b/docs/governed-inference-budgets.md @@ -0,0 +1,306 @@ +# 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 [security audit](security-audits/2026-09-08-governed-inference-budgets.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. + +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 +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. +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 +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. + +## 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. + +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: + +- 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. +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 +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. + +## Native enforcement fixture + +The existing full E2E harness runs `tests/e2e/inference-budget-enforcement.mjs` +against real routers and the private broker in its disposable Kind cluster. +It verifies the built router's manifest and platform content on **every** Kind +node, adds only a missing same-image canonical containerd alias, and checks both +canonical and literal `image:tag@manifest` references through CRI. It never +force-tags, substitutes a config ID for a manifest, or changes production image +pins or pull policies. The fixture registers its named local provider before +launch, and uses an ephemeral `CA:FALSE`, `serverAuth` TLS leaf with the exact +private broker DNS SAN. Certificate files and UID-owned fixture Secrets are +removed during cleanup; diagnostics contain only bounded, fixed stage facts. + +Readiness pins the fixture-created Task UID, its account binding, Sandbox UID, +claimed Namespace UID, and Deployment UID. Only nonterminating, Running Pods +owned through that Deployment's current ReplicaSet template, using the exact +verified image and private budget binding, can be forwarded. A legitimate Pod +roll or exited owned tunnel permits at most three reconnections within the +original 120-second deadline. Both `/healthz` and the private `/readyz` contract +must pass, followed by another ownership check. A 503 or transport error alone +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/docs/security-audits/2026-09-08-governed-inference-budgets.md b/docs/security-audits/2026-09-08-governed-inference-budgets.md new file mode 100644 index 000000000..a8398692b --- /dev/null +++ b/docs/security-audits/2026-09-08-governed-inference-budgets.md @@ -0,0 +1,319 @@ +# Security Audit — Governed inference budgets (v1) + +Date: **2026-09-08 UTC** +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/`, +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> + +## 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. + +## 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 +inference only. No claim covers compute, GPU/VM, tool/MCP, storage, networking, +all-in task spend, invoice accuracy, taxes, or exchange rates. + +## T1: New capability / attack surface? (YES) + +- A private HTTPS budget broker and durable Kubernetes accounting CRD are new + control-plane surfaces. Router identities, all Task ancestors and complete + effective authorization must be checked before any dispatch grant. +- Budget-account UIDs, signed bootstrap ownership and replay fences prevent + name reuse, arbitrary-object adoption, concurrent overspend and reset-to-zero + recovery. Accounting stays separate from credential grants and tool costs. +- The shared Task identity helper captures live UID/RV evidence; it is not an + immutable registry until a consumer verifies and persists authoritative pins. + +## T2: Security-control change? (YES) + +- Finite inference requires router-only projected Pod/audience tokens, + exact TokenReview and live identity checks, current privacy proof, dedicated + admission, a qualified immutable router image, and versioned model contracts. +- Legacy/operator opaque control tokens and the SRE API token are not broker + authorization. Pending privacy qualification never authorizes RPCs or issuance. +- Unsupported generation, opaque tunnels, and unbounded mandatory moderation + fail closed rather than bypass accounting or disable existing guardrails. + +## T3: Availability / fail-open risk? (INCREASED for opted-in finite accounts) + +- Store/API/CAS, privacy/admission, TLS, image, contract, expiry, capacity or + lineage failures intentionally deny finite sends. Uncertain accepted work + remains fully funded; conservative full-context reservations may strand + otherwise unused quota. These are explicit limitations, not zero-cost claims. +- Unbounded standalone defaults remain unchanged. Missing accounting must never + be treated as empty accounting. Pending-state rollout behavior and accepted + cancellation/retry require the planned real Kubernetes qualification. + +## Verification + +### 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 +`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 | +| 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 | +| 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 | + +No Cargo command, dependency installation, local Docker test, deployment, customer +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. + +The subsequent parent-coordinated forward merge uses exact fixed privacy/SRE +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. + +## 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 + 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. +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. + +## Earlier bounded reviewer repair candidate (historical) + +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. + +## Original unsigned authoring state (historical) + +- 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. + +## Verdict + +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. 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/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/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.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..a2626acb3 --- /dev/null +++ b/inference-router/src/inference_budget/client_tests.rs @@ -0,0 +1,416 @@ +// 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); + +#[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() { + let fixture = Fixture::for_operation(100, false, Operation::AnthropicMessages).await; + Mock::given(wiremock::matchers::method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw(wire, "text/event-stream")) + .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).set_body_raw( + super::super::anthropic_cases::complete(), + "text/event-stream", + )) + .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; + 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 { + 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 { + 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, + output_field: OutputField::Tokens, + maximum_input_tokens: 10, + maximum_output_tokens: 20, + maximum_wire_bytes: 4096, + output_bound_includes_reasoning: true, + 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 { + 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) + )); + 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 { + 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).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, + 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..338bf6e8c --- /dev/null +++ b/inference-router/src/inference_budget/dispatch.rs @@ -0,0 +1,134 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + client::{AttemptGuard, Error}, + usage, +}; +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( + upstream: &UpstreamConfig, + method: &Method, + path: &str, + body: Bytes, +) -> anyhow::Result<(Bytes, Option)> { + let Some(client) = &upstream.inference_budget else { + return Ok((body, None)); + }; + let operation = operation(method, path)?; + 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..03797f662 --- /dev/null +++ b/inference-router/src/inference_budget/mod.rs @@ -0,0 +1,13 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[cfg(test)] +mod anthropic_cases; +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/unsupported_routes_tests.rs b/inference-router/src/inference_budget/unsupported_routes_tests.rs new file mode 100644 index 000000000..d37755040 --- /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: Default::default(), + 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/inference_budget/usage.rs b/inference-router/src/inference_budget/usage.rs new file mode 100644 index 000000000..9592afa7d --- /dev/null +++ b/inference-router/src/inference_budget/usage.rs @@ -0,0 +1,365 @@ +// 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, + native_final_usage: bool, + native_content: bool, +} + +impl StreamUsage { + pub fn new(operation: Operation) -> Self { + Self { + operation, + buffer: Vec::new(), + usage: None, + terminal: false, + invalid: false, + native_final_usage: false, + native_content: 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 => 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("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_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) + || (self.operation == Operation::AnthropicMessages && !self.native_final_usage) + { + 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\",\"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); + } + + #[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 01a766fc8..e01917590 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -36,6 +36,11 @@ 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; +#[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/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..fdf038525 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, @@ -262,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( @@ -314,15 +317,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 +360,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 +441,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 +487,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 { @@ -529,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, @@ -579,6 +590,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 +657,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 +730,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..4662d0a69 --- /dev/null +++ b/shared/inference_budget/catalog.rs @@ -0,0 +1,236 @@ +// 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 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; + } + 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 { + 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::Tokens, + 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/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 new file mode 100644 index 000000000..04a517b4d --- /dev/null +++ b/shared/inference_budget/ledger.rs @@ -0,0 +1,753 @@ +// 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 + .checked_add(self.settled)? + .checked_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(); + let node = next + .nodes + .get_mut(&authority.task.uid) + .ok_or(BudgetError::Corrupt)?; + node.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()?.checked_add(maximum)?) + || path.iter().any(|uid| { + let node = &self.nodes[uid]; + node.meters + .total() + .and_then(|total| total.checked_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.checked_add(maximum)?; + for uid in path { + let node = next.nodes.get_mut(&uid).ok_or(BudgetError::Corrupt)?; + node.meters.reserved = node.meters.reserved.checked_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.checked_add(charged)?; + } else { + meters.settled = meters.settled.checked_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..39577d526 --- /dev/null +++ b/shared/inference_budget/ledger_tests.rs @@ -0,0 +1,464 @@ +// 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; + +#[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(), + 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::Tokens, + 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..674b4393e --- /dev/null +++ b/shared/inference_budget/ledger_validation.rs @@ -0,0 +1,157 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +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()?; + 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.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.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.checked_add(used)?; + } + } + 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.checked_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); + } + + // 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.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.checked_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..4de4748fa --- /dev/null +++ b/shared/inference_budget/tariff_tests.rs @@ -0,0 +1,226 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use serde_json::{Value, 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::Completion, + 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 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(); + 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::Tokens; + 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::Output; + 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..7ec71fe18 --- /dev/null +++ b/shared/inference_budget/tariffs.rs @@ -0,0 +1,248 @@ +// 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}; + +#[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 { + #[serde(rename = "MaxTokens")] + Tokens, + #[serde(rename = "MaxCompletionTokens")] + Completion, + #[serde(rename = "MaxOutputTokens")] + Output, +} + +#[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::Tokens | OutputField::Completion + ) | (Operation::AnthropicMessages, OutputField::Tokens) + | (Operation::Responses, OutputField::Output) + ); + 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 + && price.price(self.maximum_input_tokens, self.maximum_output_tokens)? + > MAX_LEDGER_INTEGER + { + return Err(BudgetError::Overflow); + } + Ok(()) + } +} + +#[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), + }) + } +} + +#[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..4e1ca8746 --- /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 checked_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/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-fixture-route.mjs b/tests/e2e/budget-fixture-route.mjs new file mode 100644 index 000000000..34f4dcb2d --- /dev/null +++ b/tests/e2e/budget-fixture-route.mjs @@ -0,0 +1,114 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import assert from "node:assert/strict"; +import { X509Certificate } from "node:crypto"; + +export const PROVIDER = "budget-fixture"; +export const ENDPOINT = "http://provider.budget-provider-fixture.svc.cluster.local:8000/v1"; +export const TLS_SERVER_EXTENSIONS = [ + "-addext", "basicConstraints=critical,CA:FALSE", + "-addext", "keyUsage=critical,digitalSignature,keyEncipherment", + "-addext", "extendedKeyUsage=serverAuth", +]; + +export function verifyFixtureCertificate(certificate) { + const leaf = new X509Certificate(certificate); + assert(!leaf.ca, "Budget fixture TLS server certificate must be an end entity, not a CA"); + assert(leaf.checkHost("kars-inference-budget.kars-system.svc", { subject: "never" }), + "Budget fixture TLS server SAN must match the private broker hostname"); + assert(leaf.keyUsage?.includes("1.3.6.1.5.5.7.3.1"), + "Budget fixture TLS certificate must explicitly authorize server authentication"); +} + +export function providerSource(namespace) { + return { apiVersion: "v1", kind: "Secret", type: "Opaque", + metadata: { name: "kars-inference-providers", namespace }, + stringData: { KARS_PROVIDER_BUDGET_FIXTURE_ENDPOINT: ENDPOINT } }; +} + +export function verifyFixturePolicy(policy) { + assert(policy?.spec?.modelPreference?.primary?.provider === PROVIDER + && policy.spec.modelPreference.primary.deployment === "fixture" + && !policy.spec.provider, "Generated fixture policy must select the registered named provider"); +} + +export function readinessFact(endpoint, response, error) { + assert(["/healthz", "/readyz"].includes(endpoint)); + const ready = endpoint === "/healthz" ? response?.status === 200 + : response?.status === 200 && response.value === "governed inference authority and model contracts available"; + return { stage: endpoint === "/healthz" ? "router-http" : "governed-readiness", + endpoint, httpStatus: response?.status ?? 0, ready, + category: error ? (error.name === "TimeoutError" || error.name === "AbortError" ? "timeout" : "transport") + : ready ? "ready" + : response?.value === "not ready \u2014 governed inference authority, provider or contracts unavailable" + ? "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", + "router-target", "router-catalog", "router-credential"], + client: ["router-budget-client", "router-broker-http", "router-contract-match"], + }, + kars_controller: { + auth: ["broker-authorization", "broker-authorization-api"], + service: ["inference_budget_exhausted", "inference_budget_capacity", "inference_budget_authority", + "inference_budget_attempt_state", "inference_budget_contract", "inference_budget_closed", + "inference_budget_frozen", "inference_budget_unavailable"], + }, +}; + +export function budgetStageFacts(log) { + const facts = new Map(); + for (const line of log.split("\n")) { + let record; + try { record = JSON.parse(line); } catch { continue; } + const [crate, module, file, ...rest] = String(record?.target).split("::"); + const fields = record?.fields; + if (module !== "inference_budget" || rest.length || record?.level !== "WARN" + || !Object.hasOwn(STAGES, crate) || !Object.hasOwn(STAGES[crate], file) + || !STAGES[crate]?.[file]?.includes(fields?.budget_stage)) continue; + const fact = { stage: fields.budget_stage }; + if (Number.isInteger(fields.source_line) && fields.source_line > 0 && fields.source_line <= 10_000) + fact.sourceLine = fields.source_line; + if (Number.isInteger(fields.http_status) && fields.http_status >= 0 && fields.http_status <= 599) + fact.httpStatus = fields.http_status; + if (fact.stage === "router-contract-match") { + for (const key of ["provider_matches", "endpoint_matches", "model_matches", "bounds_valid", "price_available"]) + if (typeof fields[key] === "boolean") fact[key] = fields[key]; + } + const key = JSON.stringify(fact); + facts.set(key, { ...fact, count: (facts.get(key)?.count ?? 0) + 1 }); + } + return [...facts.values()]; +} + +export function routerTemplateFacts(pod, replica, deployment) { + const router = pod?.spec?.containers?.find(c => c.name === "inference-router"); + const current = deployment?.spec?.template?.spec?.containers?.find(c => c.name === "inference-router"); + assert(router && current, "Real router and current Deployment template required"); + const same = key => JSON.stringify(router[key] ?? null) === JSON.stringify(current[key] ?? null); + const binding = container => container.env?.find(e => e.name === "KARS_INFERENCE_BUDGET_BINDING")?.value; + return { + stage: "router-live-template", + podOwnerUidMatches: !!replica?.metadata?.uid && pod.metadata.ownerReferences?.some( + owner => owner.controller === true && owner.kind === "ReplicaSet" && owner.uid === replica.metadata.uid), + replicaOwnerUidMatches: !!deployment?.metadata?.uid && replica.metadata.ownerReferences?.some( + owner => owner.controller === true && owner.kind === "Deployment" && owner.uid === deployment.metadata.uid), + imageMatches: same("image"), commandMatches: same("command"), argsMatch: same("args"), + environmentMatches: same("env"), environmentSourcesMatch: same("envFrom"), + securityContextMatches: same("securityContext"), volumeMountsMatch: same("volumeMounts"), + budgetBindingMatches: !!binding(router) && binding(router) === binding(current), + }; +} diff --git a/tests/e2e/budget-router-readiness.mjs b/tests/e2e/budget-router-readiness.mjs new file mode 100644 index 000000000..9b411a159 --- /dev/null +++ b/tests/e2e/budget-router-readiness.mjs @@ -0,0 +1,286 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { spawn } from "node:child_process"; +import { isDeepStrictEqual as equal } from "node:util"; +import { readinessFact, routerTemplateFacts, verifyFixturePolicy } from "./budget-fixture-route.mjs"; + +const pause = ms => new Promise(resolve => setTimeout(resolve, ms)); +const annotation = (object, suffix) => object.metadata?.annotations?.[`kars.azure.com/${suffix}`]; +const label = (object, suffix) => object.metadata?.labels?.[`kars.azure.com/${suffix}`]; + +export class FixtureIdentityError extends Error {} +export class ForwardUnavailable extends Error {} + +function requireIdentity(condition, message) { + if (!condition) throw new FixtureIdentityError(`Budget fixture identity: ${message}`); +} + +function owner(object, kind, apiVersion) { + const owners = object.metadata?.ownerReferences?.filter(ref => ref.controller === true) ?? []; + requireIdentity(owners.length === 1 && owners[0].kind === kind && owners[0].apiVersion === apiVersion + && owners[0].name && owners[0].uid, "controller owner is missing or foreign"); + return owners[0]; +} + +function ownedBy(object, parent, kind, apiVersion) { + const ref = owner(object, kind, apiVersion); + requireIdentity(ref.name === parent.metadata.name && ref.uid === parent.metadata.uid, + "controller owner UID differs"); +} + +function binding(container) { + const entries = container?.env?.filter(entry => entry.name === "KARS_INFERENCE_BUDGET_BINDING") ?? []; + requireIdentity(entries.length === 1 && typeof entries[0].value === "string" && !entries[0].valueFrom, + "private budget binding is missing or ambiguous"); + try { return JSON.parse(entries[0].value); } + catch { throw new FixtureIdentityError("Budget fixture identity: private budget binding is malformed"); } +} + +function routerContainer(spec) { + const containers = spec?.containers?.filter(container => container.name === "inference-router") ?? []; + requireIdentity(containers.length === 1, "exactly one real router is required"); + return containers[0]; +} + +function sameTemplate(replica, deployment) { + const labels = template => Object.fromEntries(Object.entries(template.metadata?.labels ?? {}) + .filter(([key]) => key !== "pod-template-hash")); + return equal(replica.spec?.template?.spec, deployment.spec.template.spec) + && equal(replica.spec?.template?.metadata?.annotations ?? {}, deployment.spec.template.metadata?.annotations ?? {}) + && equal(labels(replica.spec?.template ?? {}), labels(deployment.spec.template)); +} + +function mountsMatch(actual, template) { + const expected = template.volumeMounts ?? []; + const mounts = actual.volumeMounts ?? []; + // Match the broker's router_mounts_match rule, including API-injected SA mounts. + return expected.every(mount => mounts.some(value => equal(value, mount))) + && mounts.filter(mount => !expected.some(value => equal(value, mount))).every(mount => + mount.readOnly === true && mount.subPath == null && mount.subPathExpr == null && mount.mountPropagation == null + && ((mount.name?.startsWith("kube-api-access-") + && mount.mountPath === "/var/run/secrets/kubernetes.io/serviceaccount") + || (mount.name === "azure-identity-token" && mount.mountPath === "/var/run/secrets/azure/tokens"))); +} + +// The Task UID comes from this fixture's create response, never a name-only GET. +// Pin each later identity on first observation; missing resources may appear late, +// but a same-name replacement must never become a new readiness authority. +export function ownedRouterResolver({ task, image, read, pods, deadline }) { + const name = task?.metadata?.name; + const workspace = task?.metadata?.namespace; + const runtime = `kars-${name}`; + requireIdentity(name && workspace && task.metadata.uid && task.metadata.generation, "created Task is required"); + requireIdentity(/^kars-inference-router:e2e@sha256:[a-f0-9]{64}$/.test(image), "verified fixture image is required"); + const uids = new Map([["karstask", task.metadata.uid]]); + let taskBinding; + + function pin(kind, object, expectedName, namespace) { + requireIdentity(object.metadata?.name === expectedName && object.metadata?.namespace === namespace + && object.metadata?.uid && !object.metadata.deletionTimestamp, "resource identity or lifetime differs"); + requireIdentity(!uids.has(kind) || uids.get(kind) === object.metadata.uid, "pinned resource UID changed"); + uids.set(kind, object.metadata.uid); + } + + return async preferredUid => { + const current = await read("karstask", name, workspace, deadline); + if (!current) return null; + pin("karstask", current, name, workspace); + requireIdentity(current.metadata.generation === task.metadata.generation + && current.spec?.execution?.launch === true, "Task authority changed"); + const bound = current.status?.inferenceBudget; + if (taskBinding) requireIdentity(equal(bound, taskBinding), "Task account binding changed"); + if (current.status?.phase !== "Ready" || current.status?.observedGeneration !== current.metadata.generation + || !bound) return null; + requireIdentity(bound.taskUid === current.metadata.uid && bound.account?.uid + && bound.authorizationDigest, "Task account authority is incomplete"); + taskBinding ??= structuredClone(bound); + const policy = await read("inferencepolicy", `${name}-inference`, workspace, deadline); + if (!policy) return null; + verifyFixturePolicy(policy); + + const sandbox = await read("karssandbox", name, workspace, deadline); + if (!sandbox) return null; + pin("karssandbox", sandbox, name, workspace); + ownedBy(sandbox, current, "KarsTask", "kars.azure.com/v1alpha1"); + requireIdentity(equal(sandbox.spec?.inferenceBudgetRef, bound), "Sandbox budget authority differs"); + const namespace = await read("namespace", runtime, undefined, deadline); + if (!namespace) return null; + pin("namespace", namespace, runtime, undefined); + requireIdentity(!(namespace.metadata.ownerReferences?.length), "runtime Namespace has a foreign owner"); + let claimed = true; + for (const [suffix, expected] of [ + ["namespace-claim-version", "v1"], ["sandbox-namespace", workspace], + ["sandbox-name", name], ["sandbox-uid", sandbox.metadata.uid], + ]) { + const value = annotation(namespace, suffix); + requireIdentity(value === undefined || value === expected, "runtime Namespace claim differs"); + claimed &&= value !== undefined; + } + if (!claimed) return null; + const namespaceUid = annotation(sandbox, "namespace-uid"); + requireIdentity(namespaceUid === undefined || namespaceUid === namespace.metadata.uid, + "Sandbox runtime Namespace UID differs"); + requireIdentity(annotation(namespace, "namespace-prestage") === undefined, "runtime Namespace claim is still prestaged"); + if (!namespaceUid) return null; + const fence = label(namespace, "inference-budget"); + requireIdentity(fence === undefined || fence === "v1", "runtime budget fence differs"); + if (!fence) return null; + + const deployment = await read("deployment", name, runtime, deadline); + if (!deployment) return null; + pin("deployment", deployment, name, runtime); + requireIdentity(label(deployment, "sandbox") === name && label(deployment, "parent-namespace") === workspace + && !(deployment.metadata.ownerReferences?.length) + && equal(deployment.spec?.selector?.matchLabels, { "kars.azure.com/sandbox": name }) + && !(deployment.spec.selector.matchExpressions?.length), "Deployment workload scope differs"); + const desired = routerContainer(deployment.spec?.template?.spec); + requireIdentity(desired.image === image, "Deployment is not the exact CRI-verified image"); + const desiredBinding = binding(desired); + requireIdentity(equal(desiredBinding.task, bound) + && equal(desiredBinding.sandbox, { name, namespace: workspace, uid: sandbox.metadata.uid }) + && desiredBinding.runtimeNamespace === runtime + && desiredBinding.runtimeNamespaceUid === namespace.metadata.uid, "Deployment budget binding differs"); + if (!deployment.metadata.generation || deployment.status?.observedGeneration !== deployment.metadata.generation) + return null; + + const candidates = []; + for (const pod of await pods(runtime, name, deadline)) { + requireIdentity(pod.metadata?.namespace === runtime && pod.metadata?.uid + && label(pod, "sandbox") === name, "Pod workload scope differs"); + const replicaOwner = owner(pod, "ReplicaSet", "apps/v1"); + const replica = await read("replicaset", replicaOwner.name, runtime, deadline); + if (!replica) continue; + requireIdentity(replica.metadata?.namespace === runtime + && replica.metadata?.name === replicaOwner.name, "ReplicaSet scope differs"); + ownedBy(pod, replica, "ReplicaSet", "apps/v1"); + ownedBy(replica, deployment, "Deployment", "apps/v1"); + if (pod.metadata.deletionTimestamp || replica.metadata.deletionTimestamp || !sameTemplate(replica, deployment)) + continue; + const actual = routerContainer(pod.spec); + const facts = routerTemplateFacts(pod, replica, deployment); + requireIdentity(Object.entries(facts).every(([key, value]) => + key === "stage" || key === "volumeMountsMatch" || value === true) + && actual.image === image && equal(binding(actual), desiredBinding) && mountsMatch(actual, desired) + && pod.spec.serviceAccountName === "sandbox", "current Pod template or private binding differs"); + if (pod.status?.phase === "Running" && pod.status.containerStatuses?.some( + container => container.name === "inference-router" && container.state?.running)) { + candidates.push({ pod, runtime, facts, deploymentGeneration: deployment.metadata.generation }); + } + } + return candidates.find(candidate => candidate.pod.metadata.uid === preferredUid) + ?? (candidates.length === 1 ? candidates[0] : null); + }; +} + +export async function startForward({ context, cwd, namespace, target, port, deadline, + register = () => {}, spawnProcess = spawn, now = Date.now, sleep = pause }) { + const child = spawnProcess("kubectl", ["--context", context, "port-forward", "--address", "127.0.0.1", + "-n", namespace, target, `:${port}`], { cwd, stdio: ["ignore", "pipe", "pipe"] }); + let output = "", failed = false, stopping; + child.stdout.on("data", data => { output = (output + data).slice(-2048); }); + child.stderr.on("data", () => {}); + child.on("error", () => { failed = true; }); + const exited = () => child.exitCode !== null || child.signalCode !== null; + const handle = { + alive: () => !failed && !exited(), + stop: () => stopping ??= (async () => { + if (exited() || !child.pid) return; + child.kill("SIGTERM"); + const end = now() + 2000; + while (!exited() && now() < end) await sleep(Math.min(50, end - now())); + if (!exited()) { + child.kill("SIGKILL"); + const killed = now() + 2000; + while (!exited() && now() < killed) await sleep(Math.min(50, killed - now())); + if (!exited()) throw new Error("Budget fixture owned port-forward cleanup failed"); + } + })(), + }; + register(handle); + try { + const end = Math.min(deadline, now() + 30_000); + while (now() < end) { + if (!handle.alive()) throw new ForwardUnavailable("Budget fixture port-forward exited"); + const match = output.match(/Forwarding from 127\.0\.0\.1:(\d+) ->/); + if (match) return Object.assign(handle, { url: `http://127.0.0.1:${match[1]}` }); + await sleep(Math.min(50, end - now())); + } + throw new ForwardUnavailable("Budget fixture port-forward startup deadline"); + } catch (error) { + await handle.stop(); + throw error; + } +} + +export async function waitForOwnedRouter({ resolve, openForward, probe, deadline, + maxReconnects = 3, report = () => {}, now = Date.now, sleep = pause }) { + requireIdentity(Number.isInteger(maxReconnects) && maxReconnects >= 0 && maxReconnects <= 3 + && Number.isFinite(deadline), "readiness bounds are invalid"); + let active, forward, attempts = 0; + const lastReadiness = new Map(); + const stop = async () => { + const previous = forward; + forward = undefined; + if (previous) await previous.stop(); + }; + try { + while (now() < deadline) { + const selected = await resolve(active?.pod.metadata.uid); + if (now() >= deadline) break; + if (!selected) { + await stop(); + } else { + const changed = active?.pod.metadata.uid !== selected.pod.metadata.uid; + if (changed || !forward?.alive()) { + await stop(); + if (now() >= deadline) break; + if (attempts >= maxReconnects + 1) throw new Error("Budget fixture port-forward reconnect bound"); + active = selected; + attempts++; + report({ stage: "router-owned-target", podUid: active.pod.metadata.uid, attempt: attempts }); + try { forward = await openForward(active, deadline); } + catch (error) { if (!(error instanceof ForwardUnavailable)) throw error; } + } + // kubectl resolves a Pod name while opening the tunnel. Fence that + // asynchronous window before even issuing the first readiness probe. + const connected = forward?.alive() ? await resolve(active.pod.metadata.uid) : null; + if (connected?.pod.metadata.uid === active?.pod.metadata.uid + && connected.deploymentGeneration === selected.deploymentGeneration + && forward?.alive()) { + let ready = true; + for (const path of ["/healthz", "/readyz"]) { + if (now() >= deadline) { ready = false; break; } + let response, error; + try { response = await probe(forward.url, path, Math.min(30_000, deadline - now())); } + catch (caught) { + if (caught instanceof FixtureIdentityError) throw caught; + error = caught; + } + const fact = readinessFact(path, response, error); + const key = JSON.stringify({ ...fact, podUid: active.pod.metadata.uid }); + if (lastReadiness.get(path) !== key) { + report({ ...fact, podUid: active.pod.metadata.uid }); + lastReadiness.set(path, key); + } + if (response?.status === 401 || response?.status === 403) + throw new Error("Budget fixture readiness authorization rejected"); + ready &&= fact.ready; + } + // A successful HTTP response does not authorize a replaced workload. + const current = await resolve(active.pod.metadata.uid); + if (ready && now() < deadline && forward.alive() + && current?.pod.metadata.uid === active.pod.metadata.uid + && current.deploymentGeneration === selected.deploymentGeneration) { + return { ...current, url: forward.url }; + } + } + } + if (now() < deadline) await sleep(Math.min(500, deadline - now())); + } + throw new Error("Budget fixture deadline: owned router private readiness"); + } catch (error) { + await stop(); + throw error; + } +} diff --git a/tests/e2e/budget-workload-cases.mjs b/tests/e2e/budget-workload-cases.mjs new file mode 100644 index 000000000..c4dbaa233 --- /dev/null +++ b/tests/e2e/budget-workload-cases.mjs @@ -0,0 +1,215 @@ +// 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 { withKindApi } from "./budget-api-client.mjs"; + +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) { + await withKindApi(options, request => workloadCases(request, options.until, options, + result => console.log("BUDGET-WORKLOAD " + JSON.stringify(result)))); +} diff --git a/tests/e2e/inference-budget-api.mjs b/tests/e2e/inference-budget-api.mjs new file mode 100644 index 000000000..55a60cc97 --- /dev/null +++ b/tests/e2e/inference-budget-api.mjs @@ -0,0 +1,181 @@ +// 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"; +import { runWorkloadProof } from "./budget-workload-cases.mjs"; + +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"]); +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"); diff --git a/tests/e2e/inference-budget-enforcement.mjs b/tests/e2e/inference-budget-enforcement.mjs new file mode 100644 index 000000000..d672c5e71 --- /dev/null +++ b/tests/e2e/inference-budget-enforcement.mjs @@ -0,0 +1,341 @@ +// 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 } from "node:child_process"; +import { mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +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, + 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"; +const namespace = "kars-system"; +const source = "budget-provider-fixture"; +const endpoint = ENDPOINT; +const scratch = join(root, `.budget-kind-${process.pid}`); +const forwards = []; +const createdTasks = new Map(); +const secrets = []; +const runtimes = new Set(); +let verifiedRouterReference; + +function execute(binary, args, input, deadline = Date.now() + 120_000) { + assert(Date.now() < deadline, "Budget fixture command deadline"); + try { + return execFileSync(binary, args, { + cwd: root, encoding: "utf8", input, stdio: ["pipe", "pipe", "pipe"], + timeout: Math.max(1, Math.min(120_000, deadline - Date.now())), + }); + } 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, deadline = Date.now() + 120_000) { + const timeout = Math.max(1, Math.min(30_000, deadline - Date.now())); + return execute("kubectl", ["--context", context, `--request-timeout=${timeout}ms`, ...args], + value === undefined ? undefined : JSON.stringify(value), deadline); +} +function create(value) { + const result = JSON.parse(k(["create", "-f", "-", "-o", "json"], value)); + if (value.kind === "Secret") { + const { name, namespace, uid } = result.metadata; + secrets.push({ name, namespace, uid }); + } + return result; +} +function get(kind, name, ns = namespace) { return JSON.parse(k(["get", kind, name, "-n", ns, "-o", "json"])); } +function optionalResource(kind, name, ns, deadline) { + const output = k(["get", kind, name, ...(ns ? ["-n", ns] : []), "--ignore-not-found", "-o", "json"], + undefined, deadline); + return output.trim() ? JSON.parse(output) : null; +} +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) { + const value = await check(); + if (value) return value; + await sleep(500); + } + throw new Error(`Budget fixture deadline: ${message}`); +} + +async function portForward(ns, target, port) { + return (await startForward({ context, cwd: root, namespace: ns, target, port, + deadline: Date.now() + 30_000, register: handle => forwards.push(handle) })).url; +} + +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: PROVIDER, deployment: "fixture" }, instructions: "Fixture only" }; +} + +function task(name, tokens, usdMicros, parent, launch) { + if (launch) runtimes.add(`kars-${name}`); + const created = 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(), + }, + }); + createdTasks.set(name, created); + return created; +} + +async function router(name) { + const deadline = Date.now() + 120_000; + const resolve = ownedRouterResolver({ + task: createdTasks.get(name), image: verifiedRouterReference, deadline, read: optionalResource, + pods: (runtime, name, deadline) => JSON.parse(k(["get", "pods", "-n", runtime, + "-l", `kars.azure.com/sandbox=${name}`, "-o", "json"], undefined, deadline)).items, + }); + return waitForOwnedRouter({ + resolve, deadline, + openForward: (selected, deadline) => { + console.log("BUDGET-TEMPLATE " + JSON.stringify(selected.facts)); + return startForward({ context, cwd: root, namespace: selected.runtime, + target: `pod/${selected.pod.metadata.name}`, port: 8443, deadline, + register: handle => forwards.push(handle) }); + }, + probe: (url, path, timeout) => request(url, path, undefined, timeout), + report: fact => console.log("BUDGET-READINESS " + JSON.stringify(fact)), + }); +} + +function diagnostics() { + const report = { complete: true, sources: [] }; + const targets = [ + { runtime: namespace, selector: "app.kubernetes.io/component=controller", container: "controller" }, + ...[...runtimes].map(runtime => ({ runtime, selector: "kars.azure.com/component=sandbox", + container: "inference-router" })), + ]; + for (const { runtime, selector, container } of targets) { + try { + const pods = JSON.parse(k(["get", "pods", "-n", runtime, "-l", selector, "-o", "json"])); + for (const pod of pods.items) { + const actualContainer = container === "controller" + ? pod.spec.containers.find(c => c.name === "controller" || c.name === "kars-controller")?.name + : container; + assert(actualContainer, "Diagnostic container identity unavailable"); + const log = k(["logs", "-n", runtime, pod.metadata.name, "-c", actualContainer, + "--tail=256", "--limit-bytes=131072"]); + const source = { component: container, podUid: pod.metadata.uid, facts: budgetStageFacts(log) }; + if (container === "inference-router") { + const owner = pod.metadata.ownerReferences.find(o => o.kind === "ReplicaSet" && o.controller); + const replica = get("replicaset", owner.name, runtime); + const deploymentOwner = replica.metadata.ownerReferences.find(o => o.kind === "Deployment" && o.controller); + source.template = routerTemplateFacts(pod, replica, get("deployment", deploymentOwner.name, runtime)); + } + report.sources.push(source); + } + } catch { report.complete = false; } + } + const directory = join(root, "e2e-diag", "standalone"); + mkdirSync(directory, { recursive: true, mode: 0o700 }); + writeFileSync(join(directory, "budget-readiness-stages.json"), JSON.stringify(report, null, 2)); + console.log("BUDGET-STAGES " + JSON.stringify(report)); +} + +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 values = JSON.parse(execute("helm", ["get", "values", "kars", "--kube-context", context, "-n", namespace, "--all", "-o", "json"])); + const imageProofs = []; + const imageDirectory = join(root, "e2e-diag", "standalone"); + mkdirSync(imageDirectory, { recursive: true, mode: 0o700 }); + const image = await prepareRouterImage({ + nodes, kube: k, values, + report: proof => { + imageProofs.push(proof); + console.log("BUDGET-IMAGE " + JSON.stringify(proof)); + writeFileSync(join(imageDirectory, "router-image-preflight.json"), JSON.stringify({ proofs: imageProofs }, null, 2)); + }, + }); + const digest = image.manifestDigest; + verifiedRouterReference = image.reference; + + 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", + ...TLS_SERVER_EXTENSIONS]); + const certificate = readFileSync(join(scratch, "tls.crt")); + verifyFixtureCertificate(certificate); + 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); + // Primary-model provider names select only explicitly registered endpoints. + // OLLAMA_ENDPOINT alone does not turn informational model metadata into routing intent. + create(providerSource(namespace)); + + 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: PROVIDER, endpoint, model: "fixture", operation: "ChatCompletions", outputField: "MaxTokens", + maximumInputTokens: 10, maximumOutputTokens: 20, maximumWireBytes: 4096, + outputBoundIncludesReasoning: true, maximumPrice: { kind: "perRequest", maximumMicros: 5 } }], + }; + 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; + 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; + 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 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); + 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 { + try { + diagnostics(); + } finally { + try { + const cleanup = await Promise.allSettled(forwards.map(handle => handle.stop())); + const secretCleanup = await Promise.allSettled(secrets.map(async secret => { + const current = optionalResource("secret", secret.name, secret.namespace); + if (!current) return; + assert(current.metadata.uid === secret.uid, "Fixture Secret cleanup refuses a replacement UID"); + k(["delete", "--raw", `/api/v1/namespaces/${secret.namespace}/secrets/${secret.name}`, "-f", "-"], + { apiVersion: "v1", kind: "DeleteOptions", preconditions: { uid: secret.uid } }); + })); + assert(secretCleanup.every(result => result.status === "fulfilled"), "UID-owned fixture Secret cleanup failed"); + assert(cleanup.every(result => result.status === "fulfilled"), "Owned fixture port-forward cleanup failed"); + } finally { + rmSync(scratch, { recursive: true, force: true }); + } + } +} diff --git a/tests/e2e/kind-router-image.mjs b/tests/e2e/kind-router-image.mjs new file mode 100644 index 000000000..47ee90e22 --- /dev/null +++ b/tests/e2e/kind-router-image.mjs @@ -0,0 +1,190 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +// Kind's ctr --digests import creates import-date@digest, not repository@digest. +// CRI resolves the latter by exact reference, separately from its config-ID key. +import assert from "node:assert/strict"; +import { createHash } from "node:crypto"; +import { execFileSync } from "node:child_process"; + +export const TAG = "docker.io/library/kars-inference-router:e2e"; +export const FIXTURE = "kars-inference-router:e2e"; +const REPOSITORY = "docker.io/library/kars-inference-router"; +const DIGEST = /^sha256:[a-f0-9]{64}$/; +const MANIFESTS = new Set([ + "application/vnd.oci.image.manifest.v1+json", + "application/vnd.docker.distribution.manifest.v2+json", +]); +const INDEXES = new Set([ + "application/vnd.oci.image.index.v1+json", + "application/vnd.docker.distribution.manifest.list.v2+json", +]); +const CONFIGS = new Set([ + "application/vnd.oci.image.config.v1+json", + "application/vnd.docker.container.image.v1+json", +]); + +function command(stage, binary, args) { + try { + return execFileSync(binary, args, { + stdio: ["ignore", "pipe", "pipe"], timeout: 30_000, maxBuffer: 2 * 1024 * 1024, + }); + } catch { + throw new Error(`Budget image preflight ${stage} command failed`); + } +} + +export function imageRow(output, name) { + const rows = output.toString("utf8").trim().split("\n") + .map(line => line.trim().split(/\s+/)).filter(fields => fields[0] === name); + assert(rows.length <= 1, "Ambiguous containerd image reference"); + if (!rows.length) return null; + const [, mediaType, digest] = rows[0]; + assert(DIGEST.test(digest) && (MANIFESTS.has(mediaType) || INDEXES.has(mediaType)), + "A real manifest/index target, not a config ID, is required"); + return { mediaType, digest }; +} + +export function checkedObject(bytes, digest) { + assert(DIGEST.test(digest) && bytes.length <= 1024 * 1024, "Invalid image metadata bound"); + assert(`sha256:${createHash("sha256").update(bytes).digest("hex")}` === digest, + "Image metadata content does not match its digest"); + return JSON.parse(bytes.toString("utf8")); +} + +export function platformConfig(root, load, platform, mediaType = root.mediaType) { + let manifest = root; + assert(root.schemaVersion === 2 && (!root.mediaType || root.mediaType === mediaType), "Invalid manifest schema"); + if (INDEXES.has(mediaType)) { + const selected = root.manifests?.filter(descriptor => + descriptor.platform?.os === platform.os && descriptor.platform?.architecture === platform.architecture); + assert(selected?.length === 1, "Image platform is missing or ambiguous"); + assert(MANIFESTS.has(selected[0].mediaType), "Expected a platform manifest"); + mediaType = selected[0].mediaType; + manifest = load(selected[0].digest); + assert(manifest.schemaVersion === 2 && (!manifest.mediaType || manifest.mediaType === mediaType), + "Platform manifest type differs from its descriptor"); + } + assert(MANIFESTS.has(mediaType) && CONFIGS.has(manifest.config?.mediaType) + && DIGEST.test(manifest.config?.digest) && Array.isArray(manifest.layers) + && manifest.layers.every(layer => DIGEST.test(layer.digest)), "Invalid platform content graph"); + const config = load(manifest.config.digest); + assert(config.os === platform.os && config.architecture === platform.architecture + && Array.isArray(config.rootfs?.diff_ids) && config.rootfs.diff_ids.length === manifest.layers.length, + "Image config does not cover the node platform and layers"); + return manifest.config.digest; +} + +function completeImage(output, name, digest) { + const row = output.toString("utf8").trim().split("\n").map(line => line.trim().split(/\s+/)) + .find(fields => fields[0] === name); + assert(row && row[2] === digest && row[3] === "complete" && row.at(-1) === "true", + "Platform image content must be complete and unpacked"); +} + +export async function qualifyNodes({ nodes, platformFor, expectedConfig, run, report, pause = ms => + new Promise(resolve => setTimeout(resolve, ms)) }) { + assert(DIGEST.test(expectedConfig) && nodes.length > 0 + && new Set(nodes).size === nodes.length + && nodes.every(name => /^kars-e2e-(?:control-plane|worker)(?:\d+)?$/.test(name)), + "Only exact owned Kind nodes and the built image config may be used"); + let manifestDigest; + const proofs = []; + for (const node of nodes) { + const platform = platformFor(node); + assert(platform.os === "linux" && ["amd64", "arm64"].includes(platform.architecture), + "Unsupported Kind fixture platform"); + const ctr = (...args) => run("docker", ["exec", node, "ctr", "-n", "k8s.io", ...args]); + const cri = (...args) => JSON.parse(run("docker", ["exec", node, "crictl", ...args]).toString("utf8")); + const filter = name => `name==${JSON.stringify(name)}`; + const source = imageRow(ctr("images", "list", filter(TAG)), TAG); + assert(source, "The existing loaded router tag is missing"); + assert(source.digest !== expectedConfig, "An image config ID cannot be an enforcement manifest"); + if (manifestDigest === undefined) manifestDigest = source.digest; + assert(source.digest === manifestDigest, "Kind nodes contain different router manifests"); + const load = digest => checkedObject(ctr("content", "get", digest), digest); + const root = load(source.digest); + const configDigest = platformConfig(root, load, platform, source.mediaType); + assert(configDigest === expectedConfig, "Loaded router is not the same image built by this fixture"); + completeImage(ctr("images", "check", "--snapshotter", "overlayfs", filter(TAG)), TAG, source.digest); + const tagStatus = cri("inspecti", "--output", "json", TAG).status; + assert(tagStatus?.id === configDigest, "CRI tag does not resolve the verified platform config"); + const canonical = `${REPOSITORY}@${source.digest}`; + const exactReference = `${FIXTURE}@${source.digest}`; + const before = imageRow(ctr("images", "list", filter(canonical)), canonical); + if (before) assert(before.digest === source.digest && before.mediaType === source.mediaType, + "Refusing to overwrite a canonical alias that targets another image"); + const criBefore = cri("images", "--output", "json").images; + assert(Array.isArray(criBefore), "CRI image inventory is unavailable"); + const listed = criBefore.filter(image => image.id === configDigest); + assert(listed.length === 1 && listed[0].repoTags?.includes(TAG), "CRI loaded tag identity differs"); + const digestVisibleBefore = listed[0].repoDigests?.includes(canonical) === true; + report({ phase: "before", node, reference: exactReference, canonical, manifestDigest: source.digest, + configDigest, platform: `${platform.os}/${platform.architecture}`, + aliasPresent: before !== null, criDigestPresent: digestVisibleBefore, contentComplete: true }); + if (!before) { + // Only add a new metadata reference to this same already-verified target. + // Never force-replace a name, pull another image, or edit any content. + ctr("images", "tag", TAG, canonical); + } + const after = imageRow(ctr("images", "list", filter(canonical)), canonical); + assert(after?.digest === source.digest && after.mediaType === source.mediaType, + "Canonical alias did not retain the exact manifest target"); + completeImage(ctr("images", "check", "--snapshotter", "overlayfs", filter(canonical)), + canonical, source.digest); + let visible = false; + for (let attempt = 0; attempt < 20; attempt++) { + const images = cri("images", "--output", "json").images; + assert(Array.isArray(images), "CRI image inventory became unavailable"); + visible = images.some(image => image.id === configDigest && image.repoDigests?.includes(canonical)); + if (visible) break; + await pause(500); + } + assert(visible, "CRI did not observe the canonical same-image alias within its bound"); + for (const reference of [canonical, exactReference]) { + const status = cri("inspecti", "--output", "json", reference).status; + assert(status?.id === configDigest && status.repoDigests?.includes(canonical), + "The exact digest-qualified reference does not resolve through CRI"); + } + const proof = { phase: "verified", node, reference: exactReference, canonical, + manifestDigest: source.digest, configDigest, platform: `${platform.os}/${platform.architecture}`, + aliasCreated: before === null, criResolved: true, contentComplete: true }; + report(proof); + proofs.push(proof); + } + return { manifestDigest, reference: `${FIXTURE}@${manifestDigest}`, proofs }; +} + +export async function prepareRouterImage({ nodes, kube, values, report }) { + const deadline = Date.now() + 90_000; + if (process.env.KARS_STANDALONE_CLUSTER_UID) { + const uid = kube(["get", "namespace", "kube-system", "-o", "jsonpath={.metadata.uid}"]).trim(); + assert(uid === process.env.KARS_STANDALONE_CLUSTER_UID, "Owned standalone cluster identity changed"); + } + assert(values.inferenceRouter?.image?.repository === "kars-inference-router" + && values.inferenceRouter.image.tag === "e2e" + && !(values.controller?.extraEnv ?? []).some(entry => entry.name === "INFERENCE_ROUTER_IMAGE"), + "Only the unchanged loaded router fixture reference is allowed"); + const expectedConfig = command("built-image-identity", "docker", + ["image", "inspect", "--format", "{{.Id}}", FIXTURE]).toString("utf8").trim(); + const metadata = JSON.parse(kube(["get", "nodes", "-o", "json"])); + const byName = new Map(metadata.items.map(node => [node.metadata.name, node])); + for (const node of nodes) { + assert(byName.get(node)?.metadata?.uid, "Kind node is not in the current Kubernetes cluster"); + const cluster = command("node-ownership", "docker", ["inspect", "--format", + "{{index .Config.Labels \"io.x-k8s.kind.cluster\"}}", node]).toString("utf8").trim(); + assert(cluster === "kars-e2e", "Refusing a node outside the owned Kind fixture"); + } + return qualifyNodes({ + nodes, expectedConfig, report, + platformFor: node => ({ os: byName.get(node).status.nodeInfo.operatingSystem, + architecture: byName.get(node).status.nodeInfo.architecture }), + run: (binary, args) => { + assert(Date.now() < deadline, "Budget image preflight exceeded its total bound"); + const operation = args[2] === "ctr" ? `${args[5]}-${args[6]}` : `cri-${args[3]}`; + assert(["images-list", "images-check", "images-tag", "content-get", "cri-images", "cri-inspecti"].includes(operation), + "Unexpected image fixture operation"); + return command(operation, binary, args); + }, + }); +} diff --git a/tests/e2e/run.sh b/tests/e2e/run.sh index 9fefa1321..fb17b2a35 100755 --- a/tests/e2e/run.sh +++ b/tests/e2e/run.sh @@ -3171,6 +3171,11 @@ main() { test_managed_mcp || fail "Managed MCP lifecycle/protocol gate failed" 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 "═══════════════════════════════════════════════════════"