Skip to content

feat(routing): data_classification restricted keeps calls on self-hosted nodes - #139

Open
coyl wants to merge 6 commits into
mainfrom
feat/data-classification
Open

coyl wants to merge 6 commits into
mainfrom
feat/data-classification

Conversation

@coyl

@coyl coyl commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor

data_classification: restricted is a guarantee that a call stays on self-hosted nodes (VendorLocal, VendorOllama) with no cloud fallback. placement: nodes_only was only a request preference, and the auto-effort judge saw the prompt before placement applied.

  • Where: APIKeyPolicy.data_classification, Team.data_classification (+ Create/UpdateTeam), and per call GenerateTextRequest / ExplainRouteRequest, data_classification on /v1/chat/completions, the X-KTMS-Data-Classification header on every HTTP endpoint and x-ktms-data-classification gRPC metadata (both through the context, so media RPCs too). A call can restrict, never lift a key's or team's restriction. Migration adds the column to api_key and team.
  • Updates fail closed: SetAPIKeyPolicy and UpdateTeam keep the current value when the field is empty, so an older console or CLI cannot lift a restriction by leaving it out; "unrestricted" lifts it, and config apply sends that when the file drops the field.
  • Text: placement becomes nodes_only (nodes_first, cloud_first included; cloud_only is INVALID_ARGUMENT); no node model is FAILED_PRECONDITION no_capable_model, all nodes down UNAVAILABLE, the cloud is never tried. The auto-effort judge and the shadow evaluation judge run only on node models, else fall back / skip. The response cache is skipped, so a restricted answer cannot reach another key.
  • Restricted key or team: embeddings, Judge, EvaluateRollouts, speech and images need a node model (planMedia); SubmitBatch takes a node batch model or is refused.
  • Usage: data_classification: restricted in the event's meta; a caller's own data_classification attribute is stripped.
  • CLI / config: ktms keys mint|policy set --data-classification, ktms route explain --data-classification, DATA column in keys list (a team's restriction shows as restricted (team)), MCP mint_key/set_key_policy/explain_route, teams in the YAML file, schema regenerated with the enum.
  • Not covered, documented: LATENT_ENCODER_URL (latent guardrails, semantic cache), MCP servers, and an ollama model pointed at a hosted endpoint (nodes are vendors local and ollama).

Checked live on built binaries with two scripted vendors (cloud openai, node ollama): a restricted key, a key of a restricted team and the header all route to the node with cloud-first; with the node down a restricted key gets 503 while the cloud vendor receives nothing; an unrestricted key still falls back; cloud_only and an unknown value are 400; config plan after apply is empty.

Up to date with main after #138: /v1/messages takes the header (no body field, the Messages API has none) and answers a bad value in the Messages error shape. Part 1 of ai-process/docs#80 (hedged requests and the corrective JSON retry stay open there).

Review, second pass (with #141, docs#202 and docs#193)

Open, by severity:

  1. Conflicts with feat: pass temperature, top_p, stop and seed to the model #141. Both PRs add fields to GenerateTextRequest (17 here, 18–21 there) and touch gateway.pb.go, grpcapi/batch.go and server/openai_test.go; git merge-tree reports conflicts in those four files. Whichever lands second must be rebased and make proto re-run.
  2. Double bad_request: prefix in SubmitBatch (internal/grpcapi/batch.go). classify already returns bad_request: data_classification ..., and the batch loop wraps it as bad_request: item %q: %s, so an item with data_classification: "secret" yields bad_request: item "x": bad_request: data_classification "secret": .... feat: pass temperature, top_p, stop and seed to the model #141 trims the prefix at the same place; do the same here.
  3. Read-modify-write without a lock (internal/grpcapi/admin.go, SetAPIKeyPolicy and UpdateTeam). The current data_classification is read with GetAPIKeyByName / the team read, then written by a separate UPDATE outside a transaction. An old console sending SetAPIKeyPolicy{} (reads restricted) racing an admin's data_classification: "unrestricted" writes restricted back after the lift, or the other way round. Fix: one UPDATE with CASE/COALESCE for the keep-and-lift semantics, or SELECT ... FOR UPDATE inside the mutation.
  4. A doc comment moved onto the wrong function (internal/declarative/schema.go). keyPolicySchema and dataClassSchema were inserted between the // messageSchema describes a message as protojson reads it... comment and func messageSchema, so godoc shows the protojson description on keyPolicySchema and messageSchema has none.
  5. Comment length. AGENTS.md §6 allows 3 lines; the comment on classify (internal/grpcapi/restricted.go) is 4.

Resolved in e141aab:

  1. Conflicts with feat: pass temperature, top_p, stop and seed to the model #141: still open. Whichever PR lands second gets rebased and buf generate re-run.
  2. A batch item error now carries a single bad_request: prefix, with a test.
  3. SetAPIKeyPolicy and UpdateTeam no longer read the stored value first. A nil classification keeps it through COALESCE($n, data_classification) in the same UPDATE, so a racing lift cannot be undone. The DB-backed keep and lift tests pass against Postgres.
  4. The messageSchema doc comment is back on its function.
  5. The comment on classify is down to 3 lines.

Checked and fine: the restriction reaches the auto-effort judge, the shadow judge and the shadow model (withRestricted + nodes_only placement), both caches are skipped before lookupCache, and the HTTP stream context carries the flag.

coyl added 2 commits October 28, 2026 01:52
…ted nodes

On a request, a key's policy or a team: the chain keeps node models only with no cloud fallback,
the auto-effort and shadow judges run only on nodes, the response cache is skipped, media calls
of a restricted key need a node model and its batches are refused. Usage events carry it.
@github-actions

github-actions Bot commented Oct 28, 2026 •

Copy link
Copy Markdown

The latest Buf updates on your PR. Results from workflow Buf / buf (pull_request).

BuildFormatLintBreakingUpdated (UTC)
✅ passed✅ passed✅ passed✅ passedOct 29, 2026, 1:34 PM

coyl added 3 commits October 28, 2026 02:03
Update RPCs keep the restriction unless told "unrestricted", so older clients cannot lift it; the
header and gRPC metadata restrict every call through the context; callers cannot forge the usage
attribute; values are case-insensitive and checked in one place; a restricted batch names the
restriction only when it refused; keys list shows a team's restriction; the schema has the enum.
SetAPIKeyPolicy and UpdateTeam no longer read the stored classification
and write it back: COALESCE keeps it in the same statement, so a racing
lift is not undone. A batch item error carries one bad_request prefix,
and the messageSchema doc comment is back on its function.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant