fix(mothership): BYOK rides every worker leg — resume, copilot route, execute, title

Companion to worker batch A. The key attached only to '/api/mothership*'
sends: workflow-scoped copilot chats never pinned, and a resume landing on
a dead run became a hosted-key continuation. Resolution moves to a shared
resolveEnterpriseByokKey (entitlement-gated, revocation-fresh, fails to
hosted), applied per leg in the lifecycle loop, on child-chain resume
legs, and on title generation (which reads message content). Regenerated
protocol mirror carries the new optional fields.

Claude-Session: https://claude.ai/code/session_01CgaxNAaeD3taGdghbXn17w

Author: Sg312

apps/sim/lib/mothership/generated/protocol.ts

@@ -61,6 +61,13 @@ export interface ChatContextItem {
 export interface ResumeRequest {
   streamId: string;
   results: ResumeResult[];
+  /**
+   * Enterprise BYOK, re-resolved by sim per call (S27: context-only, zero retention).
+   * A LIVE run keeps its key inside the loop closure and ignores this; a DEAD run's
+   * continuation leg has no closure, so without it that leg would silently fall back
+   * to the hosted key mid-chat.
+   */
+  byokApiKey?: string | undefined;
 }
 
 export interface ResumeResult {
@@ -87,6 +94,8 @@ export interface SteerRequest {
 /** POST /api/generate-chat-title */
 export interface TitleRequest {
   message: string;
+  /** Enterprise BYOK: the title call reads user content, so it pins the same key (S27). */
+  byokApiKey?: string | undefined;
 }
 
 /** The 409 body for a duplicate send while a sibling instance streams (S32). */
@@ -123,6 +132,8 @@ export interface ExecuteRequest {
   integrationTools?: unknown[] | undefined;
   mothershipTools?: unknown[] | undefined;
   delegationToken?: string | undefined;
+  /** Enterprise BYOK: one-shot executions pin the customer key like chat turns (S27). */
+  byokApiKey?: string | undefined;
 }
 
 export interface ExecuteMessage {

apps/sim/lib/mothership/request/enterprise-byok.ts

@@ -0,0 +1,34 @@
+import { createLogger } from '@sim/logger'
+import { toError } from '@sim/utils/errors'
+import { getBYOKKey } from '@/lib/api-key/byok'
+import { isWorkspaceOnEnterprisePlan } from '@/lib/billing/core/subscription'
+
+const logger = createLogger('EnterpriseByok')
+
+/**
+ * Resolves the enterprise BYOK key sim-side for a mothership call (contract field
+ * `byokApiKey`, S27): the worker builds per-request provider instances from it and
+ * retains nothing. Eligibility (enterprise plan) gates resolution server-side, so a
+ * client can never assert its own eligibility; key rows are read fresh so revocation is
+ * immediate. Failures default to hosted.
+ *
+ * Every worker call that reaches a model must resolve this — the initial send, the
+ * workflow-scoped copilot send, one-shot executes, tool-resume (a dead-run continuation
+ * leg re-applies it), and title generation — or that leg silently runs on the hosted key.
+ */
+export async function resolveEnterpriseByokKey(
+  workspaceId: string | undefined
+): Promise<string | null> {
+  if (!workspaceId) return null
+  try {
+    if (!(await isWorkspaceOnEnterprisePlan(workspaceId))) return null
+    const byok = await getBYOKKey(workspaceId, 'anthropic')
+    return byok?.apiKey ?? null
+  } catch (error) {
+    logger.warn('Failed to resolve BYOK key; defaulting to hosted', {
+      workspaceId,
+      error: toError(error).message,
+    })
+    return null
+  }
+}

apps/sim/lib/mothership/request/lifecycle/run.test.ts

@@ -157,6 +157,13 @@ vi.mock('@/lib/mothership/request/tools/executor', () => ({
   pendingToolWaitBudgetMs: mockPendingToolWaitBudgetMs,
 }))
 
+const { mockResolveEnterpriseByokKey } = vi.hoisted(() => ({
+  mockResolveEnterpriseByokKey: vi.fn().mockResolvedValue(null),
+}))
+vi.mock('@/lib/mothership/request/enterprise-byok', () => ({
+  resolveEnterpriseByokKey: mockResolveEnterpriseByokKey,
+}))
+
 import {
   MothershipStreamV1CompletionStatus,
   MothershipStreamV1ToolOutcome,
@@ -428,6 +435,25 @@ describe('runCopilotLifecycle', () => {
     expect(sent).toEqual(payload)
   })
 
+  it('attaches the resolved enterprise BYOK key to the outbound payload', async () => {
+    mockResolveEnterpriseByokKey.mockResolvedValueOnce('sk-ant-enterprise-test')
+    const payload = { message: 'hi', workspaceId: 'ws-ent', messageId: 'stream-byok-attach' }
+    let capturedRequestBody = ''
+    mockRunStreamLoop.mockImplementationOnce(async (_url: string, request: RequestInit) => {
+      capturedRequestBody = String(request.body)
+    })
+
+    await runCopilotLifecycle(payload, {
+      userId: 'user-1',
+      workspaceId: 'ws-ent',
+      executionContext: { userId: 'user-1', workflowId: '', workspaceId: 'ws-ent' },
+      resolvedSecretTraceRegistry: new ResolvedSecretTraceRegistry([]),
+    })
+
+    const sent = JSON.parse(capturedRequestBody)
+    expect(sent.byokApiKey).toBe('sk-ant-enterprise-test')
+  })
+
   it('preserves large ordinary tool catalogs without scanning configured secret values', async () => {
     const registry = new ResolvedSecretTraceRegistry([
       { name: 'TOKEN', plaintext: 'catalog-secret', encryptedValue: 'ciphertext' },

apps/sim/lib/mothership/request/lifecycle/run.ts

@@ -5,15 +5,13 @@ import { toError } from '@sim/utils/errors'
 import { interruptibleSleep, sleep } from '@sim/utils/helpers'
 import { generateId } from '@sim/utils/id'
 import { omit } from '@sim/utils/object'
-import { getBYOKKey } from '@/lib/api-key/byok'
 import {
   type AttributedBillingRequestEnvelope,
   assertBillingAttributionSnapshot,
   type BillingAttributionSnapshot,
   checkAttributedUsageLimits,
   createAttributedBillingRequestEnvelope,
 } from '@/lib/billing/core/billing-attribution'
-import { isWorkspaceOnEnterprisePlan } from '@/lib/billing/core/subscription'
 import { env } from '@/lib/core/config/env'
 import { isCopilotToolPermissionsEnabled, isHosted } from '@/lib/core/config/env-flags'
 import type { AsyncCompletionSignal } from '@/lib/mothership/async-runs/lifecycle'
@@ -33,6 +31,7 @@ import { CopilotDegradedReason } from '@/lib/mothership/generated/trace-attribut
 import { getAutoAllowedTools } from '@/lib/mothership/persistence/tool-permission/auto-allow'
 import { createStreamingContext } from '@/lib/mothership/request/context/request-context'
 import { buildToolCallSummaries } from '@/lib/mothership/request/context/result'
+import { resolveEnterpriseByokKey } from '@/lib/mothership/request/enterprise-byok'
 import {
   BillingLimitError,
   CopilotBackendError,
@@ -733,11 +732,15 @@ async function driveOneChildChain(
         options.onAbortObserved?.(reason)
       },
     }
+    // Same per-leg BYOK rule as the main loop: this child-chain resume can also land on
+    // a dead run and become a hosted-key continuation without it.
+    const byokApiKey = await resolveEnterpriseByokKey(workspaceId)
     await runResumeLegWithRetry(
       `${baseURL}/api/tools/resume`,
       {
         streamId: context.messageId,
         results,
+        ...(byokApiKey ? { byokApiKey } : {}),
       },
       leg,
       execContext,
@@ -879,15 +882,15 @@ async function runCheckpointLoop(
     payload = { ...payload, workspaceId: lifecycleWorkspaceId }
   }
 
-  // Enterprise BYOK eligibility hint: set once on the initial mothership request
-  // so Go only attempts a BYOK lookup for entitled workspaces. This is only a
-  // gate — Go re-confirms entitlement authoritatively before using any key.
-  payload = await withEnterpriseByokKey(payload, route, lifecycleWorkspaceId)
-
   for (;;) {
     context.streamComplete = false
     const isResume = route === '/api/tools/resume'
 
+    // Enterprise BYOK rides EVERY leg, resume included: a resume that lands on a dead
+    // run becomes a continuation with no closure holding the key. Re-resolved per leg so
+    // revocation is immediate (key rows are read fresh; entitlement is cached).
+    payload = await withEnterpriseByokKey(payload, route, lifecycleWorkspaceId)
+
     if (isResume && isAborted(options, context)) {
       cancelPendingTools(context)
       context.awaitingAsyncContinuation = undefined
@@ -1397,30 +1400,26 @@ async function ensureHeadlessRunIdentity(input: {
 // Helpers
 
 /**
- * Resolves the enterprise BYOK key sim-side and attaches it as `byokApiKey`
- * (contract field, S27): the worker builds a per-run provider instance from it and
- * retains nothing. Eligibility (enterprise plan) gates resolution server-side, so a
- * client can never assert its own eligibility; key rows are read fresh so revocation
- * is immediate. Failures default to hosted. Mothership-only — other routes untouched.
+ * Routes whose payloads carry `byokApiKey` (see resolveEnterpriseByokKey): every
+ * model-reaching worker call, INCLUDING tool-resume — a resume that lands on a dead run
+ * becomes a continuation leg with no closure holding the key, so omitting it there
+ * silently finishes an enterprise chat on the hosted key.
  */
+const BYOK_ROUTES = [
+  '/api/mothership',
+  '/api/mothership/execute',
+  '/api/copilot',
+  '/api/tools/resume',
+]
+
 async function withEnterpriseByokKey(
   payload: Record<string, unknown>,
   route: string,
   workspaceId?: string
 ): Promise<Record<string, unknown>> {
-  if (!workspaceId || !route.startsWith('/api/mothership')) return payload
-  try {
-    if (!(await isWorkspaceOnEnterprisePlan(workspaceId))) return payload
-    const byok = await getBYOKKey(workspaceId, 'anthropic')
-    if (!byok) return payload
-    return { ...payload, byokApiKey: byok.apiKey }
-  } catch (error) {
-    logger.warn('Failed to resolve BYOK key; defaulting to hosted', {
-      workspaceId,
-      error: toError(error).message,
-    })
-    return payload
-  }
+  if (!BYOK_ROUTES.includes(route)) return payload
+  const byokApiKey = await resolveEnterpriseByokKey(workspaceId)
+  return byokApiKey ? { ...payload, byokApiKey } : payload
 }
 
 function isAborted(options: CopilotLifecycleOptions, context: StreamingContext): boolean {

apps/sim/lib/mothership/request/lifecycle/start.ts

@@ -28,6 +28,7 @@ import {
 } from '@/lib/mothership/generated/trace-attribute-values-v1'
 import { TraceAttr } from '@/lib/mothership/generated/trace-attributes-v1'
 import { TraceEvent } from '@/lib/mothership/generated/trace-events-v1'
+import { resolveEnterpriseByokKey } from '@/lib/mothership/request/enterprise-byok'
 import { mothershipRequestHeaders } from '@/lib/mothership/request/headers'
 import { finalizeStream } from '@/lib/mothership/request/lifecycle/finalize'
 import type { CopilotLifecycleOptions } from '@/lib/mothership/request/lifecycle/run'
@@ -541,11 +542,14 @@ export async function requestChatTitle(params: {
 
     const { fetchGo } = await import('@/lib/mothership/request/go/fetch')
     const mothershipBaseURL = await getMothershipBaseURL({ userId })
+    // Title reads the user's message content, so an enterprise chat pins its key here too.
+    const byokApiKey = await resolveEnterpriseByokKey(workspaceId)
     const response = await fetchGo(`${mothershipBaseURL}/api/generate-chat-title`, {
       method: 'POST',
       headers,
       body: JSON.stringify({
         message,
+        ...(byokApiKey ? { byokApiKey } : {}),
       }),
       otelContext,
       spanName: 'sim → go /api/generate-chat-title',