feat(keys): key policies, per-key spend and prompt-cache prices

CHANGELOG.md

@@ -24,6 +24,10 @@ The first stable release, paired with Keystone. The Model Switchboard. (KTMS) 1.
 - Sign-in through an access proxy: an e-mail header from `TRUSTED_PROXY_CIDRS` (oauth2-proxy, Authelia,
   Authentik, Pomerium, Tailscale) or a verified Cloudflare Access token (`CF_ACCESS_TEAM_DOMAIN`,
   `CF_ACCESS_AUD`).
+- Key policies on the Keys page: allowed models and requests and tokens per minute, set when minting
+  or later (team leads for their team's keys), and each key's spend against its caps.
+- Prompt-cache prices (cached input, cache write) and the `prompt_cache` capability in the model
+  editor.
 
 ### Changed
 - Google sign-in is optional: configure any of the methods above, alone or together. Sign-in methods

README.md

@@ -316,8 +316,10 @@ own validation.
 **Limits.** Per-user rate limits and daily token budgets, and spend caps (USD per UTC day or month) for keys
 and users.
 
-**Keys.** A wizard mints client API keys: name and team, access, optional daily and monthly spend caps; the
-secret is shown once with a copy button. Revoke keys, assign vendor credentials to clients (write-only:
+**Keys.** A wizard mints client API keys: name and team, access, optional daily and monthly spend caps and
+request and token rates; the secret is shown once with a copy button. Each key's policy (allowed models,
+requests and tokens per minute) is edited in place, and its spend shows against its caps as KTMS counts it,
+wildcard caps included; team leads see and edit their team's keys. Revoke keys, assign vendor credentials to clients (write-only:
 stored encrypted by the gateway and never shown again) and re-encrypt all credentials after adding a new
 encryption key version.
 

internal/api/api.go

