fix(db): enforce tenant ownership in session_pins and billing spend queries (#796)

db/queries/billing.sql

@@ -76,9 +76,14 @@ key_spend AS (
     -- (the debit no-ops and the app sees ErrBalanceRowMissing) we must NOT bump
     -- the key's lifetime spend, or a capped key could trip its cap with no
     -- matching ledger debit.
-    UPDATE router.model_router_api_keys
+    -- Ownership: only bump a key whose installation's external_id matches the
+    -- org being debited (#796) — a mismatched api_key_id silently no-ops.
+    UPDATE router.model_router_api_keys k
     SET spent_usd_micros = spent_usd_micros - @delta_usd_micros::bigint
-    WHERE id = sqlc.narg('api_key_id')::uuid
+    FROM router.model_router_installations i
+    WHERE k.id = sqlc.narg('api_key_id')::uuid
+      AND k.installation_id = i.id
+      AND i.external_id = @organization_id::varchar
       AND EXISTS (SELECT 1 FROM updated)
 ),
 user_month_spend AS (
@@ -88,6 +93,8 @@ user_month_spend AS (
     -- Also no-ops when the user row no longer exists (stale cached id after a
     -- cascade delete mid-request) so a dangling FK can't roll back the debit
     -- after inference was already served.
+    -- Ownership: user must belong to an installation whose external_id matches
+    -- the org being debited (#796) — a mismatched router_user_id silently no-ops.
     INSERT INTO router.model_router_user_monthly_spend (router_user_id, month, spent_usd_micros, updated_at)
     SELECT
         sqlc.narg('router_user_id')::uuid,
@@ -97,8 +104,11 @@ user_month_spend AS (
     WHERE sqlc.narg('router_user_id')::uuid IS NOT NULL
       AND EXISTS (SELECT 1 FROM updated)
       AND EXISTS (
-          SELECT 1 FROM router.model_router_users
-          WHERE id = sqlc.narg('router_user_id')::uuid
+          SELECT 1
+          FROM router.model_router_users u
+          JOIN router.model_router_installations i ON i.id = u.installation_id
+          WHERE u.id = sqlc.narg('router_user_id')::uuid
+            AND i.external_id = @organization_id::varchar
       )
     ON CONFLICT (router_user_id, month) DO UPDATE
     SET spent_usd_micros = router.model_router_user_monthly_spend.spent_usd_micros + EXCLUDED.spent_usd_micros,

db/queries/session_pins.sql

@@ -8,17 +8,21 @@
 SELECT *
 FROM router.session_pins
 WHERE session_key = @session_key::bytea
-  AND role        = @role::varchar;
+  AND role        = @role::varchar
+  AND installation_id = @installation_id::uuid;
 
 -- Upserts a pin, refreshing pinned_until on every hit (sliding TTL).
 -- turn_count increments on conflict so we can observe how many turns a
 -- single (session_key, role) lives for. installation_id is set on first
 -- insert and not touched on update — re-binding a session to a different
--- installation would indicate a bug, not a legitimate state. The
--- last_*_tokens / last_turn_ended_at columns are deliberately omitted
--- from the ON CONFLICT update set: only UpdateSessionPinUsage writes
--- them, so the at-start-of-turn refresh here cannot clobber the
--- previous turn's usage with zeros before the planner reads it.
+-- installation would indicate a bug, not a legitimate state. The ON CONFLICT
+-- DO UPDATE is gated on installation_id = EXCLUDED.installation_id so a
+-- mismatched caller silently no-ops rather than overwriting another
+-- tenant's pin. The last_*_tokens / last_turn_ended_at columns are
+-- deliberately omitted from the ON CONFLICT update set: only
+-- UpdateSessionPinUsage writes them, so the at-start-of-turn refresh
+-- here cannot clobber the previous turn's usage with zeros before the
+-- planner reads it.
 --
 -- consecutive_upstream_errors is preserved on a same-model refresh (so
 -- the two-strike eviction counter accumulates across turns of the same
@@ -80,7 +84,8 @@ ON CONFLICT (session_key, role) DO UPDATE SET
     WHEN router.session_pins.pinned_model = EXCLUDED.pinned_model
       THEN router.session_pins.consecutive_upstream_errors
     ELSE 0
-  END;
+  END
+WHERE router.session_pins.installation_id = EXCLUDED.installation_id;
 
 -- Records the previous turn's upstream token usage on an existing pin
 -- row. Fired off the request path after the upstream response
@@ -110,29 +115,34 @@ SET last_input_tokens        = @last_input_tokens::int,
       OR (@prior_served_model::varchar <> '' AND @prior_served_model::varchar <> @last_served_model::varchar),
     last_served_model        = @last_served_model::varchar
 WHERE session_key = @session_key::bytea
-  AND role        = @role::varchar;
+  AND role        = @role::varchar
+  AND installation_id = @installation_id::uuid;
 
 -- Atomically increments consecutive_upstream_errors and returns the
 -- new value. The turn loop calls this after a non-retryable upstream
 -- 4xx on a sticky-pinned turn; the returned count drives the
 -- two-strike eviction decision. Returns sql.ErrNoRows if no pin
--- exists, which the adapter maps to a no-op (pin must already be
--- evicted by another path, e.g. force-model / loop-break).
+-- exists (or installation_id mismatches), which the adapter maps to a
+-- no-op (pin must already be evicted by another path, e.g. force-model /
+-- loop-break).
 -- name: IncrementSessionPinUpstreamErrors :one
 UPDATE router.session_pins
 SET consecutive_upstream_errors = consecutive_upstream_errors + 1
 WHERE session_key = @session_key::bytea
   AND role        = @role::varchar
+  AND installation_id = @installation_id::uuid
 RETURNING consecutive_upstream_errors;
 
 -- Clears the two-strike counter after a successful turn. UPDATE
--- matches by (session_key, role); zero rows affected on missing pin
--- is a successful no-op like UpdateSessionPinUsage.
+-- matches by (session_key, role, installation_id); zero rows affected
+-- on missing pin (or ownership mismatch) is a successful no-op like
+-- UpdateSessionPinUsage.
 -- name: ResetSessionPinUpstreamErrors :exec
 UPDATE router.session_pins
 SET consecutive_upstream_errors = 0
 WHERE session_key = @session_key::bytea
   AND role        = @role::varchar
+  AND installation_id = @installation_id::uuid
   AND consecutive_upstream_errors > 0;
 
 -- Garbage-collects pins that have been expired for >24h. The 24h grace

db/queries/spend_limits.sql

@@ -10,22 +10,34 @@ WHERE organization_id = @organization_id::varchar;
 -- per-user override when an override row exists (its NULL means "explicitly
 -- uncapped"), otherwise the org-wide default; has_override distinguishes the
 -- two NULL meanings. spent is 0 when the user has no spend row this month.
+-- Spend/override subqueries require the user to belong to an installation
+-- whose external_id matches @organization_id (#796); a mismatched pair is a
+-- silent miss (spent 0, no override). Org default still resolves for the org.
 -- name: GetUserMonthlySpendAndLimit :one
 SELECT
     (EXISTS (
         SELECT 1 FROM router.model_router_user_spend_limits ovr
+        JOIN router.model_router_users u ON u.id = ovr.router_user_id
+        JOIN router.model_router_installations i ON i.id = u.installation_id
         WHERE ovr.router_user_id = @router_user_id::uuid
+          AND i.external_id = @organization_id::varchar
     ))::boolean AS has_override,
     (SELECT ovr.monthly_limit_usd_micros
      FROM router.model_router_user_spend_limits ovr
-     WHERE ovr.router_user_id = @router_user_id::uuid) AS override_limit_usd_micros,
+     JOIN router.model_router_users u ON u.id = ovr.router_user_id
+     JOIN router.model_router_installations i ON i.id = u.installation_id
+     WHERE ovr.router_user_id = @router_user_id::uuid
+       AND i.external_id = @organization_id::varchar) AS override_limit_usd_micros,
     (SELECT lim.user_monthly_limit_usd_micros
      FROM router.organization_spend_limits lim
      WHERE lim.organization_id = @organization_id::varchar) AS org_default_limit_usd_micros,
     COALESCE((
         SELECT sp.spent_usd_micros
         FROM router.model_router_user_monthly_spend sp
+        JOIN router.model_router_users u ON u.id = sp.router_user_id
+        JOIN router.model_router_installations i ON i.id = u.installation_id
         WHERE sp.router_user_id = @router_user_id::uuid
+          AND i.external_id = @organization_id::varchar
           AND sp.month = DATE_TRUNC('month', NOW() AT TIME ZONE 'utc')::date
     ), 0)::bigint AS spent_usd_micros;
 

internal/postgres/session_pin_repo.go

@@ -9,6 +9,7 @@ import (
 	"workweave/router/internal/router/sessionpin"
 	"workweave/router/internal/sqlc"
 
+	"github.com/google/uuid"
 	"github.com/jackc/pgx/v5/pgtype"
 )
 
@@ -24,11 +25,12 @@ func NewSessionPinRepo(tx sqlc.DBTX) *SessionPinRepo {
 
 var _ sessionpin.Store = (*SessionPinRepo)(nil)
 
-func (r *SessionPinRepo) Get(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) (sessionpin.Pin, bool, error) {
+func (r *SessionPinRepo) Get(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, installationID uuid.UUID) (sessionpin.Pin, bool, error) {
 	q := sqlc.New(r.tx)
 	row, err := q.GetSessionPin(ctx, sqlc.GetSessionPinParams{
-		SessionKey: sessionKey[:],
-		Role:       role,
+		SessionKey:     sessionKey[:],
+		Role:           role,
+		InstallationID: installationID,
 	})
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
@@ -56,9 +58,10 @@ func (r *SessionPinRepo) Upsert(ctx context.Context, p sessionpin.Pin) error {
 }
 
 // UpdateUsage records the previous turn's usage on the pin row. A missing
-// pin (evicted/swept/never created) is a no-op, not an error. A zero
-// EndedAt is stamped with time.Now so the column is always populated.
-func (r *SessionPinRepo) UpdateUsage(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, usage sessionpin.Usage) error {
+// pin (evicted/swept/never created) or ownership mismatch is a no-op, not
+// an error. A zero EndedAt is stamped with time.Now so the column is always
+// populated.
+func (r *SessionPinRepo) UpdateUsage(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, installationID uuid.UUID, usage sessionpin.Usage) error {
 	endedAt := usage.EndedAt
 	if endedAt.IsZero() {
 		endedAt = time.Now()
@@ -67,6 +70,7 @@ func (r *SessionPinRepo) UpdateUsage(ctx context.Context, sessionKey [sessionpin
 	return q.UpdateSessionPinUsage(ctx, sqlc.UpdateSessionPinUsageParams{
 		SessionKey:            sessionKey[:],
 		Role:                  role,
+		InstallationID:        installationID,
 		LastInputTokens:       int32(usage.InputTokens),
 		LastCachedReadTokens:  int32(usage.CachedReadTokens),
 		LastCachedWriteTokens: int32(usage.CachedWriteTokens),
@@ -79,13 +83,15 @@ func (r *SessionPinRepo) UpdateUsage(ctx context.Context, sessionKey [sessionpin
 }
 
 // IncrementUpstreamErrors atomically bumps the consecutive-error counter.
-// A missing pin (already evicted or never created) returns (0, nil): the
-// two-strike check treats it as a no-op since there's no row left to evict.
-func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) (int, error) {
+// A missing pin (already evicted, never created, or ownership mismatch)
+// returns (0, nil): the two-strike check treats it as a no-op since there's
+// no row left to evict.
+func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, installationID uuid.UUID) (int, error) {
 	q := sqlc.New(r.tx)
 	count, err := q.IncrementSessionPinUpstreamErrors(ctx, sqlc.IncrementSessionPinUpstreamErrorsParams{
-		SessionKey: sessionKey[:],
-		Role:       role,
+		SessionKey:     sessionKey[:],
+		Role:           role,
+		InstallationID: installationID,
 	})
 	if err != nil {
 		if errors.Is(err, sql.ErrNoRows) {
@@ -97,12 +103,14 @@ func (r *SessionPinRepo) IncrementUpstreamErrors(ctx context.Context, sessionKey
 }
 
 // ResetUpstreamErrors clears the consecutive-error counter after a
-// successful turn. Missing pin is a no-op, same as UpdateUsage.
-func (r *SessionPinRepo) ResetUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string) error {
+// successful turn. Missing pin / ownership mismatch is a no-op, same as
+// UpdateUsage.
+func (r *SessionPinRepo) ResetUpstreamErrors(ctx context.Context, sessionKey [sessionpin.SessionKeyLen]byte, role string, installationID uuid.UUID) error {
 	q := sqlc.New(r.tx)
 	return q.ResetSessionPinUpstreamErrors(ctx, sqlc.ResetSessionPinUpstreamErrorsParams{
-		SessionKey: sessionKey[:],
-		Role:       role,
+		SessionKey:     sessionKey[:],
+		Role:           role,
+		InstallationID: installationID,
 	})
 }