From 99cf9c06a1303c0cfe22194097aea23fd6cb4703 Mon Sep 17 00:00:00 2001 From: mintaka Date: Mon, 7 Sep 2026 15:22:25 -0400 Subject: [PATCH 1/4] feat(store): add the durable session_bindings table and its store methods (RIG-3108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RunnerHub holds the (session -> agent account, Runner) binding only in RAM, so a Server restart loses every binding and the relay resolves nothing for a session it minted moments earlier. Add the durable table and the store surface over it. Nothing reads these methods yet — demoting the hub's in-RAM maps to caches over this table is the next slice. A binding is not authorization, exactly as `agent_placements` is not: `SubscribeAgentSession` still authorizes through `agent_sessions -> agent_accounts -> channel_members` and never reads this table. The two reads it exists for are the relay resolving the account that owns an inbound session, and the delivery consumer resolving the live session of an already-authorized recipient account. UNIQUE index on `agent_account_id`, not merely an index: the in-RAM `accountSessions` map it replaces is 1:1, so a second concurrent session for one account would split that account's deliveries across two sessions and let the reverse lookup resolve whichever row Postgres happened to return. The unique index makes the double-bind unrepresentable rather than unlikely, and a rebind onto a new session id raises a unique violation mapped to `ErrConflict` instead of silently double-binding. `DeleteSessionBindingsForRunner` is `:many ... RETURNING`, not a bare DELETE. `Hub.enroll` snapshots the bindings *before* clearing them, because each cleared binding drives a presence DISCONNECTED edge and each cleared session id must be reaped from the delivery held-deliver registry; enroll emits no lifecycle frames of its own. A DELETE that returned nothing would satisfy the fail-closed invariant while silently dropping both side-effects, leaving a long-WORKING agent stuck WORKING in the projection forever. Both lookups fail closed: a miss is `ErrNotFound`, never a zero-value id with a nil error, because a zero-value account id would flow onward as a real (wrong) principal instead of stopping the call. `runner_id` is NOT NULL and rejects `''`. The `''` unknown-runner sentinel does not carry over from `agent_placements`, where a placement must outlive its Runner's attachment and the next provision self-heals it. A binding has the opposite lifetime — it exists only while a Runner is attached, and `runner_id` is the sweep key that retires it — so a binding stamped `''` could never be swept by any real Runner's re-enroll and would linger as a stale session outliving its Runner. Two deviations from the frozen record's §T4 schema, both ruled by Matt (2026-09-07) and recorded on RIG-3108: - No `0002_session_bindings.sql`. Migrations are collapsed into the squashed init per the standing 2026-08-07 ruling in that file's header ("the same reasoning folds each later migration in as it accretes"); §T4's `NNNN_` wording predates it. - `TIMESTAMPTZ`, not `bound_at_unix_ms`. The schema's split is wire exposure: `_unix_ms BIGINT` is used where the value crosses the protobuf wire as an `int64`, `TIMESTAMPTZ` for server-internal bookkeeping, and no `TIMESTAMPTZ` column name appears as an `int64` proto field. A binding is server-internal. Verified: 9 pgtest cases — round-trip both directions, fail-closed miss on both lookups, the `ON CONFLICT` rebind in place, `ErrConflict` on a second session for a bound account, `ErrInvalidArgument` on an unknown account, idempotent delete freeing the unique slot, the runner sweep returning every swept binding while leaving another Runner's binding alive, `updated_at` advancing via the trigger, and cross-tenant isolation. Red control: dropping the table from `tenant_tables` fails with `tenant B ResolveSessionAccount(sess-a) err = , want ErrNotFound — cross-tenant read leak`, caught by both the isolation test and the RLS catalog floor. Gates: `sqlc-drift`, `sqlc-vet` (every query PREPAREs against the live schema), `sql-migration-gate:check` (0 issues), `go build ./...`, `golangci-lint run ./internal/store/...` (0 issues), and the full `-tags pgtest` store suite. Co-authored-by: Matt Wilkinson --- go/internal/store/db/models.go | 9 + go/internal/store/db/querier.go | 29 ++ go/internal/store/db/session_bindings.sql.go | 117 ++++++ go/internal/store/migrations/0001_init.sql | 44 +- .../store/queries/session_bindings.sql | 46 ++ go/internal/store/rls_pgtest_test.go | 2 +- go/internal/store/session_bindings.go | 190 +++++++++ .../store/session_bindings_pgtest_test.go | 397 ++++++++++++++++++ 8 files changed, 832 insertions(+), 2 deletions(-) create mode 100644 go/internal/store/db/session_bindings.sql.go create mode 100644 go/internal/store/queries/session_bindings.sql create mode 100644 go/internal/store/session_bindings.go create mode 100644 go/internal/store/session_bindings_pgtest_test.go diff --git a/go/internal/store/db/models.go b/go/internal/store/db/models.go index 76f968c3..a0e2df50 100644 --- a/go/internal/store/db/models.go +++ b/go/internal/store/db/models.go @@ -266,6 +266,15 @@ type Secret struct { TenantID string } +type SessionBinding struct { + SessionID string + AgentAccountID string + RunnerID string + CreatedAt pgtype.Timestamptz + UpdatedAt pgtype.Timestamptz + TenantID string +} + type SystemAccount struct { AccountID string TenantID string diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 5c1aa2d5..10549153 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -91,6 +91,23 @@ type Querier interface { DeleteChannelPinReturningPosition(ctx context.Context, arg DeleteChannelPinReturningPositionParams) (int32, error) DeleteModelRegistry(ctx context.Context) error DeleteSecret(ctx context.Context, name string) (int64, error) + DeleteSessionBinding(ctx context.Context, sessionID string) error + // The reconnect sweep. Hub.enroll (internal/runnerhub/hub.go:905-957) clears + // every binding when a Runner (re-)enrolls: a reconnecting Runner has no live + // sessions, so a surviving binding would resolve a re-minted session id to a + // stale account. A durable table does not forget on reconnect, so the sweep must + // be explicit. + // + // :many with RETURNING, deliberately NOT :exec. enroll snapshots the bindings + // BEFORE clearing them (hub.go:912-928) because each cleared binding drives a + // presence DISCONNECTED edge (RIG-1569 T8) and each cleared session id must be + // reaped from the delivery held-deliver registry (RIG-1569 T3). A bare DELETE + // would satisfy the invariant while silently dropping both side-effects, leaving + // a long-WORKING agent stuck WORKING in the projection forever. RETURNING is + // what preserves them, so the returned rows are load-bearing, not diagnostic. + // A DELETE ... RETURNING takes no ORDER BY, so the Store method sorts the + // returned slice by session id to keep a sweep pass deterministic and diffable. + DeleteSessionBindingsForRunner(ctx context.Context, runnerID string) ([]DeleteSessionBindingsForRunnerRow, error) DeleteTopic(ctx context.Context, id string) error // Agent-forge-subscription / artifact-cursor queries (sqlc adoption T6, // RIG-3034). These replace the inline SQL literals in @@ -335,6 +352,16 @@ type Querier interface { // there is the NORMAL replay case, not a drop), so asserting rows-affected here // would wrongly fail an idempotent re-fire. RecordOwedMention(ctx context.Context, arg RecordOwedMentionParams) error + // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable + // (session -> agent account, Runner) binding the RunnerHub has so far held only + // in RAM. The hand-written Store methods in internal/store/session_bindings.go + // keep their signatures and map these rows into the SessionBinding domain struct + // (the AccountID newtype is done inline in the Go, as agent_placements does). + // + // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger + // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written + // `updated_at = now()` is the exact defect that convention removes. + RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error RemarkSafetyValveSuperseded(ctx context.Context, arg RemarkSafetyValveSupersededParams) error RenameTopic(ctx context.Context, arg RenameTopicParams) error RequireAgentSessionSubscriber(ctx context.Context, arg RequireAgentSessionSubscriberParams) (bool, error) @@ -368,6 +395,8 @@ type Querier interface { SeedHomeChannelMembers(ctx context.Context, arg SeedHomeChannelMembersParams) error SelfAuthoredSeqsAbove(ctx context.Context, arg SelfAuthoredSeqsAboveParams) ([]int64, error) SessionBase(ctx context.Context, sessionID string) (int64, error) + SessionBindingAccount(ctx context.Context, sessionID string) (string, error) + SessionBindingForAccount(ctx context.Context, agentAccountID string) (string, error) SessionMaxEntrySeq(ctx context.Context, sessionID string) (int64, error) SessionTranscript(ctx context.Context, sessionID string) ([]SessionTranscriptRow, error) // Agent-activity queries (sqlc adoption T5, RIG-3034). These replace the inline diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go new file mode 100644 index 00000000..3db90cef --- /dev/null +++ b/go/internal/store/db/session_bindings.sql.go @@ -0,0 +1,117 @@ +// Code generated by sqlc. DO NOT EDIT. +// versions: +// sqlc v1.31.1 +// source: session_bindings.sql + +package db + +import ( + "context" +) + +const deleteSessionBinding = `-- name: DeleteSessionBinding :exec +DELETE FROM session_bindings WHERE session_id = $1 +` + +func (q *Queries) DeleteSessionBinding(ctx context.Context, sessionID string) error { + _, err := q.db.Exec(ctx, deleteSessionBinding, sessionID) + return err +} + +const deleteSessionBindingsForRunner = `-- name: DeleteSessionBindingsForRunner :many + +DELETE FROM session_bindings + WHERE runner_id = $1 +RETURNING session_id, agent_account_id +` + +type DeleteSessionBindingsForRunnerRow struct { + SessionID string + AgentAccountID string +} + +// The reconnect sweep. Hub.enroll (internal/runnerhub/hub.go:905-957) clears +// every binding when a Runner (re-)enrolls: a reconnecting Runner has no live +// sessions, so a surviving binding would resolve a re-minted session id to a +// stale account. A durable table does not forget on reconnect, so the sweep must +// be explicit. +// +// :many with RETURNING, deliberately NOT :exec. enroll snapshots the bindings +// BEFORE clearing them (hub.go:912-928) because each cleared binding drives a +// presence DISCONNECTED edge (RIG-1569 T8) and each cleared session id must be +// reaped from the delivery held-deliver registry (RIG-1569 T3). A bare DELETE +// would satisfy the invariant while silently dropping both side-effects, leaving +// a long-WORKING agent stuck WORKING in the projection forever. RETURNING is +// what preserves them, so the returned rows are load-bearing, not diagnostic. +// A DELETE ... RETURNING takes no ORDER BY, so the Store method sorts the +// returned slice by session id to keep a sweep pass deterministic and diffable. +func (q *Queries) DeleteSessionBindingsForRunner(ctx context.Context, runnerID string) ([]DeleteSessionBindingsForRunnerRow, error) { + rows, err := q.db.Query(ctx, deleteSessionBindingsForRunner, runnerID) + if err != nil { + return nil, err + } + defer rows.Close() + var items []DeleteSessionBindingsForRunnerRow + for rows.Next() { + var i DeleteSessionBindingsForRunnerRow + if err := rows.Scan(&i.SessionID, &i.AgentAccountID); err != nil { + return nil, err + } + items = append(items, i) + } + if err := rows.Err(); err != nil { + return nil, err + } + return items, nil +} + +const recordSessionBinding = `-- name: RecordSessionBinding :exec + +INSERT INTO session_bindings (session_id, agent_account_id, runner_id) +VALUES ($1, $2, $3) +ON CONFLICT (session_id) DO UPDATE + SET agent_account_id = EXCLUDED.agent_account_id, + runner_id = EXCLUDED.runner_id +` + +type RecordSessionBindingParams struct { + SessionID string + AgentAccountID string + RunnerID string +} + +// Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable +// (session -> agent account, Runner) binding the RunnerHub has so far held only +// in RAM. The hand-written Store methods in internal/store/session_bindings.go +// keep their signatures and map these rows into the SessionBinding domain struct +// (the AccountID newtype is done inline in the Go, as agent_placements does). +// +// updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger +// (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written +// `updated_at = now()` is the exact defect that convention removes. +func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error { + _, err := q.db.Exec(ctx, recordSessionBinding, arg.SessionID, arg.AgentAccountID, arg.RunnerID) + return err +} + +const sessionBindingAccount = `-- name: SessionBindingAccount :one +SELECT agent_account_id FROM session_bindings WHERE session_id = $1 +` + +func (q *Queries) SessionBindingAccount(ctx context.Context, sessionID string) (string, error) { + row := q.db.QueryRow(ctx, sessionBindingAccount, sessionID) + var agent_account_id string + err := row.Scan(&agent_account_id) + return agent_account_id, err +} + +const sessionBindingForAccount = `-- name: SessionBindingForAccount :one +SELECT session_id FROM session_bindings WHERE agent_account_id = $1 +` + +func (q *Queries) SessionBindingForAccount(ctx context.Context, agentAccountID string) (string, error) { + row := q.db.QueryRow(ctx, sessionBindingForAccount, agentAccountID) + var session_id string + err := row.Scan(&session_id) + return session_id, err +} diff --git a/go/internal/store/migrations/0001_init.sql b/go/internal/store/migrations/0001_init.sql index b34f8919..88119683 100644 --- a/go/internal/store/migrations/0001_init.sql +++ b/go/internal/store/migrations/0001_init.sql @@ -492,6 +492,47 @@ CREATE INDEX agent_placements_runner_idx ON agent_placements (runner_id); -- owner. CREATE UNIQUE INDEX agent_placements_container_key ON agent_placements (container_name); +-- Session bindings (RIG-3108 / RIG-2861 §T4): the DURABLE (session -> agent +-- account, Runner) binding the RunnerHub has so far held only in RAM. Placement +-- (above) is where an agent RUNS; a binding is which LIVE session speaks for it, +-- so the two are siblings and neither is authorization: SubscribeAgentSession +-- still authorizes through agent_sessions -> agent_accounts -> channel_members +-- and never reads this table. +-- +-- PK on session_id, not a surrogate: a session id names exactly one binding, and +-- that is also the accountForSession read direction the relay resolves on every +-- inbound comms call. +-- +-- runner_id is deliberately NOT a FK, for the same reason agent_placements' +-- isn't: Runners are enrolled in memory under a token subject with no runners +-- table to reference. It stays NOT NULL — a binding with no Runner cannot be +-- swept by the reconnect sweep below, and an unswept binding is a stale session +-- that outlives its Runner. +CREATE TABLE session_bindings ( + session_id TEXT PRIMARY KEY, + agent_account_id TEXT NOT NULL REFERENCES agent_accounts (account_id) ON DELETE RESTRICT, + runner_id TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + tenant_id TEXT NOT NULL DEFAULT current_setting('compass.tenant_id', TRUE) +); + +-- The REVERSE lookup (SessionForAccount, runnerhub/relay_comms.go:179-184): the +-- delivery consumer holds a resolved subscriber account and needs the live +-- session to dispatch to. UNIQUE because the in-RAM accountSessions map it +-- replaces is 1:1 — an account has AT MOST ONE live session — so this index +-- makes a second concurrent session for one account UNREPRESENTABLE rather than +-- merely unlikely. A rebind of an account onto a NEW session id therefore raises +-- a unique violation (mapped to ErrConflict) instead of silently double-binding +-- and letting one account's deliveries split across two sessions. +CREATE UNIQUE INDEX session_bindings_account_key ON session_bindings (agent_account_id); + +-- The reconnect sweep deletes every binding of a re-enrolling Runner +-- (DeleteSessionBindingsForRunner), so runner_id is the read direction that +-- needs an index — the same direction, and the same reason, as +-- agent_placements_runner_idx above. Non-unique: one Runner holds many sessions. +CREATE INDEX session_bindings_runner_idx ON session_bindings (runner_id); + -- ── Agent session transcripts (two-tier store) ─────────────────────────────── -- The durable TWO-TIER transcript store (RIG-1667 T4): a Postgres HOT TAIL -- holding [latest checkpoint .. now] = the normal resume set, plus a manifest of @@ -941,7 +982,7 @@ DECLARE 'user_accounts', 'agent_accounts', 'system_accounts', 'account_handles', 'channel_groups', 'channels', 'channel_members', 'agent_workspaces', 'topics', 'messages', 'channel_pins', 'secrets', - 'agent_sessions', 'agent_placements', + 'agent_sessions', 'agent_placements', 'session_bindings', 'agent_session_transcript_entries', 'agent_session_archive_segments', 'agent_delivery_cursors', 'owed_mentions', 'agent_activity', 'agent_forge_subscriptions', 'forge_authored_artifacts', @@ -1022,6 +1063,7 @@ DECLARE updated_at_tables text[] := ARRAY[ 'secrets', 'agent_placements', + 'session_bindings', 'agent_config_bundle', 'model_registry', 'forge_repo_subscriptions' diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql new file mode 100644 index 00000000..836a853b --- /dev/null +++ b/go/internal/store/queries/session_bindings.sql @@ -0,0 +1,46 @@ +-- Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable +-- (session -> agent account, Runner) binding the RunnerHub has so far held only +-- in RAM. The hand-written Store methods in internal/store/session_bindings.go +-- keep their signatures and map these rows into the SessionBinding domain struct +-- (the AccountID newtype is done inline in the Go, as agent_placements does). +-- +-- updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger +-- (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written +-- `updated_at = now()` is the exact defect that convention removes. + +-- name: RecordSessionBinding :exec +INSERT INTO session_bindings (session_id, agent_account_id, runner_id) +VALUES ($1, $2, $3) +ON CONFLICT (session_id) DO UPDATE + SET agent_account_id = EXCLUDED.agent_account_id, + runner_id = EXCLUDED.runner_id; + +-- name: SessionBindingAccount :one +SELECT agent_account_id FROM session_bindings WHERE session_id = $1; + +-- name: SessionBindingForAccount :one +SELECT session_id FROM session_bindings WHERE agent_account_id = $1; + +-- name: DeleteSessionBinding :exec +DELETE FROM session_bindings WHERE session_id = $1; + +-- The reconnect sweep. Hub.enroll (internal/runnerhub/hub.go:905-957) clears +-- every binding when a Runner (re-)enrolls: a reconnecting Runner has no live +-- sessions, so a surviving binding would resolve a re-minted session id to a +-- stale account. A durable table does not forget on reconnect, so the sweep must +-- be explicit. +-- +-- :many with RETURNING, deliberately NOT :exec. enroll snapshots the bindings +-- BEFORE clearing them (hub.go:912-928) because each cleared binding drives a +-- presence DISCONNECTED edge (RIG-1569 T8) and each cleared session id must be +-- reaped from the delivery held-deliver registry (RIG-1569 T3). A bare DELETE +-- would satisfy the invariant while silently dropping both side-effects, leaving +-- a long-WORKING agent stuck WORKING in the projection forever. RETURNING is +-- what preserves them, so the returned rows are load-bearing, not diagnostic. +-- A DELETE ... RETURNING takes no ORDER BY, so the Store method sorts the +-- returned slice by session id to keep a sweep pass deterministic and diffable. + +-- name: DeleteSessionBindingsForRunner :many +DELETE FROM session_bindings + WHERE runner_id = $1 +RETURNING session_id, agent_account_id; diff --git a/go/internal/store/rls_pgtest_test.go b/go/internal/store/rls_pgtest_test.go index 0943cd50..1b29ed82 100644 --- a/go/internal/store/rls_pgtest_test.go +++ b/go/internal/store/rls_pgtest_test.go @@ -592,7 +592,7 @@ func TestRLSCatalogEnabledAndForced(t *testing.T) { "user_accounts", "agent_accounts", "system_accounts", "account_handles", "channel_groups", "channels", "channel_members", "agent_workspaces", "topics", "messages", "channel_pins", "secrets", - "agent_sessions", "agent_placements", + "agent_sessions", "agent_placements", "session_bindings", "agent_session_transcript_entries", "agent_session_archive_segments", "agent_delivery_cursors", "owed_mentions", "agent_activity", "agent_forge_subscriptions", "forge_authored_artifacts", diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go new file mode 100644 index 00000000..0e788fb4 --- /dev/null +++ b/go/internal/store/session_bindings.go @@ -0,0 +1,190 @@ +package store + +import ( + "cmp" + "context" + "fmt" + "slices" + + "github.com/RigelBuild/compass/go/internal/store/db" +) + +// Session bindings: the durable record of WHICH LIVE SESSION speaks for an agent +// account, and the Runner that session is attached to (RIG-3108 / RIG-2861 §T4). +// Until now the RunnerHub held this only in RAM (sessionAccounts, accountSessions), +// so a Server restart lost every binding and the relay resolved nothing for a +// session it had itself minted moments earlier. +// +// A binding is NOT authorization, exactly as agent_placements is not: +// SubscribeAgentSession authorizes through agent_sessions -> agent_accounts -> +// channel_members and never reads this table. What a binding is for is the two +// reads below — the relay resolving the account that owns an inbound session +// (accountForSession), and the delivery consumer resolving the live session of a +// recipient account it has already authorized (SessionForAccount). +// +// PR2 adds the table and these methods only. Demoting the hub's in-RAM maps to +// caches over this table is PR3; nothing in internal/runnerhub reads this yet. + +// SessionBinding is one live binding: the session, the agent account it speaks +// for, and the Runner it is attached to. Returned by +// DeleteSessionBindingsForRunner, which is the reconnect sweep — its caller +// needs every field of each binding it just removed. +type SessionBinding struct { + SessionID string + AccountID AccountID + RunnerID string +} + +// RecordSessionBinding persists the binding for a live session. It is an UPSERT +// keyed on the session, because a session id names exactly one binding: a rebind +// REPLACES it, updating agent_account_id and runner_id TOGETHER so a row can +// never pair a fresh Runner with the account from a previous attachment. +// +// updated_at is maintained by the set_updated_at() trigger, never here +// (RIG-3495) — the query file assigns it nowhere. +// +// An unknown agent_account_id is ErrInvalidArgument (the FK). +// +// An account that already holds a DIFFERENT live session is ErrConflict (the +// unique index on agent_account_id), and unlike agent_placements' container +// guard this one is a live safety property, not a latent one. The in-RAM +// accountSessions map this replaces is 1:1, so a second concurrent session for +// one account would split that account's deliveries across two sessions and let +// SessionForAccount resolve whichever row Postgres happened to return. Refusing +// the write makes the double-bind unrepresentable. The caller's remedy is to +// release the stale binding first — DeleteSessionBinding for a session it knows +// is gone, or the runner sweep on a re-enroll. +func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) error { + if sessionID == "" { + return fmt.Errorf("%w: session id is required", ErrInvalidArgument) + } + if accountID == "" { + return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + } + // Unlike agent_placements.runner_id, '' is NOT an accepted unknown-runner + // sentinel here. A placement must OUTLIVE its Runner's attachment (it is + // where the agent runs, and the next provision self-heals the sentinel), but + // a binding exists ONLY while a Runner is attached, and runner_id is the + // sweep key that retires it. A binding stamped '' could never be swept by + // any real Runner's re-enroll, so it would linger as a stale session that + // outlives its Runner — the exact leak the sweep exists to prevent. + if runnerID == "" { + return fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + } + if err := s.q.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ + SessionID: sessionID, + AgentAccountID: string(accountID), + RunnerID: runnerID, + }); err != nil { + if pgErrIs(err, pgForeignKeyViolation) { + return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) + } + if pgErrIs(err, pgUniqueViolation) { + return fmt.Errorf("%w: agent %q is already bound to another live session", ErrConflict, accountID) + } + return fmt.Errorf("store: record session binding: %w", err) + } + return nil +} + +// ResolveSessionAccount resolves the agent account a live session speaks for — +// the relay's read on every inbound comms call, where the request carries only +// the session id. +// +// An unbound session is ErrNotFound, and that is FAIL-CLOSED by design: this +// resolves the scope a comms call runs under, so a miss must never surface as an +// empty AccountID with a nil error. A zero-value account id would flow onward as +// a real (wrong) principal instead of stopping the call. +func (s *Store) ResolveSessionAccount(ctx context.Context, sessionID string) (AccountID, error) { + if sessionID == "" { + return "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) + } + accountID, err := s.q.SessionBindingAccount(ctx, sessionID) + if err != nil { + if noRows(err) { + return "", fmt.Errorf("%w: session %q is not bound", ErrNotFound, sessionID) + } + return "", fmt.Errorf("store: resolve session account: %w", err) + } + return AccountID(accountID), nil +} + +// SessionForAccount resolves the live session bound to an agent account — the +// REVERSE of ResolveSessionAccount, and the direction the delivery consumer +// needs to dispatch a deliver to an already-resolved subscriber +// (runnerhub/relay_comms.go:179-184). At most one row can answer, because the +// unique index makes the mapping 1:1. +// +// An account with no live session is ErrNotFound — never started, stopped, or +// dropped on a Runner reconnect. Fail-closed for the same reason as above: an +// empty session id with a nil error would be dispatched to as if it were a live +// session. The consumer's own contract turns this into "push nothing now, let +// the cursor sweep deliver on the recipient's next start". +func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (string, error) { + if accountID == "" { + return "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + } + sessionID, err := s.q.SessionBindingForAccount(ctx, string(accountID)) + if err != nil { + if noRows(err) { + return "", fmt.Errorf("%w: agent %q has no live session", ErrNotFound, accountID) + } + return "", fmt.Errorf("store: resolve session for account: %w", err) + } + return sessionID, nil +} + +// DeleteSessionBinding releases the binding for sessionID — the single-session +// release path a normal session end takes, freeing the account's unique slot for +// its next session. It is IDEMPOTENT: a session teardown may be retried, and a +// second release of an already-released session must succeed, so deleting an +// absent row is not an error (the same posture as DeleteAgentPlacement). +func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID string) error { + if sessionID == "" { + return fmt.Errorf("%w: session id is required", ErrInvalidArgument) + } + if err := s.q.DeleteSessionBinding(ctx, sessionID); err != nil { + return fmt.Errorf("store: delete session binding: %w", err) + } + return nil +} + +// DeleteSessionBindingsForRunner is the reconnect sweep: it releases every +// binding attached to runnerID and RETURNS the bindings it removed. Hub.enroll +// (runnerhub/hub.go:905-957) clears all bindings when a Runner (re-)enrolls, +// because a reconnecting Runner has no live sessions and a surviving binding +// would resolve a re-minted session id to a stale account. +// +// The returned slice is load-bearing, not diagnostic. Each released binding +// drives a presence DISCONNECTED edge for its account (RIG-1569 T8) and each +// released session id must be reaped from the delivery held-deliver registry +// (RIG-1569 T3); enroll emits no lifecycle frames of its own, so a caller that +// deleted without reading the rows back would leave a long-WORKING agent stuck +// WORKING in the projection forever. That is why the query is :many with +// RETURNING rather than a bare DELETE. +// +// Sorted by session id: a DELETE ... RETURNING has no ORDER BY, and a sweep pass +// should be deterministic and its logs diffable across runs. A Runner holding no +// bindings yields an empty slice, not an error — a first-ever enroll sweeps +// nothing, which is a normal outcome. +func (s *Store) DeleteSessionBindingsForRunner(ctx context.Context, runnerID string) ([]SessionBinding, error) { + if runnerID == "" { + return nil, fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + } + rows, err := s.q.DeleteSessionBindingsForRunner(ctx, runnerID) + if err != nil { + return nil, fmt.Errorf("store: delete session bindings for runner: %w", err) + } + bindings := make([]SessionBinding, 0, len(rows)) + for _, row := range rows { + bindings = append(bindings, SessionBinding{ + SessionID: row.SessionID, + AccountID: AccountID(row.AgentAccountID), + RunnerID: runnerID, + }) + } + slices.SortFunc(bindings, func(a, b SessionBinding) int { + return cmp.Compare(a.SessionID, b.SessionID) + }) + return bindings, nil +} diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go new file mode 100644 index 00000000..2576ef85 --- /dev/null +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -0,0 +1,397 @@ +//go:build pgtest + +package store + +// Session bindings: the durable (session -> agent account, Runner) record the +// RunnerHub has so far held only in RAM (RIG-3108 / RIG-2861 §T4). Four things +// must hold, and all four are DATABASE invariants rather than Go logic — the PK, +// the unique index on agent_account_id, the FK to agent_accounts, and the RLS +// policy — so all four are pgtest-backed; a mock would only re-assert the Go. +// +// A binding is SINGULAR per session (a rebind replaces it), an account holds AT +// MOST ONE live session (the unique index), a miss FAILS CLOSED (a comms call +// resolves its scope through these reads, so an empty AccountID with a nil error +// would flow onward as a real principal), and the reconnect sweep RETURNS what +// it removed (the returned rows drive presence DISCONNECTED and the held-deliver +// reap — a bare DELETE would satisfy the invariant while dropping both). + +import ( + "context" + "errors" + "testing" + "time" +) + +// bindingTimes reads a binding row's created_at/updated_at directly, so a test +// asserts the persisted timestamps rather than trusting a return value — the +// same posture as tenantOf. +func bindingTimes(t *testing.T, s *Store, sessionID string) (createdAt, updatedAt time.Time) { + t.Helper() + if err := s.pool.QueryRow(context.Background(), + "SELECT created_at, updated_at FROM session_bindings WHERE session_id = $1", sessionID, + ).Scan(&createdAt, &updatedAt); err != nil { + t.Fatalf("read timestamps of binding %q: %v", sessionID, err) + } + return createdAt, updatedAt +} + +// countBindings counts binding rows for an account, so a rebind test can prove +// the row was REPLACED rather than accumulated. +func countBindings(t *testing.T, s *Store, accountID AccountID) int { + t.Helper() + var n int + if err := s.pool.QueryRow(context.Background(), + "SELECT count(*) FROM session_bindings WHERE agent_account_id = $1", string(accountID), + ).Scan(&n); err != nil { + t.Fatalf("count bindings of %q: %v", accountID, err) + } + return n +} + +// TestRecordSessionBindingRoundTripsBothDirections pins the base contract both +// reads depend on: what the relay bound is what both directions read back — the +// account resolvable from the session id (the inbound comms-call read) and the +// session resolvable from the account (the delivery-dispatch read). +func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + t.Fatalf("RecordSessionBinding: %v", err) + } + + gotAccount, err := s.ResolveSessionAccount(ctx, "sess-1") + if err != nil { + t.Fatalf("ResolveSessionAccount: %v", err) + } + if gotAccount != agent.ID { + t.Fatalf("ResolveSessionAccount = %q, want the bound agent %q", gotAccount, agent.ID) + } + + gotSession, err := s.SessionForAccount(ctx, agent.ID) + if err != nil { + t.Fatalf("SessionForAccount: %v", err) + } + if gotSession != "sess-1" { + t.Fatalf("SessionForAccount = %q, want the bound session %q", gotSession, "sess-1") + } +} + +// TestSessionBindingLookupsFailClosed is the security-relevant assertion. Both +// reads resolve the scope a call runs under, so an unbound session (and an +// account with no live session) MUST be ErrNotFound — never a zero-value +// AccountID or session id with a nil error, which the caller would treat as a +// real principal and dispatch to. +func TestSessionBindingLookupsFailClosed(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + + account, err := s.ResolveSessionAccount(ctx, "never-bound") + if !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionAccount(never-bound) err = %v, want errors.Is(_, ErrNotFound)", err) + } + if account != "" { + t.Fatalf("ResolveSessionAccount(never-bound) = %q, want the empty AccountID — a miss must resolve no principal", account) + } + + // The reverse direction fails closed the same way: a known agent that never + // started a session resolves nothing to dispatch to. + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + session, err := s.SessionForAccount(ctx, agent.ID) + if !errors.Is(err, ErrNotFound) { + t.Fatalf("SessionForAccount(unbound agent) err = %v, want errors.Is(_, ErrNotFound)", err) + } + if session != "" { + t.Fatalf("SessionForAccount(unbound agent) = %q, want the empty session id", session) + } +} + +// TestRecordSessionBindingRebindsSessionInPlace pins the ON CONFLICT path: a +// session id names exactly ONE binding, so rebinding it onto a different Runner +// must REPLACE the row — updating account and runner TOGETHER — not add a +// second. If bindings accumulated, the reconnect sweep on the OLD Runner would +// retire a binding that has already moved, and SessionForAccount would resolve +// whichever of two rows Postgres happened to return. +func TestRecordSessionBindingRebindsSessionInPlace(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + t.Fatalf("first RecordSessionBinding: %v", err) + } + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-2"); err != nil { + t.Fatalf("rebind onto runner-2: %v", err) + } + + if n := countBindings(t, s, agent.ID); n != 1 { + t.Fatalf("bindings for the agent = %d, want exactly 1 (the rebind must replace, not accumulate)", n) + } + + // The old Runner no longer holds it: a sweep on runner-1 must retire nothing. + swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("DeleteSessionBindingsForRunner(runner-1): %v", err) + } + if len(swept) != 0 { + t.Fatalf("runner-1 still holds %+v after the session moved, want no bindings", swept) + } +} + +// TestRecordSessionBindingRejectsASecondSessionForABoundAccount proves the +// UNIQUE index on agent_account_id. The in-RAM accountSessions map this table +// replaces is 1:1, so a second concurrent session for one account would split +// that account's deliveries across two sessions and make SessionForAccount's +// answer depend on which row Postgres returned. The index refuses the write as +// ErrConflict rather than letting the double-bind land. +func TestRecordSessionBindingRejectsASecondSessionForABoundAccount(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + t.Fatalf("first binding: %v", err) + } + + err := s.RecordSessionBinding(ctx, "sess-2", agent.ID, "runner-1") + sentinelIs(t, err, ErrConflict, "a second live session for an already-bound account") + + // The refused write changed nothing: the original session still owns the + // account, in both directions. + gotSession, err := s.SessionForAccount(ctx, agent.ID) + if err != nil { + t.Fatalf("SessionForAccount after the refused bind: %v", err) + } + if gotSession != "sess-1" { + t.Fatalf("the agent resolves to session %q, want the original %q", gotSession, "sess-1") + } + if _, err := s.ResolveSessionAccount(ctx, "sess-2"); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionAccount(sess-2) err = %v, want ErrNotFound — the refused bind must not have landed", err) + } +} + +// TestRecordSessionBindingUnknownAgentIsInvalidArgument pins the FK: a binding +// for an account that is not an agent cannot land. Without it a session could +// resolve to a nonexistent principal, and every authz check downstream would be +// evaluating an account that does not exist. +func TestRecordSessionBindingUnknownAgentIsInvalidArgument(t *testing.T) { + s := newTestStore(t) + + err := s.RecordSessionBinding(t.Context(), "sess-1", "no-such-agent", "runner-1") + sentinelIs(t, err, ErrInvalidArgument, "binding for an unknown agent") +} + +// TestDeleteSessionBindingReleasesAndIsIdempotent covers the single-session +// release path: the delete removes the row (both directions stop resolving, and +// the account's unique slot is free for its next session), and a SECOND delete +// succeeds. Idempotency is load-bearing — a session teardown may be retried, and +// an error on zero rows would strand the retried teardown mid-way. +func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + t.Fatalf("RecordSessionBinding: %v", err) + } + if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("DeleteSessionBinding: %v", err) + } + + if _, err := s.ResolveSessionAccount(ctx, "sess-1"); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionAccount after delete err = %v, want ErrNotFound", err) + } + if _, err := s.SessionForAccount(ctx, agent.ID); !errors.Is(err, ErrNotFound) { + t.Fatalf("SessionForAccount after delete err = %v, want ErrNotFound", err) + } + + // The released slot is reusable: the same agent may bind a NEW session + // without hitting the unique-index conflict. + if err := s.RecordSessionBinding(ctx, "sess-2", agent.ID, "runner-1"); err != nil { + t.Fatalf("binding a new session after release: %v (want nil — the slot was freed)", err) + } + + // A second delete of an already-released session is a no-op, and so is + // deleting one that never existed. + if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { + t.Fatalf("second DeleteSessionBinding(sess-1) = %v, want nil (idempotent)", err) + } + if err := s.DeleteSessionBinding(ctx, "never-bound"); err != nil { + t.Fatalf("DeleteSessionBinding(never-bound) = %v, want nil (idempotent)", err) + } +} + +// TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding is the reconnect +// sweep itself, and the test that would catch a `:exec` regression. The returned +// rows are NOT diagnostic: each drives a presence DISCONNECTED edge for its +// account (RIG-1569 T8) and each session id must be reaped from the delivery +// held-deliver registry (RIG-1569 T3). A bare DELETE would clear the bindings — +// satisfying the invariant — while silently dropping both side-effects, leaving +// a long-WORKING agent stuck WORKING in the projection forever. +// +// It must also sweep ONLY the re-enrolling Runner's bindings: too few leaves a +// stale session resolving a re-minted id to the wrong account; too many drops +// live sessions on a healthy Runner. +func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + a := mustAgent(t, s, owner.ID, "agent-a") + b := mustAgent(t, s, owner.ID, "agent-b") + elsewhere := mustAgent(t, s, owner.ID, "agent-elsewhere") + + for _, bind := range []SessionBinding{ + {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, + {SessionID: "sess-b", AccountID: b.ID, RunnerID: "runner-1"}, + {SessionID: "sess-elsewhere", AccountID: elsewhere.ID, RunnerID: "runner-2"}, + } { + if err := s.RecordSessionBinding(ctx, bind.SessionID, bind.AccountID, bind.RunnerID); err != nil { + t.Fatalf("RecordSessionBinding(%+v): %v", bind, err) + } + } + + swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("DeleteSessionBindingsForRunner: %v", err) + } + + // The returned rows ARE the side-effect inputs, so assert them exactly — + // every swept binding, with the account each DISCONNECTED edge needs. + want := []SessionBinding{ + {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, + {SessionID: "sess-b", AccountID: b.ID, RunnerID: "runner-1"}, + } + if len(swept) != len(want) { + t.Fatalf("swept = %+v, want exactly the %d bindings on runner-1 (a :exec sweep would return none)", swept, len(want)) + } + for i, w := range want { + if swept[i] != w { + t.Fatalf("swept[%d] = %+v, want %+v (sorted by session id)", i, swept[i], w) + } + } + + // The swept bindings are actually gone in both directions. + for _, sessionID := range []string{"sess-a", "sess-b"} { + if _, err := s.ResolveSessionAccount(ctx, sessionID); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionAccount(%s) after the sweep err = %v, want ErrNotFound", sessionID, err) + } + } + + // The other Runner's binding SURVIVES: a re-enroll must not drop live + // sessions on a Runner that never reconnected. + gotAccount, err := s.ResolveSessionAccount(ctx, "sess-elsewhere") + if err != nil { + t.Fatalf("runner-2's binding did not survive the runner-1 sweep: %v", err) + } + if gotAccount != elsewhere.ID { + t.Fatalf("sess-elsewhere resolves to %q, want %q", gotAccount, elsewhere.ID) + } + + // A Runner holding no bindings sweeps nothing quietly — a first-ever enroll. + empty, err := s.DeleteSessionBindingsForRunner(ctx, "runner-never-seen") + if err != nil { + t.Fatalf("DeleteSessionBindingsForRunner(unknown) = %v, want nil (nothing to sweep is not a failure)", err) + } + if len(empty) != 0 { + t.Fatalf("unknown runner sweep = %+v, want empty", empty) + } +} + +// TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly proves the RIG-3495 +// trigger fires through a real store method: the ON CONFLICT rebind advances +// updated_at while created_at stays put. It is also the only proof the column is +// live at all — secrets.updated_at rotted precisely because no write statement +// set it, so the value could only ever equal created_at and every reader was +// reading a lie. +// +// No time.Sleep: now() is TRANSACTION time in Postgres, so the read-then-rebind +// below spans two transactions and the two values differ on their own. +func TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly(t *testing.T) { + ctx := t.Context() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + t.Fatalf("first RecordSessionBinding: %v", err) + } + createdBefore, updatedBefore := bindingTimes(t, s, "sess-1") + if !createdBefore.Equal(updatedBefore) { + t.Fatalf("on INSERT created_at = %v and updated_at = %v, want the same DEFAULT now()", createdBefore, updatedBefore) + } + + if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-2"); err != nil { + t.Fatalf("rebind: %v", err) + } + createdAfter, updatedAfter := bindingTimes(t, s, "sess-1") + + if !createdAfter.Equal(createdBefore) { + t.Fatalf("created_at moved from %v to %v across a rebind; it must record the row's birth", createdBefore, createdAfter) + } + if !updatedAfter.After(updatedBefore) { + t.Fatalf("updated_at = %v after the rebind, want later than %v — the set_updated_at trigger did not fire", updatedAfter, updatedBefore) + } +} + +// TestSessionBindingIsTenantIsolated is the RLS proof: a binding written under +// one tenant is invisible to another, and the isolation is the database policy +// rather than an application-layer WHERE — both rows live in the same physical +// table and only the per-transaction compass.tenant_id GUC differs. It matters +// more here than for most tables: these reads resolve the PRINCIPAL a comms call +// runs under, so a cross-tenant hit would run a foreign tenant's session under a +// local account. +// +// Follows the established pattern in rls_pgtest_test.go — seedTenant for a +// second tenant, WithTenant for its context, every write through the normal +// store API so rows land stamped exactly as a real request would. +func TestSessionBindingIsTenantIsolated(t *testing.T) { + s := newTestStore(t) + tenantB := seedTenant(t, s, "tenant-b") + ctxA := t.Context() // no tenant set → the bootstrap tenant + ctxB := WithTenant(t.Context(), tenantB) + + ownerA := mustUser(t, s, "owner-a") + agentA := mustAgent(t, s, ownerA.ID, "agent-a") + if err := s.RecordSessionBinding(ctxA, "sess-a", agentA.ID, "runner-1"); err != nil { + t.Fatalf("RecordSessionBinding under tenant A: %v", err) + } + + // Tenant B resolves A's session id: the row is not in B's view, so this must + // fail closed rather than hand B tenant A's account. + gotAccount, err := s.ResolveSessionAccount(ctxB, "sess-a") + if !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant B ResolveSessionAccount(sess-a) err = %v, want ErrNotFound — cross-tenant read leak", err) + } + if gotAccount != "" { + t.Fatalf("tenant B resolved tenant A's account %q from A's session — cross-tenant read leak", gotAccount) + } + + // The reverse direction leaks nothing either. + if gotSession, err := s.SessionForAccount(ctxB, agentA.ID); !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant B SessionForAccount(A's agent) = (%q, %v), want ErrNotFound — cross-tenant read leak", gotSession, err) + } + + // Nor does B's sweep of the SAME runner id retire A's binding — a shared + // runner id across tenants must not let one tenant drop another's sessions. + swept, err := s.DeleteSessionBindingsForRunner(ctxB, "runner-1") + if err != nil { + t.Fatalf("tenant B sweep of runner-1: %v", err) + } + if len(swept) != 0 { + t.Fatalf("tenant B's sweep returned %+v, want empty — it reached tenant A's bindings", swept) + } + + // Control: tenant A still sees its OWN binding, proving the policy is not + // simply hiding everything. + if gotAccount, err := s.ResolveSessionAccount(ctxA, "sess-a"); err != nil { + t.Fatalf("tenant A cannot see its OWN binding — policy over-blocks: %v", err) + } else if gotAccount != agentA.ID { + t.Fatalf("tenant A's own binding resolves to %q, want %q", gotAccount, agentA.ID) + } +} From 86e264e13ed126312b2da847bb5b2c00f0f61b5f Mon Sep 17 00:00:00 2001 From: mintaka Date: Mon, 7 Sep 2026 16:20:23 -0400 Subject: [PATCH 2/4] fix(store): key session_bindings by account, not session (RIG-3108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-1 review found that the UNIQUE index on `agent_account_id` forbids a path the hub relies on today. The in-RAM `accountSessions` map is 1:1, which is why the index looked right — but the hub re-points an account onto a NEWER session while the older one is still bound, and there is a test named for exactly that: `relay_comms_test.go:492-501` promotes `sess-old` then `sess-new` onto the same account before unbinding the stale one. Against the previous schema the second bind raised `duplicate key value violates unique constraint "session_bindings_account_key"`. `promoteSession` also returns nothing (`relay_comms.go:49`), so the next slice had nowhere to put the `ErrConflict` the index produced. The doc comment's prescribed remedy — release the stale binding first — is a delete-then-insert race with a window where the account resolves to no session at all, and the stale session id is usually unknown at promotion time because the hub reaches it only through the very map this table replaces. Key by the account instead, with the tenant folded in: PRIMARY KEY (tenant_id, agent_account_id) A re-point is now one atomic upsert on the account, which is what `promoteSession` already does to its map. `session_id` becomes a plain column with `UNIQUE INDEX session_bindings_session_key (tenant_id, session_id)`, so one session still speaks for exactly one account and a second account claiming a live session is refused — now the only reachable unique violation, remapped to `ErrConflict` accordingly. Folding `tenant_id` into both keys follows the convention the forge-coordinate tables already use (`0001_init.sql:775`, `:821`, `:854`), which the RLS block header states is deliberate so two tenants may hold the same coordinate without colliding. Previously a cross-tenant `session_id` collision surfaced as an opaque SQLSTATE 42501 RLS error rather than a mapped sentinel, because the colliding row was invisible to the inserting tenant and its `ON CONFLICT` could not fire. `RecordSessionBinding` now returns the session id it displaced, in one statement, via a CTE that reads the prior row in the same snapshot as the upsert. The next slice needs it: a displaced session must be reaped from the delivery held-deliver registry, and the previous shape gave the caller no way to learn what it had replaced. `COALESCE` to the empty string rather than NULL, because sqlc types the expression as a non-nullable `string` that pgx cannot scan a NULL into — the fresh-insert case would have failed at runtime. Also from the review: - State in the schema comment that a binding is NOT cross-checked against `agent_sessions`, which remains the authz root. There is no FK from `session_id`, deliberately, so the comment no longer implies an integrity the schema does not have. - Add the account-moves rebind test. Dropping an assignment from the upsert's SET list previously passed all 12 cases, because no test ever rebound a session onto a different account. - Seed the sweep-order fixture in reverse (`sess-z`, `sess-b`, `sess-a`). `DELETE ... RETURNING` emits physical heap order, which for freshly inserted rows equals insertion order, so the previous fixture was already sorted and the assertion could not fail if the sort were deleted. - Add a catalog-floor test for `updated_at_tables`, mirroring the one that defends `tenant_tables`: enumerate columns named `updated_at` and fail on any whose table lacks the `set_updated_at` trigger. Without it a future table could gain the column and silently miss the trigger — the rot the trigger exists to prevent. - Use `context.Background()` in both new test files, matching 19 of the 20 existing pgtest files rather than planting a second convention. - Close the sqlc comment gap so the sweep's block comment attaches to its own query instead of the preceding one's godoc. Verified: 25 pgtest cases green against Postgres 16.14. Red controls — dropping `session_id = EXCLUDED.session_id` fails the re-point test; dropping `runner_id = EXCLUDED.runner_id` fails two tests with `runner-1 still holds [...] after the re-point, want none`; removing `slices.SortFunc` fails the sweep-order test with the rows returned backwards. Every revert verified byte-identical by md5. The fresh-insert `displaced == ""` case is pinned in the suite, not just probed. `sqlc-drift`, `sql-migration-gate:check`, `go build ./...`, `gofmt`, and `golangci-lint run ./internal/store/...` all clean. Co-authored-by: Matt Wilkinson --- go/internal/store/db/models.go | 4 +- go/internal/store/db/querier.go | 29 +- go/internal/store/db/session_bindings.sql.go | 58 +++- go/internal/store/migrations/0001_init.sql | 52 +++- .../store/queries/session_bindings.sql | 46 +++- go/internal/store/session_bindings.go | 80 ++++-- .../store/session_bindings_pgtest_test.go | 252 ++++++++++++------ go/internal/store/updated_at_pgtest_test.go | 70 +++++ 8 files changed, 438 insertions(+), 153 deletions(-) diff --git a/go/internal/store/db/models.go b/go/internal/store/db/models.go index a0e2df50..b42c365b 100644 --- a/go/internal/store/db/models.go +++ b/go/internal/store/db/models.go @@ -267,12 +267,12 @@ type Secret struct { } type SessionBinding struct { - SessionID string + TenantID string AgentAccountID string + SessionID string RunnerID string CreatedAt pgtype.Timestamptz UpdatedAt pgtype.Timestamptz - TenantID string } type SystemAccount struct { diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 10549153..264d8ab2 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -355,13 +355,36 @@ type Querier interface { // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable // (session -> agent account, Runner) binding the RunnerHub has so far held only // in RAM. The hand-written Store methods in internal/store/session_bindings.go - // keep their signatures and map these rows into the SessionBinding domain struct - // (the AccountID newtype is done inline in the Go, as agent_placements does). + // map these rows into the SessionBinding domain struct (the AccountID newtype is + // done inline in the Go, as agent_placements does). + // + // No query here names tenant_id. Tenant scoping is the RLS policy's job + // (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id + // column DEFAULTs to the request GUC on insert — which is how every other query + // file here is written. // // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written // `updated_at = now()` is the exact defect that convention removes. - RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error + // The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 + // accountSessions map this replaces treats re-pointing an account at a newer + // session as an assignment, not a collision, so this is an upsert on + // (tenant_id, agent_account_id) and never refuses a re-point. + // + // It returns the session id it DISPLACED, or '' when the account held none — + // because the caller must reap that session from the delivery held-deliver + // registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING + // exists for. COALESCE'd to '' rather than left NULL so the generated signature + // is a plain string: "no displaced session" is the empty string throughout this + // package, as ResolveSessionAccount's miss is. + // + // ONE statement, deliberately. The `prev` CTE reads the pre-update row and the + // upsert writes the new one in the SAME snapshot, so no concurrent bind can slip + // between a read and a write and make the caller reap a session that is still + // live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ + // only; this targets 16), would each lose that. The upsert CTE is unreferenced + // on purpose: a data-modifying CTE always executes. + RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) RemarkSafetyValveSuperseded(ctx context.Context, arg RemarkSafetyValveSupersededParams) error RenameTopic(ctx context.Context, arg RenameTopicParams) error RequireAgentSessionSubscriber(ctx context.Context, arg RequireAgentSessionSubscriberParams) (bool, error) diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index 3db90cef..ac5317c6 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -19,7 +19,6 @@ func (q *Queries) DeleteSessionBinding(ctx context.Context, sessionID string) er } const deleteSessionBindingsForRunner = `-- name: DeleteSessionBindingsForRunner :many - DELETE FROM session_bindings WHERE runner_id = $1 RETURNING session_id, agent_account_id @@ -65,33 +64,64 @@ func (q *Queries) DeleteSessionBindingsForRunner(ctx context.Context, runnerID s return items, nil } -const recordSessionBinding = `-- name: RecordSessionBinding :exec - -INSERT INTO session_bindings (session_id, agent_account_id, runner_id) -VALUES ($1, $2, $3) -ON CONFLICT (session_id) DO UPDATE - SET agent_account_id = EXCLUDED.agent_account_id, - runner_id = EXCLUDED.runner_id +const recordSessionBinding = `-- name: RecordSessionBinding :one + +WITH prev AS ( + SELECT b.session_id FROM session_bindings b WHERE b.agent_account_id = $1 +), upsert AS ( + INSERT INTO session_bindings (agent_account_id, session_id, runner_id) + VALUES ($1, $2, $3) + ON CONFLICT (tenant_id, agent_account_id) DO UPDATE + SET session_id = EXCLUDED.session_id, + runner_id = EXCLUDED.runner_id + RETURNING session_id +) +SELECT COALESCE((SELECT session_id FROM prev), '')::text AS displaced_session_id ` type RecordSessionBindingParams struct { - SessionID string AgentAccountID string + SessionID string RunnerID string } // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable // (session -> agent account, Runner) binding the RunnerHub has so far held only // in RAM. The hand-written Store methods in internal/store/session_bindings.go -// keep their signatures and map these rows into the SessionBinding domain struct -// (the AccountID newtype is done inline in the Go, as agent_placements does). +// map these rows into the SessionBinding domain struct (the AccountID newtype is +// done inline in the Go, as agent_placements does). +// +// No query here names tenant_id. Tenant scoping is the RLS policy's job +// (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id +// column DEFAULTs to the request GUC on insert — which is how every other query +// file here is written. // // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written // `updated_at = now()` is the exact defect that convention removes. -func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error { - _, err := q.db.Exec(ctx, recordSessionBinding, arg.SessionID, arg.AgentAccountID, arg.RunnerID) - return err +// The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 +// accountSessions map this replaces treats re-pointing an account at a newer +// session as an assignment, not a collision, so this is an upsert on +// (tenant_id, agent_account_id) and never refuses a re-point. +// +// It returns the session id it DISPLACED, or ” when the account held none — +// because the caller must reap that session from the delivery held-deliver +// registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING +// exists for. COALESCE'd to ” rather than left NULL so the generated signature +// is a plain string: "no displaced session" is the empty string throughout this +// package, as ResolveSessionAccount's miss is. +// +// ONE statement, deliberately. The `prev` CTE reads the pre-update row and the +// upsert writes the new one in the SAME snapshot, so no concurrent bind can slip +// between a read and a write and make the caller reap a session that is still +// live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ +// only; this targets 16), would each lose that. The upsert CTE is unreferenced +// on purpose: a data-modifying CTE always executes. +func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) { + row := q.db.QueryRow(ctx, recordSessionBinding, arg.AgentAccountID, arg.SessionID, arg.RunnerID) + var displaced_session_id string + err := row.Scan(&displaced_session_id) + return displaced_session_id, err } const sessionBindingAccount = `-- name: SessionBindingAccount :one diff --git a/go/internal/store/migrations/0001_init.sql b/go/internal/store/migrations/0001_init.sql index 88119683..b9952cd8 100644 --- a/go/internal/store/migrations/0001_init.sql +++ b/go/internal/store/migrations/0001_init.sql @@ -499,9 +499,33 @@ CREATE UNIQUE INDEX agent_placements_container_key ON agent_placements (containe -- still authorizes through agent_sessions -> agent_accounts -> channel_members -- and never reads this table. -- --- PK on session_id, not a surrogate: a session id names exactly one binding, and --- that is also the accountForSession read direction the relay resolves on every --- inbound comms call. +-- KEYED ON THE ACCOUNT, tenant folded in: PRIMARY KEY (tenant_id, +-- agent_account_id). The account is the identity because this table replaces the +-- hub's 1:1 in-RAM accountSessions map, where re-pointing an account at a newer +-- session is an assignment, not a collision. So a re-point here is ONE upsert on +-- the account (ON CONFLICT DO UPDATE), never a conflict to be refused: the hub +-- legitimately holds two sessions for one account transiently — it promotes the +-- new session before unbinding the stale one — and RecordSessionBinding's caller +-- has nowhere to put a refusal. The NEWER binding displaces the older, and the +-- displaced session id comes back via RETURNING so the caller can reap it from +-- the held-deliver registry. Keying on the session instead would have made the +-- displacement a unique violation and the reap impossible in one statement. +-- +-- tenant_id leads the key for the reason the RLS header below states: two +-- tenants may hold the same coordinate without collision. Its declaration text +-- is character-identical to every other tenant table's, because the policy +-- compares it to the same GUC. +-- +-- session_id is UNIQUE PER TENANT (the index below) but is NOT the identity: it +-- is the accountForSession read direction the relay resolves on every inbound +-- comms call, and the uniqueness only says one session speaks for one account. +-- +-- A binding is NOT cross-checked against agent_sessions: there is deliberately +-- no FK from session_id, so a binding may name a session with no agent_sessions +-- row, or disagree with one about the owner. That is intentional — a binding is +-- independent of the session record's lifetime — and it is why agent_sessions, +-- not this table, remains the authz root. Nothing here may be read as proof a +-- session exists or as proof of who owns it. -- -- runner_id is deliberately NOT a FK, for the same reason agent_placements' -- isn't: Runners are enrolled in memory under a token subject with no runners @@ -509,23 +533,23 @@ CREATE UNIQUE INDEX agent_placements_container_key ON agent_placements (containe -- swept by the reconnect sweep below, and an unswept binding is a stale session -- that outlives its Runner. CREATE TABLE session_bindings ( - session_id TEXT PRIMARY KEY, + tenant_id TEXT NOT NULL DEFAULT current_setting('compass.tenant_id', TRUE), agent_account_id TEXT NOT NULL REFERENCES agent_accounts (account_id) ON DELETE RESTRICT, + session_id TEXT NOT NULL, runner_id TEXT NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT now(), updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), - tenant_id TEXT NOT NULL DEFAULT current_setting('compass.tenant_id', TRUE) + PRIMARY KEY (tenant_id, agent_account_id) ); --- The REVERSE lookup (SessionForAccount, runnerhub/relay_comms.go:179-184): the --- delivery consumer holds a resolved subscriber account and needs the live --- session to dispatch to. UNIQUE because the in-RAM accountSessions map it --- replaces is 1:1 — an account has AT MOST ONE live session — so this index --- makes a second concurrent session for one account UNREPRESENTABLE rather than --- merely unlikely. A rebind of an account onto a NEW session id therefore raises --- a unique violation (mapped to ErrConflict) instead of silently double-binding --- and letting one account's deliveries split across two sessions. -CREATE UNIQUE INDEX session_bindings_account_key ON session_bindings (agent_account_id); +-- The session -> account lookup (ResolveSessionAccount): the relay's read on +-- every inbound comms call, where the request carries only the session id. +-- UNIQUE so one session id speaks for exactly one account — a second account +-- claiming a live session id is refused (mapped to ErrConflict) rather than +-- letting the relay resolve whichever row Postgres happened to return. +-- Tenant-folded for the same reason the PK is: two tenants may mint the same +-- session id without colliding. +CREATE UNIQUE INDEX session_bindings_session_key ON session_bindings (tenant_id, session_id); -- The reconnect sweep deletes every binding of a re-enrolling Runner -- (DeleteSessionBindingsForRunner), so runner_id is the read direction that diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index 836a853b..386e938a 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -1,19 +1,48 @@ -- Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable -- (session -> agent account, Runner) binding the RunnerHub has so far held only -- in RAM. The hand-written Store methods in internal/store/session_bindings.go --- keep their signatures and map these rows into the SessionBinding domain struct --- (the AccountID newtype is done inline in the Go, as agent_placements does). +-- map these rows into the SessionBinding domain struct (the AccountID newtype is +-- done inline in the Go, as agent_placements does). +-- +-- No query here names tenant_id. Tenant scoping is the RLS policy's job +-- (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id +-- column DEFAULTs to the request GUC on insert — which is how every other query +-- file here is written. -- -- updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger -- (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written -- `updated_at = now()` is the exact defect that convention removes. --- name: RecordSessionBinding :exec -INSERT INTO session_bindings (session_id, agent_account_id, runner_id) -VALUES ($1, $2, $3) -ON CONFLICT (session_id) DO UPDATE - SET agent_account_id = EXCLUDED.agent_account_id, - runner_id = EXCLUDED.runner_id; +-- The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 +-- accountSessions map this replaces treats re-pointing an account at a newer +-- session as an assignment, not a collision, so this is an upsert on +-- (tenant_id, agent_account_id) and never refuses a re-point. +-- +-- It returns the session id it DISPLACED, or '' when the account held none — +-- because the caller must reap that session from the delivery held-deliver +-- registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING +-- exists for. COALESCE'd to '' rather than left NULL so the generated signature +-- is a plain string: "no displaced session" is the empty string throughout this +-- package, as ResolveSessionAccount's miss is. +-- +-- ONE statement, deliberately. The `prev` CTE reads the pre-update row and the +-- upsert writes the new one in the SAME snapshot, so no concurrent bind can slip +-- between a read and a write and make the caller reap a session that is still +-- live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ +-- only; this targets 16), would each lose that. The upsert CTE is unreferenced +-- on purpose: a data-modifying CTE always executes. +-- name: RecordSessionBinding :one +WITH prev AS ( + SELECT b.session_id FROM session_bindings b WHERE b.agent_account_id = $1 +), upsert AS ( + INSERT INTO session_bindings (agent_account_id, session_id, runner_id) + VALUES ($1, $2, $3) + ON CONFLICT (tenant_id, agent_account_id) DO UPDATE + SET session_id = EXCLUDED.session_id, + runner_id = EXCLUDED.runner_id + RETURNING session_id +) +SELECT COALESCE((SELECT session_id FROM prev), '')::text AS displaced_session_id; -- name: SessionBindingAccount :one SELECT agent_account_id FROM session_bindings WHERE session_id = $1; @@ -39,7 +68,6 @@ DELETE FROM session_bindings WHERE session_id = $1; -- what preserves them, so the returned rows are load-bearing, not diagnostic. -- A DELETE ... RETURNING takes no ORDER BY, so the Store method sorts the -- returned slice by session id to keep a sweep pass deterministic and diffable. - -- name: DeleteSessionBindingsForRunner :many DELETE FROM session_bindings WHERE runner_id = $1 diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index 0e788fb4..81458168 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -22,6 +22,11 @@ import ( // (accountForSession), and the delivery consumer resolving the live session of a // recipient account it has already authorized (SessionForAccount). // +// Nor is a binding cross-checked against agent_sessions: there is deliberately +// no FK from session_id, so a binding may name a session with no agent_sessions +// row, or disagree with one about the owner. agent_sessions remains the authz +// root, and nothing read from here may stand in for it. +// // PR2 adds the table and these methods only. Demoting the hub's in-RAM maps to // caches over this table is PR3; nothing in internal/runnerhub reads this yet. @@ -35,31 +40,42 @@ type SessionBinding struct { RunnerID string } -// RecordSessionBinding persists the binding for a live session. It is an UPSERT -// keyed on the session, because a session id names exactly one binding: a rebind -// REPLACES it, updating agent_account_id and runner_id TOGETHER so a row can -// never pair a fresh Runner with the account from a previous attachment. +// RecordSessionBinding points an agent account at the live session that speaks +// for it. It is an UPSERT keyed on the ACCOUNT, not the session: this table +// replaces the hub's 1:1 in-RAM accountSessions map, where re-pointing an +// account at a newer session is an assignment, so a re-point here is one atomic +// statement rather than a conflict to refuse. That matters concretely — the hub +// promotes a new session onto an account BEFORE unbinding the stale one, and +// promoteSession returns nothing, so it has nowhere to put a refusal. +// +// It returns the session id this bind DISPLACED — empty when the account held +// none. That value is load-bearing, not diagnostic: PR3 must reap the displaced +// session from the delivery held-deliver registry, the same side-effect +// DeleteSessionBindingsForRunner's returned rows exist for. A displaced session +// left unreaped holds deliveries for an account that has already moved on. +// +// The read of the previous session and the write of the new one are ONE +// statement (a CTE, see the query file), so no concurrent bind can land between +// them and make the caller reap a session that is once again live. // // updated_at is maintained by the set_updated_at() trigger, never here // (RIG-3495) — the query file assigns it nowhere. // // An unknown agent_account_id is ErrInvalidArgument (the FK). // -// An account that already holds a DIFFERENT live session is ErrConflict (the -// unique index on agent_account_id), and unlike agent_placements' container -// guard this one is a live safety property, not a latent one. The in-RAM -// accountSessions map this replaces is 1:1, so a second concurrent session for -// one account would split that account's deliveries across two sessions and let -// SessionForAccount resolve whichever row Postgres happened to return. Refusing -// the write makes the double-bind unrepresentable. The caller's remedy is to -// release the stale binding first — DeleteSessionBinding for a session it knows -// is gone, or the runner sweep on a re-enroll. -func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) error { +// ErrConflict now means ONE thing, and it is no longer about the account: the +// account path is an upsert and cannot conflict. The only unique index left is +// (tenant_id, session_id), so a violation means this session id is ALREADY BOUND +// TO A DIFFERENT ACCOUNT. Refusing it is what keeps ResolveSessionAccount +// single-valued — two accounts sharing a live session id would make the relay's +// answer depend on which row Postgres returned, and it resolves the principal a +// comms call runs under. +func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (string, error) { if sessionID == "" { - return fmt.Errorf("%w: session id is required", ErrInvalidArgument) + return "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) } if accountID == "" { - return fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) + return "", fmt.Errorf("%w: agent account id is required", ErrInvalidArgument) } // Unlike agent_placements.runner_id, '' is NOT an accepted unknown-runner // sentinel here. A placement must OUTLIVE its Runner's attachment (it is @@ -69,22 +85,23 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco // any real Runner's re-enroll, so it would linger as a stale session that // outlives its Runner — the exact leak the sweep exists to prevent. if runnerID == "" { - return fmt.Errorf("%w: runner id is required", ErrInvalidArgument) + return "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) } - if err := s.q.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ + displaced, err := s.q.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ SessionID: sessionID, AgentAccountID: string(accountID), RunnerID: runnerID, - }); err != nil { + }) + if err != nil { if pgErrIs(err, pgForeignKeyViolation) { - return fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) + return "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) } if pgErrIs(err, pgUniqueViolation) { - return fmt.Errorf("%w: agent %q is already bound to another live session", ErrConflict, accountID) + return "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) } - return fmt.Errorf("store: record session binding: %w", err) + return "", fmt.Errorf("store: record session binding: %w", err) } - return nil + return displaced, nil } // ResolveSessionAccount resolves the agent account a live session speaks for — @@ -112,8 +129,8 @@ func (s *Store) ResolveSessionAccount(ctx context.Context, sessionID string) (Ac // SessionForAccount resolves the live session bound to an agent account — the // REVERSE of ResolveSessionAccount, and the direction the delivery consumer // needs to dispatch a deliver to an already-resolved subscriber -// (runnerhub/relay_comms.go:179-184). At most one row can answer, because the -// unique index makes the mapping 1:1. +// (runnerhub/relay_comms.go:179-184). Exactly one row can answer, because the +// account is the table's key. // // An account with no live session is ErrNotFound — never started, stopped, or // dropped on a Runner reconnect. Fail-closed for the same reason as above: an @@ -135,10 +152,15 @@ func (s *Store) SessionForAccount(ctx context.Context, accountID AccountID) (str } // DeleteSessionBinding releases the binding for sessionID — the single-session -// release path a normal session end takes, freeing the account's unique slot for -// its next session. It is IDEMPOTENT: a session teardown may be retried, and a -// second release of an already-released session must succeed, so deleting an -// absent row is not an error (the same posture as DeleteAgentPlacement). +// release path a normal session end takes. It is IDEMPOTENT: a session teardown +// may be retried, and a second release of an already-released session must +// succeed, so deleting an absent row is not an error (the same posture as +// DeleteAgentPlacement). +// +// Note it deletes by SESSION, not by account, which is what makes it safe to +// call on a session RecordSessionBinding has already displaced: the row now +// names the newer session, so the stale release matches nothing and leaves the +// live binding alone. func (s *Store) DeleteSessionBinding(ctx context.Context, sessionID string) error { if sessionID == "" { return fmt.Errorf("%w: session id is required", ErrInvalidArgument) diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 2576ef85..21359381 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -4,16 +4,22 @@ package store // Session bindings: the durable (session -> agent account, Runner) record the // RunnerHub has so far held only in RAM (RIG-3108 / RIG-2861 §T4). Four things -// must hold, and all four are DATABASE invariants rather than Go logic — the PK, -// the unique index on agent_account_id, the FK to agent_accounts, and the RLS -// policy — so all four are pgtest-backed; a mock would only re-assert the Go. +// must hold, and all four are DATABASE invariants rather than Go logic — the +// account-keyed PK, the unique index on (tenant_id, session_id), the FK to +// agent_accounts, and the RLS policy — so all four are pgtest-backed; a mock +// would only re-assert the Go. // -// A binding is SINGULAR per session (a rebind replaces it), an account holds AT -// MOST ONE live session (the unique index), a miss FAILS CLOSED (a comms call -// resolves its scope through these reads, so an empty AccountID with a nil error -// would flow onward as a real principal), and the reconnect sweep RETURNS what -// it removed (the returned rows drive presence DISCONNECTED and the held-deliver -// reap — a bare DELETE would satisfy the invariant while dropping both). +// The table is keyed on the ACCOUNT, so re-pointing an account at a newer +// session is an UPSERT that returns what it displaced (the hub does exactly this +// and has nowhere to put a refusal), one session id speaks for at most one +// account (the unique index), a miss FAILS CLOSED (a comms call resolves its +// scope through these reads, so an empty AccountID with a nil error would flow +// onward as a real principal), and the reconnect sweep RETURNS what it removed +// (the returned rows drive presence DISCONNECTED and the held-deliver reap — a +// bare DELETE would satisfy the invariant while dropping both). +// +// context.Background is the test root (the pgtest-suite convention, sibling +// updated_at_pgtest_test.go and forge_cursors_pgtest_test.go). import ( "context" @@ -22,6 +28,17 @@ import ( "time" ) +// mustBind records a binding a test only needs to SUCCEED, returning the session +// id it displaced. Most cases below care about a later assertion, not this call. +func mustBind(t *testing.T, s *Store, ctx context.Context, sessionID string, accountID AccountID, runnerID string) string { + t.Helper() + displaced, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) + if err != nil { + t.Fatalf("RecordSessionBinding(%q, %q, %q): %v", sessionID, accountID, runnerID, err) + } + return displaced +} + // bindingTimes reads a binding row's created_at/updated_at directly, so a test // asserts the persisted timestamps rather than trusting a return value — the // same posture as tenantOf. @@ -35,7 +52,7 @@ func bindingTimes(t *testing.T, s *Store, sessionID string) (createdAt, updatedA return createdAt, updatedAt } -// countBindings counts binding rows for an account, so a rebind test can prove +// countBindings counts binding rows for an account, so a re-point test can prove // the row was REPLACED rather than accumulated. func countBindings(t *testing.T, s *Store, accountID AccountID) int { t.Helper() @@ -51,16 +68,22 @@ func countBindings(t *testing.T, s *Store, accountID AccountID) int { // TestRecordSessionBindingRoundTripsBothDirections pins the base contract both // reads depend on: what the relay bound is what both directions read back — the // account resolvable from the session id (the inbound comms-call read) and the -// session resolvable from the account (the delivery-dispatch read). +// session resolvable from the account (the delivery-dispatch read). A FIRST bind +// displaces nothing, so it must report the empty session id rather than any +// placeholder the caller would try to reap. func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { - ctx := t.Context() + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { + displaced, err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1") + if err != nil { t.Fatalf("RecordSessionBinding: %v", err) } + if displaced != "" { + t.Fatalf("first bind displaced %q, want \"\" — the account held no prior session, and a caller reaping a phantom id would clear a live registry entry", displaced) + } gotAccount, err := s.ResolveSessionAccount(ctx, "sess-1") if err != nil { @@ -85,7 +108,7 @@ func TestRecordSessionBindingRoundTripsBothDirections(t *testing.T) { // AccountID or session id with a nil error, which the caller would treat as a // real principal and dispatch to. func TestSessionBindingLookupsFailClosed(t *testing.T) { - ctx := t.Context() + ctx := context.Background() s := newTestStore(t) account, err := s.ResolveSessionAccount(ctx, "never-bound") @@ -109,23 +132,85 @@ func TestSessionBindingLookupsFailClosed(t *testing.T) { } } -// TestRecordSessionBindingRebindsSessionInPlace pins the ON CONFLICT path: a -// session id names exactly ONE binding, so rebinding it onto a different Runner -// must REPLACE the row — updating account and runner TOGETHER — not add a -// second. If bindings accumulated, the reconnect sweep on the OLD Runner would -// retire a binding that has already moved, and SessionForAccount would resolve -// whichever of two rows Postgres happened to return. -func TestRecordSessionBindingRebindsSessionInPlace(t *testing.T) { - ctx := t.Context() +// TestRecordSessionBindingRePointsAccountAndReportsDisplaced is the reason the +// table is keyed on the account. The hub promotes a NEW session onto an account +// while the old one is still bound (relay_comms_test.go's repoint case) and +// promoteSession returns nothing, so this MUST NOT be a conflict: it is one +// upsert that moves the account onto the new session and hands back the old +// session id for the caller to reap from the held-deliver registry. +// +// Both halves of the SET list are asserted here — the new session AND the new +// runner — so dropping either assignment fails: session_id via SessionForAccount +// and the reap value, runner_id via which Runner's sweep finds the binding. +func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { - t.Fatalf("first RecordSessionBinding: %v", err) + mustBind(t, s, ctx, "sess-old", agent.ID, "runner-1") + + displaced, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") + if err != nil { + t.Fatalf("re-point onto sess-new: %v (a re-point must never be refused — promoteSession has nowhere to put an error)", err) + } + if displaced != "sess-old" { + t.Fatalf("re-point displaced %q, want %q — the caller reaps this session from the held-deliver registry, so a wrong value strands or wrongly clears deliveries", displaced, "sess-old") + } + + // The account now resolves to the NEW session (the session_id assignment). + gotSession, err := s.SessionForAccount(ctx, agent.ID) + if err != nil { + t.Fatalf("SessionForAccount after the re-point: %v", err) } - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-2"); err != nil { - t.Fatalf("rebind onto runner-2: %v", err) + if gotSession != "sess-new" { + t.Fatalf("the agent resolves to session %q, want the re-pointed %q", gotSession, "sess-new") + } + + // Exactly one row: a re-point REPLACES, it does not accumulate a second + // binding whose sweep would retire a session that has already moved. + if n := countBindings(t, s, agent.ID); n != 1 { + t.Fatalf("bindings for the agent = %d, want exactly 1 (a re-point must replace, not accumulate)", n) + } + + // The displaced session id no longer resolves: it named the same row, which + // now carries the new session. + if _, err := s.ResolveSessionAccount(ctx, "sess-old"); !errors.Is(err, ErrNotFound) { + t.Fatalf("ResolveSessionAccount(sess-old) err = %v, want ErrNotFound — the displaced session must stop resolving", err) + } + + // The runner_id assignment: the binding moved to runner-2, so the OLD + // Runner's sweep finds nothing and the new one finds it. + stale, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") + if err != nil { + t.Fatalf("DeleteSessionBindingsForRunner(runner-1): %v", err) + } + if len(stale) != 0 { + t.Fatalf("runner-1 still holds %+v after the re-point, want none — runner_id was not updated with session_id", stale) + } + swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-2") + if err != nil { + t.Fatalf("DeleteSessionBindingsForRunner(runner-2): %v", err) + } + if len(swept) != 1 || swept[0].SessionID != "sess-new" { + t.Fatalf("runner-2 sweep = %+v, want exactly the sess-new binding", swept) + } +} + +// TestRecordSessionBindingRebindsSameSessionOntoANewRunner covers the other +// re-point axis: the SAME session moving Runners (a session id is stable across +// a re-attach). It must update runner_id in place rather than adding a row, and +// it displaces itself — the account's previous session IS this session. +func TestRecordSessionBindingRebindsSameSessionOntoANewRunner(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") + displaced := mustBind(t, s, ctx, "sess-1", agent.ID, "runner-2") + if displaced != "sess-1" { + t.Fatalf("re-binding the same session displaced %q, want %q — the account's previous session is this same session", displaced, "sess-1") } if n := countBindings(t, s, agent.ID); n != 1 { @@ -142,36 +227,39 @@ func TestRecordSessionBindingRebindsSessionInPlace(t *testing.T) { } } -// TestRecordSessionBindingRejectsASecondSessionForABoundAccount proves the -// UNIQUE index on agent_account_id. The in-RAM accountSessions map this table -// replaces is 1:1, so a second concurrent session for one account would split -// that account's deliveries across two sessions and make SessionForAccount's -// answer depend on which row Postgres returned. The index refuses the write as -// ErrConflict rather than letting the double-bind land. -func TestRecordSessionBindingRejectsASecondSessionForABoundAccount(t *testing.T) { - ctx := t.Context() +// TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount proves the +// UNIQUE index on (tenant_id, session_id) — the one conflict this table still +// has, now that the account path is an upsert. Two accounts sharing a live +// session id would make ResolveSessionAccount's answer depend on which row +// Postgres returned, and that read resolves the PRINCIPAL a comms call runs +// under. The index refuses the write as ErrConflict rather than letting the +// ambiguity land. +func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T) { + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") - agent := mustAgent(t, s, owner.ID, "agent") + agentA := mustAgent(t, s, owner.ID, "agent-a") + agentB := mustAgent(t, s, owner.ID, "agent-b") - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { - t.Fatalf("first binding: %v", err) - } + mustBind(t, s, ctx, "sess-1", agentA.ID, "runner-1") - err := s.RecordSessionBinding(ctx, "sess-2", agent.ID, "runner-1") - sentinelIs(t, err, ErrConflict, "a second live session for an already-bound account") + displaced, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") + sentinelIs(t, err, ErrConflict, "a session id already bound to a different agent") + if displaced != "" { + t.Fatalf("the refused bind returned displaced = %q, want \"\" — a failed write displaced nothing", displaced) + } - // The refused write changed nothing: the original session still owns the - // account, in both directions. - gotSession, err := s.SessionForAccount(ctx, agent.ID) + // The refused write changed nothing: the session still speaks for agent A, + // and agent B still has no live session. + gotAccount, err := s.ResolveSessionAccount(ctx, "sess-1") if err != nil { - t.Fatalf("SessionForAccount after the refused bind: %v", err) + t.Fatalf("ResolveSessionAccount after the refused bind: %v", err) } - if gotSession != "sess-1" { - t.Fatalf("the agent resolves to session %q, want the original %q", gotSession, "sess-1") + if gotAccount != agentA.ID { + t.Fatalf("sess-1 resolves to %q, want the original owner %q", gotAccount, agentA.ID) } - if _, err := s.ResolveSessionAccount(ctx, "sess-2"); !errors.Is(err, ErrNotFound) { - t.Fatalf("ResolveSessionAccount(sess-2) err = %v, want ErrNotFound — the refused bind must not have landed", err) + if _, err := s.SessionForAccount(ctx, agentB.ID); !errors.Is(err, ErrNotFound) { + t.Fatalf("SessionForAccount(agent-b) err = %v, want ErrNotFound — the refused bind must not have landed", err) } } @@ -182,24 +270,22 @@ func TestRecordSessionBindingRejectsASecondSessionForABoundAccount(t *testing.T) func TestRecordSessionBindingUnknownAgentIsInvalidArgument(t *testing.T) { s := newTestStore(t) - err := s.RecordSessionBinding(t.Context(), "sess-1", "no-such-agent", "runner-1") + _, err := s.RecordSessionBinding(context.Background(), "sess-1", "no-such-agent", "runner-1") sentinelIs(t, err, ErrInvalidArgument, "binding for an unknown agent") } // TestDeleteSessionBindingReleasesAndIsIdempotent covers the single-session -// release path: the delete removes the row (both directions stop resolving, and -// the account's unique slot is free for its next session), and a SECOND delete -// succeeds. Idempotency is load-bearing — a session teardown may be retried, and -// an error on zero rows would strand the retried teardown mid-way. +// release path: the delete removes the row (both directions stop resolving), and +// a SECOND delete succeeds. Idempotency is load-bearing — a session teardown may +// be retried, and an error on zero rows would strand the retried teardown +// mid-way. func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { - ctx := t.Context() + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { - t.Fatalf("RecordSessionBinding: %v", err) - } + mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } @@ -211,10 +297,10 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { t.Fatalf("SessionForAccount after delete err = %v, want ErrNotFound", err) } - // The released slot is reusable: the same agent may bind a NEW session - // without hitting the unique-index conflict. - if err := s.RecordSessionBinding(ctx, "sess-2", agent.ID, "runner-1"); err != nil { - t.Fatalf("binding a new session after release: %v (want nil — the slot was freed)", err) + // The released account is bindable again, and displaces nothing — the row is + // gone, so this is a fresh insert rather than a re-point. + if displaced := mustBind(t, s, ctx, "sess-2", agent.ID, "runner-1"); displaced != "" { + t.Fatalf("binding after a release displaced %q, want \"\" — the row was deleted, so there is nothing to reap", displaced) } // A second delete of an already-released session is a no-op, and so is @@ -238,22 +324,28 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { // It must also sweep ONLY the re-enrolling Runner's bindings: too few leaves a // stale session resolving a re-minted id to the wrong account; too many drops // live sessions on a healthy Runner. +// +// The seed order is REVERSE-SORTED on purpose (sess-z, then sess-b, then +// sess-a). DELETE ... RETURNING emits physical heap order, which for freshly +// inserted rows is insertion order — so seeding in sorted order would make the +// sort assertion below pass with slices.SortFunc removed. Seeding backwards is +// what makes the sort load-bearing. func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { - ctx := t.Context() + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") a := mustAgent(t, s, owner.ID, "agent-a") b := mustAgent(t, s, owner.ID, "agent-b") + z := mustAgent(t, s, owner.ID, "agent-z") elsewhere := mustAgent(t, s, owner.ID, "agent-elsewhere") for _, bind := range []SessionBinding{ - {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, + {SessionID: "sess-z", AccountID: z.ID, RunnerID: "runner-1"}, {SessionID: "sess-b", AccountID: b.ID, RunnerID: "runner-1"}, + {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, {SessionID: "sess-elsewhere", AccountID: elsewhere.ID, RunnerID: "runner-2"}, } { - if err := s.RecordSessionBinding(ctx, bind.SessionID, bind.AccountID, bind.RunnerID); err != nil { - t.Fatalf("RecordSessionBinding(%+v): %v", bind, err) - } + mustBind(t, s, ctx, bind.SessionID, bind.AccountID, bind.RunnerID) } swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") @@ -262,22 +354,24 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { } // The returned rows ARE the side-effect inputs, so assert them exactly — - // every swept binding, with the account each DISCONNECTED edge needs. + // every swept binding, with the account each DISCONNECTED edge needs, in + // sorted order (which is NOT the order they were seeded in). want := []SessionBinding{ {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, {SessionID: "sess-b", AccountID: b.ID, RunnerID: "runner-1"}, + {SessionID: "sess-z", AccountID: z.ID, RunnerID: "runner-1"}, } if len(swept) != len(want) { t.Fatalf("swept = %+v, want exactly the %d bindings on runner-1 (a :exec sweep would return none)", swept, len(want)) } for i, w := range want { if swept[i] != w { - t.Fatalf("swept[%d] = %+v, want %+v (sorted by session id)", i, swept[i], w) + t.Fatalf("swept[%d] = %+v, want %+v (sorted by session id; the rows were seeded in reverse order, so an unsorted sweep returns them backwards)", i, swept[i], w) } } // The swept bindings are actually gone in both directions. - for _, sessionID := range []string{"sess-a", "sess-b"} { + for _, sessionID := range []string{"sess-a", "sess-b", "sess-z"} { if _, err := s.ResolveSessionAccount(ctx, sessionID); !errors.Is(err, ErrNotFound) { t.Fatalf("ResolveSessionAccount(%s) after the sweep err = %v, want ErrNotFound", sessionID, err) } @@ -304,7 +398,7 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { } // TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly proves the RIG-3495 -// trigger fires through a real store method: the ON CONFLICT rebind advances +// trigger fires through a real store method: the ON CONFLICT re-point advances // updated_at while created_at stays put. It is also the only proof the column is // live at all — secrets.updated_at rotted precisely because no write statement // set it, so the value could only ever equal created_at and every reader was @@ -313,22 +407,18 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { // No time.Sleep: now() is TRANSACTION time in Postgres, so the read-then-rebind // below spans two transactions and the two values differ on their own. func TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly(t *testing.T) { - ctx := t.Context() + ctx := context.Background() s := newTestStore(t) owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-1"); err != nil { - t.Fatalf("first RecordSessionBinding: %v", err) - } + mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") createdBefore, updatedBefore := bindingTimes(t, s, "sess-1") if !createdBefore.Equal(updatedBefore) { t.Fatalf("on INSERT created_at = %v and updated_at = %v, want the same DEFAULT now()", createdBefore, updatedBefore) } - if err := s.RecordSessionBinding(ctx, "sess-1", agent.ID, "runner-2"); err != nil { - t.Fatalf("rebind: %v", err) - } + mustBind(t, s, ctx, "sess-1", agent.ID, "runner-2") createdAfter, updatedAfter := bindingTimes(t, s, "sess-1") if !createdAfter.Equal(createdBefore) { @@ -353,14 +443,12 @@ func TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly(t *testing.T) { func TestSessionBindingIsTenantIsolated(t *testing.T) { s := newTestStore(t) tenantB := seedTenant(t, s, "tenant-b") - ctxA := t.Context() // no tenant set → the bootstrap tenant - ctxB := WithTenant(t.Context(), tenantB) + ctxA := context.Background() // no tenant set → the bootstrap tenant + ctxB := WithTenant(context.Background(), tenantB) ownerA := mustUser(t, s, "owner-a") agentA := mustAgent(t, s, ownerA.ID, "agent-a") - if err := s.RecordSessionBinding(ctxA, "sess-a", agentA.ID, "runner-1"); err != nil { - t.Fatalf("RecordSessionBinding under tenant A: %v", err) - } + mustBind(t, s, ctxA, "sess-a", agentA.ID, "runner-1") // Tenant B resolves A's session id: the row is not in B's view, so this must // fail closed rather than hand B tenant A's account. diff --git a/go/internal/store/updated_at_pgtest_test.go b/go/internal/store/updated_at_pgtest_test.go index 83441829..0273676c 100644 --- a/go/internal/store/updated_at_pgtest_test.go +++ b/go/internal/store/updated_at_pgtest_test.go @@ -152,6 +152,76 @@ func TestSecretsUpdatedAtIsLive(t *testing.T) { } } +// TestUpdatedAtTriggerCatalogFloor is the catalog guard the three behavioural +// tests above cannot be: they each pin ONE named table, so a FUTURE table that +// declares updated_at and forgets its updated_at_tables entry is silently +// untriggered — the exact rot RIG-3495 exists to prevent, reintroduced by +// omission rather than by edit. This enumerates the live catalog instead of +// trusting a hand-maintained list, the same self-auditing posture as +// TestRLSCatalogEnabledAndForced (rls_pgtest_test.go): every table carrying an +// updated_at column must carry a set_updated_at trigger. +// +// pg_attribute rather than information_schema.columns because a dropped column +// lingers there as attisdropped, and this must not count one. The LEFT JOIN +// (not a NOT EXISTS) is deliberate: it reports every offending table in one +// run, so adding three tables and forgetting all three is one failure listing +// all three, not three successive red runs. +func TestUpdatedAtTriggerCatalogFloor(t *testing.T) { + s := newTestStore(t) + ctx := context.Background() + + rows, err := s.pool.Query(ctx, + `SELECT c.relname + FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + JOIN pg_attribute a ON a.attrelid = c.oid + LEFT JOIN pg_trigger tg + ON tg.tgrelid = c.oid AND tg.tgname = 'set_updated_at' + WHERE n.nspname = current_schema() + AND c.relkind = 'r' + AND a.attname = 'updated_at' + AND NOT a.attisdropped + AND tg.oid IS NULL + ORDER BY c.relname`) + if err != nil { + t.Fatalf("enumerate updated_at-bearing tables: %v", err) + } + defer rows.Close() + for rows.Next() { + var tbl string + if err := rows.Scan(&tbl); err != nil { + t.Fatalf("scan catalog row: %v", err) + } + t.Errorf("%s: declares updated_at but has no set_updated_at trigger — add it to updated_at_tables in 0001_init.sql, or the column can only ever equal created_at and every reader of it is reading a lie", tbl) + } + if err := rows.Err(); err != nil { + t.Fatalf("iterate catalog rows: %v", err) + } + + // Floor cross-check, mirroring TestRLSCatalogEnabledAndForced's: the + // enumeration above is silent if a table LOSES its updated_at column, which + // would drop it out of the query rather than fail it. Assert the tables that + // must carry the column still do. + for _, tbl := range []string{ + "secrets", "agent_placements", "session_bindings", + "agent_config_bundle", "model_registry", "forge_repo_subscriptions", + } { + var n int + if err := s.pool.QueryRow(ctx, + `SELECT count(*) + FROM pg_attribute a + WHERE a.attrelid = format('%I.%I', current_schema(), $1::text)::regclass + AND a.attname = 'updated_at' + AND NOT a.attisdropped`, tbl, + ).Scan(&n); err != nil { + t.Fatalf("check %s.updated_at: %v", tbl, err) + } + if n != 1 { + t.Errorf("%s: expected to carry updated_at but does not — it has dropped out of trigger coverage silently", tbl) + } + } +} + // repoSubStamps reads a forge repo subscription's (created_at, updated_at). func repoSubStamps(t *testing.T, s *Store, sub ForgeRepoSubscription) (created, updated time.Time) { t.Helper() From 7d0a9f62faf11989b81e33d4670cad27accfdf3a Mon Sep 17 00:00:00 2001 From: mintaka Date: Mon, 7 Sep 2026 17:25:05 -0400 Subject: [PATCH 3/4] fix(store): serialize session binds so the displaced session is reported once (RIG-3108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2 review reproduced a real defect against Postgres 16: the CTE upsert returned a wrong `displaced` session id under concurrency. The query claimed the `prev` read and the write shared a snapshot. They do not. This path runs at READ COMMITTED (the only `IsoLevel` in the package is an unrelated read-only snapshot at `agent_transcripts.go:321`), where `prev` uses the statement-start snapshot but `ON CONFLICT DO UPDATE` blocks on the row lock and then re-reads the latest committed version, so the write acts on a newer row than the read reported. Two interleavings, both reproduced with real concurrent sessions armed exactly as the store arms its own: - Both re-point from `sess-A`: both callers returned `displaced = sess-A`, so `sess-A` is reaped twice and `sess-B` — genuinely displaced — is reaped by nobody. - Both first-bind: the second returned `displaced = ""` having destroyed the first's binding, reporting the loss to nobody. That matters because `displaced` exists so the next slice can reap the held-deliver registry. A stranded entry holds deliveries for an account that has already moved on. Run the prior-value read and the write in one `beginTenantTx`, which arms `SET LOCAL ROLE` and the tenant GUC at BEGIN so every statement stays RLS-scoped — a raw `Begin` would run as the owner with no GUC and silently disable tenant isolation. A locked read alone is not sufficient, and this is the part worth recording: a row lock cannot serialize two binds when there is no row yet. Both first-binds read empty, and the primary key then orders only their *writes* — but `displaced` is a *read*, so the key structurally cannot cover it. Take a per-account advisory lock BEFORE the read, following the package's existing `LockOwnerCoordination`/`AcquireOwnerTreeLock` pattern, keyed on `(tenant, account)` rather than the account alone because the folded key lets two tenants hold bindings for one account id and those binds must not block each other. Both locks are load-bearing: removing either fails a test. Also from the review: - The false "same snapshot" claim appeared in four files (the query source, both generated files, and hand-written in `session_bindings.go`). All four now state the actual guarantee; the measured count of the old assertion is zero. - **The tenant fold had no test at all.** Reverting to un-folded keys left every existing assertion unchanged, so the headline half of the previous commit was unprotected. Two cases now pin it: two tenants holding the same `session_id`, and two holding the same `agent_account_id`. Both fail on the un-folded mutation. - `ErrConflict` branches on `pgConstraintName(err) == "session_bindings_session_key"` and falls through to the generic wrap otherwise. The previous comment claimed the session index was the only unique index left; the primary key is one too, and `pgErrIs` cannot tell which raised, so any other unique violation was reported as a session claim. - `SessionForAccount`'s doc said exactly one row can answer because the account is the key. True only per tenant — corrected, with the system-role carve-out stated. Under `WithSystemRole` (BYPASSRLS) the account direction is unscoped and returns an arbitrary tenant's row; a test pins that hazard rather than leaving it latent, since `WithSystemRole` is on the path the next slice takes. - A system-role write does not fail as the review first measured — it usually *succeeds* and lands an orphan with `tenant_id = ''`, because an ended `SET LOCAL` leaves a custom GUC defined-and-empty on a pooled backend, which satisfies NOT NULL. No RLS policy matches `''`, so the row is invisible to every tenant. The test pins the disjunction rather than pretending either branch is the contract. - The schema comment implied the table represents the hub's transient two-session state. It cannot: the hub keeps `sessionAccounts` (many-to-one) and `accountSessions` (1:1), and this table faithfully represents only the latter. It now says so, and says a caller needing the stale-vs-unknown distinction must keep it in RAM. - `countBindings`/`bindingTimes` read the owner pool with no RLS, so their "exactly one row" assertions were cross-tenant counts. Both now take a tenant predicate — mandatory, not decorative, since the new fold tests deliberately create a second tenant holding the same account and session ids. - `mustBind` takes `ctx` second, matching `mustTopic`; the sqlc header separator stops the file comment being absorbed into `RecordSessionBinding`'s godoc. Corrections to the previous commit message, which overstated three counts: the two touched files hold **14** test functions, not 25; **20 of 20** pgtest files use `context.Background()`, not 19; and the tenant-folded key convention is at `0001_init.sql:799`, `:845`, `:878` — the lines cited before were prose and a timestamp column. Verified: 20 pgtest cases green against Postgres 16. Red controls — removing the advisory lock fails both concurrency tests ("with no row to lock, that lock is the ONLY thing serializing two first binds"); un-folding the keys fails both new tenant tests while the pre-existing isolation test still passes, confirming the review's mutation proof. Every revert verified byte-identical by md5. `sqlc-drift`, `sql-migration-gate:check`, `go build ./...`, `gofmt`, and `golangci-lint run ./internal/store/...` all clean. Co-authored-by: Matt Wilkinson Re-pushed to trigger CI: the base PR was pushed after this head, and the workflow fires only on opened/synchronize/reopened, so no run was created for this commit. --- go/internal/store/db/querier.go | 123 ++++- go/internal/store/db/session_bindings.sql.go | 165 ++++-- go/internal/store/migrations/0001_init.sql | 39 +- .../store/queries/session_bindings.sql | 126 ++++- go/internal/store/session_bindings.go | 116 +++- .../store/session_bindings_pgtest_test.go | 517 +++++++++++++++++- 6 files changed, 937 insertions(+), 149 deletions(-) diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 264d8ab2..35a078a6 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -307,6 +307,75 @@ type Querier interface { LockChannelPolicy(ctx context.Context, id string) (LockChannelPolicyRow, error) LockOwnerCoordination(ctx context.Context, dollar_1 pgtype.Text) error LockOwnerDM(ctx context.Context, dollar_1 pgtype.Text) error + // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable + // (session -> agent account, Runner) binding the RunnerHub has so far held only + // in RAM. The hand-written Store methods in internal/store/session_bindings.go + // map these rows into the SessionBinding domain struct (the AccountID newtype is + // done inline in the Go, as agent_placements does). + // + // No query here names tenant_id, and on the REQUEST path that is complete: the + // store arms every statement with SET LOCAL ROLE compass_app + the + // compass.tenant_id GUC (tenant_tx.go), so the RLS policy (0001_init.sql) scopes + // reads to the acting tenant and the tenant_id column DEFAULTs to that GUC on + // insert. That is how every other query file here is written. + // + // It is NOT complete under WithSystemRole (tenant_tx.go), which arms the + // BYPASSRLS compass_system role and NO tenant GUC. Every query here then runs + // cross-tenant and unscoped, and each one changes meaning: + // + // * SessionBindingForAccount / SessionBindingAccount / SessionBindingForUpdate + // are :one, but with RLS gone their single predicate can match rows in + // SEVERAL tenants. pgx's QueryRow takes the first and discards the rest + // without error, so the caller gets a plausible answer from an arbitrary + // tenant. + // * DeleteSessionBindingsForRunner sweeps EVERY tenant's bindings for that + // runner id — runner ids are not tenant-unique either. + // * RecordSessionBinding does not fail closed. tenant_id DEFAULTs to + // current_setting('compass.tenant_id', TRUE); on a pooled connection that + // previously served an ARMED statement, the ended SET LOCAL leaves that + // custom GUC defined-and-EMPTY rather than undefined, so the DEFAULT + // resolves to '' and the NOT NULL is satisfied. The row lands stamped with a + // tenant that does not exist — and tenant_id here has no FK to tenants + // (accounts.tenant_id does), so nothing catches it. No RLS policy matches + // '', so that row is then invisible to every tenant and releasable by + // nothing on the request path. On a connection that never carried an armed + // statement the GUC is genuinely undefined, the DEFAULT is NULL, and the + // insert fails not-null instead — so which of the two a caller gets depends + // on the pooled connection it draws. + // + // Nothing calls these under the system role today (WithSystemRole is set at + // delivery/consumer.go and runnerhub/hub.go); a PR3 caller that wants to must + // scope them deliberately rather than inherit scoping from here. + // session_bindings_pgtest_test.go pins the observed behaviour so it is recorded + // rather than latent. + // + // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger + // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written + // `updated_at = now()` is the exact defect that convention removes. + // + // The per-account serialization the bind takes FIRST, before it reads anything. + // Auto-released at transaction end. It mirrors LockOwnerDM / LockOwnerCoordination + // / AcquireOwnerTreeLock, with a DISTINCT key domain ('binding:') so a bind never + // serializes behind a DM open, a coordination reconcile, or a reparent. hashtext + // widens the text key to the int the advisory lock takes; a hash collision across + // two accounts is a benign redundant wait, never a wrong result. + // + // It is keyed on TENANT AND ACCOUNT. Every other lock in this package keys on an + // account id alone, which is globally unique — but two tenants can legitimately + // hold bindings for one account id (the PK is folded, see 0001_init.sql), and + // those two binds are independent operations that must not block each other. + // + // Why an advisory lock and not just the row lock below: FOR UPDATE on a row that + // does NOT exist locks NOTHING, so two concurrent FIRST binds for one account + // both read no prior value and neither blocks the other. The PK does serialize + // their WRITES — the second INSERT waits on the first's uncommitted tuple and + // resolves as ON CONFLICT DO UPDATE — but by then both have already read, so both + // report "displaced nothing" while the second has in fact destroyed the first's + // live binding. That session is then reported to NOBODY and its held deliveries + // are stranded. Measured, not assumed: the PK is a write-ordering guarantee and + // the displaced value is a READ, so the PK cannot cover it. This lock is taken + // before the read, so it does. + LockSessionBindingAccount(ctx context.Context, arg LockSessionBindingAccountParams) error MarkMentionsRouted(ctx context.Context, arg MarkMentionsRoutedParams) error MergeTopicLastSeq(ctx context.Context, arg MergeTopicLastSeqParams) error MessageByID(ctx context.Context, id string) (MessageByIDRow, error) @@ -352,39 +421,20 @@ type Querier interface { // there is the NORMAL replay case, not a drop), so asserting rows-affected here // would wrongly fail an idempotent re-fire. RecordOwedMention(ctx context.Context, arg RecordOwedMentionParams) error - // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable - // (session -> agent account, Runner) binding the RunnerHub has so far held only - // in RAM. The hand-written Store methods in internal/store/session_bindings.go - // map these rows into the SessionBinding domain struct (the AccountID newtype is - // done inline in the Go, as agent_placements does). - // - // No query here names tenant_id. Tenant scoping is the RLS policy's job - // (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id - // column DEFAULTs to the request GUC on insert — which is how every other query - // file here is written. - // - // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger - // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written - // `updated_at = now()` is the exact defect that convention removes. // The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 // accountSessions map this replaces treats re-pointing an account at a newer // session as an assignment, not a collision, so this is an upsert on // (tenant_id, agent_account_id) and never refuses a re-point. // - // It returns the session id it DISPLACED, or '' when the account held none — - // because the caller must reap that session from the delivery held-deliver - // registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING - // exists for. COALESCE'd to '' rather than left NULL so the generated signature - // is a plain string: "no displaced session" is the empty string throughout this - // package, as ResolveSessionAccount's miss is. - // - // ONE statement, deliberately. The `prev` CTE reads the pre-update row and the - // upsert writes the new one in the SAME snapshot, so no concurrent bind can slip - // between a read and a write and make the caller reap a session that is still - // live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ - // only; this targets 16), would each lose that. The upsert CTE is unreferenced - // on purpose: a data-modifying CTE always executes. - RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) + // What it DISPLACED comes from SessionBindingForUpdate above, not from a + // RETURNING here: ON CONFLICT DO UPDATE's RETURNING sees the POST-update row, and + // the pre-update one is unreachable from this statement (`OLD`-aliased RETURNING + // is Postgres 18+; this targets 16). The two statements are nonetheless one + // operation, because they share a transaction and the row lock the read took — + // which is the property the caller needs. It reaps the displaced session from the + // delivery held-deliver registry, and reaping a session that is once again live + // would strand a live agent's deliveries. + RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error RemarkSafetyValveSuperseded(ctx context.Context, arg RemarkSafetyValveSupersededParams) error RenameTopic(ctx context.Context, arg RenameTopicParams) error RequireAgentSessionSubscriber(ctx context.Context, arg RequireAgentSessionSubscriberParams) (bool, error) @@ -420,6 +470,23 @@ type Querier interface { SessionBase(ctx context.Context, sessionID string) (int64, error) SessionBindingAccount(ctx context.Context, sessionID string) (string, error) SessionBindingForAccount(ctx context.Context, agentAccountID string) (string, error) + // The prior-value read of the bind, and the second of three statements the Store + // runs in ONE explicit transaction (beginTenantTx): the advisory lock above, this + // read, then the upsert below. FOR UPDATE takes a row lock on the binding this + // bind is about to overwrite, so the read and the write cannot be separated by a + // concurrent bind — that caller is already parked on the advisory lock, and would + // park here too. + // + // The single-statement form this replaces (a `prev` CTE beside the upsert) could + // not do that: under READ COMMITTED the CTE reads the statement-start snapshot + // while ON CONFLICT DO UPDATE blocks on the row lock and then re-reads the LATEST + // committed row, so the two halves saw different versions and the reported + // displacement was wrong — two concurrent re-points away from sess-A both + // reported sess-A, so the genuinely displaced sess-B was never reaped. + // + // A MISS is not an error: a first-ever bind returns pgx.ErrNoRows and the Store + // maps that to the empty displaced id. + SessionBindingForUpdate(ctx context.Context, agentAccountID string) (string, error) SessionMaxEntrySeq(ctx context.Context, sessionID string) (int64, error) SessionTranscript(ctx context.Context, sessionID string) ([]SessionTranscriptRow, error) // Agent-activity queries (sqlc adoption T5, RIG-3034). These replace the inline diff --git a/go/internal/store/db/session_bindings.sql.go b/go/internal/store/db/session_bindings.sql.go index ac5317c6..bc44b6ce 100644 --- a/go/internal/store/db/session_bindings.sql.go +++ b/go/internal/store/db/session_bindings.sql.go @@ -7,6 +7,8 @@ package db import ( "context" + + "github.com/jackc/pgx/v5/pgtype" ) const deleteSessionBinding = `-- name: DeleteSessionBinding :exec @@ -64,25 +66,14 @@ func (q *Queries) DeleteSessionBindingsForRunner(ctx context.Context, runnerID s return items, nil } -const recordSessionBinding = `-- name: RecordSessionBinding :one - -WITH prev AS ( - SELECT b.session_id FROM session_bindings b WHERE b.agent_account_id = $1 -), upsert AS ( - INSERT INTO session_bindings (agent_account_id, session_id, runner_id) - VALUES ($1, $2, $3) - ON CONFLICT (tenant_id, agent_account_id) DO UPDATE - SET session_id = EXCLUDED.session_id, - runner_id = EXCLUDED.runner_id - RETURNING session_id -) -SELECT COALESCE((SELECT session_id FROM prev), '')::text AS displaced_session_id +const lockSessionBindingAccount = `-- name: LockSessionBindingAccount :exec + +SELECT pg_advisory_xact_lock(hashtext('binding:' || $1 || ':' || $2)) ` -type RecordSessionBindingParams struct { - AgentAccountID string - SessionID string - RunnerID string +type LockSessionBindingAccountParams struct { + Column1 pgtype.Text + Column2 pgtype.Text } // Session-binding queries (RIG-3108 / RIG-2861 §T4): the durable @@ -91,37 +82,103 @@ type RecordSessionBindingParams struct { // map these rows into the SessionBinding domain struct (the AccountID newtype is // done inline in the Go, as agent_placements does). // -// No query here names tenant_id. Tenant scoping is the RLS policy's job -// (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id -// column DEFAULTs to the request GUC on insert — which is how every other query -// file here is written. +// No query here names tenant_id, and on the REQUEST path that is complete: the +// store arms every statement with SET LOCAL ROLE compass_app + the +// compass.tenant_id GUC (tenant_tx.go), so the RLS policy (0001_init.sql) scopes +// reads to the acting tenant and the tenant_id column DEFAULTs to that GUC on +// insert. That is how every other query file here is written. +// +// It is NOT complete under WithSystemRole (tenant_tx.go), which arms the +// BYPASSRLS compass_system role and NO tenant GUC. Every query here then runs +// cross-tenant and unscoped, and each one changes meaning: +// +// - SessionBindingForAccount / SessionBindingAccount / SessionBindingForUpdate +// are :one, but with RLS gone their single predicate can match rows in +// SEVERAL tenants. pgx's QueryRow takes the first and discards the rest +// without error, so the caller gets a plausible answer from an arbitrary +// tenant. +// - DeleteSessionBindingsForRunner sweeps EVERY tenant's bindings for that +// runner id — runner ids are not tenant-unique either. +// - RecordSessionBinding does not fail closed. tenant_id DEFAULTs to +// current_setting('compass.tenant_id', TRUE); on a pooled connection that +// previously served an ARMED statement, the ended SET LOCAL leaves that +// custom GUC defined-and-EMPTY rather than undefined, so the DEFAULT +// resolves to ” and the NOT NULL is satisfied. The row lands stamped with a +// tenant that does not exist — and tenant_id here has no FK to tenants +// (accounts.tenant_id does), so nothing catches it. No RLS policy matches +// ”, so that row is then invisible to every tenant and releasable by +// nothing on the request path. On a connection that never carried an armed +// statement the GUC is genuinely undefined, the DEFAULT is NULL, and the +// insert fails not-null instead — so which of the two a caller gets depends +// on the pooled connection it draws. +// +// Nothing calls these under the system role today (WithSystemRole is set at +// delivery/consumer.go and runnerhub/hub.go); a PR3 caller that wants to must +// scope them deliberately rather than inherit scoping from here. +// session_bindings_pgtest_test.go pins the observed behaviour so it is recorded +// rather than latent. // // updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger // (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written // `updated_at = now()` is the exact defect that convention removes. +// +// The per-account serialization the bind takes FIRST, before it reads anything. +// Auto-released at transaction end. It mirrors LockOwnerDM / LockOwnerCoordination +// / AcquireOwnerTreeLock, with a DISTINCT key domain ('binding:') so a bind never +// serializes behind a DM open, a coordination reconcile, or a reparent. hashtext +// widens the text key to the int the advisory lock takes; a hash collision across +// two accounts is a benign redundant wait, never a wrong result. +// +// It is keyed on TENANT AND ACCOUNT. Every other lock in this package keys on an +// account id alone, which is globally unique — but two tenants can legitimately +// hold bindings for one account id (the PK is folded, see 0001_init.sql), and +// those two binds are independent operations that must not block each other. +// +// Why an advisory lock and not just the row lock below: FOR UPDATE on a row that +// does NOT exist locks NOTHING, so two concurrent FIRST binds for one account +// both read no prior value and neither blocks the other. The PK does serialize +// their WRITES — the second INSERT waits on the first's uncommitted tuple and +// resolves as ON CONFLICT DO UPDATE — but by then both have already read, so both +// report "displaced nothing" while the second has in fact destroyed the first's +// live binding. That session is then reported to NOBODY and its held deliveries +// are stranded. Measured, not assumed: the PK is a write-ordering guarantee and +// the displaced value is a READ, so the PK cannot cover it. This lock is taken +// before the read, so it does. +func (q *Queries) LockSessionBindingAccount(ctx context.Context, arg LockSessionBindingAccountParams) error { + _, err := q.db.Exec(ctx, lockSessionBindingAccount, arg.Column1, arg.Column2) + return err +} + +const recordSessionBinding = `-- name: RecordSessionBinding :exec +INSERT INTO session_bindings (agent_account_id, session_id, runner_id) +VALUES ($1, $2, $3) +ON CONFLICT (tenant_id, agent_account_id) DO UPDATE + SET session_id = EXCLUDED.session_id, + runner_id = EXCLUDED.runner_id +` + +type RecordSessionBindingParams struct { + AgentAccountID string + SessionID string + RunnerID string +} + // The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 // accountSessions map this replaces treats re-pointing an account at a newer // session as an assignment, not a collision, so this is an upsert on // (tenant_id, agent_account_id) and never refuses a re-point. // -// It returns the session id it DISPLACED, or ” when the account held none — -// because the caller must reap that session from the delivery held-deliver -// registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING -// exists for. COALESCE'd to ” rather than left NULL so the generated signature -// is a plain string: "no displaced session" is the empty string throughout this -// package, as ResolveSessionAccount's miss is. -// -// ONE statement, deliberately. The `prev` CTE reads the pre-update row and the -// upsert writes the new one in the SAME snapshot, so no concurrent bind can slip -// between a read and a write and make the caller reap a session that is still -// live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ -// only; this targets 16), would each lose that. The upsert CTE is unreferenced -// on purpose: a data-modifying CTE always executes. -func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) (string, error) { - row := q.db.QueryRow(ctx, recordSessionBinding, arg.AgentAccountID, arg.SessionID, arg.RunnerID) - var displaced_session_id string - err := row.Scan(&displaced_session_id) - return displaced_session_id, err +// What it DISPLACED comes from SessionBindingForUpdate above, not from a +// RETURNING here: ON CONFLICT DO UPDATE's RETURNING sees the POST-update row, and +// the pre-update one is unreachable from this statement (`OLD`-aliased RETURNING +// is Postgres 18+; this targets 16). The two statements are nonetheless one +// operation, because they share a transaction and the row lock the read took — +// which is the property the caller needs. It reaps the displaced session from the +// delivery held-deliver registry, and reaping a session that is once again live +// would strand a live agent's deliveries. +func (q *Queries) RecordSessionBinding(ctx context.Context, arg RecordSessionBindingParams) error { + _, err := q.db.Exec(ctx, recordSessionBinding, arg.AgentAccountID, arg.SessionID, arg.RunnerID) + return err } const sessionBindingAccount = `-- name: SessionBindingAccount :one @@ -145,3 +202,33 @@ func (q *Queries) SessionBindingForAccount(ctx context.Context, agentAccountID s err := row.Scan(&session_id) return session_id, err } + +const sessionBindingForUpdate = `-- name: SessionBindingForUpdate :one +SELECT b.session_id + FROM session_bindings b + WHERE b.agent_account_id = $1 + FOR UPDATE +` + +// The prior-value read of the bind, and the second of three statements the Store +// runs in ONE explicit transaction (beginTenantTx): the advisory lock above, this +// read, then the upsert below. FOR UPDATE takes a row lock on the binding this +// bind is about to overwrite, so the read and the write cannot be separated by a +// concurrent bind — that caller is already parked on the advisory lock, and would +// park here too. +// +// The single-statement form this replaces (a `prev` CTE beside the upsert) could +// not do that: under READ COMMITTED the CTE reads the statement-start snapshot +// while ON CONFLICT DO UPDATE blocks on the row lock and then re-reads the LATEST +// committed row, so the two halves saw different versions and the reported +// displacement was wrong — two concurrent re-points away from sess-A both +// reported sess-A, so the genuinely displaced sess-B was never reaped. +// +// A MISS is not an error: a first-ever bind returns pgx.ErrNoRows and the Store +// maps that to the empty displaced id. +func (q *Queries) SessionBindingForUpdate(ctx context.Context, agentAccountID string) (string, error) { + row := q.db.QueryRow(ctx, sessionBindingForUpdate, agentAccountID) + var session_id string + err := row.Scan(&session_id) + return session_id, err +} diff --git a/go/internal/store/migrations/0001_init.sql b/go/internal/store/migrations/0001_init.sql index b9952cd8..0b4f7867 100644 --- a/go/internal/store/migrations/0001_init.sql +++ b/go/internal/store/migrations/0001_init.sql @@ -500,16 +500,35 @@ CREATE UNIQUE INDEX agent_placements_container_key ON agent_placements (containe -- and never reads this table. -- -- KEYED ON THE ACCOUNT, tenant folded in: PRIMARY KEY (tenant_id, --- agent_account_id). The account is the identity because this table replaces the --- hub's 1:1 in-RAM accountSessions map, where re-pointing an account at a newer --- session is an assignment, not a collision. So a re-point here is ONE upsert on --- the account (ON CONFLICT DO UPDATE), never a conflict to be refused: the hub --- legitimately holds two sessions for one account transiently — it promotes the --- new session before unbinding the stale one — and RecordSessionBinding's caller --- has nowhere to put a refusal. The NEWER binding displaces the older, and the --- displaced session id comes back via RETURNING so the caller can reap it from --- the held-deliver registry. Keying on the session instead would have made the --- displacement a unique violation and the reap impossible in one statement. +-- agent_account_id). The account is the identity because this table represents +-- the hub's 1:1 in-RAM accountSessions map (account -> its ONE live session), +-- where re-pointing an account at a newer session is an assignment, not a +-- collision. So a re-point here is an upsert on the account +-- (ON CONFLICT DO UPDATE), never a conflict to be refused: the hub legitimately +-- holds two sessions for one account transiently — it promotes the new session +-- before unbinding the stale one — and RecordSessionBinding's caller has nowhere +-- to put a refusal. Keying on the session instead would have made the +-- displacement a unique violation and the reap below impossible. +-- +-- It represents accountSessions ONLY. The hub also keeps sessionAccounts +-- (session -> account, MANY-to-one), and this table is not that map: a re-point +-- OVERWRITES the row, so the displaced session stops resolving IMMEDIATELY — +-- ResolveSessionAccount returns not-found for it the moment the newer bind +-- commits, which the suite asserts directly. There is deliberately no tombstone +-- and no history: this table answers "which session speaks for this account +-- NOW", nothing else. +-- +-- So it cannot answer STALE vs UNKNOWN. A caller that needs to tell "a session +-- this account used to hold" from "a session id we have never seen" — the hub's +-- re-point guard at runnerhub/relay_comms.go:112-116 — must keep that +-- distinction in RAM. Demoting the hub's maps to caches over this table (PR3) +-- does not change that: the guard's state has no column here to live in. +-- +-- What a bind DISPLACED is still reported to the caller, so it can reap the +-- displaced session from the held-deliver registry. It comes from the prior-value +-- read RecordSessionBinding takes under FOR UPDATE in the same transaction as the +-- write (queries/session_bindings.sql), not from a RETURNING — ON CONFLICT DO +-- UPDATE's RETURNING sees the post-update row. -- -- tenant_id leads the key for the reason the RLS header below states: two -- tenants may hold the same coordinate without collision. Its declaration text diff --git a/go/internal/store/queries/session_bindings.sql b/go/internal/store/queries/session_bindings.sql index 386e938a..126d5b87 100644 --- a/go/internal/store/queries/session_bindings.sql +++ b/go/internal/store/queries/session_bindings.sql @@ -4,45 +4,113 @@ -- map these rows into the SessionBinding domain struct (the AccountID newtype is -- done inline in the Go, as agent_placements does). -- --- No query here names tenant_id. Tenant scoping is the RLS policy's job --- (0001_init.sql) — reads see only the acting tenant's rows and the tenant_id --- column DEFAULTs to the request GUC on insert — which is how every other query --- file here is written. +-- No query here names tenant_id, and on the REQUEST path that is complete: the +-- store arms every statement with SET LOCAL ROLE compass_app + the +-- compass.tenant_id GUC (tenant_tx.go), so the RLS policy (0001_init.sql) scopes +-- reads to the acting tenant and the tenant_id column DEFAULTs to that GUC on +-- insert. That is how every other query file here is written. +-- +-- It is NOT complete under WithSystemRole (tenant_tx.go), which arms the +-- BYPASSRLS compass_system role and NO tenant GUC. Every query here then runs +-- cross-tenant and unscoped, and each one changes meaning: +-- +-- * SessionBindingForAccount / SessionBindingAccount / SessionBindingForUpdate +-- are :one, but with RLS gone their single predicate can match rows in +-- SEVERAL tenants. pgx's QueryRow takes the first and discards the rest +-- without error, so the caller gets a plausible answer from an arbitrary +-- tenant. +-- * DeleteSessionBindingsForRunner sweeps EVERY tenant's bindings for that +-- runner id — runner ids are not tenant-unique either. +-- * RecordSessionBinding does not fail closed. tenant_id DEFAULTs to +-- current_setting('compass.tenant_id', TRUE); on a pooled connection that +-- previously served an ARMED statement, the ended SET LOCAL leaves that +-- custom GUC defined-and-EMPTY rather than undefined, so the DEFAULT +-- resolves to '' and the NOT NULL is satisfied. The row lands stamped with a +-- tenant that does not exist — and tenant_id here has no FK to tenants +-- (accounts.tenant_id does), so nothing catches it. No RLS policy matches +-- '', so that row is then invisible to every tenant and releasable by +-- nothing on the request path. On a connection that never carried an armed +-- statement the GUC is genuinely undefined, the DEFAULT is NULL, and the +-- insert fails not-null instead — so which of the two a caller gets depends +-- on the pooled connection it draws. +-- +-- Nothing calls these under the system role today (WithSystemRole is set at +-- delivery/consumer.go and runnerhub/hub.go); a PR3 caller that wants to must +-- scope them deliberately rather than inherit scoping from here. +-- session_bindings_pgtest_test.go pins the observed behaviour so it is recorded +-- rather than latent. -- -- updated_at is NEVER assigned here: the set_updated_at() BEFORE UPDATE trigger -- (0001_init.sql, RIG-3495) is the one mechanism, and a hand-written -- `updated_at = now()` is the exact defect that convention removes. +-- + +-- The per-account serialization the bind takes FIRST, before it reads anything. +-- Auto-released at transaction end. It mirrors LockOwnerDM / LockOwnerCoordination +-- / AcquireOwnerTreeLock, with a DISTINCT key domain ('binding:') so a bind never +-- serializes behind a DM open, a coordination reconcile, or a reparent. hashtext +-- widens the text key to the int the advisory lock takes; a hash collision across +-- two accounts is a benign redundant wait, never a wrong result. +-- +-- It is keyed on TENANT AND ACCOUNT. Every other lock in this package keys on an +-- account id alone, which is globally unique — but two tenants can legitimately +-- hold bindings for one account id (the PK is folded, see 0001_init.sql), and +-- those two binds are independent operations that must not block each other. +-- +-- Why an advisory lock and not just the row lock below: FOR UPDATE on a row that +-- does NOT exist locks NOTHING, so two concurrent FIRST binds for one account +-- both read no prior value and neither blocks the other. The PK does serialize +-- their WRITES — the second INSERT waits on the first's uncommitted tuple and +-- resolves as ON CONFLICT DO UPDATE — but by then both have already read, so both +-- report "displaced nothing" while the second has in fact destroyed the first's +-- live binding. That session is then reported to NOBODY and its held deliveries +-- are stranded. Measured, not assumed: the PK is a write-ordering guarantee and +-- the displaced value is a READ, so the PK cannot cover it. This lock is taken +-- before the read, so it does. +-- name: LockSessionBindingAccount :exec +SELECT pg_advisory_xact_lock(hashtext('binding:' || $1 || ':' || $2)); + +-- The prior-value read of the bind, and the second of three statements the Store +-- runs in ONE explicit transaction (beginTenantTx): the advisory lock above, this +-- read, then the upsert below. FOR UPDATE takes a row lock on the binding this +-- bind is about to overwrite, so the read and the write cannot be separated by a +-- concurrent bind — that caller is already parked on the advisory lock, and would +-- park here too. +-- +-- The single-statement form this replaces (a `prev` CTE beside the upsert) could +-- not do that: under READ COMMITTED the CTE reads the statement-start snapshot +-- while ON CONFLICT DO UPDATE blocks on the row lock and then re-reads the LATEST +-- committed row, so the two halves saw different versions and the reported +-- displacement was wrong — two concurrent re-points away from sess-A both +-- reported sess-A, so the genuinely displaced sess-B was never reaped. +-- +-- A MISS is not an error: a first-ever bind returns pgx.ErrNoRows and the Store +-- maps that to the empty displaced id. +-- name: SessionBindingForUpdate :one +SELECT b.session_id + FROM session_bindings b + WHERE b.agent_account_id = $1 + FOR UPDATE; -- The bind. Keyed on the ACCOUNT (see the table comment): the hub's 1:1 -- accountSessions map this replaces treats re-pointing an account at a newer -- session as an assignment, not a collision, so this is an upsert on -- (tenant_id, agent_account_id) and never refuses a re-point. -- --- It returns the session id it DISPLACED, or '' when the account held none — --- because the caller must reap that session from the delivery held-deliver --- registry, the same side-effect DeleteSessionBindingsForRunner's RETURNING --- exists for. COALESCE'd to '' rather than left NULL so the generated signature --- is a plain string: "no displaced session" is the empty string throughout this --- package, as ResolveSessionAccount's miss is. --- --- ONE statement, deliberately. The `prev` CTE reads the pre-update row and the --- upsert writes the new one in the SAME snapshot, so no concurrent bind can slip --- between a read and a write and make the caller reap a session that is still --- live. A read-then-write from Go, or an `OLD`-aliased RETURNING (Postgres 18+ --- only; this targets 16), would each lose that. The upsert CTE is unreferenced --- on purpose: a data-modifying CTE always executes. --- name: RecordSessionBinding :one -WITH prev AS ( - SELECT b.session_id FROM session_bindings b WHERE b.agent_account_id = $1 -), upsert AS ( - INSERT INTO session_bindings (agent_account_id, session_id, runner_id) - VALUES ($1, $2, $3) - ON CONFLICT (tenant_id, agent_account_id) DO UPDATE - SET session_id = EXCLUDED.session_id, - runner_id = EXCLUDED.runner_id - RETURNING session_id -) -SELECT COALESCE((SELECT session_id FROM prev), '')::text AS displaced_session_id; +-- What it DISPLACED comes from SessionBindingForUpdate above, not from a +-- RETURNING here: ON CONFLICT DO UPDATE's RETURNING sees the POST-update row, and +-- the pre-update one is unreachable from this statement (`OLD`-aliased RETURNING +-- is Postgres 18+; this targets 16). The two statements are nonetheless one +-- operation, because they share a transaction and the row lock the read took — +-- which is the property the caller needs. It reaps the displaced session from the +-- delivery held-deliver registry, and reaping a session that is once again live +-- would strand a live agent's deliveries. +-- name: RecordSessionBinding :exec +INSERT INTO session_bindings (agent_account_id, session_id, runner_id) +VALUES ($1, $2, $3) +ON CONFLICT (tenant_id, agent_account_id) DO UPDATE + SET session_id = EXCLUDED.session_id, + runner_id = EXCLUDED.runner_id; -- name: SessionBindingAccount :one SELECT agent_account_id FROM session_bindings WHERE session_id = $1; diff --git a/go/internal/store/session_bindings.go b/go/internal/store/session_bindings.go index 81458168..a070a2e9 100644 --- a/go/internal/store/session_bindings.go +++ b/go/internal/store/session_bindings.go @@ -6,6 +6,8 @@ import ( "fmt" "slices" + "github.com/jackc/pgx/v5/pgtype" + "github.com/RigelBuild/compass/go/internal/store/db" ) @@ -43,10 +45,10 @@ type SessionBinding struct { // RecordSessionBinding points an agent account at the live session that speaks // for it. It is an UPSERT keyed on the ACCOUNT, not the session: this table // replaces the hub's 1:1 in-RAM accountSessions map, where re-pointing an -// account at a newer session is an assignment, so a re-point here is one atomic -// statement rather than a conflict to refuse. That matters concretely — the hub -// promotes a new session onto an account BEFORE unbinding the stale one, and -// promoteSession returns nothing, so it has nowhere to put a refusal. +// account at a newer session is an assignment, so a re-point here is a write +// that always lands rather than a conflict to refuse. That matters concretely — +// the hub promotes a new session onto an account BEFORE unbinding the stale one, +// and promoteSession returns nothing, so it has nowhere to put a refusal. // // It returns the session id this bind DISPLACED — empty when the account held // none. That value is load-bearing, not diagnostic: PR3 must reap the displaced @@ -54,22 +56,52 @@ type SessionBinding struct { // DeleteSessionBindingsForRunner's returned rows exist for. A displaced session // left unreaped holds deliveries for an account that has already moved on. // -// The read of the previous session and the write of the new one are ONE -// statement (a CTE, see the query file), so no concurrent bind can land between -// them and make the caller reap a session that is once again live. +// So the bind is THREE statements in one explicit transaction, and that is the +// whole reason this method opens a tx at all: a per-account advisory lock, then +// the prior-value read under FOR UPDATE, then the write. A concurrent bind for +// the same account parks on the advisory lock until this transaction commits, so +// it cannot interleave between this read and this write — the two binds +// serialize, and each caller is told the session IT actually displaced. +// +// A single statement could not give that. The `prev` CTE this replaces read the +// statement-start snapshot while the upsert's ON CONFLICT DO UPDATE blocked on +// the row lock and re-read the latest committed row — different versions, so +// under READ COMMITTED two concurrent re-points away from sess-A both reported +// sess-A and the genuinely displaced sess-B was never reaped. +// +// The advisory lock is what covers the FIRST bind, and it is not redundant with +// FOR UPDATE. FOR UPDATE on a row that does not yet exist locks NOTHING, so two +// concurrent first binds would both read no prior value; the PRIMARY KEY does +// serialize their WRITES (the second INSERT waits on the first's uncommitted +// tuple, then resolves as ON CONFLICT DO UPDATE), but both have already READ by +// then, so both would report "displaced nothing" while the second had in fact +// destroyed the first's live binding — reported to nobody, deliveries stranded. +// Measured, not assumed. The PK orders writes; the displaced value is a read, so +// only a lock taken BEFORE the read can make it correct. +// +// Both locks are kept. The advisory lock alone serializes binds for one account, +// and FOR UPDATE alone cannot cover a first bind; together the read is correct +// whether or not a row already exists, and the row lock still guards against a +// writer that reaches the row by some path not holding the advisory lock. // // updated_at is maintained by the set_updated_at() trigger, never here // (RIG-3495) — the query file assigns it nowhere. // // An unknown agent_account_id is ErrInvalidArgument (the FK). // -// ErrConflict now means ONE thing, and it is no longer about the account: the -// account path is an upsert and cannot conflict. The only unique index left is -// (tenant_id, session_id), so a violation means this session id is ALREADY BOUND -// TO A DIFFERENT ACCOUNT. Refusing it is what keeps ResolveSessionAccount -// single-valued — two accounts sharing a live session id would make the relay's -// answer depend on which row Postgres returned, and it resolves the principal a -// comms call runs under. +// ErrConflict means ONE thing, and it is not about the account: the account path +// is an upsert and cannot conflict. Both unique indexes on this table can raise +// 23505, though — the PK (tenant_id, agent_account_id) as well as +// session_bindings_session_key (tenant_id, session_id) — and the SQLSTATE alone +// does not say which did, so the mapping branches on the CONSTRAINT NAME rather +// than assuming. Only the session key is expected here (the PK is the ON CONFLICT +// arbiter, so it resolves instead of raising); a hit on it means this session id +// is ALREADY BOUND TO A DIFFERENT ACCOUNT. Refusing that keeps +// ResolveSessionAccount single-valued — two accounts sharing a live session id +// would make the relay's answer depend on which row Postgres returned, and it +// resolves the principal a comms call runs under. Any OTHER unique violation is +// unexpected, so it falls through to the generic wrap with its own message +// intact rather than being relabelled as a session collision. func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, accountID AccountID, runnerID string) (string, error) { if sessionID == "" { return "", fmt.Errorf("%w: session id is required", ErrInvalidArgument) @@ -87,20 +119,58 @@ func (s *Store) RecordSessionBinding(ctx context.Context, sessionID string, acco if runnerID == "" { return "", fmt.Errorf("%w: runner id is required", ErrInvalidArgument) } - displaced, err := s.q.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ + + // beginTenantTx, not s.pool.Begin: it arms SET LOCAL ROLE + the + // compass.tenant_id GUC at BEGIN, so every statement below is tenant-scoped + // by RLS exactly as the single-statement scopedDBTX path is. A raw Begin here + // would run as the owner with no GUC and silently disable tenant isolation. + tx, err := s.beginTenantTx(ctx) + if err != nil { + return "", fmt.Errorf("store: begin record session binding: %w", err) + } + // No-op after a successful commit; the rollback that matters is on every + // error path below, where there is nothing further to report about it. + defer func() { _ = tx.Rollback(ctx) }() + qtx := s.q.WithTx(tx) + + // FIRST, before reading anything: serialize every bind for this account. + // Keyed on (tenant, account) because two tenants may hold bindings for one + // account id and those binds are independent. Auto-released at tx end. + if err := qtx.LockSessionBindingAccount(ctx, db.LockSessionBindingAccountParams{ + Column1: pgtype.Text{String: string(s.resolveTenant(ctx)), Valid: true}, + Column2: pgtype.Text{String: string(accountID), Valid: true}, + }); err != nil { + return "", fmt.Errorf("store: lock session binding account: %w", err) + } + + displaced, err := qtx.SessionBindingForUpdate(ctx, string(accountID)) + if err != nil && !noRows(err) { + return "", fmt.Errorf("store: read prior session binding: %w", err) + } + if noRows(err) { + // No prior binding: nothing to displace. Scan left displaced at its zero + // value, which is already the empty session id this method reports for + // "the account held none". + displaced = "" + } + + if err := qtx.RecordSessionBinding(ctx, db.RecordSessionBindingParams{ SessionID: sessionID, AgentAccountID: string(accountID), RunnerID: runnerID, - }) - if err != nil { + }); err != nil { if pgErrIs(err, pgForeignKeyViolation) { return "", fmt.Errorf("%w: agent account %q does not exist", ErrInvalidArgument, accountID) } - if pgErrIs(err, pgUniqueViolation) { + if pgErrIs(err, pgUniqueViolation) && pgConstraintName(err) == "session_bindings_session_key" { return "", fmt.Errorf("%w: session %q is already bound to a different agent", ErrConflict, sessionID) } return "", fmt.Errorf("store: record session binding: %w", err) } + + if err := tx.Commit(ctx); err != nil { + return "", fmt.Errorf("store: commit record session binding: %w", err) + } return displaced, nil } @@ -129,8 +199,14 @@ func (s *Store) ResolveSessionAccount(ctx context.Context, sessionID string) (Ac // SessionForAccount resolves the live session bound to an agent account — the // REVERSE of ResolveSessionAccount, and the direction the delivery consumer // needs to dispatch a deliver to an already-resolved subscriber -// (runnerhub/relay_comms.go:179-184). Exactly one row can answer, because the -// account is the table's key. +// (runnerhub/relay_comms.go:179-184). Exactly one row can answer PER TENANT: the +// table's key is (tenant_id, agent_account_id), so the account alone is not +// unique and the query is single-valued only because RLS has already narrowed +// the visible rows to the acting tenant's. Under WithSystemRole (BYPASSRLS, no +// tenant GUC) that narrowing is gone, several tenants' rows can match, and pgx +// takes whichever comes first — so this method is a REQUEST-PATH read. Nothing +// calls it under the system role today; a PR3 caller that wants to must scope it +// itself. // // An account with no live session is ErrNotFound — never started, stopped, or // dropped on a Runner reconnect. Fail-closed for the same reason as above: an diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 21359381..5a09e31a 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -30,7 +30,8 @@ import ( // mustBind records a binding a test only needs to SUCCEED, returning the session // id it displaced. Most cases below care about a later assertion, not this call. -func mustBind(t *testing.T, s *Store, ctx context.Context, sessionID string, accountID AccountID, runnerID string) string { +// ctx is the SECOND parameter, after t, matching mustTopic (messages_test.go). +func mustBind(t *testing.T, ctx context.Context, s *Store, sessionID string, accountID AccountID, runnerID string) string { t.Helper() displaced, err := s.RecordSessionBinding(ctx, sessionID, accountID, runnerID) if err != nil { @@ -42,23 +43,34 @@ func mustBind(t *testing.T, s *Store, ctx context.Context, sessionID string, acc // bindingTimes reads a binding row's created_at/updated_at directly, so a test // asserts the persisted timestamps rather than trusting a return value — the // same posture as tenantOf. -func bindingTimes(t *testing.T, s *Store, sessionID string) (createdAt, updatedAt time.Time) { +// +// It goes through s.pool, which is the OWNER connection: no SET LOCAL ROLE, no +// GUC, so RLS does not apply and the read is cross-tenant. The tenant_id +// predicate is therefore explicit and NOT optional — without it a same-session-id +// row in another tenant would make the scan ambiguous, which is exactly the +// coexistence the tenant-fold cases below create. +func bindingTimes(t *testing.T, ctx context.Context, s *Store, sessionID string) (createdAt, updatedAt time.Time) { t.Helper() - if err := s.pool.QueryRow(context.Background(), - "SELECT created_at, updated_at FROM session_bindings WHERE session_id = $1", sessionID, + if err := s.pool.QueryRow(ctx, + "SELECT created_at, updated_at FROM session_bindings WHERE session_id = $1 AND tenant_id = $2", + sessionID, string(s.resolveTenant(ctx)), ).Scan(&createdAt, &updatedAt); err != nil { t.Fatalf("read timestamps of binding %q: %v", sessionID, err) } return createdAt, updatedAt } -// countBindings counts binding rows for an account, so a re-point test can prove -// the row was REPLACED rather than accumulated. -func countBindings(t *testing.T, s *Store, accountID AccountID) int { +// countBindings counts binding rows for an account IN ctx's tenant, so a +// re-point test can prove the row was REPLACED rather than accumulated. Same +// owner-connection caveat as bindingTimes: the tenant predicate is what keeps +// "exactly 1" a per-tenant count rather than a cross-tenant one, so a second +// tenant holding a binding for the SAME account id does not inflate it. +func countBindings(t *testing.T, ctx context.Context, s *Store, accountID AccountID) int { t.Helper() var n int - if err := s.pool.QueryRow(context.Background(), - "SELECT count(*) FROM session_bindings WHERE agent_account_id = $1", string(accountID), + if err := s.pool.QueryRow(ctx, + "SELECT count(*) FROM session_bindings WHERE agent_account_id = $1 AND tenant_id = $2", + string(accountID), string(s.resolveTenant(ctx)), ).Scan(&n); err != nil { t.Fatalf("count bindings of %q: %v", accountID, err) } @@ -148,7 +160,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - mustBind(t, s, ctx, "sess-old", agent.ID, "runner-1") + mustBind(t, ctx, s, "sess-old", agent.ID, "runner-1") displaced, err := s.RecordSessionBinding(ctx, "sess-new", agent.ID, "runner-2") if err != nil { @@ -169,7 +181,7 @@ func TestRecordSessionBindingRePointsAccountAndReportsDisplaced(t *testing.T) { // Exactly one row: a re-point REPLACES, it does not accumulate a second // binding whose sweep would retire a session that has already moved. - if n := countBindings(t, s, agent.ID); n != 1 { + if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("bindings for the agent = %d, want exactly 1 (a re-point must replace, not accumulate)", n) } @@ -207,13 +219,13 @@ func TestRecordSessionBindingRebindsSameSessionOntoANewRunner(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") - displaced := mustBind(t, s, ctx, "sess-1", agent.ID, "runner-2") + mustBind(t, ctx, s, "sess-1", agent.ID, "runner-1") + displaced := mustBind(t, ctx, s, "sess-1", agent.ID, "runner-2") if displaced != "sess-1" { t.Fatalf("re-binding the same session displaced %q, want %q — the account's previous session is this same session", displaced, "sess-1") } - if n := countBindings(t, s, agent.ID); n != 1 { + if n := countBindings(t, ctx, s, agent.ID); n != 1 { t.Fatalf("bindings for the agent = %d, want exactly 1 (the rebind must replace, not accumulate)", n) } @@ -241,7 +253,7 @@ func TestRecordSessionBindingRejectsASessionClaimedByAnotherAccount(t *testing.T agentA := mustAgent(t, s, owner.ID, "agent-a") agentB := mustAgent(t, s, owner.ID, "agent-b") - mustBind(t, s, ctx, "sess-1", agentA.ID, "runner-1") + mustBind(t, ctx, s, "sess-1", agentA.ID, "runner-1") displaced, err := s.RecordSessionBinding(ctx, "sess-1", agentB.ID, "runner-1") sentinelIs(t, err, ErrConflict, "a session id already bound to a different agent") @@ -285,7 +297,7 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") + mustBind(t, ctx, s, "sess-1", agent.ID, "runner-1") if err := s.DeleteSessionBinding(ctx, "sess-1"); err != nil { t.Fatalf("DeleteSessionBinding: %v", err) } @@ -299,7 +311,7 @@ func TestDeleteSessionBindingReleasesAndIsIdempotent(t *testing.T) { // The released account is bindable again, and displaces nothing — the row is // gone, so this is a fresh insert rather than a re-point. - if displaced := mustBind(t, s, ctx, "sess-2", agent.ID, "runner-1"); displaced != "" { + if displaced := mustBind(t, ctx, s, "sess-2", agent.ID, "runner-1"); displaced != "" { t.Fatalf("binding after a release displaced %q, want \"\" — the row was deleted, so there is nothing to reap", displaced) } @@ -345,7 +357,7 @@ func TestDeleteSessionBindingsForRunnerReturnsEverySweptBinding(t *testing.T) { {SessionID: "sess-a", AccountID: a.ID, RunnerID: "runner-1"}, {SessionID: "sess-elsewhere", AccountID: elsewhere.ID, RunnerID: "runner-2"}, } { - mustBind(t, s, ctx, bind.SessionID, bind.AccountID, bind.RunnerID) + mustBind(t, ctx, s, bind.SessionID, bind.AccountID, bind.RunnerID) } swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") @@ -412,14 +424,14 @@ func TestRecordSessionBindingTriggerAdvancesUpdatedAtOnly(t *testing.T) { owner := mustUser(t, s, "owner") agent := mustAgent(t, s, owner.ID, "agent") - mustBind(t, s, ctx, "sess-1", agent.ID, "runner-1") - createdBefore, updatedBefore := bindingTimes(t, s, "sess-1") + mustBind(t, ctx, s, "sess-1", agent.ID, "runner-1") + createdBefore, updatedBefore := bindingTimes(t, ctx, s, "sess-1") if !createdBefore.Equal(updatedBefore) { t.Fatalf("on INSERT created_at = %v and updated_at = %v, want the same DEFAULT now()", createdBefore, updatedBefore) } - mustBind(t, s, ctx, "sess-1", agent.ID, "runner-2") - createdAfter, updatedAfter := bindingTimes(t, s, "sess-1") + mustBind(t, ctx, s, "sess-1", agent.ID, "runner-2") + createdAfter, updatedAfter := bindingTimes(t, ctx, s, "sess-1") if !createdAfter.Equal(createdBefore) { t.Fatalf("created_at moved from %v to %v across a rebind; it must record the row's birth", createdBefore, createdAfter) @@ -448,7 +460,7 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { ownerA := mustUser(t, s, "owner-a") agentA := mustAgent(t, s, ownerA.ID, "agent-a") - mustBind(t, s, ctxA, "sess-a", agentA.ID, "runner-1") + mustBind(t, ctxA, s, "sess-a", agentA.ID, "runner-1") // Tenant B resolves A's session id: the row is not in B's view, so this must // fail closed rather than hand B tenant A's account. @@ -483,3 +495,462 @@ func TestSessionBindingIsTenantIsolated(t *testing.T) { t.Fatalf("tenant A's own binding resolves to %q, want %q", gotAccount, agentA.ID) } } + +// seedTenantAgent creates an owner user and an owned agent UNDER ctx's tenant, +// so a multi-tenant case can WRITE as a second tenant rather than only read as +// one. mustUser/mustAgent are hard-wired to context.Background() (the bootstrap +// tenant), which is precisely why every pre-existing multi-tenant assertion here +// could only ever write under tenant A. +func seedTenantAgent(t *testing.T, ctx context.Context, s *Store, handle string) Account { + t.Helper() + owner, err := s.CreateUser(ctx, NewUser{Handle: handle + "-owner", DisplayName: handle + "-owner"}) + if err != nil { + t.Fatalf("CreateUser(%s-owner): %v", handle, err) + } + agent, err := s.CreateAgent(ctx, owner.ID, NewAgent{Handle: handle, DisplayName: handle}) + if err != nil { + t.Fatalf("CreateAgent(%s): %v", handle, err) + } + return agent +} + +// TestSessionBindingSameSessionIDInTwoTenantsCoexist pins the SESSION half of the +// tenant fold: session_bindings_session_key is (tenant_id, session_id), not +// (session_id). Two tenants mint session ids independently — nothing coordinates +// them — so a collision between them is routine, and a global unique index would +// refuse tenant B's perfectly valid bind because tenant A happened to pick the +// same string first. That is a cross-tenant denial of service through an id +// namespace neither tenant can see. +// +// This is the case the round-1 fold had NO test for. TestSessionBindingIsTenantIsolated +// only ever WRITES under tenant A; tenant B appears solely in reads and a +// returns-nothing sweep, so reverting the index to a global UNIQUE (session_id) +// left every one of its assertions passing. This one WRITES under tenant B and +// fails on that mutation with the ErrConflict the un-folded index raises. +func TestSessionBindingSameSessionIDInTwoTenantsCoexist(t *testing.T) { + s := newTestStore(t) + tenantB := seedTenant(t, s, "tenant-b") + ctxA := context.Background() // no tenant set → the bootstrap tenant + ctxB := WithTenant(context.Background(), tenantB) + + ownerA := mustUser(t, s, "owner-a") + agentA := mustAgent(t, s, ownerA.ID, "agent-a") + agentB := seedTenantAgent(t, ctxB, s, "agent-b") + + mustBind(t, ctxA, s, "sess-shared", agentA.ID, "runner-a") + + // The load-bearing write: tenant B claims the SAME session id. It must + // SUCCEED — under a global unique index this is ErrConflict. + if displaced := mustBind(t, ctxB, s, "sess-shared", agentB.ID, "runner-b"); displaced != "" { + t.Fatalf("tenant B's first bind displaced %q, want \"\" — B's agent held no prior session, and a cross-tenant displaced id would mean B just overwrote A's row", displaced) + } + + // Each tenant resolves its OWN account from the shared session id. A single + // surviving row would make one of these two answer with the other tenant's + // account — the relay resolving a foreign principal. + gotA, err := s.ResolveSessionAccount(ctxA, "sess-shared") + if err != nil { + t.Fatalf("tenant A ResolveSessionAccount(sess-shared) after B's bind: %v (B's write clobbered A's binding)", err) + } + if gotA != agentA.ID { + t.Fatalf("tenant A resolves sess-shared to %q, want its own agent %q", gotA, agentA.ID) + } + gotB, err := s.ResolveSessionAccount(ctxB, "sess-shared") + if err != nil { + t.Fatalf("tenant B ResolveSessionAccount(sess-shared): %v", err) + } + if gotB != agentB.ID { + t.Fatalf("tenant B resolves sess-shared to %q, want its own agent %q", gotB, agentB.ID) + } + + // And the reverse direction, per tenant. + if got, err := s.SessionForAccount(ctxA, agentA.ID); err != nil || got != "sess-shared" { + t.Fatalf("tenant A SessionForAccount = (%q, %v), want (sess-shared, nil)", got, err) + } + if got, err := s.SessionForAccount(ctxB, agentB.ID); err != nil || got != "sess-shared" { + t.Fatalf("tenant B SessionForAccount = (%q, %v), want (sess-shared, nil)", got, err) + } +} + +// TestSessionBindingSameAccountIDInTwoTenantsCoexist pins the PRIMARY KEY half of +// the fold: PRIMARY KEY (tenant_id, agent_account_id), not (agent_account_id). +// Two rows for one account id, one per tenant, must COEXIST — under an un-folded +// PK the second bind is not a conflict but something worse: it resolves as +// ON CONFLICT DO UPDATE and OVERWRITES the other tenant's binding, silently +// moving a foreign tenant's account onto this tenant's session. +// +// accounts.id is a GLOBAL primary key, so the two tenants cannot own two distinct +// agent rows sharing an id; the second tenant necessarily binds the first +// tenant's account id. The FK to agent_accounts permits it — referential-integrity +// checks are not RLS-constrained — and that is the point: the DEFENCE against one +// tenant's write reaching another's binding is the tenant in the KEY, not the FK. +func TestSessionBindingSameAccountIDInTwoTenantsCoexist(t *testing.T) { + s := newTestStore(t) + tenantB := seedTenant(t, s, "tenant-b") + ctxA := context.Background() // no tenant set → the bootstrap tenant + ctxB := WithTenant(context.Background(), tenantB) + + ownerA := mustUser(t, s, "owner-a") + shared := mustAgent(t, s, ownerA.ID, "agent-shared") + + mustBind(t, ctxA, s, "sess-a", shared.ID, "runner-a") + + // The load-bearing write: tenant B binds the SAME account id to its own + // session. Under PRIMARY KEY (agent_account_id) this upserts over A's row. + if displaced := mustBind(t, ctxB, s, "sess-b", shared.ID, "runner-b"); displaced != "" { + t.Fatalf("tenant B's bind of the shared account id displaced %q, want \"\" — a non-empty value means it read (and overwrote) tenant A's row", displaced) + } + + // Both rows survive, one per tenant, each resolving its own session. + if got, err := s.SessionForAccount(ctxA, shared.ID); err != nil || got != "sess-a" { + t.Fatalf("tenant A SessionForAccount = (%q, %v), want (sess-a, nil) — B's write reached A's row", got, err) + } + if got, err := s.SessionForAccount(ctxB, shared.ID); err != nil || got != "sess-b" { + t.Fatalf("tenant B SessionForAccount = (%q, %v), want (sess-b, nil)", got, err) + } + + // Exactly one row PER TENANT — countBindings carries the tenant predicate, + // so this is 1 and 1, not a cross-tenant 2. + if n := countBindings(t, ctxA, s, shared.ID); n != 1 { + t.Fatalf("tenant A holds %d bindings for the shared account, want 1", n) + } + if n := countBindings(t, ctxB, s, shared.ID); n != 1 { + t.Fatalf("tenant B holds %d bindings for the shared account, want 1", n) + } + + // Each tenant's session resolves only in its own tenant. + if _, err := s.ResolveSessionAccount(ctxA, "sess-b"); !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant A resolved B's session err = %v, want ErrNotFound", err) + } + if _, err := s.ResolveSessionAccount(ctxB, "sess-a"); !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant B resolved A's session err = %v, want ErrNotFound", err) + } +} + +// TestSessionForAccountUnderSystemRoleIsUnscoped PINS A HAZARD rather than a +// desired property, and the assertions are written to say so. +// +// WithSystemRole arms the BYPASSRLS compass_system role with no tenant GUC +// (tenant_tx.go). SessionBindingForAccount is :one and its only predicate is the +// account id — single-valued ONLY because RLS normally narrows the visible rows +// to one tenant's. Strip that and two tenants' rows for one account id both +// match; pgx's QueryRow takes the first and discards the rest WITHOUT error, so +// the caller gets a plausible session id from an arbitrary tenant. +// +// Nothing calls this under the system role today — WithSystemRole is set at +// delivery/consumer.go:316 and runnerhub/hub.go:753,:814, i.e. on PR3's path — so +// this test exists so the behaviour is RECORDED, not discovered later. Deciding +// what to do about it (a tenant predicate, a :many + explicit refusal, a +// structural guard) is PR3's call, and this test is what will fail loudly when +// PR3 changes it. +func TestSessionForAccountUnderSystemRoleIsUnscoped(t *testing.T) { + s := newTestStore(t) + tenantB := seedTenant(t, s, "tenant-b") + ctxA := context.Background() // no tenant set → the bootstrap tenant + ctxB := WithTenant(context.Background(), tenantB) + + ownerA := mustUser(t, s, "owner-a") + shared := mustAgent(t, s, ownerA.ID, "agent-shared") + mustBind(t, ctxA, s, "sess-a", shared.ID, "runner-a") + mustBind(t, ctxB, s, "sess-b", shared.ID, "runner-b") + + // Under the system role BOTH rows are visible, so the :one read is ambiguous. + // It does NOT error — that is the hazard. It returns one of the two, and + // which one is not something the caller can control or detect. + got, err := s.SessionForAccount(WithSystemRole(context.Background()), shared.ID) + if err != nil { + t.Fatalf("SessionForAccount under the system role: %v — the current behaviour is a SILENT pick, not an error; if this now errors, PR3 changed the contract and this test must be updated deliberately", err) + } + if got != "sess-a" && got != "sess-b" { + t.Fatalf("SessionForAccount under the system role = %q, want one of the two tenants' sessions", got) + } + t.Logf("system-role SessionForAccount(%q) returned %q with two tenants holding that account id — unscoped and silently single-valued", shared.ID, got) + + // The same read on the REQUEST path is exact in both tenants: the hazard is + // the system role's missing scoping, NOT anything about the data. + if v, err := s.SessionForAccount(ctxA, shared.ID); err != nil || v != "sess-a" { + t.Fatalf("tenant A request-path SessionForAccount = (%q, %v), want (sess-a, nil)", v, err) + } + if v, err := s.SessionForAccount(ctxB, shared.ID); err != nil || v != "sess-b" { + t.Fatalf("tenant B request-path SessionForAccount = (%q, %v), want (sess-b, nil)", v, err) + } + + // A system-role WRITE is worse than a refusal: it SUCCEEDS and lands an + // ORPHAN. tenant_id DEFAULTs to current_setting('compass.tenant_id', TRUE), + // and under this role no GUC is armed — but the pooled connection has served + // armed request-path statements before, and a SET LOCAL that has ended leaves + // the custom GUC defined-and-EMPTY on that backend rather than undefined. So + // the DEFAULT resolves to '' instead of NULL, the NOT NULL is satisfied, and + // the row lands stamped with a tenant that does not exist. + // + // session_bindings.tenant_id has no FK to tenants (only accounts.tenant_id + // does), so nothing catches it. The row is then invisible to EVERY tenant — + // no RLS policy matches '' — and reachable only by another system-role read + // or the owner pool: an unswept binding no request path can see or release. + // + // It is also NOT deterministic. On a backend that never carried an armed + // statement the GUC is genuinely undefined, the DEFAULT is NULL, and the + // insert fails not-null instead. Which one a caller gets depends on the + // pooled connection it draws, so this asserts the DISJUNCTION rather than + // pretending either branch is the contract. Both are defects; PR3 owns the + // fix, and this test is what will fail when it lands one. + sysDisplaced, sysErr := s.RecordSessionBinding(WithSystemRole(context.Background()), "sess-sys", shared.ID, "runner-sys") + switch { + case sysErr != nil: + t.Logf("system-role RecordSessionBinding failed (connection had no prior armed statement, so the GUC was undefined and the DEFAULT was NULL): %v", sysErr) + default: + t.Logf("system-role RecordSessionBinding SUCCEEDED, displaced=%q — it landed an untenanted row", sysDisplaced) + var orphans int + if err := s.pool.QueryRow(ctxA, + "SELECT count(*) FROM session_bindings WHERE session_id = $1 AND tenant_id = ''", "sess-sys", + ).Scan(&orphans); err != nil { + t.Fatalf("count untenanted bindings: %v", err) + } + if orphans != 1 { + t.Fatalf("system-role write left %d rows stamped tenant_id = '', want 1 — the observed failure mode is an ORPHAN row, so if it is now stamped with a real tenant the write path grew scoping and this test must be updated deliberately", orphans) + } + // And it is invisible to every tenant: no policy matches ''. + if _, err := s.ResolveSessionAccount(ctxA, "sess-sys"); !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant A sees the untenanted binding err = %v, want ErrNotFound", err) + } + if _, err := s.ResolveSessionAccount(ctxB, "sess-sys"); !errors.Is(err, ErrNotFound) { + t.Fatalf("tenant B sees the untenanted binding err = %v, want ErrNotFound", err) + } + } +} + +// TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced is the +// regression test for the round-2 HIGH defect, and it drives the interleaving +// DETERMINISTICALLY rather than hoping a goroutine race lands on it. +// +// The setup: account X is bound to sess-A, then two callers re-point it — one to +// sess-B, one to sess-C. Exactly two sessions get displaced: sess-A by whichever +// caller wins, and the WINNER'S session by the loser. Each caller must be told +// the one IT displaced, because that id is what it reaps from the delivery +// held-deliver registry, and a session reported to nobody keeps its held +// deliveries forever. +// +// The determinism comes from a THIRD connection holding the SAME per-account +// advisory lock the bind takes first. Both callers park on it, so both are +// provably in flight before either can proceed, and releasing it starts the real +// contention. Without that the two calls would usually just serialize and the +// test would pass on a schedule that never exercised the bug. +// +// The defect this catches: with an unlocked prior-value read (the `prev` CTE this +// replaced), the second caller's read takes the statement-start snapshot while +// its write blocks on the first's row lock and then re-reads the latest committed +// row — so both callers report sess-A and the middle session is destroyed +// unreported. Removing either lock from the bind fails this test. +func TestRecordSessionBindingConcurrentRePointsReportDistinctDisplaced(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + mustBind(t, ctx, s, "sess-A", agent.ID, "runner-0") + + // A gate transaction holding the bind's per-account advisory lock, armed + // exactly as the store's own transactions are (SET LOCAL ROLE + the tenant + // GUC). The lock key is character-identical to the one in + // queries/session_bindings.sql — if that key changes, this gate stops + // blocking and the test fails loudly rather than silently going green. + gate, err := s.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin gate tx: %v", err) + } + defer func() { _ = gate.Rollback(ctx) }() // no-op once released below + if _, err := gate.Exec(ctx, "SET LOCAL ROLE "+appRole); err != nil { + t.Fatalf("gate arm role: %v", err) + } + if _, err := gate.Exec(ctx, "SELECT set_config($1, $2, true)", tenantGUC, string(s.resolveTenant(ctx))); err != nil { + t.Fatalf("gate arm guc: %v", err) + } + if _, err := gate.Exec(ctx, + "SELECT pg_advisory_xact_lock(hashtext('binding:' || $1 || ':' || $2))", + string(s.resolveTenant(ctx)), string(agent.ID), + ); err != nil { + t.Fatalf("gate advisory lock: %v", err) + } + + type outcome struct { + session string + displaced string + err error + } + results := make(chan outcome, 2) + launch := func(sessionID, runnerID string) { + go func() { + d, err := s.RecordSessionBinding(ctx, sessionID, agent.ID, runnerID) + results <- outcome{sessionID, d, err} + }() + } + launch("sess-B", "runner-1") + launch("sess-C", "runner-2") + + // Both callers must be BLOCKED on the gate's lock. If either completes now, + // the bind is not serializing per account and displaced values cannot be + // correct under concurrency. + select { + case r := <-results: + t.Fatalf("a re-point completed (session=%q displaced=%q err=%v) while the gate held the per-account advisory lock — the bind is not taking it, so two concurrent binds can read the same prior value", r.session, r.displaced, r.err) + case <-time.After(750 * time.Millisecond): + } + + if err := gate.Rollback(ctx); err != nil { + t.Fatalf("release gate: %v", err) + } + + got := make(map[string]string, 2) + for range 2 { + r := <-results + if r.err != nil { + t.Fatalf("concurrent re-point failed: %v (a re-point must never be refused)", r.err) + } + got[r.session] = r.displaced + } + + // Which caller wins the lock is not determined, so the assertion is stated + // on the SHAPE, not on a fixed pairing: the winner displaces sess-A (the + // seeded binding), and the loser displaces THE WINNER'S session, because + // that is what it actually overwrote. So {sess-A, }, one + // each. Both callers reporting the SAME id is the unlocked-read defect. + b, c := got["sess-B"], got["sess-C"] + var winner, loser string + switch { + case b == "sess-A" && c == "sess-B": + winner, loser = "sess-B", "sess-C" + case c == "sess-A" && b == "sess-C": + winner, loser = "sess-C", "sess-B" + default: + t.Fatalf("the two concurrent re-points reported displaced sess-B=%q sess-C=%q; want one of them naming sess-A and the OTHER naming the first one's session. Every displaced session must be reported to exactly one caller or its held deliveries are stranded forever; both callers naming the same id is the unlocked-read defect (the prior value must be read under a lock inside the write's transaction)", b, c) + } + t.Logf("caller %s won the lock (displaced sess-A); caller %s displaced %s", winner, loser, winner) + + // One row survives, holding the LOSER's session — it committed last. + if n := countBindings(t, ctx, s, agent.ID); n != 1 { + t.Fatalf("bindings after two concurrent re-points = %d, want 1", n) + } + live, err := s.SessionForAccount(ctx, agent.ID) + if err != nil { + t.Fatalf("SessionForAccount after the race: %v", err) + } + if live != loser { + t.Fatalf("live session is %q, want %q — the caller that committed last owns the binding", live, loser) + } + + // Both displaced sessions are gone, and between them they were reported to + // exactly the two callers: nothing was destroyed unreported. + for _, dead := range []string{"sess-A", winner} { + if _, err := s.ResolveSessionAccount(ctx, dead); !errors.Is(err, ErrNotFound) { + t.Fatalf("displaced session %q still resolves (err = %v), want ErrNotFound", dead, err) + } + } +} + +// TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession covers +// the case the ROW lock cannot cover, and it is why the bind takes a per-account +// advisory lock as well. +// +// SELECT ... FOR UPDATE on a row that does not yet exist locks NOTHING. So two +// first-ever binds for one account both read no prior value, and the row lock +// never brings them into contact. The PRIMARY KEY does serialize their WRITES — +// the second INSERT waits on the first's uncommitted tuple and resolves as +// ON CONFLICT DO UPDATE — but by then both have already READ, so with only those +// two mechanisms both callers report "displaced nothing" while the second has +// destroyed the first's live binding. That session is reported to NOBODY and its +// held deliveries are stranded: the same failure as the re-point case, reached +// by a different route. Measured, not assumed — the PK orders writes, and the +// displaced value is a read. +// +// The advisory lock is taken BEFORE the read, so it covers this: exactly one +// caller reports "" (it genuinely displaced nothing) and the other reports the +// first caller's session, which it really did overwrite and must reap. +// +// Gated the same way as the re-point case, on the same lock, so the interleaving +// is driven rather than hoped for. Removing the advisory lock from the bind makes +// this test fail with both callers reporting "". +func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *testing.T) { + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + // The gate: hold the bind's per-account advisory lock so both callers are + // provably in flight before either reads. No binding row exists yet, which + // is the entire point — there is nothing for FOR UPDATE to lock. + gate, err := s.pool.Begin(ctx) + if err != nil { + t.Fatalf("begin gate tx: %v", err) + } + defer func() { _ = gate.Rollback(ctx) }() // no-op once released below + if _, err := gate.Exec(ctx, "SET LOCAL ROLE "+appRole); err != nil { + t.Fatalf("gate arm role: %v", err) + } + if _, err := gate.Exec(ctx, "SELECT set_config($1, $2, true)", tenantGUC, string(s.resolveTenant(ctx))); err != nil { + t.Fatalf("gate arm guc: %v", err) + } + if _, err := gate.Exec(ctx, + "SELECT pg_advisory_xact_lock(hashtext('binding:' || $1 || ':' || $2))", + string(s.resolveTenant(ctx)), string(agent.ID), + ); err != nil { + t.Fatalf("gate advisory lock: %v", err) + } + + type outcome struct { + session string + displaced string + err error + } + results := make(chan outcome, 2) + for _, b := range []struct{ session, runner string }{ + {"sess-1", "runner-1"}, + {"sess-2", "runner-2"}, + } { + go func() { + d, err := s.RecordSessionBinding(ctx, b.session, agent.ID, b.runner) + results <- outcome{b.session, d, err} + }() + } + + select { + case r := <-results: + t.Fatalf("a first bind completed (session=%q displaced=%q err=%v) while the gate held the per-account advisory lock — with no row to lock, that lock is the ONLY thing serializing two first binds", r.session, r.displaced, r.err) + case <-time.After(750 * time.Millisecond): + } + + if err := gate.Rollback(ctx); err != nil { + t.Fatalf("release gate: %v", err) + } + + got := make(map[string]string, 2) + for range 2 { + r := <-results + if r.err != nil { + t.Fatalf("concurrent first bind of %q failed: %v — the PK conflict must resolve as an upsert, never surface to the caller", r.session, r.err) + } + got[r.session] = r.displaced + } + + // Which caller wins is not determined, so accept either assignment and + // reject the two broken shapes: both "" (a live session destroyed and + // reported to nobody) and both non-empty (one session reaped twice). + first, second := got["sess-1"], got["sess-2"] + var winner string + switch { + case first == "" && second == "sess-1": + winner = "sess-1" + case second == "" && first == "sess-2": + winner = "sess-2" + default: + t.Fatalf("concurrent first binds reported displaced sess-1=%q sess-2=%q; want exactly one \"\" and the other naming the session it overwrote. Both empty means a live session was destroyed and reported to NOBODY — the defect the per-account advisory lock exists to prevent, and one the PK cannot cover because it orders WRITES while the displaced value is a READ", first, second) + } + t.Logf("caller %s bound first; the other displaced it and can reap it", winner) + + if n := countBindings(t, ctx, s, agent.ID); n != 1 { + t.Fatalf("bindings after two concurrent first binds = %d, want 1 — the PK must fold them into one row", n) + } + // The winner's session is gone, and it WAS reported, so it can be reaped. + if _, err := s.ResolveSessionAccount(ctx, winner); !errors.Is(err, ErrNotFound) { + t.Fatalf("the overwritten session %q still resolves (err = %v), want ErrNotFound", winner, err) + } +} From 08ef6d3bbd6f2146c6d5e0cf478acb7701561288 Mon Sep 17 00:00:00 2001 From: mintaka Date: Mon, 7 Sep 2026 20:20:47 -0400 Subject: [PATCH 4/4] test(store): make the bind's row lock load-bearing (RIG-3108) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-3 review found that removing `FOR UPDATE` from `RecordSessionBinding`'s prior-value read leaves the entire committed suite green. I reproduced that myself before fixing it: the row lock was unprotected. **This corrects a claim in the previous commit message.** It said "Both locks are load-bearing: removing either fails a test." Half false. Removing the advisory lock does fail two tests, with the messages that commit quotes. Removing `FOR UPDATE` failed nothing. The reason no existing test could reach it: both concurrency tests pit a bind against another bind, and a bind takes the per-account advisory lock BEFORE its prior-value read, so the two callers are already serialized by the time either evaluates `SELECT ... FOR UPDATE`. The advisory lock alone is sufficient for bind-vs-bind, which makes the row lock dead weight in that interleaving. The row lock earns its place against a writer that reaches the row WITHOUT the advisory lock, and `DeleteSessionBindingsForRunner` is exactly that — the reconnect sweep is a bare `:many` DELETE and takes no advisory lock at all. So the new test races a re-point against a sweep of the account's current Runner. Both destroy the same binding, and the displaced session id drives the reap from the delivery held-deliver registry, so it must reach exactly one of them: reported to both is a double reap, reported to neither strands its held deliveries. Which side wins is deliberately not asserted — the assertion is on the shape, so it holds on every schedule. No gate transaction drives the interleaving, because the contended resource IS the lock under test and a gate holding it would beg the question. An unsynchronized loop suffices at this count: 120 iterations, measured 108/120 and 113/120 reporting the double-reap with `FOR UPDATE` removed and 0/120 with it present. A single iteration would be a ~90% test. Verified: my own red control removed `FOR UPDATE`, regenerated sqlc, and this test was the ONLY failure of the 16 — 113/120 double-reaps. Restored byte-identically (`queries/session_bindings.sql` md5 `62fb16b446847ee231c0f8a2dc2042e3`, generated `5304b52f8aca86e811ae57c746dab1c9`) and re-ran green. `sqlc-drift`, `go build ./...`, `gofmt`, and `golangci-lint run ./internal/store/...` clean. Two further count corrections to the previous message, both found by the same review: the folded-key citations `0001_init.sql:799`/`:845`/`:878` were correct at the parent commit but that commit's own edit to the file shifted them — they are `:818`/`:864`/`:897` there. And "20 pgtest cases green" was 19. Co-authored-by: Matt Wilkinson --- .../store/session_bindings_pgtest_test.go | 159 ++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/go/internal/store/session_bindings_pgtest_test.go b/go/internal/store/session_bindings_pgtest_test.go index 5a09e31a..031b044a 100644 --- a/go/internal/store/session_bindings_pgtest_test.go +++ b/go/internal/store/session_bindings_pgtest_test.go @@ -24,6 +24,8 @@ package store import ( "context" "errors" + "fmt" + "sync" "testing" "time" ) @@ -954,3 +956,160 @@ func TestRecordSessionBindingConcurrentFirstBindsReportTheDestroyedSession(t *te t.Fatalf("the overwritten session %q still resolves (err = %v), want ErrNotFound", winner, err) } } + +// TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep is what +// makes the FOR UPDATE row lock on the bind's prior-value read load-bearing, and +// it is the only case in this file that can be. +// +// The two concurrency tests above cannot reach that lock AT ALL. Both pit a bind +// against another BIND, and the bind takes the per-account advisory lock FIRST, +// before it reads — so the two callers are already serialized by the time either +// one evaluates SELECT ... FOR UPDATE. Remove FOR UPDATE and both still pass: +// the advisory lock alone is sufficient for bind-vs-bind. The row lock earns its +// place only against a writer that reaches the row WITHOUT holding the advisory +// lock, and DeleteSessionBindingsForRunner is exactly that writer — the reconnect +// sweep is a single :many DELETE and takes no advisory lock at all. +// +// So this pits a re-point against a sweep of the account's CURRENT Runner. Both +// destroy the same binding, and the displaced session id drives PR3's reap from +// the delivery held-deliver registry, so it must be handed to EXACTLY ONE of +// them: +// +// - reported to BOTH is a double reap — the bind returns sess-old AND the +// sweep returns sess-old, so PR3 reaps the same session twice. +// - reported to NEITHER strands it — its held deliveries are kept forever, +// the mirror of the defect the advisory lock covers. +// +// With the row lock, one side blocks on the other's uncommitted tuple and sees +// the committed truth: if the sweep commits first the row is gone, the bind's +// read misses, and it reports "" while the sweep reports sess-old; if the bind +// commits first it reports sess-old and the sweep — re-reading the latest +// committed row, which now names sess-new on runner-2 — matches nothing. Without +// it the bind's read runs on the statement-start snapshot and still sees +// sess-old while the sweep concurrently deletes and returns it, so both report. +// +// WHICH side wins is not asserted and must not be: the assertion is on the SHAPE +// (exactly one reporter), so it holds on every iteration regardless of schedule. +// No gate transaction drives the interleaving here — unlike the two tests above, +// the contended resource IS the row lock under test, so a gate holding it would +// beg the question. An unsynchronized loop is enough at this iteration count: +// measured 108/120 and 113/120 iterations reporting the double-reap with +// FOR UPDATE removed, and 0/120 with it present. A SINGLE iteration would be a +// ~90% test, which is why the count is 120 rather than one. +func TestRecordSessionBindingReportsDisplacedExactlyOnceAgainstARunnerSweep(t *testing.T) { + // context.Background is the test root, the pgtest-suite convention. + ctx := context.Background() + s := newTestStore(t) + owner := mustUser(t, s, "owner") + agent := mustAgent(t, s, owner.ID, "agent") + + const iterations = 120 + var doubleReported, unreported int + for i := range iterations { + old := fmt.Sprintf("sess-old-%d", i) + fresh := fmt.Sprintf("sess-new-%d", i) + + // Fresh session ids every iteration, and the seed bind RE-POINTS the + // account's one row back onto runner-1. That is what keeps each + // iteration independent AND non-vacuous: the previous iteration leaves a + // row on runner-2 (which the runner-1 sweep below would not see), so + // without this re-point iteration i+1 would race a bind against a sweep + // that had nothing to find and "exactly one reporter" would hold + // trivially. After this call there is exactly one binding, it names + // `old`, and it sits on the Runner the sweep targets. + mustBind(t, ctx, s, old, agent.ID, "runner-1") + + type bindResult struct { + displaced string + err error + } + type sweepResult struct { + swept []SessionBinding + err error + } + binds := make(chan bindResult, 1) + sweeps := make(chan sweepResult, 1) + + // Both goroutines park on `start` until each is scheduled and ready, so + // they are released together rather than sequentially. This narrows the + // window, it does not close it — the interleaving stays up to Postgres + // and the scheduler, which is why the loop runs many iterations instead + // of trusting one. + start := make(chan struct{}) + var ready sync.WaitGroup + ready.Add(2) + go func() { + ready.Done() + <-start + displaced, err := s.RecordSessionBinding(ctx, fresh, agent.ID, "runner-2") + binds <- bindResult{displaced, err} + }() + go func() { + ready.Done() + <-start + swept, err := s.DeleteSessionBindingsForRunner(ctx, "runner-1") + sweeps <- sweepResult{swept, err} + }() + ready.Wait() + close(start) + + bind, sweep := <-binds, <-sweeps + if bind.err != nil { + t.Fatalf("iteration %d: RecordSessionBinding(%q, runner-2): %v — a re-point must never be refused, whatever a concurrent sweep is doing", i, fresh, bind.err) + } + if sweep.err != nil { + t.Fatalf("iteration %d: DeleteSessionBindingsForRunner(runner-1): %v", i, sweep.err) + } + + // The bind may only ever name the session it actually displaced or "" + // (the sweep beat it to the row). Anything else is a wrong id handed to + // the reap, which is worse than either failure counted below. + if bind.displaced != old && bind.displaced != "" { + t.Fatalf("iteration %d: the bind reported displaced %q; the only correct answers are %q (it overwrote the seeded binding) or \"\" (the sweep removed it first)", i, bind.displaced, old) + } + // Likewise the sweep: the account holds one binding and it is the only + // row on runner-1, so the sweep returns that row or nothing. + if len(sweep.swept) > 1 { + t.Fatalf("iteration %d: the runner-1 sweep returned %d bindings, want at most 1 — the account holds a single binding", i, len(sweep.swept)) + } + for _, b := range sweep.swept { + if b.SessionID != old { + t.Fatalf("iteration %d: the runner-1 sweep returned session %q, want %q", i, b.SessionID, old) + } + } + + reportedByBind := bind.displaced == old + reportedBySweep := len(sweep.swept) == 1 + switch { + case reportedByBind && reportedBySweep: + doubleReported++ + case !reportedByBind && !reportedBySweep: + unreported++ + } + + // Whoever won, the account ends the iteration holding exactly the new + // binding. This is the per-iteration proof the race actually ran: a + // no-op iteration could not move the account onto `fresh`. + if n := countBindings(t, ctx, s, agent.ID); n != 1 { + t.Fatalf("iteration %d: bindings after the race = %d, want 1 — the bind must land whether or not the sweep removed the prior row", i, n) + } + live, err := s.SessionForAccount(ctx, agent.ID) + if err != nil { + t.Fatalf("iteration %d: SessionForAccount after the race: %v", i, err) + } + if live != fresh { + t.Fatalf("iteration %d: live session is %q, want %q — the bind commits last in both orderings, so its session owns the binding", i, live, fresh) + } + // And the displaced session is gone in every ordering, which is what + // makes "reported to nobody" a genuine strand rather than a deferral. + if _, err := s.ResolveSessionAccount(ctx, old); !errors.Is(err, ErrNotFound) { + t.Fatalf("iteration %d: displaced session %q still resolves (err = %v), want ErrNotFound", i, old, err) + } + } + + if doubleReported > 0 || unreported > 0 { + t.Fatalf("over %d iterations of a re-point racing a sweep of the account's Runner: %d reported the displaced session to BOTH callers (a double reap — PR3 reaps it from the held-deliver registry twice) and %d reported it to NEITHER (stranded held deliveries). Every displaced session must be reported to exactly one caller. The prior-value read must take a ROW LOCK (SELECT ... FOR UPDATE in queries/session_bindings.sql): the per-account advisory lock cannot cover this, because the sweep is a bare DELETE that never takes it", + iterations, doubleReported, unreported) + } + t.Logf("%d iterations: the displaced session was reported to exactly one caller every time", iterations) +}