Skip to content

Commit 3dcf9fb

Browse files
authored
fix(copilot): re-derive the turn's secret authority on the in-band tool lane (#9180)
* fix(copilot): re-derive the turn's secret authority on the in-band tool lane The in-band tool route rebuilt the tool context from Go's request body, which carries no secret mount policy or secret actor, so run_function/run_code read the absence as every secret acting as the request's user. - Headless runs record the authority their execution context runs under; one function now produces it for both the context and the record. - Interactive runs derive it from their saved request (assistant mode mounts nothing). - The route resolves it from the run (by stream id, user and chat) and caches it per turn; a call that cannot be tied to a run mounts no secrets. * chore(copilot): shrink ratchets and re-record the tool-registry boundary - Pass the route context in the in-band route tests (TS2554 4 -> 0). - The authority read makes secretMountPolicySchema and DEFAULT_SECRET_MOUNT_POLICY imported, so their unused-export entries go. - This change adds 2 modules to the v2 blocks/tools/connector-types routes and the credit-usage and usage-events pages (the policy module and its contract schema, reached through the run repository); the remaining +35-37 on those entries is drift already on staging against the previously recorded baseline. * improvement(copilot): hold the in-band turn context in an lru-cache
1 parent 2fe1e88 commit 3dcf9fb

10 files changed

Lines changed: 583 additions & 310 deletions

File tree

‎apps/sim/app/api/copilot/tools/execute/route.test.ts‎

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
11
import { copilotHttpMock, copilotHttpMockFns } from '@sim/testing/mocks/copilot-http.mock'
2+
import {
3+
mothershipAsyncRunsMock,
4+
mothershipAsyncRunsMockFns,
5+
} from '@sim/testing/mocks/mothership-async-runs.mock'
26
import {
37
mothershipEnvironmentContextMock,
48
mothershipEnvironmentContextMockFns,
@@ -16,6 +20,8 @@ vi.mock('@/lib/mothership/request/http', () => copilotHttpMock)
1620

1721
vi.mock('@/lib/mothership/environment-context', () => mothershipEnvironmentContextMock)
1822

23+
vi.mock('@/lib/mothership/async-runs/repository', () => mothershipAsyncRunsMock)
24+
1925
vi.mock('@/lib/mothership/tool-executor', () => ({
2026
ensureHandlersRegistered: vi.fn(),
2127
toolRequiresApprovalLane: mockToolRequiresApprovalLane,
@@ -41,6 +47,7 @@ const { mockCheckInternalApiKey } = copilotHttpMockFns
4147

4248
const mockPrepareEnvironmentContext =
4349
mothershipEnvironmentContextMockFns.mockPrepareCopilotEnvironmentContext
50+
const { mockGetRunSecretAuthority } = mothershipAsyncRunsMockFns
4451

4552
function makeRequest(body: Record<string, unknown>): Request {
4653
return new Request('http://localhost/api/copilot/tools/execute', {
@@ -81,7 +88,7 @@ describe('POST /api/copilot/tools/execute (in-band)', () => {
8188
mockPrepareEnvironmentContext.mockRejectedValue(new Error('Workspace ws-gone does not exist'))
8289
mockHandler.mockResolvedValue({ success: true, output: { content: 'sensitive' } })
8390

84-
const res = await POST(makeRequest({ ...BASE_BODY, messageId: 'msg-no-registry' }) as never)
91+
const res = await POST(makeRequest({ ...BASE_BODY, messageId: 'msg-no-registry' }) as never, {})
8592
const body = await res.json()
8693

8794
expect(mockHandler).not.toHaveBeenCalled()
@@ -112,7 +119,8 @@ describe('POST /api/copilot/tools/execute (in-band)', () => {
112119
toolName: 'generate_api_key',
113120
params: { name: 'demo' },
114121
messageId: 'msg-api-key',
115-
}) as never
122+
}) as never,
123+
{}
116124
)
117125
const body = await res.json()
118126

@@ -142,7 +150,8 @@ describe('POST /api/copilot/tools/execute (in-band)', () => {
142150
toolName: 'run_function',
143151
params: { code: 'return 1' },
144152
messageId: 'msg-gated',
145-
}) as never
153+
}) as never,
154+
{}
146155
)
147156
const body = await res.json()
148157

@@ -153,4 +162,36 @@ describe('POST /api/copilot/tools/execute (in-band)', () => {
153162
expect(body.error).toContain('checkpoint lane')
154163
})
155164
})
165+
/**
166+
* Go's body cannot carry the turn's secret authority — the turn payload never did — so a
167+
* context rebuilt from it alone read as "every secret, as the body's user" for a turn that
168+
* was admitted with less.
169+
*/
170+
describe('secret authority', () => {
171+
it("hands the tool the turn's re-derived authority, not the permissive default", async () => {
172+
const admitted = {
173+
secretMountPolicy: { secretScope: 'selected' as const, mountedSecrets: ['STRIPE_KEY'] },
174+
secretActorUserId: null,
175+
}
176+
mockGetRunSecretAuthority.mockResolvedValueOnce(admitted)
177+
mockHandler.mockResolvedValue({ success: true, output: { ran: true } })
178+
179+
await POST(
180+
makeRequest({
181+
...BASE_BODY,
182+
toolName: 'run_function',
183+
params: { code: 'return 1', secrets: ['STRIPE_KEY'] },
184+
messageId: 'msg-admitted-authority',
185+
}) as never,
186+
{}
187+
)
188+
189+
expect(mockGetRunSecretAuthority).toHaveBeenCalledWith({
190+
streamId: 'msg-admitted-authority',
191+
userId: BASE_BODY.userId,
192+
chatId: BASE_BODY.chatId,
193+
})
194+
expect(mockHandler.mock.calls[0][1]).toMatchObject(admitted)
195+
})
196+
})
156197
})

