Add byok migration script for newly hosted keys#3462
Add byok migration script for newly hosted keys#3462TheodoreSpeaks wants to merge 10 commits intofeat/mothership-copilotfrom
Conversation
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
| console.log(`Found ${rows.length} candidate blocks\n`) | ||
|
|
||
| // 2. Pre-load env vars for resolving {{VAR}} references | ||
| const personalEnvRows = await db.select().from(environment) |
There was a problem hiding this comment.
We shouldn't load all env vars into memory. Can we first find the blocks that reference these env vars, then map the env rows needed from that?
PR SummaryMedium Risk Overview The script supports Written by Cursor Bugbot for commit efb23da. This will update automatically on new commits. Configure here. |
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
Greptile SummaryThis PR adds a one-shot migration script that reads API keys stored directly on workflow blocks and inserts them into the Key points:
Confidence Score: 3/5
Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[Parse CLI args\n--map blockType=providerId\n--dry-run] --> B[Query distinct workspaceIds\nwith matching block types]
B --> C{For each workspace}
C --> D[Query all candidate blocks\nfor this workspace]
D --> E[Extract raw API key refs\nper provider]
E --> F{Has env var\nreferences?}
F -- Yes --> G[Load workspace env vars\n+ personal env vars]
F -- No --> H
G --> H[For each provider:\nresolve all keys]
H --> I{Multiple distinct\nresolved keys?}
I -- Yes --> J[Log CONFLICT\nUse first resolved key]
I -- No --> K[Use single resolved key]
J --> L{DRY_RUN?}
K --> L
L -- Yes --> M[Print preview line\nNo DB write]
L -- No --> N[encryptSecret key\nInsert into workspace_byok_keys\nonConflictDoNothing]
N --> O{rowCount check\n⚠️ always undefined}
O -- undefined === 0 false --> P[Count as inserted\n even if skipped]
O -- never reached --> Q[Count as skipped]
M --> R[Print summary]
P --> R
Q --> R
Last reviewed commit: 353ef31 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| const blocks = await db | ||
| .select({ | ||
| blockId: workflowBlocks.id, | ||
| blockName: workflowBlocks.name, | ||
| blockType: workflowBlocks.type, | ||
| subBlocks: workflowBlocks.subBlocks, | ||
| workflowId: workflow.id, |
There was a problem hiding this comment.
Hardcoded apiKey sub-block ID may silently miss keys for some providers
The script unconditionally reads API keys via subBlocks?.apiKey?.value. If any of the target block types store their credential under a differently-named sub-block key (e.g., api_key, token, secret), those keys will be silently skipped with no warning — they won't appear in Skipped (empty key) stats either, because the value simply won't be found.
Before running a live migration it is worth verifying (or documenting) that every mapped block type in BLOCK_TYPE_TO_PROVIDER uses exactly apiKey as the sub-block identifier. A quick sanity-check log line listing how many non-empty apiKey values were found per block type would make this more auditable.
|
|
||
| // Detect conflicting values | ||
| const distinctKeys = new Set(resolved.map((r) => r.key)) | ||
| if (distinctKeys.size > 1) { | ||
| stats.conflicts++ | ||
| console.log(` [CONFLICT] provider "${providerId}": ${distinctKeys.size} distinct keys`) | ||
| for (const { ref, key } of resolved) { | ||
| const display = isEnvVarReference(ref.rawValue) | ||
| ? `${ref.rawValue} -> ${maskKey(key)}` | ||
| : maskKey(ref.rawValue) | ||
| console.log(` "${ref.blockName}" in "${ref.workflowName}": ${display}`) | ||
| } | ||
| console.log(' Using first resolved key') | ||
| } | ||
|
|
||
| // Use the first resolved key | ||
| const chosen = resolved[0] | ||
|
|
||
| if (DRY_RUN) { | ||
| console.log( | ||
| ` [DRY RUN] Would insert BYOK for provider "${providerId}": ${maskKey(chosen.key)}` | ||
| ) | ||
| continue | ||
| } |
There was a problem hiding this comment.
Conflicts are silently resolved in live run; no abort option
The PR description states: "The dry run flag will detect these cases and raise them so we can look on a case-by-case basis." However, in the live run the conflict is logged as [CONFLICT] but the migration immediately proceeds to insert the first resolved key without any pause or opt-in confirmation.
For a one-shot production migration, consider adding an --abort-on-conflict flag (or making the live run abort by default when conflicts are found), so operators are not surprised when their conflict review leads to no action. At minimum, a note in the script header documenting this behavior would be helpful.
|
|
||
| const stats = { | ||
| workspacesProcessed: 0, | ||
| workspacesSkipped: 0, | ||
| conflicts: 0, | ||
| inserted: 0, | ||
| skippedExisting: 0, | ||
| errors: 0, | ||
| envVarFailures: 0, | ||
| } | ||
|
|
||
| try { | ||
| // 1. Get distinct workspace IDs that have matching blocks | ||
| const mappedBlockTypes = Object.keys(BLOCK_TYPE_TO_PROVIDER) | ||
| const agentTypes = Object.keys(TOOL_INPUT_SUBBLOCK_IDS) | ||
| const allBlockTypes = [...new Set([...mappedBlockTypes, ...agentTypes])] | ||
|
|
There was a problem hiding this comment.
Initial workspace discovery query may include false-positive workspaces
allBlockTypes is the union of mappedBlockTypes (the providers we care about) and agentTypes (agent, human_in_the_loop). This means any workspace that has any agent or HITL block — even one with no tool configured — will be included in the outer loop and generate a per-workspace log line. For large deployments this is noisy and wastes DB round-trips.
Consider restricting the initial query to only mappedBlockTypes, and querying the agent/HITL blocks only inside the workspace loop (or alongside the mapped blocks in a single per-workspace query). That way a workspace that has agents but no matching tool API keys is excluded upfront rather than found and then immediately skipped as "No API keys found".
| if ((result as any).rowCount === 0) { | ||
| console.log(` [SKIP] BYOK already exists for provider "${providerId}"`) | ||
| stats.skippedExisting++ | ||
| } else { | ||
| console.log(` [INSERT] BYOK for provider "${providerId}": ${maskKey(chosen.key)}`) | ||
| stats.inserted++ |
There was a problem hiding this comment.
rowCount does not exist on postgres-js insert result
With drizzle-orm + postgres-js (the driver used in this project), calling db.insert(...).onConflictDoNothing() without .returning() returns a postgres.js command-result object whose relevant property is count (a BigInt), not rowCount. Because (result as any).rowCount will always be undefined, the expression undefined === 0 is always false — the skip branch is never entered. Every conflict-silenced insert will be incorrectly counted as stats.inserted++ instead of stats.skippedExisting++.
The safest fix is to add .returning({ id: workspaceBYOKKeys.id }) and check whether the returned array is empty — which is how the production route.ts already handles insert results. This makes the statistics trustworthy, which is especially important for a one-shot data migration.
| for (const row of personalRows) { | ||
| personalEnvCache.set(row.userId, (row.variables as Record<string, string>) || {}) | ||
| } | ||
| } |
There was a problem hiding this comment.
All env vars loaded into memory per workspace
Low Severity
The PR reviewer specifically flagged that env vars should not all be loaded into memory. While the code now loads per-workspace rather than globally, it still does select() (all columns) from workspaceEnvironment and environment, pulling the entire variables JSON blob for each workspace and user. This includes every env var, not just the ones referenced by blocks. The referenced env var names are known from the providerKeys refs — the script could first collect the needed var names, then only decrypt those specific entries rather than loading all variables into wsEnvVars and personalEnvCache.
| @@ -0,0 +1,598 @@ | |||
| #!/usr/bin/env bun | |||
There was a problem hiding this comment.
You'd have to include this in pure SQL in an actual migration at the top. We can discuss this on Monday morning. It's hard to tell our self-hosted users to run the script is what we realized.
There was a problem hiding this comment.
On second thought, this one only applies for hosted -- so actually a script is fine here
TheodoreSpeaks
force-pushed
the
feat/migrate-byok
branch
from
f6281f9 to
62219a6
Compare
TheodoreSpeaks
force-pushed
the
feat/migrate-byok
branch
from
9df482d to
11f402b
Compare
️✅ There are no secrets present in this pull request anymore.If these secrets were true positive and are still valid, we highly recommend you to revoke them. 🦉 GitGuardian detects secrets in your source code to help developers and security teams secure the modern development process. You are seeing this because you or someone else with access to this repository has authorized GitGuardian to scan your pull request. |
There was a problem hiding this comment.
Cursor Bugbot has reviewed your changes and found 1 potential issue.
Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.
| } | ||
|
|
||
| console.log(` [${index}/${total}] Done with workspace ${workspaceId}\n`) | ||
| return { stats, shouldWriteWorkspaceId: DRY_RUN } |
There was a problem hiding this comment.
Dry run writes workspace IDs even without insertable keys
Low Severity
shouldWriteWorkspaceId: DRY_RUN is unconditionally true for every processed workspace during a dry run, even when no keys actually resolved successfully (e.g., all env var lookups failed or all were skipped as non-owner personal keys). This causes the output file to include workspace IDs that have no valid keys to migrate. The live run then wastes time re-processing those workspaces with no result, and the dry-run summary inaccurately reports the count of "workspace IDs (with keys)."
2 participants
Summary
Currently, hosted keys will override users' currently provided api keys. Instead, we need to create a script to migrate their provided api keys to the BYOK section so that users continue to use their existing keys.
This handles:
Note: BYOK only allows for a single api key. This could cause different behavior if a user has multiple api keys provided for different blocks for the same service. The dry run flag will detect these cases and raise them so we can look on a case-by-case basis.
Example usage:
Type of Change
Testing
Tested with existing api key, migration moves it to byok uses key
Tested with existing env key, migration moves to byok and uses key
Tested with agent subblock, migration moves to byok and uses key
Tested with multiple same keys, migration moves to byok
Tested with different keys, dry run flags as potential issue.
Checklist
Screenshots/Videos