Skip to content

Commit c1f8e3f

Browse files
committed
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.
1 parent 2d0ac79 commit c1f8e3f

7 files changed

Lines changed: 327 additions & 32 deletions

File tree

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

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

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

Lines changed: 35 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import { type NextRequest, NextResponse } from 'next/server'
44
import { copilotToolExecuteInternalBodySchema } from '@/lib/api/contracts/copilot'
55
import { validationErrorResponse } from '@/lib/api/server'
66
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
7+
import { getRunSecretAuthority } from '@/lib/mothership/async-runs/repository'
78
import { toolResultForModel } from '@/lib/mothership/chat/sim-key-redaction'
89
import { prepareCopilotEnvironmentContext } from '@/lib/mothership/environment-context'
910
import { MothershipStreamV1ToolOutcome } from '@/lib/mothership/generated/mothership-stream-v1'
@@ -19,6 +20,7 @@ import {
1920
} from '@/lib/mothership/request/tools/resolved-secret-result'
2021
import { handleResourceSideEffects } from '@/lib/mothership/request/tools/resources'
2122
import type { ToolCallResult } from '@/lib/mothership/request/types'
23+
import { NO_SECRET_AUTHORITY, type TurnSecretAuthority } from '@/lib/mothership/secret-mount-policy'
2224
import { ensureHandlersRegistered, toolRequiresApprovalLane } from '@/lib/mothership/tool-executor'
2325
import { executeTool } from '@/lib/mothership/tool-executor/executor'
2426
import { TOOL_EFFECT_PHASE } from '@/lib/mothership/tool-executor/types'
@@ -32,43 +34,58 @@ const logger = createLogger('CopilotToolExecuteInternalAPI')
3234
* egress projection, exactly as the request-lifecycle registry accumulates
3335
* across a turn. Keyed by the turn (messageId) so one background lane shares
3436
* one registry; TTL-evicted since nothing signals turn end on this route.
37+
*
38+
* The turn's secret authority, which Go's body cannot carry, is cached beside it
39+
* (see {@link getRunSecretAuthority}).
3540
*/
3641
const TURN_REGISTRY_TTL_MS = 10 * 60 * 1000
3742
const TURN_REGISTRY_CACHE_MAX = 256
3843
const turnRegistryCache = new Map<
3944
string,
40-
{ registry: ResolvedSecretTraceRegistry; expiresAt: number }
45+
{
46+
registry: ResolvedSecretTraceRegistry
47+
secretAuthority: TurnSecretAuthority
48+
expiresAt: number
49+
}
4150
>()
4251