‎apps/sim/app/api/copilot/tools/execute/route.ts‎

Lines changed: 34 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
import { createLogger } from '@sim/logger'
22
import { getErrorMessage } from '@sim/utils/errors'
3+
import { LRUCache } from 'lru-cache'
34
import { type NextRequest, NextResponse } from 'next/server'
45
import { copilotToolExecuteInternalBodySchema } from '@/lib/api/contracts/copilot'
56
import { validationErrorResponse } from '@/lib/api/server'
67
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
8+
import { getRunSecretAuthority } from '@/lib/mothership/async-runs/repository'
79
import { toolResultForModel } from '@/lib/mothership/chat/sim-key-redaction'
810
import { prepareCopilotEnvironmentContext } from '@/lib/mothership/environment-context'
911
import { MothershipStreamV1ToolOutcome } from '@/lib/mothership/generated/mothership-stream-v1'
@@ -19,6 +21,7 @@ import {
1921
} from '@/lib/mothership/request/tools/resolved-secret-result'
2022
import { handleResourceSideEffects } from '@/lib/mothership/request/tools/resources'
2123
import type { ToolCallResult } from '@/lib/mothership/request/types'
24+
import { NO_SECRET_AUTHORITY, type TurnSecretAuthority } from '@/lib/mothership/secret-mount-policy'
2225
import { ensureHandlersRegistered, toolRequiresApprovalLane } from '@/lib/mothership/tool-executor'
2326
import { executeTool } from '@/lib/mothership/tool-executor/executor'
2427
import { TOOL_EFFECT_PHASE } from '@/lib/mothership/tool-executor/types'
@@ -31,44 +34,38 @@ const logger = createLogger('CopilotToolExecuteInternalAPI')
3134
* provenance activated by one tool call must stay visible to the next call's
3235
* egress projection, exactly as the request-lifecycle registry accumulates
3336
* across a turn. Keyed by the turn (messageId) so one background lane shares
34-
* one registry; TTL-evicted since nothing signals turn end on this route.
37+
* one registry; idle-evicted since nothing signals turn end on this route.
38+
*
39+
* The turn's secret authority, which Go's body cannot carry, is cached beside it
40+
* (see {@link getRunSecretAuthority}).
3541
*/
36-
const TURN_REGISTRY_TTL_MS = 10 * 60 * 1000
37-
const TURN_REGISTRY_CACHE_MAX = 256
38-
const turnRegistryCache = new Map<
42+
const turnContextCache = new LRUCache<
3943
string,
40-
{ registry: ResolvedSecretTraceRegistry; expiresAt: number }
41-
>()
44+
{ registry: ResolvedSecretTraceRegistry; secretAuthority: TurnSecretAuthority }
45+
>({ max: 256, ttl: 10 * 60 * 1000, updateAgeOnGet: true })
4246

