feat(schema): per-org agent id PK — close two-orgs-same-id footgun#750
Merged
Conversation
Phase C of the per-org PK migration: swap `agents_pkey` from a global `(id)` to a composite `(organization_id, id)` and widen every per-(agent, …) UNIQUE on the FK children with `organization_id`. The 6 single-column FKs into `agents(id)` are rebuilt as composite FKs into `agents(organization_id, id)`. NOT NULL is set on the org column on the 5 child tables that Phase A backfilled. Phase B (storage refactor) plumbs `organization_id` through every INSERT/UPDATE/DELETE/SELECT touching these tables: - postgres-stores.ts: `saveMetadata` UPSERTs on the composite key, the cross-org rejection branch is gone; connection / channel-binding / grant / user-agent reads + writes scope by `tryGetOrgId()` (HTTP request paths) or accept an explicit `organizationId` (worker-spawn paths). `getAgentOrganizationId(agentId)` deleted — ambiguous post-PK- swap; chat-instance-manager reads the org from `connection.organizationId` instead. - grant-store.ts, user-agents-store.ts, channels/binding-service.ts: optional `organizationId` parameter on every method, falls back to ALS. - preview/slack.ts: `agent_channel_bindings` INSERTs include `organization_id` from the claim payload / preview-connection org. - base-provider-module.ts: `readOrgSharedProviderKey` keys directly off `agent_secrets.(organization_id, name)` instead of joining through `agents WHERE id = agentId` (which returns multiple rows post-swap). `ProviderCredentialContext` gains `organizationId`, plumbed from `MessagePayload.organizationId` at worker spawn. - StoredConnection / PlatformConnection / MessagePayload carry `organizationId`; the row mapper reads `agent_connections.organization_id`. The embedded-schema-patches mirror is idempotent — it detects the composite PK via `pg_get_constraintdef` and skips when already applied. The migration's down path will fail if two orgs end up sharing an agent id (by design — that's the whole point). Test: `agents-per-org-pk.test.ts` creates two agents with id `shared-id` in two different orgs, verifies they coexist, and proves grants + user-agent associations are independent per-org.
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (4)
📝 WalkthroughWalkthroughThreads organizationId through deployment grant syncs and interaction-triggered grants, updates GrantStore call sites to accept org-scoped arguments, and adjusts tests to pass the organization id. ChangesAgents Per-Organization Grant & Deployment Scoping
Sequence DiagramsequenceDiagram
participant Client
participant BaseDeploymentManager
participant GrantStore
participant DB
Client->>BaseDeploymentManager: deployment message (includes organizationId)
BaseDeploymentManager->>GrantStore: grant(pattern, expiresAt, denied?, organizationId)
GrantStore->>DB: INSERT/SELECT on grants WHERE organization_id = organizationId
DB-->>GrantStore: org-scoped rows/result
GrantStore-->>BaseDeploymentManager: confirmation
BaseDeploymentManager->>Client: response
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested labels
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
ESLint skipped: no ESLint configuration detected in root package.json. To enable, add Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. Comment |
|
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 288bcf31a5
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+135
to
+138
| WHERE agent_id = ${agentId} | ||
| AND kind = ${kind} | ||
| AND pattern = ANY(${pgTextArray(candidates)}::text[]) | ||
| AND (expires_at IS NULL OR expires_at > now()) |
There was a problem hiding this comment.
When two organizations now share the same agentId, worker egress checks can hit this no-org fallback because the HTTP proxy calls proxyGrantStore.hasGrant(agentId, hostname) with only the worker-token agentId and no org. In that path this query reads grants from every org with the same agent id, so org B can inherit org A's allow/deny grant for domains or MCP tools; require/pass an organization id instead of falling back to agent_id alone.
Useful? React with 👍 / 👎.
Comment on lines
+64
to
+67
| if (!orgId) { | ||
| throw new Error( | ||
| "GrantStore.grant requires organizationId (explicit or via orgContext)" | ||
| ); |
There was a problem hiding this comment.
For worker spawns that include networkConfig.allowedDomains, preApprovedTools, or Nix packages, BaseDeploymentManager.storeDeploymentConfigs() / syncNetworkConfigGrants() call grantStore.grant(agentId, ...) outside the request orgContext; the org is only present on MessagePayload. This new throw therefore bubbles out of generateEnvironmentVariables() and prevents those agents from starting, so the deployment path needs to pass messageData.organizationId (and do the same for revoke).
Useful? React with 👍 / 👎.
Comment on lines
425
to
+427
| async saveConnection(connection) { | ||
| const sql = getDb(); | ||
| const orgId = getOrgId(); |
There was a problem hiding this comment.
Connection updates can run outside an HTTP org context: during boot, ChatInstanceManager.initialize() loads stored rows and, if an active connection fails to start, calls connectionStore.updateConnection() to mark it errored. updateConnection() has already loaded connection.organizationId, but saveConnection() ignores it and calls getOrgId(), so a bad stored connection can throw Organization context not available and skip/abort startup error handling; use connection.organizationId ?? getOrgId() here.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
packages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.ts (1)
1-177: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick winMove this test to the repository-standard test directory.
This test is placed in
packages/server/src/lobu/stores/__tests__/, but repo policy requires package tests underpackages/*/src/__tests__. Please relocate it (and update relative imports) to keep test discovery/layout consistent.As per coding guidelines,
packages/*/src/**/*.{ts,tsx}: "TypeScript sources must live inpackages/*/src, tests must live inpackages/*/src/__tests__".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.ts` around lines 1 - 177, Move the test file (the one with the describe title "agents per-org PK — same agent id in two orgs") into the package-standard tests folder (package src/__tests__) so it follows repo test layout; after moving, update all relative imports referenced in the file — specifically adjust imports for cleanupTestDatabase, getTestDb, createTestOrganization, orgContext, createPostgresAgentConfigStore, and createPostgresAgentAccessStore — so they point to the correct locations from the new __tests__ location and ensure the test name and exported symbols remain unchanged.packages/server/src/preview/slack.ts (1)
520-529:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winMissing
organizationIdargument inupsertBindingcall causes compilation failure.The
upsertBindingfunction now requires 6 arguments includingorganizationId, but line 528 only passes 5. This matches the pipeline error TS2554.The fix should retrieve
organization_idfrom the existing ownership check query and pass it toupsertBinding.🐛 Proposed fix
- const owned = await sql<{ one: number }>` - SELECT 1 AS one + const owned = await sql<{ organization_id: string }>` + SELECT a.organization_id FROM agents a JOIN "member" m ON m."organizationId" = a.organization_id WHERE a.id = ${agentId} AND m."userId" = ${lobuUserId} LIMIT 1 `; if (owned.length === 0) return { status: 'forbidden' }; - await sql.begin((tx) => upsertBinding(tx, platform, channelId, teamId, agentId)); + const organizationId = owned[0]!.organization_id; + await sql.begin((tx) => upsertBinding(tx, platform, channelId, teamId, agentId, organizationId)); return { status: 'bound' };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/preview/slack.ts` around lines 520 - 529, The ownership query that populates `owned` does not select the agent's organization_id, but `upsertBinding(tx, platform, channelId, teamId, agentId)` now requires an `organizationId` argument; change the SELECT in the `owned` query to also return a.id's organization (e.g., SELECT a.organization_id AS organization_id), extract that value from the `owned` result (e.g., owned[0].organization_id) and pass it as the sixth argument to `upsertBinding` inside the `sql.begin` call (keeping references to `upsertBinding`, `sql.begin`, `platform`, `channelId`, `teamId`, `agentId`, and `lobuUserId`).
🧹 Nitpick comments (2)
packages/server/src/gateway/connections/message-handler-bridge.ts (1)
334-348: ⚡ Quick winSilent skip of agent-user association when
organizationIdis missing.When
connection.organizationIdis undefined, the agent-user association is not recorded and no log message indicates this occurred. While this behavior may be correct (sinceagent_userslikely requiresorganization_idper the composite PK), the silent skip makes it difficult to detect when connections are operating without org context.Consider adding a debug or warning log when the association is skipped:
📋 Proposed logging addition
const userAgentsStore = this.services.getUserAgentsStore(); if (userAgentsStore && this.connection.organizationId) { try { await userAgentsStore.addAgent( platform, userId, agentId, this.connection.organizationId ); } catch (error) { logger.warn( { agentId, userId, error: String(error) }, "Failed to record agent_users association" ); } +} else if (userAgentsStore && !this.connection.organizationId) { + logger.debug( + { agentId, userId, connectionId: this.connection.id }, + "Skipped agent_users association (connection has no organizationId)" + ); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/gateway/connections/message-handler-bridge.ts` around lines 334 - 348, The code silently skips recording the agent-user association when this.connection.organizationId is undefined; update the block around userAgentsStore.addAgent in message-handler-bridge.ts to log that the association was skipped when this.connection.organizationId is missing (use logger.debug or logger.warn with context { agentId, userId, organizationId: this.connection.organizationId }) and keep the existing try/catch for addAgent unchanged so failures are still logged by the existing logger.warn; reference the userAgentsStore variable, the addAgent call, and this.connection.organizationId to locate where to add the conditional log.packages/server/src/gateway/channels/binding-service.ts (1)
48-76: ⚡ Quick winConsider refactoring nested ternary conditionals for readability.
The nested ternary logic correctly handles the four combinations of
teamIdandorgId, but deeply nested conditionals are difficult to maintain and verify. Consider refactoring to a more explicit structure:♻️ Proposed refactor using explicit conditionals
async getBinding( platform: string, channelId: string, teamId?: string, organizationId?: string ): Promise<ChannelBinding | null> { const sql = getDb(); const orgId = organizationId ?? tryGetOrgId(); - const rows = teamId - ? orgId - ? await sql` - SELECT * FROM agent_channel_bindings - WHERE organization_id = ${orgId} - AND platform = ${platform} - AND channel_id = ${channelId} - AND team_id = ${teamId} - ` - : await sql` - SELECT * FROM agent_channel_bindings - WHERE platform = ${platform} - AND channel_id = ${channelId} - AND team_id = ${teamId} - ` - : orgId - ? await sql` - SELECT * FROM agent_channel_bindings - WHERE organization_id = ${orgId} - AND platform = ${platform} - AND channel_id = ${channelId} - AND team_id IS NULL - ` - : await sql` - SELECT * FROM agent_channel_bindings - WHERE platform = ${platform} - AND channel_id = ${channelId} - AND team_id IS NULL - `; + let rows; + if (orgId && teamId) { + rows = await sql` + SELECT * FROM agent_channel_bindings + WHERE organization_id = ${orgId} + AND platform = ${platform} + AND channel_id = ${channelId} + AND team_id = ${teamId} + `; + } else if (orgId) { + rows = await sql` + SELECT * FROM agent_channel_bindings + WHERE organization_id = ${orgId} + AND platform = ${platform} + AND channel_id = ${channelId} + AND team_id IS NULL + `; + } else if (teamId) { + rows = await sql` + SELECT * FROM agent_channel_bindings + WHERE platform = ${platform} + AND channel_id = ${channelId} + AND team_id = ${teamId} + `; + } else { + rows = await sql` + SELECT * FROM agent_channel_bindings + WHERE platform = ${platform} + AND channel_id = ${channelId} + AND team_id IS NULL + `; + } if (rows.length === 0) return null; return rowToBinding(rows[0]); }🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/gateway/channels/binding-service.ts` around lines 48 - 76, The nested ternary that assigns rows is hard to read; replace it with an explicit conditional structure (if/else or switch) or build the SQL dynamically by composing WHERE clauses so the four cases are clear: check teamId and orgId separately and call sql`SELECT * FROM agent_channel_bindings WHERE platform = ${platform} AND channel_id = ${channelId}` then add `AND team_id = ${teamId}` or `AND team_id IS NULL` and optionally `AND organization_id = ${orgId}` as needed; modify the assignment to the rows variable and keep references to the same identifiers (rows, sql, teamId, orgId, platform, channelId, agent_channel_bindings).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@db/migrations/20260516120000_agents_per_org_pk_swap.sql`:
- Line 16: Remove the explicit transaction wrappers in the PK-swap migration by
deleting the standalone "BEGIN;" and matching "COMMIT;" statements (they appear
around the primary-key swap steps in the
20260516120000_agents_per_org_pk_swap.sql migration), leaving the SQL to run
within the runner-managed transaction; do the same for the other occurrences
noted (lines containing "BEGIN;" / "COMMIT;" at the other blocks) so the
migration doesn't trigger UNSAFE_TRANSACTION.
In `@packages/server/src/db/embedded-schema-patches.ts`:
- Around line 571-699: The scheduled-jobs embedded patch must be updated so it
doesn't try to add the old single-column scheduled_jobs_agent_fkey against
public.agents(id) when the agents PK has been swapped to (organization_id,id);
in the scheduled-jobs patch (the branch that adds/drops
scheduled_jobs_agent_fkey) either remove that single-column FK branch entirely
or guard it by first checking for the new composite FK/constraint (e.g. look for
scheduled_jobs_org_agent_fkey or verify agents has PRIMARY KEY on
(organization_id,id) / agents_pkey) and only run the legacy single-column logic
when that composite FK/PK is not present. Ensure you reference and gate on the
constraint names scheduled_jobs_agent_fkey and scheduled_jobs_org_agent_fkey
(and agents_pkey) so the embedded boot won’t attempt to add the incompatible
single-column FK before the Phase C PK swap runs.
- Around line 578-592: The current completion check that inspects pkDef/def for
the composite PK on the agents table can falsely succeed if the PK swap became
visible but subsequent steps (adding FK/child constraints on agent_grants,
agent_users, grants) failed; update the check in the block using
sql.unsafe/pkDef/def to also verify at least one of the new child constraints
exists (e.g., look for the expected constraint names or FK definitions on
agent_grants, agent_users, grants) before returning, or alternatively make each
step (the PK swap and each ADD CONSTRAINT on agent_grants/agent_users/grants)
idempotent so the function can safely re-run without short-circuiting on
agents_pkey alone.
In `@packages/server/src/gateway/infrastructure/queue/queue-producer.ts`:
- Around line 33-36: The queued payload type in queue-producer.ts makes
organizationId optional—change the payload/interface declaration so
organizationId is required (remove the optional marker on organizationId) so all
enqueued jobs include org context; update any call sites that construct the
payload (e.g., where enqueue/produce functions are called) to pass
organizationId, and introduce a separate test-only helper/type (e.g., a
TestQueuedPayload or a factory) to allow omission in tests rather than keeping
the production type lax. Reference the organizationId property and the
payload/interface used by the producer to locate and update the code.
---
Outside diff comments:
In `@packages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.ts`:
- Around line 1-177: Move the test file (the one with the describe title "agents
per-org PK — same agent id in two orgs") into the package-standard tests folder
(package src/__tests__) so it follows repo test layout; after moving, update all
relative imports referenced in the file — specifically adjust imports for
cleanupTestDatabase, getTestDb, createTestOrganization, orgContext,
createPostgresAgentConfigStore, and createPostgresAgentAccessStore — so they
point to the correct locations from the new __tests__ location and ensure the
test name and exported symbols remain unchanged.
In `@packages/server/src/preview/slack.ts`:
- Around line 520-529: The ownership query that populates `owned` does not
select the agent's organization_id, but `upsertBinding(tx, platform, channelId,
teamId, agentId)` now requires an `organizationId` argument; change the SELECT
in the `owned` query to also return a.id's organization (e.g., SELECT
a.organization_id AS organization_id), extract that value from the `owned`
result (e.g., owned[0].organization_id) and pass it as the sixth argument to
`upsertBinding` inside the `sql.begin` call (keeping references to
`upsertBinding`, `sql.begin`, `platform`, `channelId`, `teamId`, `agentId`, and
`lobuUserId`).
---
Nitpick comments:
In `@packages/server/src/gateway/channels/binding-service.ts`:
- Around line 48-76: The nested ternary that assigns rows is hard to read;
replace it with an explicit conditional structure (if/else or switch) or build
the SQL dynamically by composing WHERE clauses so the four cases are clear:
check teamId and orgId separately and call sql`SELECT * FROM
agent_channel_bindings WHERE platform = ${platform} AND channel_id =
${channelId}` then add `AND team_id = ${teamId}` or `AND team_id IS NULL` and
optionally `AND organization_id = ${orgId}` as needed; modify the assignment to
the rows variable and keep references to the same identifiers (rows, sql,
teamId, orgId, platform, channelId, agent_channel_bindings).
In `@packages/server/src/gateway/connections/message-handler-bridge.ts`:
- Around line 334-348: The code silently skips recording the agent-user
association when this.connection.organizationId is undefined; update the block
around userAgentsStore.addAgent in message-handler-bridge.ts to log that the
association was skipped when this.connection.organizationId is missing (use
logger.debug or logger.warn with context { agentId, userId, organizationId:
this.connection.organizationId }) and keep the existing try/catch for addAgent
unchanged so failures are still logged by the existing logger.warn; reference
the userAgentsStore variable, the addAgent call, and
this.connection.organizationId to locate where to add the conditional log.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e564870c-8d0c-4eb2-901b-d2cdd171f225
📒 Files selected for processing (20)
db/migrations/20260516120000_agents_per_org_pk_swap.sqldb/schema.sqlpackages/core/src/agent-store.tspackages/server/src/db/embedded-schema-patches.tspackages/server/src/gateway/auth/base-provider-module.tspackages/server/src/gateway/auth/user-agents-store.tspackages/server/src/gateway/channels/binding-service.tspackages/server/src/gateway/connections/chat-instance-manager.tspackages/server/src/gateway/connections/message-handler-bridge.tspackages/server/src/gateway/connections/types.tspackages/server/src/gateway/embedded.tspackages/server/src/gateway/infrastructure/queue/queue-producer.tspackages/server/src/gateway/orchestration/base-deployment-manager.tspackages/server/src/gateway/permissions/grant-store.tspackages/server/src/gateway/services/platform-helpers.tspackages/server/src/lobu/agent-routes.tspackages/server/src/lobu/client-routes.tspackages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.tspackages/server/src/lobu/stores/postgres-stores.tspackages/server/src/preview/slack.ts
💤 Files with no reviewable changes (1)
- packages/server/src/lobu/agent-routes.ts
Comment on lines
+571
to
+699
| { | ||
| // Mirrors db/migrations/20260516120000_agents_per_org_pk_swap.sql. | ||
| // Detects whether the swap has already happened by reading the current | ||
| // PK definition on `agents`; skips silently when the composite PK is | ||
| // already in place. | ||
| id: 'agents-per-org-pk-phase-c', | ||
| apply: async (sql) => { | ||
| const pkDef = (await sql.unsafe(` | ||
| SELECT pg_get_constraintdef(c.oid) AS def | ||
| FROM pg_constraint c | ||
| JOIN pg_class t ON t.oid = c.conrelid | ||
| JOIN pg_namespace n ON n.oid = t.relnamespace | ||
| WHERE n.nspname = 'public' | ||
| AND t.relname = 'agents' | ||
| AND c.contype = 'p' | ||
| LIMIT 1 | ||
| `)) as Array<{ def: string }>; | ||
| const def = pkDef[0]?.def ?? ''; | ||
| if (def.includes('organization_id') && def.includes('id')) { | ||
| // Composite PK already in place — nothing to do. | ||
| return; | ||
| } | ||
|
|
||
| // Backfill any stragglers and drop orphans. | ||
| for (const t of [ | ||
| 'agent_grants', | ||
| 'agent_connections', | ||
| 'agent_users', | ||
| 'agent_channel_bindings', | ||
| 'grants', | ||
| ]) { | ||
| await sql.unsafe(` | ||
| UPDATE public.${t} c | ||
| SET organization_id = a.organization_id | ||
| FROM public.agents a | ||
| WHERE c.organization_id IS NULL AND c.agent_id = a.id | ||
| `); | ||
| await sql.unsafe(` | ||
| DELETE FROM public.${t} WHERE organization_id IS NULL | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.${t} ALTER COLUMN organization_id SET NOT NULL | ||
| `); | ||
| } | ||
|
|
||
| // Drop legacy single-column FKs. | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_grants DROP CONSTRAINT IF EXISTS agent_grants_agent_id_fkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_connections DROP CONSTRAINT IF EXISTS agent_connections_agent_id_fkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_users DROP CONSTRAINT IF EXISTS agent_users_agent_id_fkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_channel_bindings DROP CONSTRAINT IF EXISTS agent_channel_bindings_agent_id_fkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.grants DROP CONSTRAINT IF EXISTS grants_agent_id_fkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.scheduled_jobs DROP CONSTRAINT IF EXISTS scheduled_jobs_agent_fkey` | ||
| ); | ||
|
|
||
| // Drop legacy uniques/PKs scoped to bare agent_id. | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_grants DROP CONSTRAINT IF EXISTS agent_grants_agent_id_pattern_key` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agent_users DROP CONSTRAINT IF EXISTS agent_users_pkey` | ||
| ); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.grants DROP CONSTRAINT IF EXISTS grants_pkey` | ||
| ); | ||
|
|
||
| // Swap PK on agents. | ||
| await sql.unsafe(`ALTER TABLE public.agents DROP CONSTRAINT IF EXISTS agents_pkey`); | ||
| await sql.unsafe( | ||
| `ALTER TABLE public.agents ADD CONSTRAINT agents_pkey PRIMARY KEY (organization_id, id)` | ||
| ); | ||
|
|
||
| // Re-add per-org-scoped uniques. | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_grants | ||
| ADD CONSTRAINT agent_grants_org_agent_pattern_key UNIQUE (organization_id, agent_id, pattern) | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_users | ||
| ADD CONSTRAINT agent_users_pkey PRIMARY KEY (organization_id, agent_id, platform, user_id) | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.grants | ||
| ADD CONSTRAINT grants_pkey PRIMARY KEY (organization_id, agent_id, kind, pattern) | ||
| `); | ||
|
|
||
| // Re-add composite FKs into agents(organization_id, id). | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_grants | ||
| ADD CONSTRAINT agent_grants_org_agent_fkey | ||
| FOREIGN KEY (organization_id, agent_id) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_connections | ||
| ADD CONSTRAINT agent_connections_org_agent_fkey | ||
| FOREIGN KEY (organization_id, agent_id) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_users | ||
| ADD CONSTRAINT agent_users_org_agent_fkey | ||
| FOREIGN KEY (organization_id, agent_id) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.agent_channel_bindings | ||
| ADD CONSTRAINT agent_channel_bindings_org_agent_fkey | ||
| FOREIGN KEY (organization_id, agent_id) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.grants | ||
| ADD CONSTRAINT grants_org_agent_fkey | ||
| FOREIGN KEY (organization_id, agent_id) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| await sql.unsafe(` | ||
| ALTER TABLE public.scheduled_jobs | ||
| ADD CONSTRAINT scheduled_jobs_org_agent_fkey | ||
| FOREIGN KEY (organization_id, created_by_agent) REFERENCES public.agents(organization_id, id) ON DELETE CASCADE | ||
| `); | ||
| }, | ||
| }, |
There was a problem hiding this comment.
Update the legacy scheduled-jobs embedded patch before landing this PK swap.
This change makes public.agents(id) non-unique, but the earlier scheduled-jobs patch in this same file still replays an scheduled_jobs_agent_fkey against public.agents(id) whenever that legacy constraint is absent. On embedded boots, that branch can now fail before Phase C runs. Gate the old patch on the new composite FK or remove the legacy single-column FK branch entirely.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/server/src/db/embedded-schema-patches.ts` around lines 571 - 699,
The scheduled-jobs embedded patch must be updated so it doesn't try to add the
old single-column scheduled_jobs_agent_fkey against public.agents(id) when the
agents PK has been swapped to (organization_id,id); in the scheduled-jobs patch
(the branch that adds/drops scheduled_jobs_agent_fkey) either remove that
single-column FK branch entirely or guard it by first checking for the new
composite FK/constraint (e.g. look for scheduled_jobs_org_agent_fkey or verify
agents has PRIMARY KEY on (organization_id,id) / agents_pkey) and only run the
legacy single-column logic when that composite FK/PK is not present. Ensure you
reference and gate on the constraint names scheduled_jobs_agent_fkey and
scheduled_jobs_org_agent_fkey (and agents_pkey) so the embedded boot won’t
attempt to add the incompatible single-column FK before the Phase C PK swap
runs.
Comment on lines
+578
to
+592
| const pkDef = (await sql.unsafe(` | ||
| SELECT pg_get_constraintdef(c.oid) AS def | ||
| FROM pg_constraint c | ||
| JOIN pg_class t ON t.oid = c.conrelid | ||
| JOIN pg_namespace n ON n.oid = t.relnamespace | ||
| WHERE n.nspname = 'public' | ||
| AND t.relname = 'agents' | ||
| AND c.contype = 'p' | ||
| LIMIT 1 | ||
| `)) as Array<{ def: string }>; | ||
| const def = pkDef[0]?.def ?? ''; | ||
| if (def.includes('organization_id') && def.includes('id')) { | ||
| // Composite PK already in place — nothing to do. | ||
| return; | ||
| } |
There was a problem hiding this comment.
Don't use agents_pkey alone as the completion check.
This returns as soon as the PK swap is visible, even if a later ADD CONSTRAINT fails. A partial run would then no-op on the next boot and leave agent_grants/agent_users/grants/FKs missing. Check for at least one of the new child constraints as well, or make each step independently idempotent before returning.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/server/src/db/embedded-schema-patches.ts` around lines 578 - 592,
The current completion check that inspects pkDef/def for the composite PK on the
agents table can falsely succeed if the PK swap became visible but subsequent
steps (adding FK/child constraints on agent_grants, agent_users, grants) failed;
update the check in the block using sql.unsafe/pkDef/def to also verify at least
one of the new child constraints exists (e.g., look for the expected constraint
names or FK definitions on agent_grants, agent_users, grants) before returning,
or alternatively make each step (the PK swap and each ADD CONSTRAINT on
agent_grants/agent_users/grants) idempotent so the function can safely re-run
without short-circuiting on agents_pkey alone.
Comment on lines
+33
to
+36
| // Organization id of the agent. Plumbed through so child queries (grants, | ||
| // user-agents, channel-bindings, secrets) can scope by org — agent ids | ||
| // are per-org-unique, so `agent_id = ?` alone is ambiguous. | ||
| organizationId?: string; |
There was a problem hiding this comment.
Make organizationId required on queued payloads.
Once agentId is only unique within an organization, async consumers cannot safely recover org context from the queue job if this field is omitted. Keeping it optional means any producer that misses the field can enqueue an ambiguous job. Make it required here, and keep any test-only laxness behind a separate helper/type instead.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@packages/server/src/gateway/infrastructure/queue/queue-producer.ts` around
lines 33 - 36, The queued payload type in queue-producer.ts makes organizationId
optional—change the payload/interface declaration so organizationId is required
(remove the optional marker on organizationId) so all enqueued jobs include org
context; update any call sites that construct the payload (e.g., where
enqueue/produce functions are called) to pass organizationId, and introduce a
separate test-only helper/type (e.g., a TestQueuedPayload or a factory) to allow
omission in tests rather than keeping the production type lax. Reference the
organizationId property and the payload/interface used by the producer to locate
and update the code.
…ization_id, id)
- dbmate auto-wraps each migration in a transaction; the inner BEGIN/COMMIT
caused 'UNSAFE_TRANSACTION' / 'unexpected transaction status idle' on the
postgres-driver-backed test runners (bun:test, vitest) and the dbmate up CI
job. Removed the explicit transaction blocks from both migrate:up and
migrate:down — dbmate's wrapping transaction is sufficient.
- agent-routes.ts: `POST /api/<org>/agents` was using `ON CONFLICT (id)`,
which no longer exists post-swap (the PK is now composite). Widened to
`ON CONFLICT (organization_id, id)`. The cross-org rejection branch is
gone — collisions are no longer possible at the schema level.
- agent-routes.ts: provider-api-key endpoint reads orgId from the request
context (`c.get('organizationId')`) instead of the deleted
`getAgentOrganizationId` helper.
- agent-routes.ts: `GET /api/<org>/agents/:agentId` connection-count query
now scopes by organization_id.
- preview/slack.ts: `bindChatToAgentForOwner` resolves the agent's
organization_id from the membership join and passes it through to
upsertBinding (the helper now requires it for the INSERT).
The agent_connections.organization_id column is now NOT NULL post-Phase-C; the slack-preview test was seeding a row with the legacy 9-column INSERT that omitted it.
There was a problem hiding this comment.
🧹 Nitpick comments (1)
packages/server/src/preview/slack.ts (1)
427-442: 💤 Low valueConsider extracting to
upsertBindingto avoid duplication.This INSERT/ON CONFLICT logic duplicates
upsertBinding(). For theteamId-null branch (lines 435-442), the DELETE-then-INSERT is not transactional here, whereasupsertBindingexpects a transaction handle. Wrapping insql.begin()and callingupsertBindingwould consolidate the logic and ensure atomicity.♻️ Proposed refactor
const { platform, teamId, channelId } = args; - if (teamId) { - await sql` - INSERT INTO agent_channel_bindings (organization_id, agent_id, platform, channel_id, team_id, created_at) - VALUES (${org.organizationId}, ${target.id}, ${platform}, ${channelId}, ${teamId}, now()) - ON CONFLICT (platform, channel_id, team_id) DO UPDATE SET - agent_id = EXCLUDED.agent_id, - organization_id = EXCLUDED.organization_id - `; - } else { - await sql` - DELETE FROM agent_channel_bindings - WHERE platform = ${platform} AND channel_id = ${channelId} AND team_id IS NULL - `; - await sql` - INSERT INTO agent_channel_bindings (organization_id, agent_id, platform, channel_id, team_id, created_at) - VALUES (${org.organizationId}, ${target.id}, ${platform}, ${channelId}, NULL, now()) - `; - } + await sql.begin((tx) => + upsertBinding(tx, platform, channelId, teamId, target.id, org.organizationId) + ); return { status: 'bound', agentId: target.id };🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/preview/slack.ts` around lines 427 - 442, The DELETE-then-INSERT branch duplicates the INSERT/ON CONFLICT logic already implemented in upsertBinding and is not atomic; refactor the teamId-null branch to run inside a transaction (sql.begin()) and call upsertBinding(transaction, { organizationId: org.organizationId, agentId: target.id, platform, channelId, teamId: null }) instead of separate DELETE and INSERT, so the upsert is consolidated and executed atomically against the agent_channel_bindings table.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@packages/server/src/preview/slack.ts`:
- Around line 427-442: The DELETE-then-INSERT branch duplicates the INSERT/ON
CONFLICT logic already implemented in upsertBinding and is not atomic; refactor
the teamId-null branch to run inside a transaction (sql.begin()) and call
upsertBinding(transaction, { organizationId: org.organizationId, agentId:
target.id, platform, channelId, teamId: null }) instead of separate DELETE and
INSERT, so the upsert is consolidated and executed atomically against the
agent_channel_bindings table.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 729f48e1-1977-4dcf-9fba-3f055be173f9
📒 Files selected for processing (3)
db/migrations/20260516120000_agents_per_org_pk_swap.sqlpackages/server/src/lobu/agent-routes.tspackages/server/src/preview/slack.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- db/migrations/20260516120000_agents_per_org_pk_swap.sql
…_connections INSERT
After the per-org PK swap, grant-store / user-agents-store / channel-binding-service writes require an organization_id (explicit or via AsyncLocalStorage). Update the tests + a few callers that lacked org context:
- grant-store.test.ts: every test body now runs inside orgContext.run({organizationId: 'test-org'}) — matches the seedAgentRow helper's default org.
- base-deployment-grants.test.ts: buildPayload includes organizationId in the MessagePayload, and base-deployment-manager.syncNetworkConfigGrants threads messageData.organizationId through to grantStore.grant/.revoke.
- mcp-proxy-edge-cases.test.ts: explicit organizationId passed to grantStore.grant calls.
- agent-routes.ts (PUT /platforms/by-stable-id): the agent_connections INSERT now writes organization_id from the request context.
- agent-routes-apply.test.ts: 'cross-org collision still returns 409' rewritten to 'cross-org create succeeds' — the whole point of the PK swap is that two orgs can independently own an agent with the same id. Also fixed the seedAgent ON CONFLICT to use the composite key.
- helpers/db-setup.ts seedAgentRow: ON CONFLICT uses (organization_id, id).
…oyment-manager + interaction-bridge storeDeploymentConfigs (network domains, pre-approved tools, Nix cache, npm registry domains) and registerActionHandlers (slack/telegram tool-approval click handler) now thread organization_id from messageData / connection through to grantStore.grant. Also passes orgId on the mcp-proxy.test.ts 'allows with grant' test.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
packages/server/src/lobu/agent-routes.ts (1)
379-395:⚠️ Potential issue | 🔴 Critical | ⚡ Quick winCritical: Cross-org data leakage in agent statistics queries.
After the composite PK migration, agent IDs can be duplicated across organizations. The JOIN conditions in these queries (
JOIN agents a ON a.id = c.agent_id) match onidalone, creating ambiguous joins when multiple orgs have agents with the same ID. TheWHERE a.organization_id = ${orgId}clause filters the agent side but not the connection side, causing connections from other orgs to be counted for the current org's agent.Example:
- Org A has agent "foo" with 1 connection
- Org B has agent "foo" with 2 connections
- Query for org A joins all 3 connections to both agents, filters to org A's agent, but counts all 3 connections → incorrect count of 3 instead of 1
The same issue affects:
connCounts(lines 379-385)activeConnCounts(lines 388-395)userCounts(lines 419-425)platformRows(lines 432-438)Line 573 correctly filters both
agent_idANDorganization_iddirectly onagent_connections— apply the same pattern here.🔒 Proposed fix to prevent cross-org data leakage
const connCounts = await sql` SELECT c.agent_id, count(*)::int as count FROM agent_connections c JOIN agents a ON a.id = c.agent_id - WHERE a.organization_id = ${orgId} + WHERE a.organization_id = ${orgId} AND c.organization_id = ${orgId} GROUP BY c.agent_id `; const countMap = new Map(connCounts.map((r: any) => [r.agent_id, r.count])); const activeConnCounts = await sql` SELECT c.agent_id, count(*)::int as count FROM agent_connections c JOIN agents a ON a.id = c.agent_id - WHERE a.organization_id = ${orgId} + WHERE a.organization_id = ${orgId} AND c.organization_id = ${orgId} AND c.status = 'active' GROUP BY c.agent_id `;Apply similar fixes to
userCounts(addAND u.organization_id = ${orgId}) andplatformRows(addAND c.organization_id = ${orgId}).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/server/src/lobu/agent-routes.ts` around lines 379 - 395, The SELECTs building connCounts and activeConnCounts (and similarly userCounts and platformRows) join agent_connections c to agents a using only a.id = c.agent_id which allows cross-org matches; update those queries to restrict connections to the current org by adding c.organization_id = ${orgId} (or include a.organization_id = ${orgId} in the JOIN condition) so the WHERE/JOIN filters both sides (apply the same change to the userCounts query by adding AND u.organization_id = ${orgId} and to platformRows by adding AND c.organization_id = ${orgId}); modify the SQL around the connCounts, activeConnCounts, userCounts, and platformRows symbols accordingly to prevent cross-org leakage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@packages/server/src/lobu/agent-routes.ts`:
- Around line 379-395: The SELECTs building connCounts and activeConnCounts (and
similarly userCounts and platformRows) join agent_connections c to agents a
using only a.id = c.agent_id which allows cross-org matches; update those
queries to restrict connections to the current org by adding c.organization_id =
${orgId} (or include a.organization_id = ${orgId} in the JOIN condition) so the
WHERE/JOIN filters both sides (apply the same change to the userCounts query by
adding AND u.organization_id = ${orgId} and to platformRows by adding AND
c.organization_id = ${orgId}); modify the SQL around the connCounts,
activeConnCounts, userCounts, and platformRows symbols accordingly to prevent
cross-org leakage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 0d0d808c-f2c7-4765-ae58-c2498f2d81c4
📒 Files selected for processing (7)
packages/server/src/gateway/__tests__/base-deployment-grants.test.tspackages/server/src/gateway/__tests__/grant-store.test.tspackages/server/src/gateway/__tests__/helpers/db-setup.tspackages/server/src/gateway/__tests__/mcp-proxy-edge-cases.test.tspackages/server/src/gateway/orchestration/base-deployment-manager.tspackages/server/src/lobu/__tests__/agent-routes-apply.test.tspackages/server/src/lobu/agent-routes.ts
✅ Files skipped from review due to trivial changes (1)
- packages/server/src/gateway/tests/mcp-proxy-edge-cases.test.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- packages/server/src/gateway/orchestration/base-deployment-manager.ts
This was referenced May 18, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Closes a multi-month-old footgun where two organizations cannot share an agent ID — a stale
food-orderingin one org silently blocked another org from using the same name.This is Phase B + C of the per-org PK migration (Phase A landed in #734, the parallel-UNIQUE rollback in #747):
Schema (Phase C)
20260516120000_agents_per_org_pk_swap.sql:agents_pkeyswapped from(id)to(organization_id, id).agents(id)rebuilt as composite FKs intoagents(organization_id, id)— coversagent_grants,agent_connections,agent_users,agent_channel_bindings,grants, andscheduled_jobs.created_by_agent.agent_users_pkey→(organization_id, agent_id, platform, user_id).agent_grants_org_agent_pattern_keyUNIQUE →(organization_id, agent_id, pattern).grants_pkey→(organization_id, agent_id, kind, pattern).organization_idset NOT NULL on the 5 child tables Phase A backfilled.embedded-schema-patches.ts(agents-per-org-pk-phase-c) — detects viapg_get_constraintdefand skips when already applied.db/schema.sqlregenerated + normalized;schema_migrations.versionpatched back tovarchar(128).Storage (Phase B)
postgres-stores.ts:saveMetadataUPSERTs on the composite key, the cross-org rejection branch is gone. All connection / channel-binding / grant / user-agent reads + writes scope bytryGetOrgId()(HTTP request paths) or accept an explicitorganizationId(worker-spawn paths).getAgentOrganizationId(agentId)deleted — ambiguous post-PK-swap. Sole caller (chat-instance-manager.startInstance) now reads the org fromconnection.organizationIdinstead.grant-store.ts,user-agents-store.ts,channels/binding-service.ts: optionalorganizationIdparameter on every method, falls back to ALS viatryGetOrgId(). Writes require an org id (explicit or from ALS); reads degrade gracefully when none is available (legacy callers).preview/slack.ts:agent_channel_bindingsINSERTs includeorganization_idfrom the claim payload / preview-connection org.base-provider-module.ts:readOrgSharedProviderKeynow keys directly offagent_secrets.(organization_id, name)instead of joining throughagents WHERE id = agentId(which returns multiple rows post-swap).ProviderCredentialContextgainsorganizationId, plumbed fromMessagePayload.organizationIdat worker spawn.StoredConnection/PlatformConnection/MessagePayloadcarryorganizationId; the row mapper readsagent_connections.organization_id.Test
packages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.tscreates two agents with idshared-idin two different orgs and verifies:getMetadatareturns the right one per org.listAgentsis org-scoped — neither org sees the other's row.agent_grantsandagent_usersfor the same(agent_id, …)triple in two orgs don't collide.hasGrant/ownsAgentscope to the active org context.Test plan
bun run typecheckcleanbun run format:checkcleanbun test packages/server/src/lobu/stores/__tests__/agents-per-org-pk.test.ts(new test) — 2/2 passbun test packages/server/src/lobu/stores/__tests__/postgres-agent-config-store.test.ts(existing) — 3/3 pass, no regressionbun test packages/server/src/gateway/permissions src/gateway/auth src/gateway/channels— 44/46 pass, 2 pre-existing encryption-key-rotation failures unrelated to this PRbun test packages/cli src/commands/_lib/apply— 115/115 passbun test packages/core— 706/706 passmake build-packages— server.bundle.mjs + start-local.bundle.mjs build cleandbmate upagainst fresh PG; verified composite PK + 6 composite FKs landed; manualINSERT shared-idinto two orgs succeeds🤖 Generated with Claude Code
Summary by CodeRabbit
Features
Refactor
Chores
Tests