43-
async function getTurnEgressRegistry(
52+
async function getTurnToolContext(
4453
userId: string,
4554
workspaceId: string | undefined,
55+
chatId: string | undefined,
4656
messageId: string | undefined,
4757
requestMode?: string,
4858
organizationId?: string
49-
): Promise<ResolvedSecretTraceRegistry> {
50-
const key = `${userId}\u0000${workspaceId ?? ''}\u0000${organizationId ?? ''}\u0000${messageId ?? ''}\u0000${requestMode ?? ''}`
59+
): Promise<{ registry: ResolvedSecretTraceRegistry; secretAuthority: TurnSecretAuthority }> {
60+
const key = `${userId}\u0000${workspaceId ?? ''}\u0000${organizationId ?? ''}\u0000${chatId ?? ''}\u0000${messageId ?? ''}\u0000${requestMode ?? ''}`
5161
const now = Date.now()
5262
const hit = turnRegistryCache.get(key)
5363
if (hit && hit.expiresAt > now) {
5464
hit.expiresAt = now + TURN_REGISTRY_TTL_MS
55-
return hit.registry
65+
return hit
5666
}
57-
const environmentContext = await prepareCopilotEnvironmentContext(userId, workspaceId, {
58-
includeSecrets: requestMode !== 'assistant',
59-
})
67+
const [environmentContext, secretAuthority] = await Promise.all([
68+
prepareCopilotEnvironmentContext(userId, workspaceId, {
69+
includeSecrets: requestMode !== 'assistant',
70+
}),
71+
chatId && messageId
72+
? getRunSecretAuthority({ streamId: messageId, userId, chatId })
73+
: NO_SECRET_AUTHORITY,
74+
])
6075
for (const [cachedKey, cached] of turnRegistryCache) {
6176
if (cached.expiresAt <= now) turnRegistryCache.delete(cachedKey)
6277
}
6378
if (turnRegistryCache.size >= TURN_REGISTRY_CACHE_MAX) {
6479
const oldest = turnRegistryCache.keys().next().value
6580
if (oldest !== undefined) turnRegistryCache.delete(oldest)
6681
}
67-
turnRegistryCache.set(key, {
82+
const entry = {
6883
registry: environmentContext.resolvedSecretTraceRegistry,
84+
secretAuthority,
6985
expiresAt: now + TURN_REGISTRY_TTL_MS,
70-
})
71-
return environmentContext.resolvedSecretTraceRegistry
86+
}
87+
turnRegistryCache.set(key, entry)
88+
return entry
7289
}
7390

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

154171
let toolRegistry: ResolvedSecretTraceRegistry
155172
let turnRegistry: ResolvedSecretTraceRegistry
173+
let secretAuthority: TurnSecretAuthority
156174
try {
157-
turnRegistry = await getTurnEgressRegistry(
175+
const turn = await getTurnToolContext(
158176
userId,
159177
workspaceId,
178+
chatId,
160179
messageId,
161180
requestMode,
162181
organizationId
163182
)
183+
turnRegistry = turn.registry
184+
secretAuthority = turn.secretAuthority
164185
toolRegistry = turnRegistry.forkForInputPaths([])
165186
} catch (err) {
166187
/**
@@ -218,6 +239,8 @@ export const POST = withRouteHandler((request: NextRequest) =>
218239
assistantSearch,
219240
targetWorkspaceId,
220241
resolvedSecretTraceRegistry: toolRegistry,
242+
secretMountPolicy: secretAuthority.secretMountPolicy,
243+
secretActorUserId: secretAuthority.secretActorUserId,
221244
})
222245
const projection = inspectToolResultForCopilot(result, toolRegistry, toolName)
223246
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) {
Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
/**
2+
* Re-derives a turn's secret authority from real `copilot_runs` rows in a temporary PostgreSQL
3+
* table. Set TEST_DATABASE_URL to a local PostgreSQL database.
4+
*/
5+
import { generateId } from '@sim/utils/id'
6+
import { drizzle, type PostgresJsDatabase } from 'drizzle-orm/postgres-js'
7+
import postgres from 'postgres'
8+
import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'
9+
10+
const { database, databaseUrl } = await vi.hoisted(async () => {
11+
const { readTestDatabaseUrl } = await import('@sim/db/testing/test-infrastructure')
12+
return {
13+
databaseUrl: readTestDatabaseUrl(),
14+
database: { current: undefined as PostgresJsDatabase | undefined },
15+
}
16+
})
17+
18+
vi.mock('@sim/db', () => ({
19+
db: {
20+
select: (...args: unknown[]) => {
21+
if (!database.current) throw new Error('PostgreSQL test database is not initialized')
22+
return Reflect.apply(database.current.select, database.current, args)
23+
},
24+
},
25+
}))
26+
vi.mock('@/lib/mothership/request/otel', () => ({ markSpanForError: vi.fn() }))
27+
28+
import { getRunSecretAuthority } from '@/lib/mothership/async-runs/repository'
29+
import { NO_SECRET_AUTHORITY } from '@/lib/mothership/secret-mount-policy'
30+
31+
const connection = postgres(databaseUrl, { max: 1 })
32+
const userId = 'user-secret-authority'
33+
const chatId = generateId()
34+
35+
async function insertRun(streamId: string, requestContext: unknown, owner = userId) {
36+
await connection`
37+
INSERT INTO copilot_runs (execution_id, chat_id, user_id, stream_id, request_context)
38+
VALUES (${generateId()}, ${chatId}, ${owner}, ${streamId},
39+
${JSON.stringify(requestContext)}::jsonb)
40+
`
41+
}
42+
43+
function interactiveContext(mode?: string) {
44+
return {
45+
controllerToken: 'lease',
46+
recovery: {
47+
kind: 'interactive_stream',
48+
request: { message: 'hi', userId, chatId, ...(mode ? { mode } : {}) },
49+
goRoute: '/api/copilot',
50+
clientToolPickupExpected: true,
51+
},
52+
}
53+
}
54+
55+
afterAll(async () => {
56+
await connection.end()
57+
})
58+
59+
describe('getRunSecretAuthority with PostgreSQL', () => {
60+
beforeAll(async () => {
61+
database.current = drizzle(connection)
62+
await connection.unsafe(`
63+
CREATE TEMP TABLE copilot_runs
64+
(LIKE public.copilot_runs INCLUDING DEFAULTS INCLUDING CONSTRAINTS INCLUDING INDEXES)
65+
`)
66+
})
67+
68+
beforeEach(async () => {
69+
await connection.unsafe('TRUNCATE pg_temp.copilot_runs')
70+
})
71+
72+
it("returns a headless run's recorded authority, including a refused actor", async () => {
73+
await insertRun('stream-inbox', {
74+
source: 'headless_lifecycle',
75+
secretAuthority: {
76+
secretMountPolicy: { secretScope: 'selected', mountedSecrets: ['STRIPE_KEY'] },
77+
secretActorUserId: null,
78+
},
79+
})
80+
81+
await expect(
82+
getRunSecretAuthority({ streamId: 'stream-inbox', userId, chatId })
83+
).resolves.toEqual({
84+
secretMountPolicy: { secretScope: 'selected', mountedSecrets: ['STRIPE_KEY'] },
85+
secretActorUserId: null,
86+
})
87+
})
88+
89+
it('gives an interactive agent turn every secret, as the turn user', async () => {
90+
await insertRun('stream-agent', interactiveContext())
91+
await insertRun('stream-build', interactiveContext('build'))
92+
93+
for (const streamId of ['stream-agent', 'stream-build']) {
94+
await expect(getRunSecretAuthority({ streamId, userId, chatId })).resolves.toEqual({
95+
secretMountPolicy: { secretScope: 'all', mountedSecrets: [] },
96+
secretActorUserId: userId,
97+
})
98+
}
99+
})
100+
101+
it('gives an interactive assistant turn no secrets', async () => {
102+
await insertRun('stream-assistant', interactiveContext('assistant'))
103+
104+
await expect(
105+
getRunSecretAuthority({ streamId: 'stream-assistant', userId, chatId })
106+
).resolves.toEqual(NO_SECRET_AUTHORITY)
107+
})
108+
109+
it('gives a headless run that recorded nothing no secrets', async () => {
110+
await insertRun('stream-legacy', { source: 'headless_lifecycle' })
111+
await insertRun('stream-empty', {})
112+
113+
for (const streamId of ['stream-legacy', 'stream-empty']) {
114+
await expect(getRunSecretAuthority({ streamId, userId, chatId })).resolves.toEqual(
115+
NO_SECRET_AUTHORITY
116+
)
117+
}
118+
})
119+
120+
it('refuses a malformed recorded authority instead of reading it as the permissive default', async () => {
121+
await insertRun('stream-no-scope', {
122+
secretAuthority: { secretMountPolicy: {}, secretActorUserId: userId },
123+
})
124+
await insertRun('stream-bad-actor', {
125+
secretAuthority: {
126+
secretMountPolicy: { secretScope: 'all', mountedSecrets: [] },
127+
secretActorUserId: 42,
128+
},
129+
})
130+
131+
for (const streamId of ['stream-no-scope', 'stream-bad-actor']) {
132+
await expect(getRunSecretAuthority({ streamId, userId, chatId })).resolves.toEqual(
133+
NO_SECRET_AUTHORITY
134+
)
135+
}
136+
})
137+
138+
it("never lends a run's authority to another user, another chat, or an unknown turn", async () => {
139+
await insertRun('stream-owned', interactiveContext())
140+
141+
const lookups = [
142+
{ streamId: 'stream-owned', userId: 'someone-else', chatId },
143+
{ streamId: 'stream-owned', userId, chatId: generateId() },
144+
{ streamId: 'stream-owned', userId, chatId: 'not-a-uuid' },
145+
{ streamId: 'stream-missing', userId, chatId },
146+
]
147+
for (const lookup of lookups) {
148+
await expect(getRunSecretAuthority(lookup)).resolves.toEqual(NO_SECRET_AUTHORITY)
149+
}
150+
})
151+
})

0 commit comments

Comments
 (0)