43-
async function getTurnEgressRegistry(
47+
async function getTurnToolContext(
4448
userId: string,
4549
workspaceId: string | undefined,
50+
chatId: string | undefined,
4651
messageId: string | undefined,
4752
requestMode?: string,
4853
organizationId?: string
49-
): Promise<ResolvedSecretTraceRegistry> {
50-
const key = `${userId}\u0000${workspaceId ?? ''}\u0000${organizationId ?? ''}\u0000${messageId ?? ''}\u0000${requestMode ?? ''}`
51-
const now = Date.now()
52-
const hit = turnRegistryCache.get(key)
53-
if (hit && hit.expiresAt > now) {
54-
hit.expiresAt = now + TURN_REGISTRY_TTL_MS
55-
return hit.registry
56-
}
57-
const environmentContext = await prepareCopilotEnvironmentContext(userId, workspaceId, {
58-
includeSecrets: requestMode !== 'assistant',
59-
})
60-
for (const [cachedKey, cached] of turnRegistryCache) {
61-
if (cached.expiresAt <= now) turnRegistryCache.delete(cachedKey)
62-
}
63-
if (turnRegistryCache.size >= TURN_REGISTRY_CACHE_MAX) {
64-
const oldest = turnRegistryCache.keys().next().value
65-
if (oldest !== undefined) turnRegistryCache.delete(oldest)
66-
}
67-
turnRegistryCache.set(key, {
68-
registry: environmentContext.resolvedSecretTraceRegistry,
69-
expiresAt: now + TURN_REGISTRY_TTL_MS,
70-
})
71-
return environmentContext.resolvedSecretTraceRegistry
54+
): Promise<{ registry: ResolvedSecretTraceRegistry; secretAuthority: TurnSecretAuthority }> {
55+
const key = `${userId}\u0000${workspaceId ?? ''}\u0000${organizationId ?? ''}\u0000${chatId ?? ''}\u0000${messageId ?? ''}\u0000${requestMode ?? ''}`
56+
const hit = turnContextCache.get(key)
57+
if (hit) return hit
58+
const [environmentContext, secretAuthority] = await Promise.all([
59+
prepareCopilotEnvironmentContext(userId, workspaceId, {
60+
includeSecrets: requestMode !== 'assistant',
61+
}),
62+
chatId && messageId
63+
? getRunSecretAuthority({ streamId: messageId, userId, chatId })
64+
: NO_SECRET_AUTHORITY,
65+
])
66+
const entry = { registry: environmentContext.resolvedSecretTraceRegistry, secretAuthority }
67+
turnContextCache.set(key, entry)
68+
return entry
7269
}
7370

7471
// POST /api/copilot/tools/execute — internal (Go → Sim) in-band execution of
@@ -153,14 +150,18 @@ export const POST = withRouteHandler((request: NextRequest) =>
153150

154151
let toolRegistry: ResolvedSecretTraceRegistry
155152
let turnRegistry: ResolvedSecretTraceRegistry
153+
let secretAuthority: TurnSecretAuthority
156154
try {
157-
turnRegistry = await getTurnEgressRegistry(
155+
const turn = await getTurnToolContext(
158156
userId,
159157
workspaceId,
158+
chatId,
160159
messageId,
161160
requestMode,
162161
organizationId
163162
)
163+
turnRegistry = turn.registry
164+
secretAuthority = turn.secretAuthority
164165
toolRegistry = turnRegistry.forkForInputPaths([])
165166
} catch (err) {
166167
/**
@@ -218,6 +219,8 @@ export const POST = withRouteHandler((request: NextRequest) =>
218219
assistantSearch,
219220
targetWorkspaceId,
220221
resolvedSecretTraceRegistry: toolRegistry,
222+
secretMountPolicy: secretAuthority.secretMountPolicy,
223+
secretActorUserId: secretAuthority.secretActorUserId,
221224
})
222225
const projection = inspectToolResultForCopilot(result, toolRegistry, toolName)
223226
const projected = projection.result

‎apps/sim/lib/mothership/async-runs/repository.ts‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,8 @@ import {
3030
type SQL,
3131
sql,
3232
} from 'drizzle-orm'
33+
import { z } from 'zod'
34+
import { secretMountPolicySchema } from '@/lib/api/contracts/secret-mount-policy'
3335
import { type ResourceOwner, resourceScopeFromOwner } from '@/lib/core/resource-scope'
3436
import { acquireAdvisoryXactLock } from '@/lib/db/advisory-locks'
3537
import type { DbOrTx, DbTransaction } from '@/lib/db/types'
@@ -65,6 +67,11 @@ import {
6567
import type { BillingAdmission } from '@/lib/mothership/request/lifecycle/recovery-config'
6668
import { markSpanForError } from '@/lib/mothership/request/otel'
6769
import { sealClientToolSettlement } from '@/lib/mothership/request/tools/client-completion-seal.server'
70+
import {
71+
DEFAULT_SECRET_MOUNT_POLICY,
72+
NO_SECRET_AUTHORITY,
73+
type TurnSecretAuthority,
74+
} from '@/lib/mothership/secret-mount-policy'
6875
import {
6976
type DesktopToolClaimOwner,
7077
isDesktopToolCall,
@@ -573,6 +580,49 @@ export async function recordRunBillingAdmission(
573580
)
574581
}
575582

583+
const recordedSecretAuthoritySchema = z.object({
584+
secretMountPolicy: secretMountPolicySchema,
585+
secretActorUserId: z.string().nullable(),
586+
})
587+
588+
/**
589+
* The secret authority a turn was admitted with, for a lane that rebuilds its tool context from
590+
* Go's request instead of inheriting the turn's. Headless runs record it at admission; an
591+
* interactive chat narrows it only in assistant mode, so its saved request decides. Anything
592+
* else — no run, another user's or chat's run, an unrecorded headless run — mounts nothing.
593+
*/
594+
export async function getRunSecretAuthority(input: {
595+
streamId: string
596+
userId: string
597+
chatId: string
598+
}): Promise<TurnSecretAuthority> {
599+
const [run] = await db
600+
.select({
601+
chatId: copilotRuns.chatId,
602+
recorded: sql<unknown>`${copilotRuns.requestContext}->'secretAuthority'`,
603+
recoveryKind: sql<string | null>`${copilotRuns.requestContext}->'recovery'->>'kind'`,
604+
requestMode: sql<
605+
string | null
606+
>`COALESCE(${copilotRuns.requestContext}->'recovery'->>'requestMode', ${copilotRuns.requestContext}->'recovery'->'request'->>'mode')`,
607+
})
608+
.from(copilotRuns)
609+
.where(and(eq(copilotRuns.streamId, input.streamId), eq(copilotRuns.userId, input.userId)))
610+
.limit(1)
611+
if (!run || run.chatId !== input.chatId) return NO_SECRET_AUTHORITY
612+
613+
if (run.recorded !== null) {
614+
const recorded = recordedSecretAuthoritySchema.safeParse(run.recorded)
615+
return recorded.success ? recorded.data : NO_SECRET_AUTHORITY
616+
}
617+
if (run.recoveryKind === 'interactive_stream' && run.requestMode !== 'assistant') {
618+
return {
619+
secretMountPolicy: { ...DEFAULT_SECRET_MOUNT_POLICY },
620+
secretActorUserId: input.userId,
621+
}
622+
}
623+
return NO_SECRET_AUTHORITY
624+
}
625+
576626
// Un-instrumented: called from a 4 Hz resume poll; per-call spans
577627
// swamped traces. Use Prom histograms if latency visibility is needed.
578628
export async function getLatestRunForStream(streamId: string, userId?: string) {

0 commit comments

Comments
 (0)