@@ -110,6 +110,9 @@ func (s *Server) Handler() http.Handler {
 		api("GET /api/proxy/keys", team, s.proxyListKeys)
 		api("POST /api/proxy/keys", team, s.proxyMintKey)
 		api("DELETE /api/proxy/keys/{name}", team, s.proxyRevokeKey)
+		api("PUT /api/proxy/keys/{name}/policy", team, s.proxySetKeyPolicy)
+		api("GET /api/proxy/spend", team, s.proxySpend)
+		api("GET /api/proxy/key-spend", team, s.proxyKeysSpend)
 		api("GET /api/proxy/nodes", globalRead, s.proxyListNodes)
 		api("POST /api/proxy/nodes/{id}/drain", adminOnly, s.proxyDrainNode)
 		api("GET /api/proxy/node-keys", globalRead, s.proxyListNodeKeys)

internal/api/api_test.go

@@ -33,6 +33,8 @@ type fakeAdmin struct {
 	lastMint         *pb.MintAPIKeyRequest
 	lastRevoke       string
 	spendLimits      []*pb.SpendLimitSpec
+	lastPolicy       *pb.SetAPIKeyPolicyRequest
+	lastSpend        *pb.GetSpendRequest
 }
 
 type fakeData struct {

internal/api/keypolicy_test.go

@@ -0,0 +1,96 @@
+package api
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"strings"
+	"sync"
+	"testing"
+
+	"google.golang.org/grpc"
+
+	pb "github.com/ai-process/ktms-admin/internal/gen/ktms/v1"
+)
+
+var spendMu sync.Mutex
+
+func (f *fakeAdmin) SetAPIKeyPolicy(_ context.Context, r *pb.SetAPIKeyPolicyRequest, _ ...grpc.CallOption) (*pb.SetAPIKeyPolicyResponse, error) {
+	f.lastPolicy = r
+	return &pb.SetAPIKeyPolicyResponse{}, nil
+}
+
+func (f *fakeAdmin) GetSpend(_ context.Context, r *pb.GetSpendRequest, _ ...grpc.CallOption) (*pb.GetSpendResponse, error) {
+	spendMu.Lock() // key-spend calls it concurrently
+	f.lastSpend = r
+	spendMu.Unlock()
+	return &pb.GetSpendResponse{Windows: []*pb.SpendWindow{{Period: "day", SpentNanousd: 1e9, LimitNanousd: 5e9, RemainingNanousd: 4e9, Metered: true}}}, nil
+}
+
+func TestKeyPolicyAccess(t *testing.T) {
+	h, fa := tenantServer(t)
+	lead, dev, admin := sessionCookie(t, "lead@corp.com"), sessionCookie(t, "dev@corp.com"), sessionCookie(t, "boss@x.com")
+	body := `{"allowedModels":["m1"],"rpmLimit":60,"tpmLimit":"1000"}`
+
+	if w := call(h, "PUT", "/api/proxy/keys/billing-service/policy", body, lead, csrf); w.Code != http.StatusForbidden || fa.lastPolicy != nil {
+		t.Fatalf("lead on another team's key = %d, want 403 and no change", w.Code)
+	}
+	if w := call(h, "PUT", "/api/proxy/keys/search-service/policy", body, dev, csrf); w.Code != http.StatusForbidden || fa.lastPolicy != nil {
+		t.Fatalf("member = %d, want 403", w.Code)
+	}
+	if w := call(h, "PUT", "/api/proxy/keys/search-service/policy", body, lead, csrf); w.Code != http.StatusOK {
+		t.Fatalf("lead on own key = %d %s", w.Code, w.Body)
+	}
+	if p := fa.lastPolicy; p.GetName() != "search-service" || p.GetPolicy().GetRpmLimit() != 60 || p.GetPolicy().GetTpmLimit() != 1000 || len(p.GetPolicy().GetAllowedModels()) != 1 {
+		t.Fatalf("policy sent = %v", p)
+	}
+	if w := call(h, "PUT", "/api/proxy/keys/billing-service/policy", body, admin, csrf); w.Code != http.StatusOK {
+		t.Fatalf("admin = %d %s", w.Code, w.Body)
+	}
+}
+
+func TestSpendAccess(t *testing.T) {
+	h, fa := tenantServer(t)
+	dev, viewer := sessionCookie(t, "dev@corp.com"), sessionCookie(t, "p@view.com")
+	for path, want := range map[string]int{
+		"/api/proxy/spend?subject_type=key&subject_id=search-service":  http.StatusOK,
+		"/api/proxy/spend?subject_type=team&subject_id=team-1":         http.StatusOK,
+		"/api/proxy/spend?subject_type=key&subject_id=billing-service": http.StatusForbidden,
+		"/api/proxy/spend?subject_type=user&subject_id=u1":             http.StatusForbidden,
+		"/api/proxy/spend?subject_type=team&subject_id=team-2":         http.StatusForbidden,
+	} {
+		if w := call(h, "GET", path, "", dev, nil); w.Code != want {
+			t.Errorf("member GET %s = %d, want %d", path, w.Code, want)
+		}
+	}
+	w := call(h, "GET", "/api/proxy/spend?subject_type=user&subject_id=u1", "", viewer, nil)
+	if w.Code != http.StatusOK || !strings.Contains(w.Body.String(), `"remaining_nanousd":"4000000000"`) || fa.lastSpend.GetSubjectId() != "u1" {
+		t.Fatalf("viewer = %d %s", w.Code, w.Body)
+	}
+}
+
+func TestKeySpendIsScopedAndBounded(t *testing.T) {
+	h, fa := tenantServer(t)
+	fa.keys = append(fa.keys, &pb.APIKeyInfo{Name: "old-search", TeamId: "team-1", Disabled: true})
+	var got struct {
+		Keys map[string]struct {
+			Windows []struct {
+				Period string `json:"period"`
+			} `json:"windows"`
+		} `json:"keys"`
+	}
+	w := call(h, "GET", "/api/proxy/key-spend", "", sessionCookie(t, "lead@corp.com"), nil)
+	if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil || w.Code != http.StatusOK {
+		t.Fatalf("lead = %d %s", w.Code, w.Body)
+	}
+	if _, ok := got.Keys["search-service"]; !ok || len(got.Keys) != 1 {
+		t.Fatalf("lead sees %v, want only the team's active key", got.Keys)
+	}
+	if got.Keys["search-service"].Windows[0].Period != "day" {
+		t.Fatalf("windows = %v", got.Keys["search-service"].Windows)
+	}
+	w = call(h, "GET", "/api/proxy/key-spend", "", sessionCookie(t, "boss@x.com"), nil)
+	if err := json.Unmarshal(w.Body.Bytes(), &got); err != nil || len(got.Keys) != 3 {
+		t.Fatalf("admin sees %v (%v), want the three active keys", got.Keys, err)
+	}
+}

internal/api/proxy.go

@@ -97,15 +97,11 @@ func (s *Server) proxyRevokeKey(w http.ResponseWriter, r *http.Request, u *auth.
 	ctx, cancel := context.WithTimeout(r.Context(), rpcTimeout)
 	defer cancel()
 	if !u.IsSuperAdmin() {
-		keysResp, err := s.admin.ListAPIKeys(ctx, &pb.ListAPIKeysRequest{})
+		isOwnTeam, err := s.ownsKey(ctx, u, name)
 		if err != nil {
 			backendErr(w, err)
 			return
 		}
-		isOwnTeam := false
-		for _, k := range keysOfTeam(keysResp.Keys, u.TeamID) {
-			isOwnTeam = isOwnTeam || k.Name == name
-		}
 		if !isOwnTeam {
 			writeErr(w, http.StatusForbidden, "forbidden", "cannot revoke keys belonging to other teams")
 			return
@@ -118,6 +114,81 @@ func (s *Server) proxyRevokeKey(w http.ResponseWriter, r *http.Request, u *auth.
 	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
 }
 
+// ownsKey reports whether a user who is not a super admin may manage the key:
+// only keys of their own team.
+func (s *Server) ownsKey(ctx context.Context, u *auth.User, name string) (bool, error) {
+	keysResp, err := s.admin.ListAPIKeys(ctx, &pb.ListAPIKeysRequest{})
+	if err != nil {
+		return false, err
+	}
+	for _, k := range keysOfTeam(keysResp.Keys, u.TeamID) {
+		if k.Name == name {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+
+// proxySetKeyPolicy replaces a key's allowed models and rates. A team lead may
+// set them on the team's keys, as they may mint keys without any: a key's policy
+// only narrows its team's models and budget.
+func (s *Server) proxySetKeyPolicy(w http.ResponseWriter, r *http.Request, u *auth.User) {
+	var policy pb.APIKeyPolicy
+	if err := readProto(r, &policy); err != nil {
+		backendErr(w, badBody(err))
+		return
+	}
+	name := r.PathValue("name")
+	ctx, cancel := context.WithTimeout(r.Context(), rpcTimeout)
+	defer cancel()
+	if !u.IsSuperAdmin() {
+		ok, err := s.ownsKey(ctx, u, name)
+		if err != nil {
+			backendErr(w, err)
+			return
+		}
+		if !ok {
+			writeErr(w, http.StatusForbidden, "forbidden", "cannot change keys belonging to other teams")
+			return
+		}
+	}
+	if _, err := s.admin.SetAPIKeyPolicy(ctx, &pb.SetAPIKeyPolicyRequest{Name: name, Policy: &policy}); err != nil {
+		backendErr(w, err)
+		return
+	}
+	writeJSON(w, http.StatusOK, map[string]bool{"ok": true})
+}
+
+// proxySpend reports a subject's spend as KTMS enforces it. Team members see
+// their own team's keys and their team; everything else needs global read.
+func (s *Server) proxySpend(w http.ResponseWriter, r *http.Request, u *auth.User) {
+	q := r.URL.Query()
+	req := &pb.GetSpendRequest{SubjectType: q.Get("subject_type"), SubjectId: q.Get("subject_id")}
+	ctx, cancel := context.WithTimeout(r.Context(), rpcTimeout)
+	defer cancel()
+	if !u.CanReadGlobal() {
+		allowed := req.SubjectType == "team" && req.SubjectId == u.TeamID && u.TeamID != ""
+		if req.SubjectType == "key" {
+			ok, err := s.ownsKey(ctx, u, req.SubjectId)
+			if err != nil {
+				backendErr(w, err)
+				return
+			}
+			allowed = ok
+		}
+		if !allowed {
+			writeErr(w, http.StatusForbidden, "forbidden", "only your team's spend is visible to you")
+			return
+		}
+	}
+	resp, err := s.admin.GetSpend(ctx, req)
+	if err != nil {
+		backendErr(w, err)
+		return
+	}
+	writeProto(w, resp)
+}
+
 func hasField(m map[string]json.RawMessage, keys ...string) bool {
 	for _, k := range keys {
 		if _, ok := m[k]; ok {
@@ -916,6 +987,54 @@ func countersDown(err error) bool {
 	return c == codes.Unavailable || c == codes.DeadlineExceeded
 }
 
+// proxyKeysSpend reads the spend of every active key the user may see in one
+// browser request, at most maxSpendReads at a time against the gateway. Caps
+// come from GetSpend itself, so a wildcard cap counts like a key's own.
+func (s *Server) proxyKeysSpend(w http.ResponseWriter, r *http.Request, u *auth.User) {
+	ctx, cancel := context.WithTimeout(r.Context(), rpcTimeout)
+	defer cancel()
+	resp, err := s.admin.ListAPIKeys(ctx, &pb.ListAPIKeysRequest{})
+	if err != nil {
+		backendErr(w, err)
+		return
+	}
+	keys := resp.Keys
+	if !u.CanReadGlobal() {
+		keys = keysOfTeam(keys, u.TeamID)
+	}
+	out := map[string]any{}
+	sem := make(chan struct{}, maxSpendReads)
+	var (
+		wg sync.WaitGroup
+		mu sync.Mutex
+	)
+	for _, k := range keys {
+		if k.Disabled {
+			continue
+		}
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			sem <- struct{}{}
+			defer func() { <-sem }()
+			spend, err := s.admin.GetSpend(ctx, &pb.GetSpendRequest{SubjectType: "key", SubjectId: k.Name})
+			var row any = map[string]any{"error": "unavailable"}
+			if err == nil {
+				if v, encErr := protoValue(spend); encErr == nil {
+					row = v
+				}
+			} else {
+				log.Warn().Err(err).Str("key", k.Name).Msg("key spend unavailable")
+			}
+			mu.Lock()
+			out[k.Name] = row
+			mu.Unlock()
+		}()
+	}
+	wg.Wait()
+	writeJSON(w, http.StatusOK, map[string]any{"keys": out})
+}
+
 // maxSpendReads bounds the GetTeamSpend calls one Usage page load can cause.
 const maxSpendReads = 8