diff --git a/packages/session-ingest-contracts/src/cloud-agent-worktree-scope.ts b/packages/session-ingest-contracts/src/cloud-agent-worktree-scope.ts new file mode 100644 index 00000000000..db4ebd25e66 --- /dev/null +++ b/packages/session-ingest-contracts/src/cloud-agent-worktree-scope.ts @@ -0,0 +1,123 @@ +import { z } from 'zod'; + +import { containedKiloSessionIdSchema } from './cloud-agent-session-scope'; +import { cloudAgentWorktreeIdSchema } from './rpc-contract'; + +export const cloudAgentSessionScopeProtocolVersionV2 = '2' as const; + +export const cloudAgentSessionScopeV2Headers = { + worktreeId: 'X-Kilo-Cloud-Agent-Worktree', + expectedOrganization: 'X-Kilo-Expected-Organization', + memberRootKiloSessionIds: 'X-Kilo-Member-Root-Sessions', +} as const; + +export const CLOUD_AGENT_SESSION_SCOPE_PERSONAL_ORGANIZATION = 'personal' as const; + +export const MAX_CLOUD_AGENT_SESSION_SCOPE_MEMBER_ROOTS = 32; +export const MAX_CLOUD_AGENT_SESSION_SCOPE_HEADER_BYTES = 4096; + +export const cloudAgentSessionScopeV2AssertionSchema = z.object({ + protocolVersion: z.literal(cloudAgentSessionScopeProtocolVersionV2), + worktreeId: cloudAgentWorktreeIdSchema, + expectedOrganization: z.union([ + z.literal(CLOUD_AGENT_SESSION_SCOPE_PERSONAL_ORGANIZATION), + z.uuid(), + ]), + memberRootKiloSessionIds: z + .array(containedKiloSessionIdSchema) + .min(1) + .max(MAX_CLOUD_AGENT_SESSION_SCOPE_MEMBER_ROOTS), +}); + +export type CloudAgentSessionScopeV2Assertion = z.infer< + typeof cloudAgentSessionScopeV2AssertionSchema +>; + +export type CloudAgentSessionScopeV2 = { + worktreeId: string; + organizationId: string | null; + memberRootKiloSessionIds: string[]; +}; + +export type CloudAgentSessionScopeV2HeaderInput = { + worktreeId?: string; + expectedOrganization?: string; + memberRootKiloSessionIds?: string; +}; + +export type CloudAgentSessionScopeV2ParseResult = + | { kind: 'ok'; scope: CloudAgentSessionScopeV2 } + | { kind: 'invalid' }; + +function headerValue(value: string | undefined): string { + return value ?? ''; +} + +export function measureCloudAgentSessionScopeV2HeaderBytes( + input: CloudAgentSessionScopeV2HeaderInput +): number { + const encoder = new TextEncoder(); + return ( + encoder.encode(headerValue(input.worktreeId)).length + + encoder.encode(headerValue(input.expectedOrganization)).length + + encoder.encode(headerValue(input.memberRootKiloSessionIds)).length + ); +} + +export function encodeCloudAgentSessionScopeV2Headers( + scope: CloudAgentSessionScopeV2 +): Required { + return { + worktreeId: scope.worktreeId, + expectedOrganization: scope.organizationId ?? CLOUD_AGENT_SESSION_SCOPE_PERSONAL_ORGANIZATION, + memberRootKiloSessionIds: scope.memberRootKiloSessionIds.join(','), + }; +} + +export function parseCloudAgentSessionScopeV2Headers( + input: CloudAgentSessionScopeV2HeaderInput +): CloudAgentSessionScopeV2ParseResult { + const { worktreeId, expectedOrganization, memberRootKiloSessionIds } = input; + if ( + worktreeId === undefined || + expectedOrganization === undefined || + memberRootKiloSessionIds === undefined + ) { + return { kind: 'invalid' }; + } + if ( + measureCloudAgentSessionScopeV2HeaderBytes(input) > MAX_CLOUD_AGENT_SESSION_SCOPE_HEADER_BYTES + ) { + return { kind: 'invalid' }; + } + + const rawMemberRoots = memberRootKiloSessionIds.split(','); + if (rawMemberRoots.some(value => value.trim().length === 0)) return { kind: 'invalid' }; + const normalizedMemberRoots = [...new Set(rawMemberRoots.map(value => value.trim()))]; + if ( + normalizedMemberRoots.length === 0 || + normalizedMemberRoots.length > MAX_CLOUD_AGENT_SESSION_SCOPE_MEMBER_ROOTS + ) { + return { kind: 'invalid' }; + } + + const parsed = cloudAgentSessionScopeV2AssertionSchema.safeParse({ + protocolVersion: cloudAgentSessionScopeProtocolVersionV2, + worktreeId, + expectedOrganization, + memberRootKiloSessionIds: normalizedMemberRoots, + }); + if (!parsed.success) return { kind: 'invalid' }; + + return { + kind: 'ok', + scope: { + worktreeId: parsed.data.worktreeId, + organizationId: + parsed.data.expectedOrganization === CLOUD_AGENT_SESSION_SCOPE_PERSONAL_ORGANIZATION + ? null + : parsed.data.expectedOrganization, + memberRootKiloSessionIds: parsed.data.memberRootKiloSessionIds, + }, + }; +} diff --git a/packages/session-ingest-contracts/src/index.ts b/packages/session-ingest-contracts/src/index.ts index 8cce0f52698..d8bbf02e60d 100644 --- a/packages/session-ingest-contracts/src/index.ts +++ b/packages/session-ingest-contracts/src/index.ts @@ -1,6 +1,7 @@ export * from './rpc-contract'; export * from './session-live-summary'; export * from './cloud-agent-session-scope'; +export * from './cloud-agent-worktree-scope'; export * from './public-worktree-projection'; export const DEFAULT_SESSION_TITLE_PATTERN = diff --git a/services/session-ingest/package.json b/services/session-ingest/package.json index 37e05a2f032..0581283d1e0 100644 --- a/services/session-ingest/package.json +++ b/services/session-ingest/package.json @@ -7,10 +7,11 @@ "cf-typegen": "wrangler types --include-runtime=false --env-interface CloudflareBindings", "test": "vitest run", "test:watch": "vitest", + "test:postgres": "vitest run --config vitest.postgres.config.ts", "test:integration": "vitest run --config vitest.workers.config.ts", "test:integration:watch": "vitest --config vitest.workers.config.ts", "typecheck": "tsgo --noEmit", - "lint": "pnpm -w exec oxlint --config .oxlintrc.json services/session-ingest/src" + "lint": "pnpm -w exec oxlint --config .oxlintrc.json services/session-ingest/src services/session-ingest/test/postgres" }, "dependencies": { "@kilocode/db": "workspace:*", diff --git a/services/session-ingest/src/routes/api.ts b/services/session-ingest/src/routes/api.ts index 6171ef2d4d3..1560a849794 100644 --- a/services/session-ingest/src/routes/api.ts +++ b/services/session-ingest/src/routes/api.ts @@ -23,8 +23,8 @@ import { createConnectionTicket } from '../dos/connection-ticket'; import { getSessionExport } from '../services/session-export'; import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; import { handleDirectIngestRequest } from '../ingest/direct-ingest'; -import { isDefaultSessionTitle } from '../ingest/default-session-title'; import { resolveAccessibleKiloSession } from '../services/session-access'; +import { applyReportedSessionTitle, reportSessionTitleSchema } from '../services/session-title'; import { signSessionShareToken } from '../services/session-share-token'; import { canCreateCliSessionForUser } from '../services/user-session-admission'; import { isWorktreeSessionDeleting } from '../services/worktree-deletion'; @@ -820,11 +820,6 @@ api.post('/session/:sessionId/unshare', async c => { return c.json({ success: true }, 200); }); -const reportSessionTitleSchema = z.object({ - title: z.string().trim().min(1).max(200), - generated: z.boolean(), -}); - api.post('/session/:sessionId/rename-notify', async c => { const rawSessionId = c.req.param('sessionId'); const parsed = sessionIdSchema.safeParse(rawSessionId); @@ -867,66 +862,26 @@ api.post('/session/:sessionId/title', async c => { const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); const kiloUserId = c.get('user_id'); - const { title, generated } = bodyParse.data; - - const existingRows = await db - .select() - .from(cli_sessions_v2) - .where( - and(eq(cli_sessions_v2.session_id, parsed.data), eq(cli_sessions_v2.kilo_user_id, kiloUserId)) - ) - .limit(1); - const existing = existingRows[0]; - if (!existing) { + const result = await applyReportedSessionTitle(db, { + sessionId: parsed.data, + kiloUserId, + input: bodyParse.data, + }); + if (result.kind === 'not_found') { return c.json({ success: false, error: 'session_not_found' }, 404); } - - // generated: true applies only over NULL or default-pattern titles (Decision 8). - // generated: false is last-write-wins. - if (generated && !isDefaultSessionTitle(existing.title)) { + if (result.kind === 'refused') { return c.json({ applied: false }, 200); } - // Ownership filter shared by both paths. For generated writes, also CAS on the - // title observed at read time so a concurrent rename (or other non-default write) - // between select and update cannot be clobbered — Decision 8 must hold at write time. - const ownership = and( - eq(cli_sessions_v2.session_id, parsed.data), - eq(cli_sessions_v2.kilo_user_id, kiloUserId) - ); - const writeWhere = generated - ? and( - ownership, - existing.title == null - ? isNull(cli_sessions_v2.title) - : eq(cli_sessions_v2.title, existing.title) - ) - : ownership; - - // Preserve updated_at so title-only writes do not reorder activity lists - // (same as the web rename path). - const [updatedRow] = await db - .update(cli_sessions_v2) - .set({ title, updated_at: existing.updated_at }) - .where(writeWhere) - .returning(); - - if (!updatedRow) { - // Generated CAS miss (intervening rename / non-default write) → refused, not 404. - if (generated) { - return c.json({ applied: false }, 200); - } - return c.json({ success: false, error: 'session_not_found' }, 404); - } - - const session = mapSessionEventRow(updatedRow); + const session = mapSessionEventRow(result.row); notifyUserSessionEventFromContext(c, kiloUserId, { type: 'session.updated', data: { source: 'v2', session, changedAt: session.updatedAt }, }); - return c.json({ title: session.title ?? title, applied: true }, 200); + return c.json({ title: session.title ?? bodyParse.data.title, applied: true }, 200); }); api.get('/sessions/active', async c => { diff --git a/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts b/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts index 9686b1dd473..ea6ca80e264 100644 --- a/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts +++ b/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts @@ -12,6 +12,12 @@ vi.mock('@kilocode/db/client', () => ({ getWorkerDb: vi.fn() })); vi.mock('../dos/SessionAccessCacheDO', () => ({ getSessionAccessCacheDO: vi.fn() })); vi.mock('../ingest/direct-ingest', () => ({ handleDirectIngestRequest: vi.fn() })); vi.mock('../services/session-access', () => ({ resolveAccessibleKiloSession: vi.fn() })); +vi.mock('../services/cloud-agent-worktree-scope', async importOriginal => ({ + ...(await importOriginal()), + resolveScopedChildSession: vi.fn(), + resolveScopedCreateParent: vi.fn(), +})); +vi.mock('../services/session-export', () => ({ getSessionExport: vi.fn() })); vi.mock('../services/user-session-admission', () => ({ canCreateCliSessionForUser: vi.fn(), })); @@ -27,9 +33,18 @@ vi.mock('@kilocode/worker-utils', async importOriginal => ({ })); import { getWorkerDb } from '@kilocode/db/client'; -import { cloudAgentSessionScopeHeaders } from '@kilocode/session-ingest-contracts'; +import { + cloudAgentSessionScopeHeaders, + cloudAgentSessionScopeV2Headers, +} from '@kilocode/session-ingest-contracts'; import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; import { handleDirectIngestRequest } from '../ingest/direct-ingest'; +import { + resolveScopedChildSession, + resolveScopedCreateParent, +} from '../services/cloud-agent-worktree-scope'; +import type * as CloudAgentWorktreeScopeModule from '../services/cloud-agent-worktree-scope'; +import { getSessionExport } from '../services/session-export'; import { resolveAccessibleKiloSession } from '../services/session-access'; import { canCreateCliSessionForUser } from '../services/user-session-admission'; import { cloudAgentSessionScopeApi } from './cloud-agent-session-scope'; @@ -58,6 +73,10 @@ function assertionHeaders() { }; } +function queryOf(condition: SQL) { + return new PgDialect().sqlToQuery(condition); +} + function persistedRow(overrides: Record = {}) { return { session_id: childSessionId, @@ -81,6 +100,8 @@ function persistedRow(overrides: Record = {}) { function makeDb(selectResults: unknown[][], insertResult: unknown[], updateResult: unknown[] = []) { const updateSets: unknown[] = []; + const updateConditions: SQL[] = []; + const forConditions: SQL[] = []; const insertedValues: unknown[] = []; const selectConditions: SQL[] = []; @@ -91,7 +112,11 @@ function makeDb(selectResults: unknown[][], insertResult: unknown[], updateResul return select; }), limit: vi.fn(() => select), - for: vi.fn(async () => selectResults.shift() ?? []), + for: vi.fn(async () => { + const condition = selectConditions.at(-1); + if (condition) forConditions.push(condition); + return selectResults.shift() ?? []; + }), then: (resolve: (rows: unknown[]) => unknown) => resolve(selectResults.shift() ?? []), }; const update = { @@ -99,7 +124,10 @@ function makeDb(selectResults: unknown[][], insertResult: unknown[], updateResul updateSets.push(values); return update; }), - where: vi.fn(() => update), + where: vi.fn((condition: SQL) => { + updateConditions.push(condition); + return update; + }), returning: vi.fn(async () => updateResult), then: vi.fn(resolve => resolve(undefined)), }; @@ -118,8 +146,10 @@ function makeDb(selectResults: unknown[][], insertResult: unknown[], updateResul }; const db = { transaction: vi.fn(async callback => callback(tx)), + select: vi.fn(() => select), + update: vi.fn(() => update), }; - return { db, insertedValues, updateSets, selectConditions }; + return { db, insertedValues, updateSets, updateConditions, forConditions, selectConditions }; } describe('Cloud Agent session scope routes', () => { @@ -149,7 +179,7 @@ describe('Cloud Agent session scope routes', () => { expect(workerUtils.hasOrganizationAccess).not.toHaveBeenCalled(); }); - it('rejects an incomplete or invalid session scope assertion before database access', async () => { + it('rejects an incomplete protocol v2 assertion before database access', async () => { const { db } = makeDb([], []); vi.mocked(getWorkerDb).mockReturnValue(db as never); const headers = assertionHeaders(); @@ -165,6 +195,7 @@ describe('Cloud Agent session scope routes', () => { ); expect(response.status).toBe(400); + expect(getWorkerDb).not.toHaveBeenCalled(); expect(db.transaction).not.toHaveBeenCalled(); }); @@ -682,3 +713,541 @@ describe('Cloud Agent session scope routes', () => { } ); }); + +describe('Cloud Agent session scope v2 routes', () => { + const v2WorktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; + const v2OrganizationId = '33333333-3333-4333-8333-333333333333'; + const memberRootId = rootSessionId; + + function v2Headers(overrides: Record = {}) { + return { + 'Content-Type': 'application/json', + [cloudAgentSessionScopeHeaders.protocolVersion]: '2', + [cloudAgentSessionScopeV2Headers.worktreeId]: v2WorktreeId, + [cloudAgentSessionScopeV2Headers.expectedOrganization]: v2OrganizationId, + [cloudAgentSessionScopeV2Headers.memberRootKiloSessionIds]: memberRootId, + ...overrides, + }; + } + + function createRequest(headers: Record, parentSessionId = memberRootId) { + return new Request('http://local/session', { + method: 'POST', + headers, + body: JSON.stringify({ sessionId: childSessionId, parentSessionId }), + }); + } + + function rootParentRow() { + return { + sessionId: memberRootId, + parentSessionId: null, + organizationId: v2OrganizationId, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: null, + cloudAgentWorktreeId: v2WorktreeId, + }; + } + + function expectedLockSql(organizationId: string | null): string { + return organizationId === null + ? `("cli_sessions_v2"."session_id" = $1 and ("cli_sessions_v2"."kilo_user_id" = $2 and "cli_sessions_v2"."organization_id" is null and "cli_sessions_v2"."cloud_agent_worktree_id" = $3) and "cli_sessions_v2"."parent_session_id" is null)` + : `("cli_sessions_v2"."session_id" = $1 and ("cli_sessions_v2"."kilo_user_id" = $2 and "cli_sessions_v2"."organization_id" = $3 and "cli_sessions_v2"."cloud_agent_worktree_id" = $4) and "cli_sessions_v2"."parent_session_id" is null)`; + } + + function expectedTitleWhereSql(organizationId: string | null): string { + return organizationId === null + ? `("cli_sessions_v2"."session_id" = $1 and "cli_sessions_v2"."kilo_user_id" = $2 and (("cli_sessions_v2"."kilo_user_id" = $3 and "cli_sessions_v2"."organization_id" is null and "cli_sessions_v2"."cloud_agent_worktree_id" = $4) and "cli_sessions_v2"."cloud_agent_session_scope_id" = $5))` + : `("cli_sessions_v2"."session_id" = $1 and "cli_sessions_v2"."kilo_user_id" = $2 and (("cli_sessions_v2"."kilo_user_id" = $3 and "cli_sessions_v2"."organization_id" = $4 and "cli_sessions_v2"."cloud_agent_worktree_id" = $5) and "cli_sessions_v2"."cloud_agent_session_scope_id" = $6))`; + } + + beforeEach(() => { + vi.resetAllMocks(); + vi.mocked(canCreateCliSessionForUser).mockResolvedValue(true); + }); + + it.each([ + ['worktree id', cloudAgentSessionScopeV2Headers.worktreeId], + ['expected organization', cloudAgentSessionScopeV2Headers.expectedOrganization], + ['member roots', cloudAgentSessionScopeV2Headers.memberRootKiloSessionIds], + ])('rejects protocol v2 missing the %s before database access', async (_label, header) => { + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const headers = v2Headers(); + delete headers[header]; + + const response = await makeApp().fetch(createRequest(headers), env); + + expect(response.status).toBe(400); + expect(getWorkerDb).not.toHaveBeenCalled(); + }); + + it('rejects protocol v2 with more than 32 member roots before database access', async () => { + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const roots = Array.from( + { length: 33 }, + (_, index) => `ses_${String(index).padStart(26, '0')}` + ); + + const response = await makeApp().fetch( + createRequest( + v2Headers({ [cloudAgentSessionScopeV2Headers.memberRootKiloSessionIds]: roots.join(',') }) + ), + env + ); + + expect(response.status).toBe(400); + expect(getWorkerDb).not.toHaveBeenCalled(); + }); + + it('rejects a member-root header over 4096 bytes that de-duplicates to one root before database access', async () => { + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const repeated = Array.from({ length: 140 }, () => memberRootId).join(','); + expect(repeated.length).toBeGreaterThan(4096); + + const response = await makeApp().fetch( + createRequest( + v2Headers({ [cloudAgentSessionScopeV2Headers.memberRootKiloSessionIds]: repeated }) + ), + env + ); + + expect(response.status).toBe(400); + expect(getWorkerDb).not.toHaveBeenCalled(); + }); + + it('ignores an oversized Cloud Agent session header when the v2 assertion fields are valid', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ kind: 'denied' }); + const { db } = makeDb([[]], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + createRequest( + v2Headers({ + [cloudAgentSessionScopeHeaders.cloudAgentSessionId]: `workspace_${'a'.repeat(5000)}`, + }) + ), + env + ); + + expect(response.status).toBe(409); + expect(getWorkerDb).toHaveBeenCalled(); + }); + + it('rejects protocol v1 assertions on the scoped export and title routes', async () => { + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const exportResponse = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/export`, { + method: 'GET', + headers: assertionHeaders(), + }), + env + ); + const titleResponse = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/title`, { + method: 'POST', + headers: assertionHeaders(), + body: JSON.stringify({ title: 'Renamed', generated: false }), + }), + env + ); + + expect(exportResponse.status).toBe(400); + expect(titleResponse.status).toBe(400); + expect(getWorkerDb).not.toHaveBeenCalled(); + }); + + it('accepts a complete protocol v2 assertion and applies the v2 predicate', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ kind: 'denied' }); + const { db } = makeDb([[]], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch(createRequest(v2Headers()), env); + + expect(response.status).toBe(409); + expect(resolveScopedCreateParent).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + parentSessionId: memberRootId, + scope: expect.objectContaining({ + worktreeId: v2WorktreeId, + organizationId: v2OrganizationId, + memberRootKiloSessionIds: [memberRootId], + }), + }) + ); + }); + + it('creates a child under a member-root parent and copies the root scope id', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ + kind: 'ok', + parent: rootParentRow(), + memberRoot: { + sessionId: memberRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: null, + }, + memberRootScopeId: 'workspace_root', + }); + const created = persistedRow({ + parent_session_id: memberRootId, + organization_id: v2OrganizationId, + cloud_agent_session_scope_id: 'workspace_root', + cloud_agent_worktree_id: v2WorktreeId, + }); + const { db, insertedValues, forConditions } = makeDb( + [[{ sessionId: memberRootId }], []], + [created] + ); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const putValidated = vi.fn(async () => undefined); + vi.mocked(getSessionAccessCacheDO).mockReturnValue({ putValidated } as never); + + const response = await makeApp().fetch(createRequest(v2Headers()), env); + + expect(response.status).toBe(200); + expect(forConditions).toHaveLength(1); + expect(queryOf(forConditions[0]).sql).toBe(expectedLockSql(v2OrganizationId)); + expect(queryOf(forConditions[0]).params).toEqual([ + memberRootId, + 'usr_test', + v2OrganizationId, + v2WorktreeId, + ]); + expect(insertedValues).toContainEqual( + expect.objectContaining({ + session_id: childSessionId, + parent_session_id: memberRootId, + organization_id: v2OrganizationId, + cloud_agent_session_scope_id: 'workspace_root', + cloud_agent_worktree_id: v2WorktreeId, + cloud_agent_session_id: null, + }) + ); + expect(putValidated).toHaveBeenCalledWith({ + sessionId: childSessionId, + organizationId: v2OrganizationId, + cloudAgentSessionScopeId: 'workspace_root', + }); + }); + + it('locks the member root before checking worktree deletion', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ + kind: 'ok', + parent: rootParentRow(), + memberRoot: { + sessionId: memberRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: 'workspace_root', + }, + memberRootScopeId: 'workspace_root', + }); + const { db, insertedValues, forConditions, selectConditions } = makeDb( + [[{ sessionId: memberRootId }], [{ started: '2026-01-01T00:00:00.000Z' }]], + [] + ); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch(createRequest(v2Headers()), env); + + expect(response.status).toBe(404); + expect(insertedValues).toHaveLength(0); + expect(forConditions).toHaveLength(1); + expect(queryOf(forConditions[0]).sql).toBe(expectedLockSql(v2OrganizationId)); + const lockIndex = selectConditions.indexOf(forConditions[0]); + const deletionIndex = selectConditions.findIndex(condition => + queryOf(condition).sql.includes('"cloud_agent_worktrees"') + ); + expect(lockIndex).toBeGreaterThanOrEqual(0); + expect(deletionIndex).toBeGreaterThan(lockIndex); + }); + + it('rejects an unmatched existing orphan even with the trusted lineage header', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ + kind: 'ok', + parent: rootParentRow(), + memberRoot: { + sessionId: memberRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: 'workspace_root', + }, + memberRootScopeId: 'workspace_root', + }); + const orphan = persistedRow({ + parent_session_id: null, + organization_id: null, + cloud_agent_session_scope_id: null, + cloud_agent_worktree_id: null, + }); + const { db, updateSets, updateConditions } = makeDb( + [[{ sessionId: memberRootId }], [], [orphan]], + [], + [] + ); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + createRequest(v2Headers({ [cloudAgentSessionScopeHeaders.trustedLineage]: '1' })), + env + ); + + expect(response.status).toBe(409); + expect(updateSets).toHaveLength(0); + expect(updateConditions).toHaveLength(0); + }); + + it('copies the actual parent child scope, not the serving member', async () => { + const parentChildId = 'ses_55555555555555555555555555'; + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ + kind: 'ok', + parent: { + sessionId: parentChildId, + parentSessionId: memberRootId, + organizationId: v2OrganizationId, + cloudAgentSessionId: null, + cloudAgentSessionScopeId: 'workspace_A', + cloudAgentWorktreeId: v2WorktreeId, + }, + memberRoot: { + sessionId: memberRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_A', + cloudAgentSessionScopeId: 'workspace_A', + }, + memberRootScopeId: 'workspace_A', + }); + const created = persistedRow({ + parent_session_id: parentChildId, + organization_id: v2OrganizationId, + cloud_agent_session_scope_id: 'workspace_A', + cloud_agent_worktree_id: v2WorktreeId, + }); + const { db, insertedValues } = makeDb([[{ sessionId: memberRootId }], []], [created]); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + const putValidated = vi.fn(async () => undefined); + vi.mocked(getSessionAccessCacheDO).mockReturnValue({ putValidated } as never); + + const response = await makeApp().fetch( + createRequest( + v2Headers({ + [cloudAgentSessionScopeHeaders.cloudAgentSessionId]: 'workspace_B', + }), + parentChildId + ), + env + ); + + expect(response.status).toBe(200); + expect(resolveScopedCreateParent).toHaveBeenCalledWith( + expect.anything(), + expect.objectContaining({ + scope: expect.objectContaining({ memberRootKiloSessionIds: [memberRootId] }), + }) + ); + expect(insertedValues).toContainEqual( + expect.objectContaining({ + session_id: childSessionId, + parent_session_id: parentChildId, + organization_id: v2OrganizationId, + cloud_agent_session_scope_id: 'workspace_A', + cloud_agent_worktree_id: v2WorktreeId, + }) + ); + }); + + it('denies v2 create when the parent is outside the asserted scope', async () => { + vi.mocked(resolveScopedCreateParent).mockResolvedValue({ kind: 'denied' }); + const { db, insertedValues } = makeDb([[]], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch(createRequest(v2Headers()), env); + + expect(response.status).toBe(409); + expect(insertedValues).toHaveLength(0); + }); + + it('allows v2 ingest only after the child predicate passes', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ + kind: 'ok', + rootScopeId: 'workspace_root', + }); + vi.mocked(handleDirectIngestRequest).mockResolvedValue({ + status: 200, + body: { success: true }, + } as never); + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/ingest?v=2`, { + method: 'POST', + headers: v2Headers(), + body: JSON.stringify({ data: [] }), + }), + env + ); + + expect(response.status).toBe(200); + expect(handleDirectIngestRequest).toHaveBeenCalledWith( + expect.objectContaining({ sessionId: childSessionId, kiloUserId: 'usr_test' }) + ); + }); + + it('denies v2 ingest when the child predicate fails', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ kind: 'denied' }); + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/ingest?v=2`, { + method: 'POST', + headers: v2Headers(), + body: JSON.stringify({ data: [] }), + }), + env + ); + + expect(response.status).toBe(404); + expect(handleDirectIngestRequest).not.toHaveBeenCalled(); + }); + + it('exports a child only after the child predicate passes', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ + kind: 'ok', + rootScopeId: 'workspace_root', + }); + vi.mocked(getSessionExport).mockResolvedValue({ + status: 'ok', + stream: new Blob(['{"events":[]}']).stream(), + }); + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/export`, { + method: 'GET', + headers: v2Headers(), + }), + env + ); + + expect(response.status).toBe(200); + expect(getSessionExport).toHaveBeenCalledWith(env, childSessionId, 'usr_test'); + }); + + it('denies v2 export when the child predicate fails without reading the export', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ kind: 'denied' }); + const { db } = makeDb([], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/export`, { + method: 'GET', + headers: v2Headers(), + }), + env + ); + + expect(response.status).toBe(404); + expect(getSessionExport).not.toHaveBeenCalled(); + }); + + it('renames a child only after the child predicate passes', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ + kind: 'ok', + rootScopeId: 'workspace_root', + }); + const existing = persistedRow({ + title: null, + cloud_agent_worktree_id: v2WorktreeId, + organization_id: v2OrganizationId, + }); + const { db, updateSets, updateConditions } = makeDb( + [[existing]], + [], + [{ ...existing, title: 'Renamed' }] + ); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/title`, { + method: 'POST', + headers: v2Headers(), + body: JSON.stringify({ title: 'Renamed', generated: false }), + }), + env + ); + + expect(response.status).toBe(200); + expect(updateSets).toContainEqual({ title: 'Renamed', updated_at: existing.updated_at }); + expect(updateConditions).toHaveLength(1); + expect(queryOf(updateConditions[0]).sql).toBe(expectedTitleWhereSql(v2OrganizationId)); + expect(queryOf(updateConditions[0]).params).toEqual([ + childSessionId, + 'usr_test', + 'usr_test', + v2OrganizationId, + v2WorktreeId, + 'workspace_root', + ]); + }); + + it('scopes a personal title update by owner, worktree, and root scope id', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ + kind: 'ok', + rootScopeId: 'workspace_root', + }); + const existing = persistedRow({ + title: null, + organization_id: null, + cloud_agent_worktree_id: v2WorktreeId, + }); + const { db, updateConditions } = makeDb([[existing]], [], [{ ...existing, title: 'Renamed' }]); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/title`, { + method: 'POST', + headers: v2Headers({ [cloudAgentSessionScopeV2Headers.expectedOrganization]: 'personal' }), + body: JSON.stringify({ title: 'Renamed', generated: false }), + }), + env + ); + + expect(response.status).toBe(200); + expect(updateConditions).toHaveLength(1); + expect(queryOf(updateConditions[0]).sql).toBe(expectedTitleWhereSql(null)); + expect(queryOf(updateConditions[0]).params).toEqual([ + childSessionId, + 'usr_test', + 'usr_test', + v2WorktreeId, + 'workspace_root', + ]); + }); + + it('denies v2 title when the child predicate fails', async () => { + vi.mocked(resolveScopedChildSession).mockResolvedValue({ kind: 'denied' }); + const { db, updateSets } = makeDb([], [], []); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request(`http://local/session/${childSessionId}/title`, { + method: 'POST', + headers: v2Headers(), + body: JSON.stringify({ title: 'Renamed', generated: false }), + }), + env + ); + + expect(response.status).toBe(404); + expect(updateSets).toHaveLength(0); + }); +}); diff --git a/services/session-ingest/src/routes/cloud-agent-session-scope.ts b/services/session-ingest/src/routes/cloud-agent-session-scope.ts index 57cc7ce0ce4..f0c4ff23d9b 100644 --- a/services/session-ingest/src/routes/cloud-agent-session-scope.ts +++ b/services/session-ingest/src/routes/cloud-agent-session-scope.ts @@ -7,14 +7,28 @@ import { cli_sessions_v2 } from '@kilocode/db/schema'; import { cloudAgentSessionScopeAssertionSchema, cloudAgentSessionScopeHeaders, + cloudAgentSessionScopeProtocolVersion, + cloudAgentSessionScopeProtocolVersionV2, + cloudAgentSessionScopeV2Headers, containedKiloSessionIdSchema, + parseCloudAgentSessionScopeV2Headers, + type CloudAgentSessionScopeAssertion, + type CloudAgentSessionScopeV2, } from '@kilocode/session-ingest-contracts'; import { hasOrganizationAccess, zodJsonValidator, withDORetry } from '@kilocode/worker-utils'; import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; import { handleDirectIngestRequest } from '../ingest/direct-ingest'; import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; +import { + resolveScopedChildSession, + resolveScopedCreateParent, + scopedSessionCondition, + worktreeScopeCondition, +} from '../services/cloud-agent-worktree-scope'; +import { getSessionExport } from '../services/session-export'; import { resolveAccessibleKiloSession } from '../services/session-access'; +import { applyReportedSessionTitle, reportSessionTitleSchema } from '../services/session-title'; import { canCreateCliSessionForUser } from '../services/user-session-admission'; import { isWorktreeDeleting } from '../services/worktree-deletion'; import type { ApiContext } from './api'; @@ -28,14 +42,35 @@ const createScopedSessionSchema = z const ingestVersionSchema = z.coerce.number().int().nonnegative().catch(0); +type ScopedCreateBody = z.infer; + export const cloudAgentSessionScopeApi = new Hono(); -function getSessionScopeAssertion(c: Context) { - return cloudAgentSessionScopeAssertionSchema.safeParse({ - cloudAgentSessionId: c.req.header(cloudAgentSessionScopeHeaders.cloudAgentSessionId), - rootKiloSessionId: c.req.header(cloudAgentSessionScopeHeaders.rootKiloSessionId), - protocolVersion: c.req.header(cloudAgentSessionScopeHeaders.protocolVersion), +type ScopedRequestAssertion = + | { kind: 'v1'; data: CloudAgentSessionScopeAssertion } + | { kind: 'v2'; data: CloudAgentSessionScopeV2 } + | { kind: 'invalid' }; + +function getScopeAssertion(c: Context): ScopedRequestAssertion { + const protocolVersion = c.req.header(cloudAgentSessionScopeHeaders.protocolVersion); + if (protocolVersion === cloudAgentSessionScopeProtocolVersion) { + const parsed = cloudAgentSessionScopeAssertionSchema.safeParse({ + cloudAgentSessionId: c.req.header(cloudAgentSessionScopeHeaders.cloudAgentSessionId), + rootKiloSessionId: c.req.header(cloudAgentSessionScopeHeaders.rootKiloSessionId), + protocolVersion, + }); + return parsed.success ? { kind: 'v1', data: parsed.data } : { kind: 'invalid' }; + } + if (protocolVersion !== cloudAgentSessionScopeProtocolVersionV2) return { kind: 'invalid' }; + + const parsed = parseCloudAgentSessionScopeV2Headers({ + worktreeId: c.req.header(cloudAgentSessionScopeV2Headers.worktreeId), + expectedOrganization: c.req.header(cloudAgentSessionScopeV2Headers.expectedOrganization), + memberRootKiloSessionIds: c.req.header( + cloudAgentSessionScopeV2Headers.memberRootKiloSessionIds + ), }); + return parsed.kind === 'ok' ? { kind: 'v2', data: parsed.scope } : { kind: 'invalid' }; } function getRequestBodyStream(request: Request): ReadableStream { @@ -59,14 +94,16 @@ function getOptionalExecutionContext(c: Context): ExecutionContext | } } -cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionSchema), async c => { - const assertion = getSessionScopeAssertion(c); - if (!assertion.success) { - return c.json({ success: false, error: 'Invalid Cloud Agent session scope assertion' }, 400); - } +function invalidAssertion(c: Context) { + return c.json({ success: false, error: 'Invalid Cloud Agent session scope assertion' }, 400); +} - const body = c.req.valid('json'); - if (body.sessionId === assertion.data.rootKiloSessionId) { +async function handleV1Create( + c: Context, + assertion: CloudAgentSessionScopeAssertion, + body: ScopedCreateBody +) { + if (body.sessionId === assertion.rootKiloSessionId) { return c.json({ success: false, error: 'Root session must use the public route' }, 400); } @@ -87,13 +124,13 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS .from(cli_sessions_v2) .where( and( - eq(cli_sessions_v2.session_id, assertion.data.rootKiloSessionId), + eq(cli_sessions_v2.session_id, assertion.rootKiloSessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId), - eq(cli_sessions_v2.cloud_agent_session_id, assertion.data.cloudAgentSessionId), + eq(cli_sessions_v2.cloud_agent_session_id, assertion.cloudAgentSessionId), isNull(cli_sessions_v2.parent_session_id), or( isNull(cli_sessions_v2.cloud_agent_session_scope_id), - eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.data.cloudAgentSessionId) + eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.cloudAgentSessionId) ) ) ) @@ -116,19 +153,19 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS if (root.cloudAgentSessionScopeId === null) { await tx .update(cli_sessions_v2) - .set({ cloud_agent_session_scope_id: assertion.data.cloudAgentSessionId }) + .set({ cloud_agent_session_scope_id: assertion.cloudAgentSessionId }) .where( and( - eq(cli_sessions_v2.session_id, assertion.data.rootKiloSessionId), + eq(cli_sessions_v2.session_id, assertion.rootKiloSessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId), isNull(cli_sessions_v2.cloud_agent_session_scope_id) ) ); } - const parentSessionId = body.parentSessionId ?? assertion.data.rootKiloSessionId; + const parentSessionId = body.parentSessionId ?? assertion.rootKiloSessionId; if (parentSessionId === body.sessionId) return { status: 'conflict' } as const; - if (parentSessionId !== assertion.data.rootKiloSessionId) { + if (parentSessionId !== assertion.rootKiloSessionId) { const [parent] = await tx .select({ sessionId: cli_sessions_v2.session_id }) .from(cli_sessions_v2) @@ -136,7 +173,7 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS and( eq(cli_sessions_v2.session_id, parentSessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId), - eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.data.cloudAgentSessionId), + eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.cloudAgentSessionId), root.organizationId === null ? isNull(cli_sessions_v2.organization_id) : eq(cli_sessions_v2.organization_id, root.organizationId), @@ -157,7 +194,7 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS parent_session_id: parentSessionId, organization_id: root.organizationId, cloud_agent_session_id: null, - cloud_agent_session_scope_id: assertion.data.cloudAgentSessionId, + cloud_agent_session_scope_id: assertion.cloudAgentSessionId, cloud_agent_worktree_id: root.worktreeId, }) .onConflictDoNothing({ @@ -194,7 +231,7 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS .update(cli_sessions_v2) .set({ parent_session_id: parentSessionId, - cloud_agent_session_scope_id: assertion.data.cloudAgentSessionId, + cloud_agent_session_scope_id: assertion.cloudAgentSessionId, cloud_agent_worktree_id: root.worktreeId ?? null, ...(existing.organization_id !== root.organizationId ? { organization_id: root.organizationId } @@ -216,7 +253,7 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS } if ( existing.parent_session_id === null || - existing.cloud_agent_session_scope_id !== assertion.data.cloudAgentSessionId || + existing.cloud_agent_session_scope_id !== assertion.cloudAgentSessionId || existing.organization_id !== root.organizationId || (existing.cloud_agent_worktree_id != null && existing.cloud_agent_worktree_id !== (root.worktreeId ?? null)) @@ -240,7 +277,7 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS and( eq(cli_sessions_v2.session_id, body.sessionId), eq(cli_sessions_v2.kilo_user_id, kiloUserId), - eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.data.cloudAgentSessionId), + eq(cli_sessions_v2.cloud_agent_session_scope_id, assertion.cloudAgentSessionId), isNull(cli_sessions_v2.cloud_agent_worktree_id) ) ) @@ -252,40 +289,219 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS return { status: 'existing', row: existing } as const; }); + return respondToV1Create(c, assertion, body.sessionId, result); +} + +async function handleV2Create( + c: Context, + assertion: CloudAgentSessionScopeV2, + body: ScopedCreateBody +) { + const parentSessionId = body.parentSessionId; + if (parentSessionId === undefined) { + return c.json({ success: false, error: 'parentSessionId is required' }, 400); + } + if ( + parentSessionId === body.sessionId || + assertion.memberRootKiloSessionIds.includes(body.sessionId) + ) { + return c.json({ success: false, error: 'session_scope_conflict' }, 409); + } + + const kiloUserId = c.get('user_id'); + const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); + const result = await db.transaction(async tx => { + if (!(await canCreateCliSessionForUser(tx, kiloUserId))) { + return { status: 'user_not_admitted' } as const; + } + + const parentResolution = await resolveScopedCreateParent(tx, { + kiloUserId, + parentSessionId, + scope: assertion, + }); + if (parentResolution.kind !== 'ok') return { status: 'parent_not_found' } as const; + const { parent, memberRoot, memberRootScopeId } = parentResolution; + const childScopeId = parent.cloudAgentSessionScopeId ?? memberRootScopeId; + + const [lockedRoot] = await tx + .select({ sessionId: cli_sessions_v2.session_id }) + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.session_id, memberRoot.sessionId), + worktreeScopeCondition(kiloUserId, assertion), + isNull(cli_sessions_v2.parent_session_id) + ) + ) + .limit(1) + .for('update'); + if (!lockedRoot) return { status: 'root_not_found' } as const; + + if (await isWorktreeDeleting(tx, assertion.worktreeId)) { + return { status: 'worktree_deleting' } as const; + } + + if (memberRoot.cloudAgentSessionScopeId === null && memberRoot.cloudAgentSessionId !== null) { + await tx + .update(cli_sessions_v2) + .set({ cloud_agent_session_scope_id: memberRoot.cloudAgentSessionId }) + .where( + and( + eq(cli_sessions_v2.session_id, memberRoot.sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId), + isNull(cli_sessions_v2.cloud_agent_session_scope_id) + ) + ); + } + + const [created] = await tx + .insert(cli_sessions_v2) + .values({ + session_id: body.sessionId, + kilo_user_id: kiloUserId, + parent_session_id: parentSessionId, + organization_id: parent.organizationId, + cloud_agent_session_id: null, + cloud_agent_session_scope_id: childScopeId, + cloud_agent_worktree_id: assertion.worktreeId, + }) + .onConflictDoNothing({ + target: [cli_sessions_v2.session_id, cli_sessions_v2.kilo_user_id], + }) + .returning(); + + if (created) return { status: 'created', row: created, scopeId: childScopeId } as const; + + const [existing] = await tx + .select() + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.session_id, body.sessionId), + eq(cli_sessions_v2.kilo_user_id, kiloUserId) + ) + ) + .limit(1) + .for('update'); + + if (!existing || existing.cloud_agent_session_id !== null) { + return { status: 'conflict' } as const; + } + + if ( + existing.parent_session_id === null || + existing.cloud_agent_session_scope_id !== childScopeId || + existing.organization_id !== parent.organizationId || + existing.cloud_agent_worktree_id !== assertion.worktreeId + ) { + return { status: 'conflict' } as const; + } + return { status: 'existing', row: existing, scopeId: childScopeId } as const; + }); + if (result.status === 'user_not_admitted') { return c.json({ success: false, error: 'User account not found' }, 403); } - if (result.status === 'root_not_found') { + if (result.status === 'root_not_found' || result.status === 'worktree_deleting') { return c.json({ success: false, error: 'session_not_found' }, 404); } - if (result.status === 'conflict') { + if (result.status === 'parent_not_found' || result.status === 'conflict') { return c.json({ success: false, error: 'session_scope_conflict' }, 409); } + await warmScopedSessionAccessCache(c, { + kiloUserId, + sessionId: body.sessionId, + organizationId: result.row.organization_id, + cloudAgentSessionScopeId: result.scopeId, + }); + + if (result.status === 'created') { + const session = mapSessionEventRow(result.row); + notifyUserSessionEvent( + c.env, + kiloUserId, + { + type: 'session.created', + data: { source: 'v2', session, changedAt: session.updatedAt }, + }, + getOptionalExecutionContext(c) + ); + } + + return c.json( + { + id: body.sessionId, + ingestPath: `/api/session/${body.sessionId}/ingest`, + }, + 200 + ); +} + +async function warmScopedSessionAccessCache( + c: Context, + params: { + kiloUserId: string; + sessionId: string; + organizationId: string | null; + cloudAgentSessionScopeId: string | null; + } +): Promise { try { await withDORetry( - () => getSessionAccessCacheDO(c.env, { kiloUserId }), + () => getSessionAccessCacheDO(c.env, { kiloUserId: params.kiloUserId }), sessionCache => sessionCache.putValidated({ - sessionId: body.sessionId, - organizationId: result.row.organization_id, - cloudAgentSessionScopeId: assertion.data.cloudAgentSessionId, + sessionId: params.sessionId, + organizationId: params.organizationId, + cloudAgentSessionScopeId: params.cloudAgentSessionScopeId, }), 'SessionAccessCacheDO.putValidated' ); } catch (error) { console.error('Failed to warm session access cache after scoped session bootstrap', { - kiloUserId, - sessionId: body.sessionId, + kiloUserId: params.kiloUserId, + sessionId: params.sessionId, error: error instanceof Error ? error.message : String(error), }); } +} + +async function respondToV1Create( + c: Context, + assertion: CloudAgentSessionScopeAssertion, + sessionId: string, + result: + | { status: 'user_not_admitted' } + | { status: 'root_not_found' } + | { status: 'conflict' } + | { status: 'created'; row: typeof cli_sessions_v2.$inferSelect } + | { status: 'adopted'; row: typeof cli_sessions_v2.$inferSelect } + | { status: 'existing'; row: typeof cli_sessions_v2.$inferSelect } +) { + if (result.status === 'user_not_admitted') { + return c.json({ success: false, error: 'User account not found' }, 403); + } + if (result.status === 'root_not_found') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + if (result.status === 'conflict') { + return c.json({ success: false, error: 'session_scope_conflict' }, 409); + } + + await warmScopedSessionAccessCache(c, { + kiloUserId: c.get('user_id'), + sessionId, + organizationId: result.row.organization_id, + cloudAgentSessionScopeId: assertion.cloudAgentSessionId, + }); if (result.status === 'created' || result.status === 'adopted') { const session = mapSessionEventRow(result.row); notifyUserSessionEvent( c.env, - kiloUserId, + c.get('user_id'), { type: result.status === 'created' ? 'session.created' : 'session.updated', data: { source: 'v2', session, changedAt: session.updatedAt }, @@ -296,21 +512,16 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS return c.json( { - id: body.sessionId, - ingestPath: `/api/session/${body.sessionId}/ingest`, + id: sessionId, + ingestPath: `/api/session/${sessionId}/ingest`, }, 200 ); -}); - -cloudAgentSessionScopeApi.post('/session/:sessionId/ingest', async c => { - const assertion = getSessionScopeAssertion(c); - if (!assertion.success) { - return c.json({ success: false, error: 'Invalid Cloud Agent session scope assertion' }, 400); - } +} +async function handleV1Ingest(c: Context, assertion: CloudAgentSessionScopeAssertion) { const sessionId = containedKiloSessionIdSchema.safeParse(c.req.param('sessionId')); - if (!sessionId.success || sessionId.data === assertion.data.rootKiloSessionId) { + if (!sessionId.success || sessionId.data === assertion.rootKiloSessionId) { return c.json({ success: false, error: 'Invalid child sessionId' }, 400); } @@ -318,7 +529,7 @@ cloudAgentSessionScopeApi.post('/session/:sessionId/ingest', async c => { const accessibleSession = await resolveAccessibleKiloSession(c.env, { kiloUserId, kiloSessionId: sessionId.data, - expectedCloudAgentSessionScopeId: assertion.data.cloudAgentSessionId, + expectedCloudAgentSessionScopeId: assertion.cloudAgentSessionId, }); if (!accessibleSession) { return c.json({ success: false, error: 'session_not_found' }, 404); @@ -337,4 +548,138 @@ cloudAgentSessionScopeApi.post('/session/:sessionId/ingest', async c => { }); return c.json(result.body, result.status); +} + +async function handleV2Ingest(c: Context, assertion: CloudAgentSessionScopeV2) { + const sessionId = containedKiloSessionIdSchema.safeParse(c.req.param('sessionId')); + if (!sessionId.success) { + return c.json({ success: false, error: 'Invalid child sessionId' }, 400); + } + + const kiloUserId = c.get('user_id'); + const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); + const resolution = await resolveScopedChildSession(db, { + kiloUserId, + sessionId: sessionId.data, + scope: assertion, + }); + if (resolution.kind !== 'ok') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + + const result = await handleDirectIngestRequest({ + env: c.env, + body: getRequestBodyStream(c.req.raw), + contentLength: c.req.header('content-length'), + kiloUserId, + sessionId: sessionId.data, + ingestVersion: ingestVersionSchema.parse(c.req.query('v') ?? 0), + ingestedAt: Date.now(), + ingestRequestId: crypto.randomUUID(), + executionContext: getOptionalExecutionContext(c), + }); + + return c.json(result.body, result.status); +} + +cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionSchema), async c => { + const assertion = getScopeAssertion(c); + if (assertion.kind === 'invalid') return invalidAssertion(c); + const body = c.req.valid('json'); + if (assertion.kind === 'v1') return handleV1Create(c, assertion.data, body); + return handleV2Create(c, assertion.data, body); +}); + +cloudAgentSessionScopeApi.post('/session/:sessionId/ingest', async c => { + const assertion = getScopeAssertion(c); + if (assertion.kind === 'invalid') return invalidAssertion(c); + if (assertion.kind === 'v1') return handleV1Ingest(c, assertion.data); + return handleV2Ingest(c, assertion.data); +}); + +cloudAgentSessionScopeApi.get('/session/:sessionId/export', async c => { + const assertion = getScopeAssertion(c); + if (assertion.kind !== 'v2') { + return c.json({ success: false, error: 'Cloud Agent session scope protocol v2 required' }, 400); + } + + const sessionId = containedKiloSessionIdSchema.safeParse(c.req.param('sessionId')); + if (!sessionId.success) { + return c.json({ success: false, error: 'Invalid sessionId' }, 400); + } + + const kiloUserId = c.get('user_id'); + const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); + const resolution = await resolveScopedChildSession(db, { + kiloUserId, + sessionId: sessionId.data, + scope: assertion.data, + }); + if (resolution.kind !== 'ok') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + + const exported = await getSessionExport(c.env, sessionId.data, kiloUserId); + if (exported.status !== 'ok') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + return c.body(exported.stream, 200, { + 'content-type': 'application/json; charset=utf-8', + }); +}); + +cloudAgentSessionScopeApi.post('/session/:sessionId/title', async c => { + const assertion = getScopeAssertion(c); + if (assertion.kind !== 'v2') { + return c.json({ success: false, error: 'Cloud Agent session scope protocol v2 required' }, 400); + } + + const sessionId = containedKiloSessionIdSchema.safeParse(c.req.param('sessionId')); + if (!sessionId.success) { + return c.json({ success: false, error: 'Invalid sessionId' }, 400); + } + + const bodyParse = reportSessionTitleSchema.safeParse(await c.req.json().catch(() => null)); + if (!bodyParse.success) { + return c.json({ success: false, error: 'Invalid body', issues: bodyParse.error.issues }, 400); + } + + const kiloUserId = c.get('user_id'); + const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); + const resolution = await resolveScopedChildSession(db, { + kiloUserId, + sessionId: sessionId.data, + scope: assertion.data, + }); + if (resolution.kind !== 'ok') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + + const scope = scopedSessionCondition(kiloUserId, assertion.data, resolution.rootScopeId); + + const result = await applyReportedSessionTitle(db, { + sessionId: sessionId.data, + kiloUserId, + input: bodyParse.data, + scope, + }); + if (result.kind === 'not_found') { + return c.json({ success: false, error: 'session_not_found' }, 404); + } + if (result.kind === 'refused') { + return c.json({ applied: false }, 200); + } + + const session = mapSessionEventRow(result.row); + notifyUserSessionEvent( + c.env, + kiloUserId, + { + type: 'session.updated', + data: { source: 'v2', session, changedAt: session.updatedAt }, + }, + getOptionalExecutionContext(c) + ); + + return c.json({ title: session.title ?? bodyParse.data.title, applied: true }, 200); }); diff --git a/services/session-ingest/src/services/cloud-agent-worktree-scope.test.ts b/services/session-ingest/src/services/cloud-agent-worktree-scope.test.ts new file mode 100644 index 00000000000..d9a2ba78786 --- /dev/null +++ b/services/session-ingest/src/services/cloud-agent-worktree-scope.test.ts @@ -0,0 +1,433 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as WorkerUtils from '@kilocode/worker-utils'; +import type { SQL } from 'drizzle-orm'; +import { PgDialect } from 'drizzle-orm/pg-core'; + +const workerUtils = vi.hoisted(() => ({ hasOrganizationAccess: vi.fn() })); +vi.mock('@kilocode/worker-utils', async importOriginal => ({ + ...(await importOriginal()), + hasOrganizationAccess: workerUtils.hasOrganizationAccess, +})); + +import type { CloudAgentSessionScopeV2 } from '@kilocode/session-ingest-contracts'; +import { + encodeCloudAgentSessionScopeV2Headers, + parseCloudAgentSessionScopeV2Headers, +} from '@kilocode/session-ingest-contracts'; +import { + CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH, + resolveScopedChildSession, + resolveScopedCreateParent, + type AncestryNode, +} from './cloud-agent-worktree-scope'; + +const userId = 'usr_test'; +const worktreeId = 'worktree_11111111-1111-4111-8111-111111111111'; +const organizationId = '33333333-3333-4333-8333-333333333333'; +const memberRootId = 'ses_11111111111111111111111111'; +const otherRootId = 'ses_22222222222222222222222222'; +const childId = 'ses_33333333333333333333333333'; +const grandchildId = 'ses_44444444444444444444444444'; + +const scope: CloudAgentSessionScopeV2 = { + worktreeId, + organizationId, + memberRootKiloSessionIds: [memberRootId], +}; + +function node(overrides: Partial & Pick): AncestryNode { + return { + parentSessionId: null, + cloudAgentSessionId: null, + cloudAgentSessionScopeId: null, + ...overrides, + }; +} + +function childRow(overrides: Record = {}) { + return { + sessionId: childId, + parentSessionId: memberRootId, + organizationId, + cloudAgentSessionId: null, + cloudAgentSessionScopeId: 'workspace_scope', + cloudAgentWorktreeId: worktreeId, + ...overrides, + }; +} + +function makeScopeDb() { + const selectConditions: SQL[] = []; + const selectResults: unknown[][] = []; + const ancestryQueries: SQL[] = []; + + function makeSelect() { + const chain = { + from: () => chain, + where: (condition: SQL) => { + selectConditions.push(condition); + return chain; + }, + limit: () => chain, + for: () => chain, + orderBy: () => chain, + then: (resolve: (rows: unknown[]) => unknown) => resolve(selectResults.shift() ?? []), + }; + return chain; + } + + return { + db: { + select: vi.fn(makeSelect), + execute: vi.fn(async (query: SQL) => { + ancestryQueries.push(query); + return { rows: selectResults.shift() ?? [] }; + }), + } as never, + selectConditions, + ancestryQueries, + pushRow: (row: unknown | null) => selectResults.push(row === null ? [] : [row]), + pushRows: (...rows: unknown[]) => selectResults.push(rows), + }; +} + +function queryOf(condition: SQL) { + return new PgDialect().sqlToQuery(condition); +} + +function expectedScopedLookupSql(organizationId: string | null): string { + return organizationId === null + ? `("cli_sessions_v2"."session_id" = $1 and ("cli_sessions_v2"."kilo_user_id" = $2 and "cli_sessions_v2"."organization_id" is null and "cli_sessions_v2"."cloud_agent_worktree_id" = $3))` + : `("cli_sessions_v2"."session_id" = $1 and ("cli_sessions_v2"."kilo_user_id" = $2 and "cli_sessions_v2"."organization_id" = $3 and "cli_sessions_v2"."cloud_agent_worktree_id" = $4))`; +} + +describe('resolveScopedChildSession', () => { + beforeEach(() => { + vi.resetAllMocks(); + workerUtils.hasOrganizationAccess.mockResolvedValue(true); + }); + + it('scopes the target lookup by owner, asserted organization, and worktree', async () => { + const fake = makeScopeDb(); + fake.pushRow(null); + + await resolveScopedChildSession(fake.db, { kiloUserId: userId, sessionId: childId, scope }); + + expect(fake.selectConditions).toHaveLength(1); + const query = queryOf(fake.selectConditions[0]); + expect(query.sql).toBe(expectedScopedLookupSql(organizationId)); + expect(query.params).toEqual([childId, userId, organizationId, worktreeId]); + }); + + it('encodes a personal asserted organization as IS NULL, not any organization', async () => { + const fake = makeScopeDb(); + fake.pushRow(null); + + await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope: { ...scope, organizationId: null }, + }); + + const query = queryOf(fake.selectConditions[0]); + expect(query.sql).toBe(expectedScopedLookupSql(null)); + expect(query.params).toEqual([childId, userId, worktreeId]); + }); + + it('denies a root row on the scoped path without walking ancestry', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow({ parentSessionId: null })); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope, + }); + expect(result.kind).toBe('denied'); + expect(fake.selectConditions).toHaveLength(1); + }); + + it('denies before reading the row when the asserted organization has no current membership', async () => { + workerUtils.hasOrganizationAccess.mockResolvedValue(false); + const fake = makeScopeDb(); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope, + }); + expect(result.kind).toBe('denied'); + expect(fake.selectConditions).toHaveLength(0); + }); + + it('authorizes a child whose bounded walk reaches an asserted member root', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow({ cloudAgentSessionScopeId: 'workspace_root' })); + fake.pushRow( + node({ + sessionId: memberRootId, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: null, + }) + ); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope, + }); + expect(result.kind).toBe('ok'); + if (result.kind !== 'ok') return; + expect(result.rootScopeId).toBe('workspace_root'); + expect(fake.selectConditions).toHaveLength(1); + expect(fake.ancestryQueries).toHaveLength(1); + }); + + it('authorizes a grandchild through an intermediate child', async () => { + const fake = makeScopeDb(); + fake.pushRow( + childRow({ + sessionId: grandchildId, + parentSessionId: childId, + cloudAgentSessionScopeId: 'workspace_scope', + }) + ); + fake.pushRow( + node({ + sessionId: memberRootId, + cloudAgentSessionId: 'workspace_scope', + cloudAgentSessionScopeId: 'workspace_scope', + }) + ); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: grandchildId, + scope, + }); + expect(result.kind).toBe('ok'); + }); + + it('denies when the walk reaches a root that is not an asserted member', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow()); + fake.pushRow(node({ sessionId: otherRootId, cloudAgentSessionId: 'workspace_other' })); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope, + }); + expect(result.kind).toBe('denied'); + }); + + it('denies when the child scope id does not match the reached member root scope', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow({ cloudAgentSessionScopeId: 'workspace_stale' })); + fake.pushRow( + node({ + sessionId: memberRootId, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: 'workspace_root', + }) + ); + const result = await resolveScopedChildSession(fake.db, { + kiloUserId: userId, + sessionId: childId, + scope, + }); + expect(result.kind).toBe('denied'); + }); +}); + +describe('resolveScopedCreateParent', () => { + beforeEach(() => { + vi.resetAllMocks(); + workerUtils.hasOrganizationAccess.mockResolvedValue(true); + }); + + it('resolves a member-root parent and uses the root cloud agent session id as scope', async () => { + const fake = makeScopeDb(); + fake.pushRow( + childRow({ + sessionId: memberRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_root', + cloudAgentSessionScopeId: null, + }) + ); + const result = await resolveScopedCreateParent(fake.db, { + kiloUserId: userId, + parentSessionId: memberRootId, + scope, + }); + expect(result.kind).toBe('ok'); + if (result.kind !== 'ok') return; + expect(result.memberRootScopeId).toBe('workspace_root'); + expect(result.parent.cloudAgentSessionScopeId).toBeNull(); + }); + + it('resolves a child parent only when its walk reaches an asserted member root', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow({ sessionId: childId, cloudAgentSessionScopeId: 'workspace_scope' })); + fake.pushRow( + node({ + sessionId: memberRootId, + cloudAgentSessionId: 'workspace_scope', + cloudAgentSessionScopeId: 'workspace_scope', + }) + ); + const result = await resolveScopedCreateParent(fake.db, { + kiloUserId: userId, + parentSessionId: childId, + scope, + }); + expect(result.kind).toBe('ok'); + if (result.kind !== 'ok') return; + expect(result.memberRoot.sessionId).toBe(memberRootId); + }); + + it('denies a root parent that is not an asserted member root', async () => { + const fake = makeScopeDb(); + fake.pushRow( + childRow({ + sessionId: otherRootId, + parentSessionId: null, + cloudAgentSessionId: 'workspace_other', + }) + ); + const result = await resolveScopedCreateParent(fake.db, { + kiloUserId: userId, + parentSessionId: otherRootId, + scope, + }); + expect(result.kind).toBe('denied'); + }); + + it('denies when the scoped parent lookup returns no owned row (cross-worktree, org, or user)', async () => { + const fake = makeScopeDb(); + fake.pushRow(null); + const result = await resolveScopedCreateParent(fake.db, { + kiloUserId: userId, + parentSessionId: childId, + scope, + }); + expect(result.kind).toBe('denied'); + const query = queryOf(fake.selectConditions[0]); + expect(query.sql).toBe(expectedScopedLookupSql(organizationId)); + expect(query.params).toEqual([childId, userId, organizationId, worktreeId]); + }); + + it('scopes every ancestry hop by owner, asserted organization, and worktree', async () => { + const fake = makeScopeDb(); + fake.pushRow(childRow({ sessionId: grandchildId, parentSessionId: childId })); + fake.pushRow( + node({ + sessionId: memberRootId, + cloudAgentSessionId: 'workspace_scope', + cloudAgentSessionScopeId: 'workspace_scope', + }) + ); + await resolveScopedCreateParent(fake.db, { + kiloUserId: userId, + parentSessionId: grandchildId, + scope, + }); + expect(fake.selectConditions).toHaveLength(1); + expect(fake.ancestryQueries).toHaveLength(1); + const query = queryOf(fake.ancestryQueries[0]); + expect(query.sql).toContain('WITH RECURSIVE ancestry'); + expect(query.sql).toContain('parent.kilo_user_id ='); + expect(query.sql).toContain('parent.organization_id ='); + expect(query.sql).toContain('parent.cloud_agent_worktree_id ='); + expect(query.sql).toContain('NOT (parent.session_id = ANY(ancestry.path))'); + expect(query.params).toEqual([ + childId, + userId, + organizationId, + worktreeId, + userId, + organizationId, + worktreeId, + CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH, + ]); + }); +}); + +describe('Cloud Agent session scope v2 header codec', () => { + it('round-trips an organization scope', () => { + const input: CloudAgentSessionScopeV2 = { + worktreeId, + organizationId, + memberRootKiloSessionIds: [memberRootId, otherRootId], + }; + expect( + parseCloudAgentSessionScopeV2Headers(encodeCloudAgentSessionScopeV2Headers(input)) + ).toEqual({ kind: 'ok', scope: input }); + }); + + it('round-trips a personal scope as the explicit personal sentinel', () => { + const input: CloudAgentSessionScopeV2 = { + worktreeId, + organizationId: null, + memberRootKiloSessionIds: [memberRootId], + }; + const headers = encodeCloudAgentSessionScopeV2Headers(input); + expect(headers.expectedOrganization).toBe('personal'); + expect(parseCloudAgentSessionScopeV2Headers(headers)).toEqual({ kind: 'ok', scope: input }); + }); + + it('de-duplicates repeated member roots', () => { + const parsed = parseCloudAgentSessionScopeV2Headers({ + worktreeId, + expectedOrganization: organizationId, + memberRootKiloSessionIds: `${memberRootId},${memberRootId}`, + }); + expect(parsed).toEqual({ + kind: 'ok', + scope: { + worktreeId, + organizationId, + memberRootKiloSessionIds: [memberRootId], + }, + }); + }); + + it.each([ + [ + 'a missing worktree id', + { expectedOrganization: organizationId, memberRootKiloSessionIds: memberRootId }, + ], + ['a missing organization', { worktreeId, memberRootKiloSessionIds: memberRootId }], + ['missing member roots', { worktreeId, expectedOrganization: organizationId }], + [ + 'a capitalized personal sentinel', + { worktreeId, expectedOrganization: 'Personal', memberRootKiloSessionIds: memberRootId }, + ], + [ + 'an empty organization', + { worktreeId, expectedOrganization: '', memberRootKiloSessionIds: memberRootId }, + ], + [ + 'a non-uuid organization', + { worktreeId, expectedOrganization: 'org-a', memberRootKiloSessionIds: memberRootId }, + ], + [ + 'an empty member-root entry', + { worktreeId, expectedOrganization: organizationId, memberRootKiloSessionIds: 'a,,b' }, + ], + [ + 'an empty member-root list', + { worktreeId, expectedOrganization: organizationId, memberRootKiloSessionIds: '' }, + ], + ])('rejects %s', (_label, input) => { + expect(parseCloudAgentSessionScopeV2Headers(input)).toEqual({ kind: 'invalid' }); + }); + + it('rejects a header block over 4096 bytes even when it de-duplicates to one member root', () => { + const repeated = Array.from({ length: 140 }, () => memberRootId).join(','); + expect(repeated.length).toBeGreaterThan(4096); + expect( + parseCloudAgentSessionScopeV2Headers({ + worktreeId, + expectedOrganization: organizationId, + memberRootKiloSessionIds: repeated, + }) + ).toEqual({ kind: 'invalid' }); + }); +}); diff --git a/services/session-ingest/src/services/cloud-agent-worktree-scope.ts b/services/session-ingest/src/services/cloud-agent-worktree-scope.ts new file mode 100644 index 00000000000..ab281a4ed44 --- /dev/null +++ b/services/session-ingest/src/services/cloud-agent-worktree-scope.ts @@ -0,0 +1,187 @@ +import { and, eq, isNull, sql, type SQL } from 'drizzle-orm'; +import { z } from 'zod'; + +import type { WorkerDb } from '@kilocode/db/client'; +import { cli_sessions_v2 } from '@kilocode/db/schema'; +import type { CloudAgentSessionScopeV2 } from '@kilocode/session-ingest-contracts'; +import { hasOrganizationAccess } from '@kilocode/worker-utils'; + +export const CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH = 32; + +type WorktreeScopeDb = Pick; + +export type AncestryNode = { + sessionId: string; + parentSessionId: string | null; + cloudAgentSessionId: string | null; + cloudAgentSessionScopeId: string | null; +}; + +export type ScopedSessionRow = AncestryNode & { + organizationId: string | null; + cloudAgentWorktreeId: string | null; +}; + +const SCOPED_SESSION_COLUMNS = { + sessionId: cli_sessions_v2.session_id, + parentSessionId: cli_sessions_v2.parent_session_id, + organizationId: cli_sessions_v2.organization_id, + cloudAgentSessionId: cli_sessions_v2.cloud_agent_session_id, + cloudAgentSessionScopeId: cli_sessions_v2.cloud_agent_session_scope_id, + cloudAgentWorktreeId: cli_sessions_v2.cloud_agent_worktree_id, +} as const; + +export function worktreeScopeCondition( + kiloUserId: string, + scope: CloudAgentSessionScopeV2 +): SQL | undefined { + return and( + eq(cli_sessions_v2.kilo_user_id, kiloUserId), + scope.organizationId === null + ? isNull(cli_sessions_v2.organization_id) + : eq(cli_sessions_v2.organization_id, scope.organizationId), + eq(cli_sessions_v2.cloud_agent_worktree_id, scope.worktreeId) + ); +} + +export function scopedSessionCondition( + kiloUserId: string, + scope: CloudAgentSessionScopeV2, + rootScopeId: string +): SQL | undefined { + return and( + worktreeScopeCondition(kiloUserId, scope), + eq(cli_sessions_v2.cloud_agent_session_scope_id, rootScopeId) + ); +} + +const ancestryRootSchema = z.object({ + sessionId: z.string(), + parentSessionId: z.null(), + cloudAgentSessionId: z.string().nullable(), + cloudAgentSessionScopeId: z.string().nullable(), +}); + +async function resolveScopedAncestryRoot( + db: WorktreeScopeDb, + kiloUserId: string, + scope: CloudAgentSessionScopeV2, + firstParentSessionId: string +): Promise { + // Keep every hop inside the asserted scope, including intermediate children. + const result = await db.execute(sql` + WITH RECURSIVE ancestry AS ( + SELECT session_id, parent_session_id, cloud_agent_session_id, cloud_agent_session_scope_id, + 1 AS depth, ARRAY[session_id] AS path + FROM ${cli_sessions_v2} + WHERE ${and( + eq(cli_sessions_v2.session_id, firstParentSessionId), + worktreeScopeCondition(kiloUserId, scope) + )} + UNION ALL + SELECT parent.session_id, parent.parent_session_id, + parent.cloud_agent_session_id, parent.cloud_agent_session_scope_id, + ancestry.depth + 1, ancestry.path || parent.session_id + FROM ${cli_sessions_v2} parent + INNER JOIN ancestry ON parent.session_id = ancestry.parent_session_id + WHERE parent.kilo_user_id = ${kiloUserId} + AND ${scope.organizationId === null ? sql`parent.organization_id IS NULL` : sql`parent.organization_id = ${scope.organizationId}`} + AND parent.cloud_agent_worktree_id = ${scope.worktreeId} + AND ancestry.depth < ${CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH} + AND NOT (parent.session_id = ANY(ancestry.path)) + ) + SELECT session_id AS "sessionId", parent_session_id AS "parentSessionId", + cloud_agent_session_id AS "cloudAgentSessionId", + cloud_agent_session_scope_id AS "cloudAgentSessionScopeId" + FROM ancestry WHERE parent_session_id IS NULL + `); + const row = result.rows[0]; + if (row === undefined) return null; + const root = ancestryRootSchema.parse(row); + return scope.memberRootKiloSessionIds.includes(root.sessionId) ? root : null; +} + +export type ScopedChildSessionResolution = { kind: 'denied' } | { kind: 'ok'; rootScopeId: string }; + +export async function resolveScopedChildSession( + db: WorktreeScopeDb, + params: { kiloUserId: string; sessionId: string; scope: CloudAgentSessionScopeV2 } +): Promise { + const { kiloUserId, sessionId, scope } = params; + if ( + scope.organizationId !== null && + !(await hasOrganizationAccess(db, { kiloUserId, organizationId: scope.organizationId })) + ) { + return { kind: 'denied' }; + } + + const [row] = await db + .select(SCOPED_SESSION_COLUMNS) + .from(cli_sessions_v2) + .where( + and(eq(cli_sessions_v2.session_id, sessionId), worktreeScopeCondition(kiloUserId, scope)) + ) + .limit(1); + if (!row || row.parentSessionId === null || row.cloudAgentSessionScopeId === null) { + return { kind: 'denied' }; + } + + const memberRoot = await resolveScopedAncestryRoot(db, kiloUserId, scope, row.parentSessionId); + if (!memberRoot) return { kind: 'denied' }; + + const rootScopeId = memberRoot.cloudAgentSessionScopeId ?? memberRoot.cloudAgentSessionId; + if (rootScopeId === null || row.cloudAgentSessionScopeId !== rootScopeId) { + return { kind: 'denied' }; + } + return { kind: 'ok', rootScopeId }; +} + +export type ScopedCreateParentResolution = + | { kind: 'denied' } + | { kind: 'ok'; parent: ScopedSessionRow; memberRoot: AncestryNode; memberRootScopeId: string }; + +export async function resolveScopedCreateParent( + db: WorktreeScopeDb, + params: { kiloUserId: string; parentSessionId: string; scope: CloudAgentSessionScopeV2 } +): Promise { + const { kiloUserId, parentSessionId, scope } = params; + if ( + scope.organizationId !== null && + !(await hasOrganizationAccess(db, { kiloUserId, organizationId: scope.organizationId })) + ) { + return { kind: 'denied' }; + } + + const [parent] = await db + .select(SCOPED_SESSION_COLUMNS) + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.session_id, parentSessionId), + worktreeScopeCondition(kiloUserId, scope) + ) + ) + .limit(1); + if (!parent) return { kind: 'denied' }; + + const memberRoots = new Set(scope.memberRootKiloSessionIds); + let memberRoot: AncestryNode; + if (parent.parentSessionId === null) { + if (!memberRoots.has(parent.sessionId)) return { kind: 'denied' }; + memberRoot = parent; + } else { + const root = await resolveScopedAncestryRoot(db, kiloUserId, scope, parent.parentSessionId); + if (!root) return { kind: 'denied' }; + memberRoot = root; + } + + const memberRootScopeId = memberRoot.cloudAgentSessionScopeId ?? memberRoot.cloudAgentSessionId; + if (memberRootScopeId === null) return { kind: 'denied' }; + if ( + parent.cloudAgentSessionScopeId !== null && + parent.cloudAgentSessionScopeId !== memberRootScopeId + ) { + return { kind: 'denied' }; + } + return { kind: 'ok', parent, memberRoot, memberRootScopeId }; +} diff --git a/services/session-ingest/src/services/session-title.ts b/services/session-ingest/src/services/session-title.ts new file mode 100644 index 00000000000..f363012ce10 --- /dev/null +++ b/services/session-ingest/src/services/session-title.ts @@ -0,0 +1,59 @@ +import { and, eq, isNull, type SQL } from 'drizzle-orm'; +import { z } from 'zod'; + +import type { WorkerDb } from '@kilocode/db/client'; +import { cli_sessions_v2 } from '@kilocode/db/schema'; + +import { isDefaultSessionTitle } from '../ingest/default-session-title'; + +export const reportSessionTitleSchema = z.object({ + title: z.string().trim().min(1).max(200), + generated: z.boolean(), +}); + +export type ReportedSessionTitleInput = z.infer; + +export type ReportedSessionTitleResult = + | { kind: 'not_found' } + | { kind: 'refused' } + | { kind: 'applied'; row: typeof cli_sessions_v2.$inferSelect }; + +export async function applyReportedSessionTitle( + db: Pick, + params: { + sessionId: string; + kiloUserId: string; + input: ReportedSessionTitleInput; + scope?: SQL; + } +): Promise { + const ownership = and( + eq(cli_sessions_v2.session_id, params.sessionId), + eq(cli_sessions_v2.kilo_user_id, params.kiloUserId), + params.scope + ); + + const [existing] = await db.select().from(cli_sessions_v2).where(ownership).limit(1); + if (!existing) return { kind: 'not_found' }; + + const { title, generated } = params.input; + if (generated && !isDefaultSessionTitle(existing.title)) return { kind: 'refused' }; + + const writeWhere = generated + ? and( + ownership, + existing.title == null + ? isNull(cli_sessions_v2.title) + : eq(cli_sessions_v2.title, existing.title) + ) + : ownership; + + const [updatedRow] = await db + .update(cli_sessions_v2) + .set({ title, updated_at: existing.updated_at }) + .where(writeWhere) + .returning(); + + if (!updatedRow) return generated ? { kind: 'refused' } : { kind: 'not_found' }; + return { kind: 'applied', row: updatedRow }; +} diff --git a/services/session-ingest/test/postgres/cloud-agent-worktree-scope.test.ts b/services/session-ingest/test/postgres/cloud-agent-worktree-scope.test.ts new file mode 100644 index 00000000000..9ac164a3f2c --- /dev/null +++ b/services/session-ingest/test/postgres/cloud-agent-worktree-scope.test.ts @@ -0,0 +1,197 @@ +import { afterAll, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest'; +import { createDrizzleClient } from '@kilocode/db/client'; +import type { CloudAgentSessionScopeV2 } from '@kilocode/session-ingest-contracts'; +import { sql } from 'drizzle-orm'; +import { + CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH, + resolveScopedChildSession, + resolveScopedCreateParent, +} from '../../src/services/cloud-agent-worktree-scope'; + +vi.mock('@kilocode/worker-utils', () => ({ hasOrganizationAccess: async () => true })); + +const userId = 'oauth/scope-test'; +const organizationId = '33333333-3333-4333-8333-333333333333'; +const scope: CloudAgentSessionScopeV2 = { + worktreeId: 'worktree_test', + organizationId: null, + memberRootKiloSessionIds: ['root'], +}; +let client: ReturnType; + +beforeAll(async () => { + const connectionString = process.env.POSTGRES_URL; + if (!connectionString) throw new Error('Set POSTGRES_URL to a local test database'); + if (!['localhost', '127.0.0.1', '[::1]'].includes(new URL(connectionString).hostname)) { + throw new Error('The scope tests require a local PostgreSQL database'); + } + client = createDrizzleClient({ connectionString, poolConfig: { max: 1 }, ssl: false }); + // A session-local table tests the real queries without changing shared data. + await client.db.execute(sql` + CREATE TEMP TABLE cli_sessions_v2 ( + session_id text NOT NULL, + kilo_user_id text NOT NULL, + parent_session_id text, + organization_id uuid, + cloud_agent_session_id text, + cloud_agent_session_scope_id text, + cloud_agent_worktree_id text, + PRIMARY KEY (session_id, kilo_user_id) + ) + `); +}); + +afterAll(async () => { + if (client) await client.pool.end(); +}); + +beforeEach(async () => { + vi.restoreAllMocks(); + await client.db.execute(sql`TRUNCATE pg_temp.cli_sessions_v2`); +}); + +async function insertSession( + sessionId: string, + parentSessionId: string | null, + overrides: { + userId?: string; + organizationId?: string | null; + worktreeId?: string; + rootScopeId?: string | null; + } = {} +) { + const rootScopeId = + overrides.rootScopeId === undefined ? 'workspace_root' : overrides.rootScopeId; + await client.db.execute(sql` + INSERT INTO pg_temp.cli_sessions_v2 ( + session_id, kilo_user_id, parent_session_id, organization_id, + cloud_agent_session_id, cloud_agent_session_scope_id, cloud_agent_worktree_id + ) VALUES ( + ${sessionId}, ${overrides.userId ?? userId}, ${parentSessionId}, + ${overrides.organizationId ?? null}, ${parentSessionId === null ? 'workspace_root' : null}, + ${rootScopeId}, ${overrides.worktreeId ?? scope.worktreeId} + ) + `); +} + +function resolveChild(assertedScope = scope) { + return resolveScopedChildSession(client.db, { + kiloUserId: userId, + sessionId: 'child', + scope: assertedScope, + }); +} + +describe('scoped ancestry PostgreSQL queries', () => { + it('resolves a child and a create parent through the same bounded ancestry query', async () => { + await insertSession('root', null, { rootScopeId: null }); + await insertSession('parent', 'root'); + await insertSession('child', 'parent'); + expect(await resolveChild()).toEqual({ kind: 'ok', rootScopeId: 'workspace_root' }); + expect( + await resolveScopedCreateParent(client.db, { + kiloUserId: userId, + parentSessionId: 'parent', + scope, + }) + ).toMatchObject({ + kind: 'ok', + memberRootScopeId: 'workspace_root', + memberRoot: { sessionId: 'root' }, + }); + }); + + it.each([CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH, CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH + 1])( + 'bounds %i ancestors with one ancestry round trip', + async depth => { + await insertSession('root', null); + let parentId = 'root'; + for (let index = 1; index < depth; index += 1) { + const sessionId = `ancestor_${index}`; + await insertSession(sessionId, parentId); + parentId = sessionId; + } + await insertSession('child', parentId); + const select = vi.spyOn(client.db, 'select'); + const execute = vi.spyOn(client.db, 'execute'); + const result = await resolveChild(); + expect(select).toHaveBeenCalledTimes(1); + expect(execute).toHaveBeenCalledTimes(1); + expect(result).toEqual( + depth === CLOUD_AGENT_SCOPE_MAX_ANCESTRY_DEPTH + ? { kind: 'ok', rootScopeId: 'workspace_root' } + : { kind: 'denied' } + ); + } + ); + + it.each(['self', 'cycle', 'missing', 'unlisted-root', 'scope-mismatch'])( + 'denies %s ancestry', + async problem => { + await insertSession('root', null); + await insertSession( + 'parent', + problem === 'self' + ? 'parent' + : problem === 'cycle' + ? 'child' + : problem === 'missing' + ? 'absent' + : 'root' + ); + await insertSession( + 'child', + 'parent', + problem === 'scope-mismatch' ? { rootScopeId: 'workspace_other' } : {} + ); + expect( + await resolveChild( + problem === 'unlisted-root' ? { ...scope, memberRootKiloSessionIds: ['other'] } : scope + ) + ).toEqual({ kind: 'denied' }); + } + ); + + it.each([{ userId: 'oauth/another-user' }, { organizationId }, { worktreeId: 'worktree_other' }])( + 'denies an intermediate ancestor outside the asserted scope: %j', + async overrides => { + await insertSession('root', null); + await insertSession('parent', 'root', overrides); + await insertSession('child', 'parent'); + expect(await resolveChild()).toEqual({ kind: 'denied' }); + expect( + await resolveScopedCreateParent(client.db, { + kiloUserId: userId, + parentSessionId: 'child', + scope, + }) + ).toEqual({ kind: 'denied' }); + } + ); + + it('requires the asserted organization on every hop', async () => { + await insertSession('root', null, { organizationId }); + await insertSession('parent', 'root'); + await insertSession('child', 'parent', { organizationId }); + const orgScope = { ...scope, organizationId }; + expect(await resolveChild(orgScope)).toEqual({ kind: 'denied' }); + await client.db.execute( + sql`UPDATE pg_temp.cli_sessions_v2 SET organization_id = ${organizationId} WHERE session_id = 'parent'` + ); + expect(await resolveChild(orgScope)).toEqual({ kind: 'ok', rootScopeId: 'workspace_root' }); + }); + + it('denies a root on the child path but permits it as a create parent', async () => { + await insertSession('root', null, { rootScopeId: null }); + expect( + await resolveScopedChildSession(client.db, { kiloUserId: userId, sessionId: 'root', scope }) + ).toEqual({ kind: 'denied' }); + expect( + await resolveScopedCreateParent(client.db, { + kiloUserId: userId, + parentSessionId: 'root', + scope, + }) + ).toMatchObject({ kind: 'ok', memberRootScopeId: 'workspace_root' }); + }); +}); diff --git a/services/session-ingest/tsconfig.json b/services/session-ingest/tsconfig.json index caee557fbab..446addc881f 100644 --- a/services/session-ingest/tsconfig.json +++ b/services/session-ingest/tsconfig.json @@ -10,5 +10,5 @@ "jsxImportSource": "hono/jsx", "types": ["./worker-configuration.d.ts", "node", "@cloudflare/workers-types", "vitest/globals"] }, - "include": ["worker-configuration.d.ts", "src/**/*.ts"] + "include": ["worker-configuration.d.ts", "src/**/*.ts", "test/postgres/**/*.ts"] } diff --git a/services/session-ingest/vitest.postgres.config.ts b/services/session-ingest/vitest.postgres.config.ts new file mode 100644 index 00000000000..23c30c3516f --- /dev/null +++ b/services/session-ingest/vitest.postgres.config.ts @@ -0,0 +1,9 @@ +import { defineConfig } from 'vitest/config'; + +export default defineConfig({ + test: { + name: 'postgres', + environment: 'node', + include: ['test/postgres/**/*.test.ts'], + }, +});