diff --git a/go/internal/store/db/accounts.sql.go b/go/internal/store/db/accounts.sql.go index 080f85096..f9bf73ad3 100644 --- a/go/internal/store/db/accounts.sql.go +++ b/go/internal/store/db/accounts.sql.go @@ -11,6 +11,19 @@ import ( "github.com/jackc/pgx/v5/pgtype" ) +const accountTenant = `-- name: AccountTenant :one +SELECT tenant_id FROM accounts WHERE id = $1 +` + +// The tenant an account belongs to. A system-role caller (the delivery loop) reads +// it to act for that account under its own tenant. +func (q *Queries) AccountTenant(ctx context.Context, id string) (string, error) { + row := q.db.QueryRow(ctx, accountTenant, id) + var tenant_id string + err := row.Scan(&tenant_id) + return tenant_id, err +} + const accountVisibleTo = `-- name: AccountVisibleTo :one SELECT EXISTS ( SELECT 1 diff --git a/go/internal/store/db/querier.go b/go/internal/store/db/querier.go index 0ed8a52e8..2674d7269 100644 --- a/go/internal/store/db/querier.go +++ b/go/internal/store/db/querier.go @@ -11,6 +11,9 @@ import ( ) type Querier interface { + // The tenant an account belongs to. A system-role caller (the delivery loop) reads + // it to act for that account under its own tenant. + AccountTenant(ctx context.Context, id string) (string, error) AccountVisibleTo(ctx context.Context, arg AccountVisibleToParams) (bool, error) AcquireOwnerTreeLock(ctx context.Context, hashtext string) error ActivityFor(ctx context.Context, dollar_1 []string) ([]ActivityForRow, error) diff --git a/go/internal/store/queries/accounts.sql b/go/internal/store/queries/accounts.sql index 7222a2767..185aefafc 100644 --- a/go/internal/store/queries/accounts.sql +++ b/go/internal/store/queries/accounts.sql @@ -209,3 +209,8 @@ SELECT EXISTS ( ) AND a.id = $2 ); + +-- The tenant an account belongs to. A system-role caller (the delivery loop) reads +-- it to act for that account under its own tenant. +-- name: AccountTenant :one +SELECT tenant_id FROM accounts WHERE id = $1; diff --git a/go/internal/store/tenant.go b/go/internal/store/tenant.go index 09c520cf8..20923e80a 100644 --- a/go/internal/store/tenant.go +++ b/go/internal/store/tenant.go @@ -70,3 +70,16 @@ func (s *Store) resolveTenant(ctx context.Context) TenantID { func (s *Store) EffectiveTenant(ctx context.Context) TenantID { return s.resolveTenant(ctx) } + +// AccountTenant returns the tenant an account belongs to. ErrNotFound when the +// account is not visible to ctx (call it under the system role to see all). +func (s *Store) AccountTenant(ctx context.Context, account AccountID) (TenantID, error) { + t, err := s.q.AccountTenant(ctx, string(account)) + if err != nil { + if noRows(err) { + return "", fmt.Errorf("%w: account %q", ErrNotFound, account) + } + return "", fmt.Errorf("store: account tenant: %w", err) + } + return TenantID(t), nil +} diff --git a/go/internal/store/tenant_tx.go b/go/internal/store/tenant_tx.go index 3bf135da6..5888d4603 100644 --- a/go/internal/store/tenant_tx.go +++ b/go/internal/store/tenant_tx.go @@ -49,6 +49,12 @@ func WithSystemRole(ctx context.Context) context.Context { return context.WithValue(ctx, systemRoleKey{}, true) } +// WithoutSystemRole clears the system-role mark, so store calls on the returned +// ctx are tenant-scoped again. Pair it with WithTenant. +func WithoutSystemRole(ctx context.Context) context.Context { + return context.WithValue(ctx, systemRoleKey{}, false) +} + // isSystemRole reports whether ctx is the cross-tenant system path. func isSystemRole(ctx context.Context) bool { v, _ := ctx.Value(systemRoleKey{}).(bool) diff --git a/go/server/lifecycle.go b/go/server/lifecycle.go index f0e3d821f..78c243b6f 100644 --- a/go/server/lifecycle.go +++ b/go/server/lifecycle.go @@ -78,6 +78,10 @@ var _ runnerhub.LifecycleCaller = (*lifecycleService)(nil) // caller and no RequireAgentSessionSubscriber); no prior session → a fresh // hub.Start; no placement → a logged no-op. func (l *lifecycleService) WakeAgent(ctx context.Context, agent store.AccountID) { + ctx, ok := l.wakeCtx(ctx, agent) + if !ok { + return + } // 1. Not-live pre-check (cost control): a live agent is already awake, so // there is nothing to resume. No-op, no log line — a wake is only an attempt // against an OFFLINE agent. @@ -628,3 +632,18 @@ func (l *lifecycleService) freshStart(ctx context.Context, agent store.AccountID } return wakeOutcomeFreshStarted } + +// wakeCtx scopes a wake to the agent's tenant. The delivery loop wakes under the +// system role, which sets no tenant, so rows the wake writes (agent_sessions, +// bindings) would land with an empty tenant_id, invisible to every tenant. +func (l *lifecycleService) wakeCtx(ctx context.Context, agent store.AccountID) (context.Context, bool) { + if !store.IsSystemRole(ctx) { + return ctx, true + } + tenant, err := l.store.AccountTenant(ctx, agent) + if err != nil { + slog.ErrorContext(ctx, "agent wake: resolving tenant failed", "agent_account_id", agent, "error", err) + return ctx, false + } + return store.WithTenant(store.WithoutSystemRole(ctx), tenant), true +} diff --git a/go/server/offline_mention_e2e_pgtest_test.go b/go/server/offline_mention_e2e_pgtest_test.go index 11f3a9b51..f6bfa5b24 100644 --- a/go/server/offline_mention_e2e_pgtest_test.go +++ b/go/server/offline_mention_e2e_pgtest_test.go @@ -498,12 +498,10 @@ func TestMentionBeforeRunnerEnrollsWakesOnAttach(t *testing.T) { // The wake fires during attach, before the harness clears its probe bookkeeping, // so assert the durable session the fresh start records, not the command log. - // Read under the system role, as the wake path does: a wake from the delivery - // loop records its session without a tenant. w.runner = attachFakeRunner(t, w.store, w.hub, false) deadline = timeAfter() for { - if _, ok, err := w.store.LatestSessionForAccount(store.WithSystemRole(w.ctx), agent.ID); err != nil { + if _, ok, err := w.store.LatestSessionForAccount(w.ctx, agent.ID); err != nil { t.Fatalf("LatestSessionForAccount: %v", err) } else if ok { return