From d60ca6754d08c5eae8387c82885e41c7bbde83da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 01:22:17 +0200 Subject: [PATCH 01/23] feat(auth): verify current pepper and active account on session and cloud-agent JWT auth --- .../worker-utils/src/kilo-token-auth.test.ts | 83 ++++++++ packages/worker-utils/src/kilo-token-auth.ts | 28 ++- services/cloud-agent-next/src/auth.test.ts | 103 +++++++++- services/cloud-agent-next/src/auth.ts | 25 ++- .../src/middleware/auth.test.ts | 12 +- .../cloud-agent-next/src/middleware/auth.ts | 5 +- services/cloud-agent-next/src/server.test.ts | 14 ++ services/cloud-agent-next/src/server.ts | 5 +- .../src/middleware/kilo-jwt-auth.test.ts | 194 ++++++++++++++---- .../src/middleware/kilo-jwt-auth.ts | 64 +----- 10 files changed, 413 insertions(+), 120 deletions(-) diff --git a/packages/worker-utils/src/kilo-token-auth.test.ts b/packages/worker-utils/src/kilo-token-auth.test.ts index a9e8b12918c..989ad111272 100644 --- a/packages/worker-utils/src/kilo-token-auth.test.ts +++ b/packages/worker-utils/src/kilo-token-auth.test.ts @@ -88,6 +88,13 @@ describe('verifyKiloBearerAgainstCurrentPepper', () => { await expect(verifyToken(token)).resolves.toBeNull(); }); + it('rejects a user token with a null pepper when the stored pepper is non-null', async () => { + userResultByUserId.set('user-xyz-789', { pepper: 'pepper-current', blockedReason: null }); + const { token } = await signToken({ pepper: null, tokenSource: 'kilo-chat' }); + + await expect(verifyToken(token)).resolves.toBeNull(); + }); + it('rejects tokens for blocked users (pepper matches, but blocked_reason is set)', async () => { userResultByUserId.set('user-xyz-789', { pepper: 'pepper-current', @@ -115,6 +122,82 @@ describe('verifyKiloBearerAgainstCurrentPepper', () => { await expect(verifyToken(token)).resolves.toEqual({ userId: 'user-xyz-789' }); }); + + it('succeeds when a token with env is verified without workerEnv', async () => { + const { token } = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { get: async () => TEST_JWT_SECRET }, + connectionString: 'postgres://test', + getUserPepper, + }) + ).resolves.toEqual({ userId: 'user-xyz-789' }); + }); + + it('fails when token env does not match a provided workerEnv', async () => { + const { token } = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { get: async () => TEST_JWT_SECRET }, + workerEnv: 'staging', + connectionString: 'postgres://test', + getUserPepper, + }) + ).resolves.toBeNull(); + }); +}); + +describe('internal service tokens (no apiTokenPepper, no env)', () => { + beforeEach(() => { + clearSecretCacheForTest(); + userResultByUserId.clear(); + userResultByUserId.set('user-xyz-789', { pepper: 'pepper-current', blockedReason: null }); + }); + + async function signInternalToken(): Promise { + const now = Math.floor(Date.now() / 1000); + return new SignJWT({ + version: 3, + kiloUserId: 'user-xyz-789', + }) + .setProtectedHeader({ alg: 'HS256' }) + .setIssuedAt(now) + .setExpirationTime(now + 3600) + .sign(new TextEncoder().encode(TEST_JWT_SECRET)); + } + + function verifyInternalToken(token: string, workerEnv?: string) { + return verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { get: async () => TEST_JWT_SECRET }, + ...(workerEnv ? { workerEnv } : {}), + connectionString: 'postgres://test', + getUserPepper, + }); + } + + it('succeeds when workerEnv is omitted and the user is active', async () => { + const token = await signInternalToken(); + + await expect(verifyInternalToken(token)).resolves.toEqual({ userId: 'user-xyz-789' }); + }); + + it('fails when workerEnv is provided', async () => { + const token = await signInternalToken(); + + await expect(verifyInternalToken(token, 'production')).resolves.toBeNull(); + }); + + it('fails when blocked_reason is set', async () => { + userResultByUserId.set('user-xyz-789', { pepper: null, blockedReason: 'manual block' }); + const token = await signInternalToken(); + + await expect(verifyInternalToken(token)).resolves.toBeNull(); + }); }); describe('C15 deviceSessionId compatibility', () => { diff --git a/packages/worker-utils/src/kilo-token-auth.ts b/packages/worker-utils/src/kilo-token-auth.ts index 6f4f8b9cd40..1c52b51c676 100644 --- a/packages/worker-utils/src/kilo-token-auth.ts +++ b/packages/worker-utils/src/kilo-token-auth.ts @@ -7,6 +7,7 @@ import { verifyKiloToken } from './kilo-token'; export type KiloBearerAuthResult = { userId: string; + botId?: string; }; export type KiloSecretBinding = { @@ -40,8 +41,8 @@ export async function findKiloUserPepper( export async function verifyKiloBearerAgainstCurrentPepper(params: { token: string | null; - nextAuthSecret: KiloSecretBinding; - workerEnv: string; + nextAuthSecret: KiloSecretBinding | string; + workerEnv?: string; connectionString: string; getUserPepper?: GetKiloUserPepper; }): Promise { @@ -50,9 +51,15 @@ export async function verifyKiloBearerAgainstCurrentPepper(params: { const getUserPepper = params.getUserPepper ?? findKiloUserPepper; try { - const secret = await getCachedSecret(params.nextAuthSecret, 'NEXTAUTH_SECRET'); + const secret = + typeof params.nextAuthSecret === 'string' + ? params.nextAuthSecret + : await getCachedSecret(params.nextAuthSecret, 'NEXTAUTH_SECRET'); const payload = await verifyKiloToken(params.token, secret); - if (payload.env !== params.workerEnv) { + + // Env check is skipped only when the caller does not pass a workerEnv. + // When workerEnv is set, a token without env (or with a mismatched env) fails. + if (params.workerEnv && payload.env !== params.workerEnv) { return null; } @@ -61,16 +68,21 @@ export async function verifyKiloBearerAgainstCurrentPepper(params: { return null; } - const tokenPepper = payload.apiTokenPepper ?? null; - if (result.pepper !== tokenPepper) { + if (result.blockedReason !== null) { return null; } - if (result.blockedReason !== null) { + // Pepper equality is skipped only when the claim is absent (internal + // service tokens). A present claim — string or null — is always compared. + if (payload.apiTokenPepper !== undefined && result.pepper !== payload.apiTokenPepper) { return null; } - return { userId: payload.kiloUserId }; + const authResult: KiloBearerAuthResult = { userId: payload.kiloUserId }; + if (payload.botId) { + authResult.botId = payload.botId; + } + return authResult; } catch { return null; } diff --git a/services/cloud-agent-next/src/auth.test.ts b/services/cloud-agent-next/src/auth.test.ts index 98b17afa10b..f5bf2d11970 100644 --- a/services/cloud-agent-next/src/auth.test.ts +++ b/services/cloud-agent-next/src/auth.test.ts @@ -1,13 +1,36 @@ -import { describe, expect, it } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import jwt from 'jsonwebtoken'; +import { signKiloToken } from '@kilocode/worker-utils'; import { mintWrapperDispatchTicket, + validateKiloToken, validateStreamTicket, validateWrapperDispatchTicket, type WrapperDispatchTicketClaims, } from './auth.js'; const secret = 'test-secret'; +const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; + +const userRowByUserId = vi.hoisted( + () => new Map() +); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => { + const row = userRowByUserId.get('user-1'); + if (!row) return []; + return [{ api_token_pepper: row.pepper, blocked_reason: row.blockedReason }]; + }, + }), + }), + }), + }), +})); describe('validateStreamTicket', () => { it('returns a configuration error when NEXTAUTH_SECRET is missing', () => { @@ -183,3 +206,81 @@ describe('validateWrapperDispatchTicket', () => { }); }); }); + +describe('validateKiloToken', () => { + beforeEach(() => { + userRowByUserId.clear(); + }); + + it('returns a configuration error when NEXTAUTH_SECRET is missing', async () => { + await expect( + validateKiloToken('Bearer token', { secret: null, connectionString: 'postgres://test' }) + ).resolves.toEqual({ + success: false, + error: 'NEXTAUTH_SECRET is not configured on the worker', + }); + }); + + it('returns an error when the Authorization header is missing or malformed', async () => { + await expect( + validateKiloToken(null, { secret: TEST_JWT_SECRET, connectionString: 'postgres://test' }) + ).resolves.toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); + }); + + it('accepts a token for an active account with the current pepper', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toMatchObject({ success: true, userId: 'user-1', token }); + }); + + it('rejects a token with a stale pepper', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-stale', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); + }); + + it('rejects a token for a blocked account', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: 'manual block' }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); + }); +}); diff --git a/services/cloud-agent-next/src/auth.ts b/services/cloud-agent-next/src/auth.ts index ab1b7687c32..622bbd5e99c 100644 --- a/services/cloud-agent-next/src/auth.ts +++ b/services/cloud-agent-next/src/auth.ts @@ -1,5 +1,6 @@ import jwt from 'jsonwebtoken'; import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; +import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; type StreamTicketPayload = { type: 'stream_ticket'; @@ -60,11 +61,16 @@ export async function resolveSecret( export async function validateKiloToken( authHeader: string | null, - secret: string | null | undefined + params: { + secret: string | null | undefined; + connectionString: string; + workerEnv?: string; + } ): Promise< | { success: true; userId: string; token: string; botId?: string } | { success: false; error: string } > { + const { secret, connectionString, workerEnv } = params; if (!secret) { return { success: false, error: 'NEXTAUTH_SECRET is not configured on the worker' }; } @@ -74,13 +80,18 @@ export async function validateKiloToken( return { success: false, error: 'Missing or malformed Authorization header' }; } - try { - const payload = await verifyKiloToken(token, secret); - return { success: true, userId: payload.kiloUserId, token, botId: payload.botId }; - } catch (err) { - const message = err instanceof Error ? err.message : 'JWT verification failed'; - return { success: false, error: message }; + const auth = await verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: secret, + connectionString, + ...(workerEnv ? { workerEnv } : {}), + }); + + if (!auth) { + return { success: false, error: 'Invalid or expired token' }; } + + return { success: true, userId: auth.userId, token, botId: auth.botId }; } export function validateStreamTicket( diff --git a/services/cloud-agent-next/src/middleware/auth.test.ts b/services/cloud-agent-next/src/middleware/auth.test.ts index 5c44638084f..d7ee71075f5 100644 --- a/services/cloud-agent-next/src/middleware/auth.test.ts +++ b/services/cloud-agent-next/src/middleware/auth.test.ts @@ -51,7 +51,7 @@ describe('authMiddleware', () => { const response = await app.fetch( new Request('https://worker.test/trpc/send', { method: 'POST' }), - { NEXTAUTH_SECRET: 'secret' } as Env + { NEXTAUTH_SECRET: 'secret', HYPERDRIVE: { connectionString: 'postgres://test' } } as Env ); const body: any = await response.json(); @@ -83,11 +83,17 @@ describe('authMiddleware', () => { method: 'POST', headers: { Authorization: 'Bearer token-1' }, }), - { NEXTAUTH_SECRET: secretBinding } as unknown as Env + { + NEXTAUTH_SECRET: secretBinding, + HYPERDRIVE: { connectionString: 'postgres://test' }, + } as unknown as Env ); expect(response.status).toBe(200); expect(secretBinding.get).toHaveBeenCalledOnce(); - expect(validateKiloToken).toHaveBeenCalledWith('Bearer token-1', 'secret-from-store'); + expect(validateKiloToken).toHaveBeenCalledWith('Bearer token-1', { + secret: 'secret-from-store', + connectionString: 'postgres://test', + }); }); }); diff --git a/services/cloud-agent-next/src/middleware/auth.ts b/services/cloud-agent-next/src/middleware/auth.ts index c64b24fa50d..6277a6aa227 100644 --- a/services/cloud-agent-next/src/middleware/auth.ts +++ b/services/cloud-agent-next/src/middleware/auth.ts @@ -10,7 +10,10 @@ export const authMiddleware = createMiddleware( async (c: Context, next: Next) => { const authHeader = c.req.header('authorization'); const nextAuthSecret = await resolveSecret(c.env.NEXTAUTH_SECRET); - const result = await validateKiloToken(authHeader ?? null, nextAuthSecret); + const result = await validateKiloToken(authHeader ?? null, { + secret: nextAuthSecret, + connectionString: c.env.HYPERDRIVE.connectionString, + }); if (!result.success) { logger.withFields({ error: result.error }).warn('Authentication failed'); diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index d3001c7d7c7..0d35349e7c0 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -93,6 +93,18 @@ vi.mock('./persistence/CloudAgentSession.js', () => ({ CloudAgentSession: class CloudAgentSession {}, })); +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [{ api_token_pepper: null, blocked_reason: null }], + }), + }), + }), + }), +})); + const { default: worker } = await import('./server.js'); const secret = 'test-secret'; @@ -102,6 +114,7 @@ type MockEnv = { Sandbox: unknown; SandboxSmall: unknown; WS_ALLOWED_ORIGINS?: string; + HYPERDRIVE: { connectionString: string }; CLOUD_AGENT_SESSION: { idFromName: ReturnType; get: ReturnType; @@ -117,6 +130,7 @@ function createEnv(): MockEnv { NEXTAUTH_SECRET: secret, Sandbox: {}, SandboxSmall: {}, + HYPERDRIVE: { connectionString: 'postgres://test' }, CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn(), diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index a98e2a6f554..533a3128100 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -240,7 +240,10 @@ async function routeToUserKiloFacade( async function routeAuthenticatedKiloFacade(c: Context): Promise { const nextAuthSecret = await resolveSecret(c.env.NEXTAUTH_SECRET); - const authResult = await validateKiloToken(c.req.header('Authorization') ?? null, nextAuthSecret); + const authResult = await validateKiloToken(c.req.header('Authorization') ?? null, { + secret: nextAuthSecret, + connectionString: c.env.HYPERDRIVE.connectionString, + }); if (!authResult.success) { return c.text(authResult.error, 401); } diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts index 52c440012e2..7bb9d52eb56 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts @@ -1,92 +1,196 @@ import { Hono } from 'hono'; import { SignJWT } from 'jose'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { clearSecretCacheForTest, signKiloToken } from '@kilocode/worker-utils'; import { kiloJwtAuthMiddleware } from './kilo-jwt-auth'; +const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; + +const userRowByUserId = vi.hoisted( + () => new Map() +); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => { + const row = userRowByUserId.get('usr_123'); + if (!row) return []; + return [{ api_token_pepper: row.pepper, blocked_reason: row.blockedReason }]; + }, + }), + }), + }), + }), +})); + type TestEnv = { NEXTAUTH_SECRET_PROD: { get: () => Promise; }; - USER_EXISTS_CACHE: { - get: (key: string) => Promise; - put: (key: string, value: string, options?: { expirationTtl: number }) => Promise; - }; HYPERDRIVE: { connectionString: string; }; }; -function makeEnv(secret: string, opts?: { cachedUserState?: '1' | '0' | null }): TestEnv { +function makeEnv(secret: string): TestEnv { return { NEXTAUTH_SECRET_PROD: { get: async () => secret, }, - USER_EXISTS_CACHE: { - get: async () => opts?.cachedUserState ?? null, - put: async () => {}, - }, - // Not used when the KV cache returns a hit. - HYPERDRIVE: { connectionString: '' }, + HYPERDRIVE: { connectionString: 'postgres://test' }, }; } -async function sign(payload: Record, secret: string): Promise { - return new SignJWT(payload) +function makeApp() { + const app = new Hono<{ Bindings: TestEnv; Variables: { user_id: string } }>(); + app.use('/api/*', kiloJwtAuthMiddleware); + app.get('/api/me', c => c.json({ user_id: c.get('user_id') })); + app.get('/api/user/cli', c => c.json({ user_id: c.get('user_id') })); + app.get('/api/user/web', c => c.json({ user_id: c.get('user_id') })); + return app; +} + +async function signUserToken(pepper: string): Promise { + const { token } = await signKiloToken({ + userId: 'usr_123', + pepper, + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + extra: { tokenSource: 'kilo-chat' }, + }); + return token; +} + +async function signInternalToken(): Promise { + const now = Math.floor(Date.now() / 1000); + return new SignJWT({ version: 3, kiloUserId: 'usr_123' }) .setProtectedHeader({ alg: 'HS256' }) - .setIssuedAt() - .setExpirationTime('1h') - .sign(new TextEncoder().encode(secret)); + .setIssuedAt(now) + .setExpirationTime(now + 3600) + .sign(new TextEncoder().encode(TEST_JWT_SECRET)); } describe('kiloJwtAuthMiddleware', () => { + beforeEach(() => { + clearSecretCacheForTest(); + userRowByUserId.clear(); + }); + it('rejects missing Authorization header', async () => { - const app = new Hono<{ Bindings: TestEnv; Variables: { user_id: string } }>(); - app.use('/api/*', kiloJwtAuthMiddleware); - app.get('/api/me', c => c.json({ user_id: c.get('user_id') })); + const res = await makeApp().fetch(new Request('http://local/api/me'), makeEnv(TEST_JWT_SECRET)); + expect(res.status).toBe(401); + }); + + it('accepts a token with the current pepper for an active account', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signUserToken('pepper-current'); + + const res = await makeApp().fetch( + new Request('http://local/api/me', { + headers: { Authorization: `Bearer ${token}` }, + }), + makeEnv(TEST_JWT_SECRET) + ); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: 'usr_123' }); + }); + + it('rejects a stale pepper even when the user exists', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signUserToken('pepper-stale'); + + const res = await makeApp().fetch( + new Request('http://local/api/me', { + headers: { Authorization: `Bearer ${token}` }, + }), + makeEnv(TEST_JWT_SECRET) + ); - const res = await app.fetch(new Request('http://local/api/me'), makeEnv('secret')); expect(res.status).toBe(401); + expect(await res.json()).toEqual({ success: false, error: 'Invalid or expired token' }); }); - it('accepts valid v3 token when user exists in cache', async () => { - const secret = 'test-secret'; - const token = await sign({ kiloUserId: 'usr_123', version: 3 }, secret); + it('rejects a blocked account', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: 'manual block' }); + const token = await signUserToken('pepper-current'); - const app = new Hono<{ Bindings: TestEnv; Variables: { user_id: string } }>(); - app.use('/api/*', kiloJwtAuthMiddleware); - app.get('/api/me', c => c.json({ user_id: c.get('user_id') })); + const res = await makeApp().fetch( + new Request('http://local/api/me', { + headers: { Authorization: `Bearer ${token}` }, + }), + makeEnv(TEST_JWT_SECRET) + ); - const res = await app.fetch( + expect(res.status).toBe(401); + expect(await res.json()).toEqual({ success: false, error: 'Invalid or expired token' }); + }); + + it('accepts an internal token (no pepper) for an active account', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signInternalToken(); + + const res = await makeApp().fetch( new Request('http://local/api/me', { - headers: { - Authorization: `Bearer ${token}`, - }, + headers: { Authorization: `Bearer ${token}` }, }), - makeEnv(secret, { cachedUserState: '1' }) + makeEnv(TEST_JWT_SECRET) ); expect(res.status).toBe(200); expect(await res.json()).toEqual({ user_id: 'usr_123' }); }); - it('rejects valid v3 token when user is cached as not-found', async () => { - const secret = 'test-secret'; - const token = await sign({ kiloUserId: 'deleted_user', version: 3 }, secret); - - const app = new Hono<{ Bindings: TestEnv; Variables: { user_id: string } }>(); - app.use('/api/*', kiloJwtAuthMiddleware); - app.get('/api/me', c => c.json({ user_id: c.get('user_id') })); + it('rejects an internal token (no pepper) for a blocked account', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: 'manual block' }); + const token = await signInternalToken(); - const res = await app.fetch( + const res = await makeApp().fetch( new Request('http://local/api/me', { - headers: { - Authorization: `Bearer ${token}`, - }, + headers: { Authorization: `Bearer ${token}` }, + }), + makeEnv(TEST_JWT_SECRET) + ); + + expect(res.status).toBe(401); + expect(await res.json()).toEqual({ success: false, error: 'Invalid or expired token' }); + }); + + it('reads ?token= on a websocket upgrade to /api/user/cli', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signUserToken('pepper-current'); + + const res = await makeApp().fetch( + new Request(`http://local/api/user/cli?token=${token}`, { + headers: { Upgrade: 'websocket' }, }), - makeEnv(secret, { cachedUserState: '0' }) + makeEnv(TEST_JWT_SECRET) ); - expect(res.status).toBe(403); - expect(await res.json()).toEqual({ success: false, error: 'User account not found' }); + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: 'usr_123' }); + }); + + it('does not read ?token= on a websocket upgrade to /api/user/web', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signUserToken('pepper-current'); + + const res = await makeApp().fetch( + new Request(`http://local/api/user/web?token=${token}`, { + headers: { Upgrade: 'websocket' }, + }), + makeEnv(TEST_JWT_SECRET) + ); + + expect(res.status).toBe(401); + expect(await res.json()).toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); }); }); diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.ts index b893e7e6954..88cead9d786 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.ts @@ -1,48 +1,9 @@ import { createMiddleware } from 'hono/factory'; -import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; -import { eq } from 'drizzle-orm'; -import { getWorkerDb } from '@kilocode/db/client'; -import { kilocode_users } from '@kilocode/db/schema'; +import { extractBearerToken } from '@kilocode/worker-utils'; +import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; import type { Env } from '../env'; -const USER_EXISTS_TTL_SECONDS = 24 * 60 * 60; // 24h -const USER_NOT_FOUND_TTL_SECONDS = 5 * 60; // 5m - -/** - * Check whether a user exists, using KV as a cache in front of Postgres. - * Positive results are cached for 24h. Negative results are cached for 5m - * to rate-limit DB hits from deleted/nonexistent users with valid tokens. - */ -async function userExists(env: Env, userId: string): Promise { - const cacheKey = `user-exists:${userId}`; - - const cached = await env.USER_EXISTS_CACHE.get(cacheKey); - if (cached === '1') { - return true; - } - if (cached === '0') { - return false; - } - - const db = getWorkerDb(env.HYPERDRIVE.connectionString); - const rows = await db - .select({ id: kilocode_users.id }) - .from(kilocode_users) - .where(eq(kilocode_users.id, userId)) - .limit(1); - - const row = rows[0]; - - if (!row) { - void env.USER_EXISTS_CACHE.put(cacheKey, '0', { expirationTtl: USER_NOT_FOUND_TTL_SECONDS }); - return false; - } - - void env.USER_EXISTS_CACHE.put(cacheKey, '1', { expirationTtl: USER_EXISTS_TTL_SECONDS }); - return true; -} - export const kiloJwtAuthMiddleware = createMiddleware<{ Bindings: Env; Variables: { @@ -50,7 +11,7 @@ export const kiloJwtAuthMiddleware = createMiddleware<{ }; }>(async (c, next) => { let token = extractBearerToken(c.req.header('Authorization')); - if (!token && c.req.header('Upgrade') === 'websocket') { + if (!token && c.req.header('Upgrade') === 'websocket' && c.req.path !== '/api/user/web') { token = c.req.query('token') ?? null; } @@ -58,21 +19,16 @@ export const kiloJwtAuthMiddleware = createMiddleware<{ return c.json({ success: false, error: 'Missing or malformed Authorization header' }, 401); } - const secret = await c.env.NEXTAUTH_SECRET_PROD.get(); + const auth = await verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: c.env.NEXTAUTH_SECRET_PROD, + connectionString: c.env.HYPERDRIVE.connectionString, + }); - let kiloUserId: string; - try { - const payload = await verifyKiloToken(token, secret); - kiloUserId = payload.kiloUserId; - } catch { + if (!auth) { return c.json({ success: false, error: 'Invalid or expired token' }, 401); } - const exists = await userExists(c.env, kiloUserId); - if (!exists) { - return c.json({ success: false, error: 'User account not found' }, 403); - } - - c.set('user_id', kiloUserId); + c.set('user_id', auth.userId); return next(); }); From 6ef326265c035f9fea3fa0b72d1981f5ca15d059 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 01:53:55 +0200 Subject: [PATCH 02/23] feat(session-ingest): mint and consume one-use web tickets for /api/user/web --- .../src/dos/connection-ticket-do.test.ts | 117 +++++++++++++++++ .../src/dos/connection-ticket-do.ts | 53 ++++++++ services/session-ingest/src/index.ts | 1 + .../src/middleware/kilo-jwt-auth.test.ts | 118 +++++++++++++++++- .../src/middleware/kilo-jwt-auth.ts | 22 +++- .../session-ingest/src/routes/api.test.ts | 37 ++++++ services/session-ingest/src/routes/api.ts | 19 ++- .../session-ingest/worker-configuration.d.ts | 5 +- services/session-ingest/wrangler.jsonc | 8 ++ services/session-ingest/wrangler.test.jsonc | 8 ++ 10 files changed, 382 insertions(+), 6 deletions(-) create mode 100644 services/session-ingest/src/dos/connection-ticket-do.test.ts create mode 100644 services/session-ingest/src/dos/connection-ticket-do.ts diff --git a/services/session-ingest/src/dos/connection-ticket-do.test.ts b/services/session-ingest/src/dos/connection-ticket-do.test.ts new file mode 100644 index 00000000000..472289c3a13 --- /dev/null +++ b/services/session-ingest/src/dos/connection-ticket-do.test.ts @@ -0,0 +1,117 @@ +import { describe, expect, it, vi } from 'vitest'; + +// Mock cloudflare:workers before importing ConnectionTicketDO. +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + ctx: unknown; + env: unknown; + constructor(state: unknown, env: unknown) { + this.ctx = state; + this.env = env; + } + }, +})); + +import { ConnectionTicketDO } from './connection-ticket-do'; + +function makeStorage() { + const store = new Map(); + let alarm: number | null = null; + + const storage = { + put: async (key: string, value: unknown) => { + store.set(key, value); + }, + get: async (key: string): Promise => store.get(key) as T | undefined, + delete: async (key: string) => { + store.delete(key); + }, + setAlarm: async (at: number) => { + alarm = at; + }, + deleteAlarm: async () => { + alarm = null; + }, + transaction: async ( + fn: (txn: { + get: (key: string) => Promise; + delete: (key: string) => Promise; + }) => Promise + ): Promise => { + const txn = { + get: async (key: string): Promise => store.get(key) as T2 | undefined, + delete: async (key: string) => { + store.delete(key); + }, + }; + return fn(txn); + }, + }; + + return { store, storage, getAlarm: () => alarm }; +} + +function makeDO() { + const { store, storage, getAlarm } = makeStorage(); + const instance = new ConnectionTicketDO({ storage } as never, {} as never); + return { instance, store, getAlarm }; +} + +describe('ConnectionTicketDO', () => { + it('mints a ticket by storing state and arming an alarm', async () => { + const { instance, store, getAlarm } = makeDO(); + const expiresAt = Date.now() + 60_000; + + await instance.mint({ userId: 'usr_1', expiresAt }); + + expect(store.get('ticket')).toEqual({ userId: 'usr_1', expiresAt }); + expect(getAlarm()).toBe(expiresAt); + }); + + it('consumes a fresh ticket once and returns the userId', async () => { + const { instance } = makeDO(); + await instance.mint({ userId: 'usr_1', expiresAt: Date.now() + 60_000 }); + + await expect(instance.consume()).resolves.toEqual({ userId: 'usr_1' }); + }); + + it('rejects a replay of an already-consumed ticket', async () => { + const { instance } = makeDO(); + await instance.mint({ userId: 'usr_1', expiresAt: Date.now() + 60_000 }); + + await expect(instance.consume()).resolves.toEqual({ userId: 'usr_1' }); + await expect(instance.consume()).resolves.toBeNull(); + }); + + it('rejects an expired ticket', async () => { + const { instance } = makeDO(); + await instance.mint({ userId: 'usr_1', expiresAt: Date.now() - 1 }); + + await expect(instance.consume()).resolves.toBeNull(); + }); + + it('deletes ticket storage and alarm after a successful consume', async () => { + const { instance, store, getAlarm } = makeDO(); + const expiresAt = Date.now() + 60_000; + await instance.mint({ userId: 'usr_1', expiresAt }); + + expect(store.get('ticket')).toEqual({ userId: 'usr_1', expiresAt }); + expect(getAlarm()).toBe(expiresAt); + + await expect(instance.consume()).resolves.toEqual({ userId: 'usr_1' }); + + expect(store.get('ticket')).toBeUndefined(); + expect(getAlarm()).toBeNull(); + }); + + it('deletes unconsumed expired ticket storage when the alarm runs', async () => { + const { instance, store } = makeDO(); + await instance.mint({ userId: 'usr_1', expiresAt: Date.now() + 60_000 }); + + // Simulate expiry before the alarm fires. + store.set('ticket', { userId: 'usr_1', expiresAt: Date.now() - 1 }); + await instance.alarm(); + + expect(store.get('ticket')).toBeUndefined(); + }); +}); diff --git a/services/session-ingest/src/dos/connection-ticket-do.ts b/services/session-ingest/src/dos/connection-ticket-do.ts new file mode 100644 index 00000000000..de3c1440884 --- /dev/null +++ b/services/session-ingest/src/dos/connection-ticket-do.ts @@ -0,0 +1,53 @@ +import { DurableObject } from 'cloudflare:workers'; +import { z } from 'zod'; + +import type { Env } from '../env'; + +const ticketStateSchema = z.object({ + userId: z.string().min(1), + expiresAt: z.number().int(), +}); + +export const connectionTicketConsumeResponseSchema = z.object({ + userId: z.string().min(1), +}); + +type TicketState = z.infer; +export type TicketMintRequest = TicketState; +export type ConnectionTicketConsumeResponse = z.infer; + +export class ConnectionTicketDO extends DurableObject { + async mint(input: TicketMintRequest): Promise { + await this.ctx.storage.put('ticket', input); + await this.ctx.storage.setAlarm(input.expiresAt); + } + + async consume(): Promise { + const userId = await this.ctx.storage.transaction(async txn => { + const stored = await txn.get('ticket'); + const parsed = ticketStateSchema.safeParse(stored); + if (!parsed.success || parsed.data.expiresAt <= Date.now()) { + await txn.delete('ticket'); + return null; + } + + await txn.delete('ticket'); + return parsed.data.userId; + }); + + if (!userId) { + return null; + } + + await this.ctx.storage.deleteAlarm(); + return { userId } satisfies ConnectionTicketConsumeResponse; + } + + async alarm(): Promise { + const stored = await this.ctx.storage.get('ticket'); + const parsed = ticketStateSchema.safeParse(stored); + if (!parsed.success || parsed.data.expiresAt <= Date.now()) { + await this.ctx.storage.delete('ticket'); + } + } +} diff --git a/services/session-ingest/src/index.ts b/services/session-ingest/src/index.ts index 351772bb26f..008a06de1db 100644 --- a/services/session-ingest/src/index.ts +++ b/services/session-ingest/src/index.ts @@ -1,6 +1,7 @@ export { SessionIngestDO } from './dos/SessionIngestDO'; export { SessionAccessCacheDO } from './dos/SessionAccessCacheDO'; export { UserConnectionDO } from './dos/UserConnectionDO'; +export { ConnectionTicketDO } from './dos/connection-ticket-do'; export { SessionIngestRPC } from './session-ingest-rpc'; export { app } from './app'; diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts index 7bb9d52eb56..4ff8257cdd3 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts @@ -27,6 +27,41 @@ vi.mock('@kilocode/db/client', () => ({ }), })); +type TicketStore = { + tickets: Map; + mint: (userId: string, expiresAt: number) => string; + namespace: { + idFromName: (name: string) => string; + get: (id: string) => { consume: () => Promise<{ userId: string } | null> }; + }; +}; + +function makeTicketStore(): TicketStore { + const tickets = new Map(); + return { + tickets, + mint(userId: string, expiresAt: number): string { + const ticket = crypto.randomUUID(); + tickets.set(ticket, { userId, expiresAt }); + return ticket; + }, + namespace: { + idFromName: (name: string) => name, + get: (id: string) => ({ + consume: async () => { + const entry = tickets.get(id); + if (!entry || entry.expiresAt <= Date.now()) { + tickets.delete(id); + return null; + } + tickets.delete(id); + return { userId: entry.userId }; + }, + }), + }, + }; +} + type TestEnv = { NEXTAUTH_SECRET_PROD: { get: () => Promise; @@ -34,14 +69,16 @@ type TestEnv = { HYPERDRIVE: { connectionString: string; }; + CONNECTION_TICKET_DO: TicketStore['namespace']; }; -function makeEnv(secret: string): TestEnv { +function makeEnv(secret: string, ticketStore: TicketStore = makeTicketStore()): TestEnv { return { NEXTAUTH_SECRET_PROD: { get: async () => secret, }, HYPERDRIVE: { connectionString: 'postgres://test' }, + CONNECTION_TICKET_DO: ticketStore.namespace, }; } @@ -190,7 +227,84 @@ describe('kiloJwtAuthMiddleware', () => { expect(res.status).toBe(401); expect(await res.json()).toEqual({ success: false, - error: 'Missing or malformed Authorization header', + error: 'Missing or invalid ticket', + }); + }); + + it('rejects a raw bearer on a websocket upgrade to /api/user/web', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signUserToken('pepper-current'); + + const res = await makeApp().fetch( + new Request('http://local/api/user/web', { + headers: { Upgrade: 'websocket', Authorization: `Bearer ${token}` }, + }), + makeEnv(TEST_JWT_SECRET) + ); + + expect(res.status).toBe(401); + expect(await res.json()).toEqual({ + success: false, + error: 'Missing or invalid ticket', + }); + }); + + it('accepts a fresh ticket on a websocket upgrade to /api/user/web', async () => { + const ticketStore = makeTicketStore(); + const ticket = ticketStore.mint('usr_123', Date.now() + 60_000); + + const res = await makeApp().fetch( + new Request(`http://local/api/user/web?ticket=${ticket}`, { + headers: { Upgrade: 'websocket' }, + }), + makeEnv(TEST_JWT_SECRET, ticketStore) + ); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: 'usr_123' }); + }); + + it('rejects a replay of the same ticket on /api/user/web', async () => { + const ticketStore = makeTicketStore(); + const ticket = ticketStore.mint('usr_123', Date.now() + 60_000); + + const first = await makeApp().fetch( + new Request(`http://local/api/user/web?ticket=${ticket}`, { + headers: { Upgrade: 'websocket' }, + }), + makeEnv(TEST_JWT_SECRET, ticketStore) + ); + expect(first.status).toBe(200); + + const replay = await makeApp().fetch( + new Request(`http://local/api/user/web?ticket=${ticket}`, { + headers: { Upgrade: 'websocket' }, + }), + makeEnv(TEST_JWT_SECRET, ticketStore) + ); + + expect(replay.status).toBe(401); + expect(await replay.json()).toEqual({ + success: false, + error: 'Invalid or expired ticket', + }); + }); + + it('rejects an expired ticket on /api/user/web', async () => { + const ticketStore = makeTicketStore(); + const ticket = ticketStore.mint('usr_123', Date.now() - 1); + + const res = await makeApp().fetch( + new Request(`http://local/api/user/web?ticket=${ticket}`, { + headers: { Upgrade: 'websocket' }, + }), + makeEnv(TEST_JWT_SECRET, ticketStore) + ); + + expect(res.status).toBe(401); + expect(await res.json()).toEqual({ + success: false, + error: 'Invalid or expired ticket', }); }); }); diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.ts index 88cead9d786..d4841ce542e 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.ts @@ -10,8 +10,28 @@ export const kiloJwtAuthMiddleware = createMiddleware<{ user_id: string; }; }>(async (c, next) => { + // One-use web ticket path. The /api/user/web websocket upgrade consumes an + // opaque ticket minted by POST /api/user/web-ticket. This branch runs before + // the JWT path so a missing ticket is not mistaken for a missing JWT, and so + // the Authorization header / ?token= query are never read on this path. + if (c.req.header('Upgrade') === 'websocket' && c.req.path === '/api/user/web') { + const ticket = c.req.query('ticket'); + if (!ticket) { + return c.json({ success: false, error: 'Missing or invalid ticket' }, 401); + } + + const stub = c.env.CONNECTION_TICKET_DO.get(c.env.CONNECTION_TICKET_DO.idFromName(ticket)); + const consumed = await stub.consume(); + if (!consumed) { + return c.json({ success: false, error: 'Invalid or expired ticket' }, 401); + } + + c.set('user_id', consumed.userId); + return next(); + } + let token = extractBearerToken(c.req.header('Authorization')); - if (!token && c.req.header('Upgrade') === 'websocket' && c.req.path !== '/api/user/web') { + if (!token && c.req.header('Upgrade') === 'websocket') { token = c.req.query('token') ?? null; } diff --git a/services/session-ingest/src/routes/api.test.ts b/services/session-ingest/src/routes/api.test.ts index c908fa2c09c..303527f2ece 100644 --- a/services/session-ingest/src/routes/api.test.ts +++ b/services/session-ingest/src/routes/api.test.ts @@ -56,6 +56,10 @@ type TestBindings = { sendSessionReadyNotification: ReturnType; sendCloudAgentSessionNotification: ReturnType; }; + CONNECTION_TICKET_DO: { + get: ReturnType; + idFromName: ReturnType; + }; DIRECT_INGEST_PERCENT: string; DIRECT_INGEST_USER_IDS: string; DIRECT_INGEST_MAX_BYTES: string; @@ -73,6 +77,10 @@ function makeTestEnv(overrides: Partial = {}): TestBindings { sendSessionReadyNotification: vi.fn(async () => ({ dispatched: true })), sendCloudAgentSessionNotification: vi.fn(async () => ({ dispatched: true })), }, + CONNECTION_TICKET_DO: { + get: vi.fn(), + idFromName: vi.fn(), + }, DIRECT_INGEST_PERCENT: '0', DIRECT_INGEST_USER_IDS: '', DIRECT_INGEST_MAX_BYTES: '4194304', @@ -1867,6 +1875,35 @@ describe('api routes', () => { expect(forwardedUrl.searchParams.get('connectionId')).toBe('viewer-1'); }); + it('POST /user/web-ticket mints seconds for the body and milliseconds for the DO store', async () => { + const mint = vi.fn(async () => undefined); + const ticketStub = { mint }; + const env = makeTestEnv(); + env.CONNECTION_TICKET_DO.idFromName.mockReturnValue('ticket-do-id'); + env.CONNECTION_TICKET_DO.get.mockReturnValue(ticketStub); + + const now = vi.spyOn(Date, 'now').mockReturnValue(1_700_000_000_000); + + const app = makeApiApp(); + const res = await app.fetch( + new Request('http://local/user/web-ticket', { method: 'POST' }), + env + ); + now.mockRestore(); + + expect(res.status).toBe(200); + const body = await res.json(); + expect(body).toEqual({ + ticket: expect.any(String), + expiresAt: 1_700_000_060, + }); + expect(Number.isInteger(body.expiresAt)).toBe(true); + expect(mint).toHaveBeenCalledWith({ + userId: 'usr_test', + expiresAt: 1_700_000_060_000, + }); + }); + // ------------------------------------------------------------------------- // GET /api/instances/active (W3) // ------------------------------------------------------------------------- diff --git a/services/session-ingest/src/routes/api.ts b/services/session-ingest/src/routes/api.ts index 9fce023974f..e96931bde1a 100644 --- a/services/session-ingest/src/routes/api.ts +++ b/services/session-ingest/src/routes/api.ts @@ -718,7 +718,7 @@ api.get('/user/cli', async c => { return stub.fetch(new Request(wsUrl.toString(), c.req.raw)); }); -// Web UI connects to /api/user/web without userId in the path — userId comes from the JWT. +// Web UI connects to /api/user/web without userId in the path — the websocket upgrade takes userId from the consumed ticket. api.get('/user/web', async c => { if (c.req.header('Upgrade') !== 'websocket') { return c.json({ success: false, error: 'Expected WebSocket upgrade' }, 426); @@ -731,3 +731,20 @@ api.get('/user/web', async c => { return stub.fetch(new Request(wsUrl.toString(), c.req.raw)); }); + +// Web UI mints a one-use ticket before opening /api/user/web. The ticket is +// opaque (UUID); the DO stores { userId, expiresAt } keyed by the ticket and +// the websocket upgrade consumes it exactly once. +api.post('/user/web-ticket', async c => { + const kiloUserId = c.get('user_id'); + const ticket = crypto.randomUUID(); + const doExpiresAt = Date.now() + 60_000; + + const stub = c.env.CONNECTION_TICKET_DO.get(c.env.CONNECTION_TICKET_DO.idFromName(ticket)); + await stub.mint({ + userId: kiloUserId, + expiresAt: doExpiresAt, + }); + + return c.json({ ticket, expiresAt: Math.floor(doExpiresAt / 1000) }, 200); +}); diff --git a/services/session-ingest/worker-configuration.d.ts b/services/session-ingest/worker-configuration.d.ts index 3666ac876c2..6db48f3e86b 100644 --- a/services/session-ingest/worker-configuration.d.ts +++ b/services/session-ingest/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types --include-runtime=false --env-interface CloudflareBindings` (hash: b51b7ed6e348b311085e4656eecfd9d5) +// Generated by Wrangler by running `wrangler types --include-runtime=false --env-interface=CloudflareBindings` (hash: ee4846975bd8a5b95782c924bda4d9c0) interface __BaseEnv_CloudflareBindings { USER_EXISTS_CACHE: KVNamespace; SESSION_INGEST_R2: R2Bucket; @@ -13,13 +13,14 @@ interface __BaseEnv_CloudflareBindings { SESSION_INGEST_DO: DurableObjectNamespace; SESSION_ACCESS_CACHE_DO: DurableObjectNamespace; USER_CONNECTION_DO: DurableObjectNamespace; + CONNECTION_TICKET_DO: DurableObjectNamespace; O11Y: Fetcher /* o11y */; NOTIFICATIONS: Service /* entrypoint NotificationsService from notifications */; } declare namespace Cloudflare { interface GlobalProps { mainModule: typeof import("./src/index"); - durableNamespaces: "SessionIngestDO" | "SessionAccessCacheDO" | "UserConnectionDO"; + durableNamespaces: "SessionIngestDO" | "SessionAccessCacheDO" | "UserConnectionDO" | "ConnectionTicketDO"; } interface Env extends __BaseEnv_CloudflareBindings {} } diff --git a/services/session-ingest/wrangler.jsonc b/services/session-ingest/wrangler.jsonc index a95108b86bf..424a44479d8 100644 --- a/services/session-ingest/wrangler.jsonc +++ b/services/session-ingest/wrangler.jsonc @@ -63,6 +63,10 @@ "name": "USER_CONNECTION_DO", "class_name": "UserConnectionDO", }, + { + "name": "CONNECTION_TICKET_DO", + "class_name": "ConnectionTicketDO", + }, ], }, "migrations": [ @@ -74,6 +78,10 @@ "tag": "v2", "new_sqlite_classes": ["UserConnectionDO"], }, + { + "tag": "v3", + "new_sqlite_classes": ["ConnectionTicketDO"], + }, ], "r2_buckets": [ { diff --git a/services/session-ingest/wrangler.test.jsonc b/services/session-ingest/wrangler.test.jsonc index 8fc85d75610..e076a4cca40 100644 --- a/services/session-ingest/wrangler.test.jsonc +++ b/services/session-ingest/wrangler.test.jsonc @@ -24,6 +24,10 @@ "name": "SESSION_ACCESS_CACHE_DO", "class_name": "SessionAccessCacheDO", }, + { + "name": "CONNECTION_TICKET_DO", + "class_name": "ConnectionTicketDO", + }, ], }, @@ -32,6 +36,10 @@ "tag": "v1", "new_sqlite_classes": ["SessionIngestDO", "SessionAccessCacheDO"], }, + { + "tag": "v2", + "new_sqlite_classes": ["ConnectionTicketDO"], + }, ], "r2_buckets": [ From 21c12cf16b094cc05e219d5c78fe07320ad01526 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 02:18:11 +0200 Subject: [PATCH 03/23] feat(active-sessions): mint one-use web tickets via getToken mutation --- .../entrypoints/sidepanel/agents-provider.tsx | 2 +- apps/extension/tests/e2e/agents-fixture.ts | 4 +- .../user-web-connection-provider.test.ts | 83 +++++++++++++++++++ .../agents/user-web-connection-provider.tsx | 2 +- .../cloud-agent-next/CloudAgentProvider.tsx | 2 +- .../routers/active-sessions-router.test.ts | 73 ++++++++++++++++ .../web/src/routers/active-sessions-router.ts | 58 ++++++++++++- .../src/user-web-connection.test.ts | 22 ++--- .../src/user-web-connection.ts | 2 +- 9 files changed, 230 insertions(+), 18 deletions(-) create mode 100644 apps/mobile/src/components/agents/user-web-connection-provider.test.ts diff --git a/apps/extension/entrypoints/sidepanel/agents-provider.tsx b/apps/extension/entrypoints/sidepanel/agents-provider.tsx index 0ea5f627afa..630d4f3b77c 100644 --- a/apps/extension/entrypoints/sidepanel/agents-provider.tsx +++ b/apps/extension/entrypoints/sidepanel/agents-provider.tsx @@ -53,7 +53,7 @@ export const ExtensionAgentsProvider = ({ const userWebConnection = createUserWebConnection({ getAuthToken: async () => { - const tokenResult = await trpcClient.activeSessions.getToken.query(); + const tokenResult = await trpcClient.activeSessions.getToken.mutate(); return tokenResult.token; }, lifecycleHooks: createBrowserLifecycleHooks(), diff --git a/apps/extension/tests/e2e/agents-fixture.ts b/apps/extension/tests/e2e/agents-fixture.ts index e26ecb8b28c..062de2ea46d 100644 --- a/apps/extension/tests/e2e/agents-fixture.ts +++ b/apps/extension/tests/e2e/agents-fixture.ts @@ -519,7 +519,9 @@ export const mockAgentsApi = async ( } if (proc === 'activeSessions.getToken') { - return { result: { data: { token: 'mock-ingest-token' } } }; + return { + result: { data: { token: 'mock-ingest-token', expiresAt: 1_700_000_000 } }, + }; } if ( diff --git a/apps/mobile/src/components/agents/user-web-connection-provider.test.ts b/apps/mobile/src/components/agents/user-web-connection-provider.test.ts new file mode 100644 index 00000000000..40dd35d969e --- /dev/null +++ b/apps/mobile/src/components/agents/user-web-connection-provider.test.ts @@ -0,0 +1,83 @@ +/* eslint-disable typescript-eslint/no-deprecated, react/no-children-prop -- react-test-renderer is the DOM-free renderer used to mount React/RN trees under vitest (same pattern as user-web-connection-provider.mounted.test.tsx); `children` must be passed as a prop because this is a .ts file with no JSX */ +import { createElement } from 'react'; +import TestRenderer, { act } from 'react-test-renderer'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { UserWebConnectionProvider } from './user-web-connection-provider'; + +type AuthConfig = { getAuthToken: () => Promise }; + +const mocks = vi.hoisted(() => ({ + mutate: vi.fn(async () => { + await Promise.resolve(); + return { token: 'ticket-1', expiresAt: 1_700_000_060 }; + }), + query: vi.fn(), + capturedConfig: null as AuthConfig | null, +})); + +vi.mock('@kilocode/cloud-agent-sdk/user-web-connection', () => ({ + createUserWebConnection: (config: AuthConfig) => { + mocks.capturedConfig = config; + return { retain: () => vi.fn() }; + }, +})); + +vi.mock('@/lib/config', () => ({ + SESSION_INGEST_WS_URL: 'wss://ingest.example.com', +})); + +vi.mock('@/lib/user-web-connection-lifecycle', () => ({ + createNativeUserWebConnectionLifecycleHooks: () => ({}), +})); + +vi.mock('@/lib/trpc', () => ({ + trpcClient: { + activeSessions: { + getToken: { + mutate: mocks.mutate, + query: mocks.query, + }, + }, + }, +})); + +describe('UserWebConnectionProvider', () => { + beforeEach(() => { + mocks.capturedConfig = null; + mocks.mutate.mockClear(); + mocks.query.mockClear(); + }); + + it('mints the ingest ticket via the getToken mutation (not query)', async () => { + const rendererRef: { current: TestRenderer.ReactTestRenderer | undefined } = { + current: undefined, + }; + await act(async () => { + await Promise.resolve(); + rendererRef.current = TestRenderer.create( + createElement(UserWebConnectionProvider, { children: null }) + ); + }); + + const config = mocks.capturedConfig; + if (!config) { + throw new Error('createUserWebConnection was not called'); + } + + const token = await config.getAuthToken(); + + expect(token).toBe('ticket-1'); + expect(mocks.mutate).toHaveBeenCalledTimes(1); + expect(mocks.query).not.toHaveBeenCalled(); + + const renderer = rendererRef.current; + if (!renderer) { + throw new Error('renderer was not created'); + } + await act(async () => { + await Promise.resolve(); + renderer.unmount(); + }); + }); +}); diff --git a/apps/mobile/src/components/agents/user-web-connection-provider.tsx b/apps/mobile/src/components/agents/user-web-connection-provider.tsx index 06c05f40da7..34557c857bd 100644 --- a/apps/mobile/src/components/agents/user-web-connection-provider.tsx +++ b/apps/mobile/src/components/agents/user-web-connection-provider.tsx @@ -25,7 +25,7 @@ export function UserWebConnectionProvider({ children }: Readonly { - const result = await trpcClient.activeSessions.getToken.query(); + const result = await trpcClient.activeSessions.getToken.mutate(); return result.token; }, lifecycleHooks: createNativeUserWebConnectionLifecycleHooks(), diff --git a/apps/web/src/components/cloud-agent-next/CloudAgentProvider.tsx b/apps/web/src/components/cloud-agent-next/CloudAgentProvider.tsx index 7c921b0c476..385dd34b097 100644 --- a/apps/web/src/components/cloud-agent-next/CloudAgentProvider.tsx +++ b/apps/web/src/components/cloud-agent-next/CloudAgentProvider.tsx @@ -69,7 +69,7 @@ export function CloudAgentProvider({ children, organizationId }: CloudAgentProvi sharedConnectionRef.current = createUserWebConnection({ websocketUrl: `${SESSION_INGEST_WS_URL}/api/user/web`, getAuthToken: async () => { - const result = await trpcClient.activeSessions.getToken.query(); + const result = await trpcClient.activeSessions.getToken.mutate(); return result.token; }, lifecycleHooks: createBrowserLifecycleHooks(), diff --git a/apps/web/src/routers/active-sessions-router.test.ts b/apps/web/src/routers/active-sessions-router.test.ts index 382ab7a79c1..b7e0acf93bf 100644 --- a/apps/web/src/routers/active-sessions-router.test.ts +++ b/apps/web/src/routers/active-sessions-router.test.ts @@ -72,6 +72,79 @@ describe('active-sessions-router', () => { jest.restoreAllMocks(); }); + describe('getToken', () => { + it('mints a one-use ticket and returns { token, expiresAt }', async () => { + const fetchSpy = jest.spyOn(global, 'fetch').mockResolvedValue( + new Response(JSON.stringify({ ticket: 'ticket-abc', expiresAt: 1_700_000_060 }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) + ); + + const caller = await createCallerForUser(regularUser.id); + const result = await caller.activeSessions.getToken(); + + expect(result).toEqual({ token: 'ticket-abc', expiresAt: 1_700_000_060 }); + + const [calledUrl, init] = fetchSpy.mock.calls[0] as unknown as [string, RequestInit]; + expect(calledUrl).toBe('https://test-ingest.example.com/api/user/web-ticket'); + expect(init.method).toBe('POST'); + expect((init.headers as Record).Authorization).toContain('Bearer '); + }); + + it('throws PRECONDITION_FAILED when the worker returns a non-2xx response', async () => { + jest + .spyOn(global, 'fetch') + .mockResolvedValue(new Response('upstream failed', { status: 502 })); + + const caller = await createCallerForUser(regularUser.id); + const rejection = caller.activeSessions.getToken(); + await expect(rejection).rejects.toBeInstanceOf(TRPCError); + try { + await rejection; + } catch (err) { + if (!(err instanceof TRPCError)) throw err; + expect(err.code).toBe('PRECONDITION_FAILED'); + expect(err.message).toBe('Session ingest is not configured'); + } + }); + + it('throws PRECONDITION_FAILED when fetch itself fails (network error)', async () => { + jest.spyOn(global, 'fetch').mockRejectedValue(new Error('socket hang up')); + + const caller = await createCallerForUser(regularUser.id); + const rejection = caller.activeSessions.getToken(); + await expect(rejection).rejects.toBeInstanceOf(TRPCError); + try { + await rejection; + } catch (err) { + if (!(err instanceof TRPCError)) throw err; + expect(err.code).toBe('PRECONDITION_FAILED'); + expect(err.message).toBe('Session ingest is not configured'); + } + }); + + it('throws PRECONDITION_FAILED when the worker returns a malformed 200 body', async () => { + jest.spyOn(global, 'fetch').mockResolvedValue( + new Response(JSON.stringify({ ticket: '', expiresAt: 'not-a-number' }), { + status: 200, + headers: { 'Content-Type': 'application/json' }, + }) + ); + + const caller = await createCallerForUser(regularUser.id); + const rejection = caller.activeSessions.getToken(); + await expect(rejection).rejects.toBeInstanceOf(TRPCError); + try { + await rejection; + } catch (err) { + if (!(err instanceof TRPCError)) throw err; + expect(err.code).toBe('PRECONDITION_FAILED'); + expect(err.message).toBe('Invalid ticket response from session ingest'); + } + }); + }); + describe('listInstances', () => { it('returns the instances from the worker when the upstream call succeeds', async () => { const fetchSpy = jest.spyOn(global, 'fetch').mockResolvedValue( diff --git a/apps/web/src/routers/active-sessions-router.ts b/apps/web/src/routers/active-sessions-router.ts index 376e7ae5acb..0a2868998de 100644 --- a/apps/web/src/routers/active-sessions-router.ts +++ b/apps/web/src/routers/active-sessions-router.ts @@ -65,6 +65,15 @@ const connectedInstancesResponseSchema = z.object({ instances: z.array(connectedInstanceSchema), }); +/** + * Session Ingest `/api/user/web-ticket` mint response. Parsed at runtime so a + * malformed 200 fails the mutation instead of returning undefined fields. + */ +const webTicketResponseSchema = z.object({ + ticket: z.string().min(1), + expiresAt: z.number(), +}); + /** * A live session as this router returns it: the worker's wire row plus the * fields enriched from `cli_sessions_v2`. @@ -220,9 +229,54 @@ function throwOrgContextFailure(error: unknown): never { } export const activeSessionsRouter = createTRPCRouter({ - getToken: baseProcedure.query(async ({ ctx }) => { + /** + * Mint a one-use web ticket from Session Ingest for the authenticated user. + * The returned `token` is the opaque ticket; `expiresAt` is the Unix-seconds + * expiry from the worker body. A missing worker URL or a non-2xx mint + * response fails fast with PRECONDITION_FAILED rather than hanging. + */ + getToken: baseProcedure.mutation(async ({ ctx }) => { + if (!SESSION_INGEST_WORKER_URL) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Session ingest is not configured', + }); + } + const token = generateInternalServiceToken(ctx.user.id); - return { token }; + const url = `${SESSION_INGEST_WORKER_URL}/api/user/web-ticket`; + + let response: Response; + try { + response = await fetch(url, { + method: 'POST', + headers: { Authorization: `Bearer ${token}` }, + }); + } catch (error) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Session ingest is not configured', + cause: error, + }); + } + + if (!response.ok) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Session ingest is not configured', + }); + } + + const raw = await response.json(); + const parsed = webTicketResponseSchema.safeParse(raw); + if (!parsed.success) { + throw new TRPCError({ + code: 'PRECONDITION_FAILED', + message: 'Invalid ticket response from session ingest', + cause: parsed.error, + }); + } + return { token: parsed.data.ticket, expiresAt: parsed.data.expiresAt }; }), list: baseProcedure.input(listInputSchema).query(async ({ ctx, input }) => { diff --git a/packages/cloud-agent-sdk/src/user-web-connection.test.ts b/packages/cloud-agent-sdk/src/user-web-connection.test.ts index 88abcad83a8..3b4fed409bb 100644 --- a/packages/cloud-agent-sdk/src/user-web-connection.test.ts +++ b/packages/cloud-agent-sdk/src/user-web-connection.test.ts @@ -116,20 +116,20 @@ describe('createUserWebConnection', () => { const release = client.retain(); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?source=web&token=token+with+spaces&connectionId=uuid-1` + `${WS_URL}?source=web&ticket=token+with+spaces&connectionId=uuid-1` ); release(); const secondRelease = client.retain(); expect(webSocketConstructor).toHaveBeenLastCalledWith( - `${WS_URL}?source=web&token=token+with+spaces&connectionId=uuid-1` + `${WS_URL}?source=web&ticket=token+with+spaces&connectionId=uuid-1` ); secondRelease(); const other = createUserWebConnection({ websocketUrl: WS_URL, getAuthToken: () => 'token' }); const otherRelease = other.retain(); expect(webSocketConstructor).toHaveBeenLastCalledWith( - `${WS_URL}?token=token&connectionId=uuid-2` + `${WS_URL}?ticket=token&connectionId=uuid-2` ); otherRelease(); client.destroy(); @@ -203,7 +203,7 @@ describe('createUserWebConnection', () => { expect(sockets[0].close).toHaveBeenCalledTimes(1); expect(getAuthToken).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenLastCalledWith( - `${WS_URL}?token=new-token&connectionId=uuid-1` + `${WS_URL}?ticket=new-token&connectionId=uuid-1` ); open(sockets[1]); expect(sockets[1].send).toHaveBeenCalledWith( @@ -463,7 +463,7 @@ describe('createUserWebConnection', () => { client.connect(); expect(webSocketConstructor).toHaveBeenCalledTimes(1); - expect(webSocketConstructor).toHaveBeenCalledWith(`${WS_URL}?token=token&connectionId=uuid-1`); + expect(webSocketConstructor).toHaveBeenCalledWith(`${WS_URL}?ticket=token&connectionId=uuid-1`); client.destroy(); }); @@ -485,7 +485,7 @@ describe('createUserWebConnection', () => { expect(webSocketConstructor).toHaveBeenCalledTimes(1); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=async-token&connectionId=uuid-1` + `${WS_URL}?ticket=async-token&connectionId=uuid-1` ); client.destroy(); }); @@ -552,7 +552,7 @@ describe('createUserWebConnection', () => { expect(getAuthToken).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=recovered-token&connectionId=uuid-1` + `${WS_URL}?ticket=recovered-token&connectionId=uuid-1` ); release(); client.destroy(); @@ -590,7 +590,7 @@ describe('createUserWebConnection', () => { expect(getAuthToken).toHaveBeenCalledTimes(10); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=eventual-token&connectionId=uuid-1` + `${WS_URL}?ticket=eventual-token&connectionId=uuid-1` ); release(); client.destroy(); @@ -634,7 +634,7 @@ describe('createUserWebConnection', () => { expect(getAuthToken).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=online-token&connectionId=uuid-1` + `${WS_URL}?ticket=online-token&connectionId=uuid-1` ); expect(removeOnline).toHaveBeenCalledTimes(1); open(); @@ -682,7 +682,7 @@ describe('createUserWebConnection', () => { expect(getAuthToken).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=restored-token&connectionId=uuid-1` + `${WS_URL}?ticket=restored-token&connectionId=uuid-1` ); release(); client.destroy(); @@ -722,7 +722,7 @@ describe('createUserWebConnection', () => { expect(getAuthToken).toHaveBeenCalledTimes(2); expect(webSocketConstructor).toHaveBeenCalledWith( - `${WS_URL}?token=foreground-token&connectionId=uuid-1` + `${WS_URL}?ticket=foreground-token&connectionId=uuid-1` ); release(); client.destroy(); diff --git a/packages/cloud-agent-sdk/src/user-web-connection.ts b/packages/cloud-agent-sdk/src/user-web-connection.ts index 7973d438f8b..98c28e82607 100644 --- a/packages/cloud-agent-sdk/src/user-web-connection.ts +++ b/packages/cloud-agent-sdk/src/user-web-connection.ts @@ -416,7 +416,7 @@ function createUserWebConnection( function buildUrl(): string { const url = new URL(config.websocketUrl); - url.searchParams.set('token', token); + url.searchParams.set('ticket', token); url.searchParams.set('connectionId', connectionId); return url.toString(); } From a9a00a1b4fc4f85a28c6ba739a27509158141258 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 02:18:45 +0200 Subject: [PATCH 04/23] test(mobile): assert getToken mints the ingest ticket as a mutation --- ...r-web-connection-provider.mounted.test.tsx | 85 +++++++++++++++++++ 1 file changed, 85 insertions(+) create mode 100644 apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx diff --git a/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx b/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx new file mode 100644 index 00000000000..2e8c1e36021 --- /dev/null +++ b/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx @@ -0,0 +1,85 @@ +/* eslint-disable typescript-eslint/no-deprecated -- react-test-renderer is the DOM-free renderer used to mount React/RN trees under vitest (same pattern as src/test/render-with-providers.tsx) */ +import { createElement } from 'react'; +import TestRenderer, { act } from 'react-test-renderer'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { UserWebConnectionProvider } from './user-web-connection-provider'; + +type AuthConfig = { getAuthToken: () => Promise }; + +const mocks = vi.hoisted(() => ({ + mutate: vi.fn(async () => ({ token: 'ticket-1', expiresAt: 1_700_000_060 })), + query: vi.fn(), + createUserWebConnection: vi.fn(), + capturedConfig: null as AuthConfig | null, +})); + +vi.mock('@kilocode/cloud-agent-sdk/user-web-connection', () => ({ + createUserWebConnection: (config: AuthConfig) => { + mocks.capturedConfig = config; + mocks.createUserWebConnection(config); + return { + retain: () => () => {}, + connect: () => {}, + disconnect: () => {}, + destroy: () => {}, + isConnected: () => false, + onConnectionChange: () => () => {}, + subscribeToCliSession: () => () => {}, + sendCommand: async () => ({}), + sendCommandToConnection: async () => ({}), + onCliEvent: () => () => {}, + onSystemEvent: () => () => {}, + onReconnect: () => () => {}, + onSessionEvent: () => () => {}, + }; + }, +})); + +vi.mock('@/lib/config', () => ({ + SESSION_INGEST_WS_URL: 'wss://ingest.example.com', +})); + +vi.mock('@/lib/user-web-connection-lifecycle', () => ({ + createNativeUserWebConnectionLifecycleHooks: () => ({}), +})); + +vi.mock('@/lib/trpc', () => ({ + trpcClient: { + activeSessions: { + getToken: { + mutate: mocks.mutate, + query: mocks.query, + }, + }, + }, +})); + +describe('UserWebConnectionProvider', () => { + beforeEach(() => { + mocks.capturedConfig = null; + mocks.mutate.mockClear(); + mocks.query.mockClear(); + mocks.createUserWebConnection.mockClear(); + }); + + it('mints the ingest ticket via the getToken mutation', async () => { + let renderer: TestRenderer.ReactTestRenderer | undefined; + await act(async () => { + renderer = TestRenderer.create(createElement(UserWebConnectionProvider, { children: null })); + }); + + expect(mocks.createUserWebConnection).toHaveBeenCalledTimes(1); + expect(mocks.capturedConfig).not.toBeNull(); + + const token = await mocks.capturedConfig!.getAuthToken(); + + expect(token).toBe('ticket-1'); + expect(mocks.mutate).toHaveBeenCalledTimes(1); + expect(mocks.query).not.toHaveBeenCalled(); + + await act(async () => { + renderer!.unmount(); + }); + }); +}); From 6a0218827f7b3a95c3e95ecdde33983a5aea6374 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 03:08:15 +0200 Subject: [PATCH 05/23] feat(cloud-agent): bind stream tickets to audience and consume nonce once --- apps/web/src/lib/cloud-agent/stream-ticket.ts | 2 + services/cloud-agent-next/src/auth.test.ts | 45 ++- services/cloud-agent-next/src/auth.ts | 11 +- services/cloud-agent-next/src/index.ts | 1 + .../CloudAgentSession-stream-ticket.test.ts | 181 ++++++++++++ .../src/persistence/CloudAgentSession.ts | 4 +- .../src/persistence/StreamTicketNonceDO.ts | 33 +++ .../src/server-stream-ticket.test.ts | 269 ++++++++++++++++++ services/cloud-agent-next/src/server.test.ts | 23 +- services/cloud-agent-next/src/server.ts | 29 +- services/cloud-agent-next/src/types.ts | 3 + .../worker-configuration.d.ts | 2 + services/cloud-agent-next/wrangler.jsonc | 12 + services/cloud-agent-next/wrangler.test.jsonc | 8 + 14 files changed, 608 insertions(+), 15 deletions(-) create mode 100644 services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts create mode 100644 services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts create mode 100644 services/cloud-agent-next/src/server-stream-ticket.test.ts diff --git a/apps/web/src/lib/cloud-agent/stream-ticket.ts b/apps/web/src/lib/cloud-agent/stream-ticket.ts index 06acc937806..b5800aa11f8 100644 --- a/apps/web/src/lib/cloud-agent/stream-ticket.ts +++ b/apps/web/src/lib/cloud-agent/stream-ticket.ts @@ -16,6 +16,7 @@ export function signStreamTicket( expiresInSeconds = 60 ): { ticket: string; expiresAt: number } { const expiresAt = Math.floor(Date.now() / 1000) + expiresInSeconds; + const audience = payload.purpose === 'terminal' ? 'cloud-agent-terminal' : 'cloud-agent-stream'; const ticket = jwt.sign( { @@ -27,6 +28,7 @@ export function signStreamTicket( { algorithm: 'HS256', expiresIn: expiresInSeconds, + audience, } ); diff --git a/services/cloud-agent-next/src/auth.test.ts b/services/cloud-agent-next/src/auth.test.ts index f5bf2d11970..4c07fd7d1b6 100644 --- a/services/cloud-agent-next/src/auth.test.ts +++ b/services/cloud-agent-next/src/auth.test.ts @@ -34,7 +34,7 @@ vi.mock('@kilocode/db/client', () => ({ describe('validateStreamTicket', () => { it('returns a configuration error when NEXTAUTH_SECRET is missing', () => { - expect(validateStreamTicket('ticket', null)).toEqual({ + expect(validateStreamTicket('ticket', null, 'cloud-agent-stream')).toEqual({ success: false, error: 'NEXTAUTH_SECRET is not configured on the worker', }); @@ -48,10 +48,10 @@ describe('validateStreamTicket', () => { cloudAgentSessionId: 'session-1', }, secret, - { algorithm: 'HS256', expiresIn: -1 } + { algorithm: 'HS256', expiresIn: -1, audience: 'cloud-agent-stream' } ); - expect(validateStreamTicket(ticket, secret)).toEqual({ + expect(validateStreamTicket(ticket, secret, 'cloud-agent-stream')).toEqual({ success: false, error: 'Ticket expired', }); @@ -65,10 +65,10 @@ describe('validateStreamTicket', () => { cloudAgentSessionId: 'session-1', }, secret, - { algorithm: 'HS256', expiresIn: '1 minute' } + { algorithm: 'HS256', expiresIn: '1 minute', audience: 'cloud-agent-stream' } ); - expect(validateStreamTicket(ticket, secret)).toMatchObject({ + expect(validateStreamTicket(ticket, secret, 'cloud-agent-stream')).toMatchObject({ success: true, payload: { type: 'stream_ticket', @@ -77,6 +77,41 @@ describe('validateStreamTicket', () => { }, }); }); + + it('rejects a stream ticket presented with the terminal audience', () => { + const ticket = jwt.sign( + { + type: 'stream_ticket', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + }, + secret, + { algorithm: 'HS256', expiresIn: '1 minute', audience: 'cloud-agent-stream' } + ); + + expect(validateStreamTicket(ticket, secret, 'cloud-agent-terminal')).toEqual({ + success: false, + error: 'Invalid ticket audience', + }); + }); + + it('rejects a terminal ticket presented with the stream audience', () => { + const ticket = jwt.sign( + { + type: 'stream_ticket', + purpose: 'terminal', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + }, + secret, + { algorithm: 'HS256', expiresIn: '1 minute', audience: 'cloud-agent-terminal' } + ); + + expect(validateStreamTicket(ticket, secret, 'cloud-agent-stream')).toEqual({ + success: false, + error: 'Invalid ticket audience', + }); + }); }); const wrapperDispatchTicketClaims: WrapperDispatchTicketClaims = { diff --git a/services/cloud-agent-next/src/auth.ts b/services/cloud-agent-next/src/auth.ts index 622bbd5e99c..261174f56f2 100644 --- a/services/cloud-agent-next/src/auth.ts +++ b/services/cloud-agent-next/src/auth.ts @@ -12,8 +12,12 @@ type StreamTicketPayload = { organizationId?: string; ptyId?: string; nonce?: string; + aud?: string; }; +export const STREAM_TICKET_AUDIENCE = 'cloud-agent-stream'; +export const TERMINAL_TICKET_AUDIENCE = 'cloud-agent-terminal'; + export type WrapperDispatchTicketClaims = { type: 'wrapper_dispatch_ticket'; userId: string; @@ -96,7 +100,8 @@ export async function validateKiloToken( export function validateStreamTicket( ticket: string | null, - secret: string | null | undefined + secret: string | null | undefined, + expectedAudience: string ): { success: true; payload: StreamTicketPayload } | { success: false; error: string } { if (!ticket) { return { success: false, error: 'Missing stream ticket' }; @@ -114,6 +119,10 @@ export function validateStreamTicket( return { success: false, error: 'Invalid ticket type' }; } + if (payload.aud !== expectedAudience) { + return { success: false, error: 'Invalid ticket audience' }; + } + return { success: true, payload }; } catch (error) { if (error instanceof jwt.TokenExpiredError) { diff --git a/services/cloud-agent-next/src/index.ts b/services/cloud-agent-next/src/index.ts index 34a9406791d..d9081e6b9e9 100644 --- a/services/cloud-agent-next/src/index.ts +++ b/services/cloud-agent-next/src/index.ts @@ -10,4 +10,5 @@ export { ContainerProxy, } from './sandbox-outbound.js'; export { CloudAgentSession } from './persistence/CloudAgentSession.js'; +export { StreamTicketNonceDO } from './persistence/StreamTicketNonceDO.js'; export { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts new file mode 100644 index 00000000000..c6eb348c925 --- /dev/null +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts @@ -0,0 +1,181 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; + +const { createStreamHandlerMock, createAgentRuntimeMock } = vi.hoisted(() => ({ + createStreamHandlerMock: vi.fn(), + createAgentRuntimeMock: vi.fn(), +})); + +vi.mock('@cloudflare/sandbox', () => ({ + Sandbox: class Sandbox {}, + getSandbox: vi.fn(), + ContainerProxy: class ContainerProxy {}, +})); + +vi.mock('@cloudflare/containers', () => ({})); + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + ctx: unknown; + env: unknown; + constructor(ctx: unknown, env: unknown) { + this.ctx = ctx; + this.env = env; + } + }, +})); + +vi.mock('../logger.js', () => { + const logger = { + setTags: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + withFields: vi.fn(), + }; + logger.withFields.mockReturnValue(logger); + return { + logger, + withLogTags: async (_tags: unknown, fn: () => Promise) => fn(), + WithLogTags: () => (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) => + descriptor, + }; +}); + +vi.mock('drizzle-orm/durable-sqlite', () => ({ + drizzle: vi.fn(() => ({})), +})); + +vi.mock('drizzle-orm/durable-sqlite/migrator', () => ({ + migrate: vi.fn(), +})); + +vi.mock('../../drizzle/migrations', () => ({ + default: { journal: {}, migrations: {} }, +})); + +vi.mock('../session/queries/index.js', () => ({ + createExecutionQueries: vi.fn(() => ({})), + createEventQueries: vi.fn(() => ({})), + createLeaseQueries: vi.fn(() => ({})), +})); + +vi.mock('../websocket/stream.js', () => ({ + createStreamHandler: createStreamHandlerMock, + getConnectedStreamClientCount: vi.fn(() => 0), +})); + +vi.mock('../session/agent-runtime.js', () => ({ + createAgentRuntime: createAgentRuntimeMock, +})); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [{ api_token_pepper: null, blocked_reason: null }], + }), + }), + }), + }), +})); + +const { CloudAgentSession } = await import('./CloudAgentSession.js'); + +const secret = 'test-secret'; + +type MockStorage = { + get: ReturnType; + sql: unknown; + getAlarm: ReturnType; + setAlarm: ReturnType; +}; + +function createCtx() { + const storage: MockStorage = { + get: vi.fn().mockResolvedValue(undefined), + sql: {}, + getAlarm: vi.fn().mockResolvedValue(null), + setAlarm: vi.fn().mockResolvedValue(undefined), + }; + const ctx = { + id: { name: 'user-1:session-1' }, + storage, + blockConcurrencyWhile: vi.fn().mockResolvedValue(undefined), + }; + return ctx; +} + +function createEnv() { + return { + NEXTAUTH_SECRET: secret, + WS_ALLOWED_ORIGINS: '', + }; +} + +function signTicket(audience: string, purpose?: 'stream' | 'terminal'): string { + return jwt.sign( + { + type: 'stream_ticket', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + ...(purpose ? { purpose } : {}), + nonce: 'nonce-1', + }, + secret, + { algorithm: 'HS256', expiresIn: 60, audience } + ); +} + +function streamRequest(ticket: string): Request { + return new Request( + `http://worker.test/stream?cloudAgentSessionId=session-1&ticket=${encodeURIComponent(ticket)}`, + { headers: { Upgrade: 'websocket' } } + ); +} + +beforeEach(() => { + createStreamHandlerMock.mockReset(); + createAgentRuntimeMock.mockReset(); +}); + +describe('CloudAgentSession /stream ticket audience', () => { + it('accepts a valid stream-audience ticket without consuming the nonce', async () => { + const handleStreamRequest = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); + createStreamHandlerMock.mockReturnValue({ handleStreamRequest }); + createAgentRuntimeMock.mockReturnValue({ requestSnapshot: vi.fn() }); + + const env = { + ...createEnv(), + // A nonce DO whose methods throw: the DO path must never touch it. + STREAM_TICKET_NONCE_DO: { + idFromName: vi.fn(() => { + throw new Error('consume must not be called on the DO path'); + }), + get: vi.fn(() => { + throw new Error('consume must not be called on the DO path'); + }), + }, + }; + + const doInstance = new CloudAgentSession(createCtx() as never, env as never); + const response = await doInstance.fetch(streamRequest(signTicket('cloud-agent-stream'))); + + expect(response.status).toBe(200); + expect(handleStreamRequest).toHaveBeenCalledOnce(); + }); + + it('rejects a terminal-audience ticket on the DO /stream path', async () => { + const env = createEnv(); + const doInstance = new CloudAgentSession(createCtx() as never, env as never); + + const response = await doInstance.fetch( + streamRequest(signTicket('cloud-agent-terminal', 'terminal')) + ); + + expect(response.status).toBe(401); + await expect(response.text()).resolves.toBe('Invalid ticket audience'); + expect(createStreamHandlerMock).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index e6ae60ae797..99bb7d1b52f 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -96,7 +96,7 @@ import { recordSharedSandboxFailover } from '../shared-sandbox-route.js'; import { nextMetadataAfterAdmittedAgentModel } from './persist-admitted-agent-model.js'; import { dispatchedKilocodeModelId } from './model-utils.js'; -import { resolveSecret, validateStreamTicket } from '../auth.js'; +import { resolveSecret, validateStreamTicket, STREAM_TICKET_AUDIENCE } from '../auth.js'; import { isAllowedStreamWebSocketOrigin } from './ws-origin.js'; import { resolveTerminalWrapperClient, type TerminalWrapperClient } from '../terminal/access.js'; import type { WrapperPty } from '../kilo/wrapper-client.js'; @@ -985,7 +985,7 @@ export class CloudAgentSession extends DurableObject { } const nextAuthSecret = await resolveSecret(this.env.NEXTAUTH_SECRET); - const authResult = validateStreamTicket(ticket, nextAuthSecret); + const authResult = validateStreamTicket(ticket, nextAuthSecret, STREAM_TICKET_AUDIENCE); if (!authResult.success) { return new Response(authResult.error, { status: 401 }); } diff --git a/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts b/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts new file mode 100644 index 00000000000..b94064f8dda --- /dev/null +++ b/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts @@ -0,0 +1,33 @@ +import { DurableObject } from 'cloudflare:workers'; +import type { Env } from '../types.js'; + +/** + * One-time nonce consumer for stream/terminal tickets. + * + * The nonce is an unguessable UUID minted inside a signed JWT. This DO makes + * the first use of a nonce the only successful one: `consume()` atomically + * marks the nonce used and returns true, or returns false on any replay. + */ +export class StreamTicketNonceDO extends DurableObject { + async consume(): Promise { + return this.ctx.storage.transaction(async txn => { + const used = await txn.get('used'); + if (used !== undefined) { + return false; + } + await txn.put('used', { usedAt: Date.now() }); + return true; + }); + } +} + +/** + * The only consume call site for stream/terminal ticket nonces. Kept separate + * from `validateStreamTicket` so signature/type/expiry/audience validation + * never consumes the nonce. + */ +export async function consumeStreamTicketNonce(env: Env, nonce: string): Promise { + const doId = env.STREAM_TICKET_NONCE_DO.idFromName(nonce); + const stub = env.STREAM_TICKET_NONCE_DO.get(doId); + return stub.consume(); +} diff --git a/services/cloud-agent-next/src/server-stream-ticket.test.ts b/services/cloud-agent-next/src/server-stream-ticket.test.ts new file mode 100644 index 00000000000..0d5863c53aa --- /dev/null +++ b/services/cloud-agent-next/src/server-stream-ticket.test.ts @@ -0,0 +1,269 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; +import type { Env } from './types.js'; + +const { + getRunningTerminalClientMock, + consumeCloudAgentReportBatchMock, + removeExpiredCloudAgentReportDataMock, + requireCurrentSessionAccessMock, +} = vi.hoisted(() => ({ + getRunningTerminalClientMock: vi.fn(), + consumeCloudAgentReportBatchMock: vi.fn().mockResolvedValue(undefined), + removeExpiredCloudAgentReportDataMock: vi.fn().mockResolvedValue(undefined), + requireCurrentSessionAccessMock: vi.fn(), +})); + +vi.mock('./logger.js', () => { + const logger = { + setTags: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + withFields: vi.fn(), + }; + logger.withFields.mockReturnValue(logger); + + return { + logger, + withLogTags: async (_tags: unknown, fn: () => Promise) => fn(), + WithLogTags: () => (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) => + descriptor, + }; +}); + +vi.mock('@cloudflare/sandbox', () => ({ + Sandbox: class Sandbox {}, + getSandbox: vi.fn(), +})); + +vi.mock('./agent-sandbox/factory.js', () => ({ + createAgentSandbox: vi.fn(() => ({ + getRunningTerminalClient: getRunningTerminalClientMock, + })), +})); + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + constructor(_state: unknown, _env: unknown) {} + }, +})); + +vi.mock('./router.js', () => ({ + appRouter: {}, +})); + +vi.mock('./callbacks/index.js', () => ({ + createCallbackQueueConsumer: vi.fn(), +})); + +vi.mock('./telemetry/report-consumer.js', () => ({ + CLOUD_AGENT_REPORT_QUEUE_NAMES: new Set([ + 'cloud-agent-next-report-queue', + 'cloud-agent-next-report-queue-dev', + 'cloud-agent-next-report-queue-test', + ]), + consumeCloudAgentReportBatch: consumeCloudAgentReportBatchMock, + removeExpiredCloudAgentReportData: removeExpiredCloudAgentReportDataMock, +})); + +vi.mock('./middleware/auth.js', () => ({ + authMiddleware: vi.fn(), +})); + +vi.mock('./middleware/balance.js', () => ({ + balanceMiddleware: vi.fn(), +})); + +vi.mock('./session-access.js', () => ({ + requireCurrentSessionAccess: requireCurrentSessionAccessMock, + projectSessionAccessHttpError: (error: unknown) => + new Response( + error instanceof Error && 'code' in error && error.code === 'FORBIDDEN' + ? 'Session access denied' + : 'Session access is temporarily unavailable', + { + status: error instanceof Error && 'code' in error && error.code === 'FORBIDDEN' ? 403 : 503, + } + ), +})); + +vi.mock('./persistence/CloudAgentSession.js', () => ({ + CloudAgentSession: class CloudAgentSession {}, +})); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [{ api_token_pepper: null, blocked_reason: null }], + }), + }), + }), + }), +})); + +const { default: worker } = await import('./server.js'); + +const secret = 'test-secret'; + +type MockEnv = { + NEXTAUTH_SECRET: string; + Sandbox: unknown; + SandboxSmall: unknown; + WS_ALLOWED_ORIGINS?: string; + HYPERDRIVE: { connectionString: string }; + CLOUD_AGENT_SESSION: { + idFromName: ReturnType; + get: ReturnType; + }; + STREAM_TICKET_NONCE_DO: { + idFromName: ReturnType; + get: ReturnType; + }; +}; + +function createEnv(): MockEnv { + return { + NEXTAUTH_SECRET: secret, + Sandbox: {}, + SandboxSmall: {}, + HYPERDRIVE: { connectionString: 'postgres://test' }, + CLOUD_AGENT_SESSION: { + idFromName: vi.fn(), + get: vi.fn(), + }, + STREAM_TICKET_NONCE_DO: { + idFromName: vi.fn(), + get: vi.fn(), + }, + }; +} + +/** + * Stateful in-memory nonce store that mirrors StreamTicketNonceDO's + * consume-on-first-use semantics: the first consume returns true, any replay + * returns false. + */ +function installNonceStore(env: MockEnv): Set { + const used = new Set(); + env.STREAM_TICKET_NONCE_DO.idFromName.mockImplementation((nonce: string) => nonce); + env.STREAM_TICKET_NONCE_DO.get.mockImplementation((nonce: string) => ({ + consume: async () => { + if (used.has(nonce)) return false; + used.add(nonce); + return true; + }, + })); + return used; +} + +function fetchWorker(request: Request, env: MockEnv): Promise | Response { + return worker.fetch(request, env as unknown as Env, {} as ExecutionContext); +} + +function signStreamTicket(overrides: Record = {}): string { + return jwt.sign( + { + type: 'stream_ticket', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + nonce: 'nonce-1', + ...overrides, + }, + secret, + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-stream' } + ); +} + +function streamRequest(ticket: string): Request { + return new Request( + `http://worker.test/stream?cloudAgentSessionId=session-1&ticket=${encodeURIComponent(ticket)}`, + { headers: { Upgrade: 'websocket' } } + ); +} + +beforeEach(() => { + getRunningTerminalClientMock.mockReset(); + consumeCloudAgentReportBatchMock.mockClear(); + removeExpiredCloudAgentReportDataMock.mockClear(); + requireCurrentSessionAccessMock.mockReset().mockResolvedValue({ + kiloSessionId: 'ses_12345678901234567890123456', + organizationId: null, + }); +}); + +describe('server /stream ticket nonce consume', () => { + it('accepts the first use of a stream ticket and forwards to the session DO', async () => { + const env = createEnv(); + installNonceStore(env); + const doFetch = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ fetch: doFetch }); + const ticket = signStreamTicket(); + + const response = await fetchWorker(streamRequest(ticket), env); + + expect(response.status).toBe(200); + expect(env.STREAM_TICKET_NONCE_DO.idFromName).toHaveBeenCalledWith('nonce-1'); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('user-1:session-1'); + expect(doFetch).toHaveBeenCalledOnce(); + }); + + it('rejects an immediate replay of the same nonce on the worker path', async () => { + const env = createEnv(); + installNonceStore(env); + const doFetch = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ fetch: doFetch }); + const ticket = signStreamTicket(); + + const first = await fetchWorker(streamRequest(ticket), env); + expect(first.status).toBe(200); + + const replay = await fetchWorker(streamRequest(ticket), env); + expect(replay.status).toBe(401); + await expect(replay.text()).resolves.toBe('Ticket nonce already used'); + expect(doFetch).toHaveBeenCalledTimes(1); + }); + + it('rejects a terminal-audience ticket on /stream', async () => { + const env = createEnv(); + installNonceStore(env); + const ticket = jwt.sign( + { + type: 'stream_ticket', + purpose: 'terminal', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + nonce: 'nonce-1', + }, + secret, + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } + ); + + const response = await fetchWorker(streamRequest(ticket), env); + + expect(response.status).toBe(401); + await expect(response.text()).resolves.toBe('Invalid ticket audience'); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + expect(env.STREAM_TICKET_NONCE_DO.idFromName).not.toHaveBeenCalled(); + }); + + it('rejects a stream ticket missing a nonce before forwarding', async () => { + const env = createEnv(); + installNonceStore(env); + const doFetch = vi.fn(); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ fetch: doFetch }); + const ticket = signStreamTicket({ nonce: undefined }); + + const response = await fetchWorker(streamRequest(ticket), env); + + expect(response.status).toBe(401); + await expect(response.text()).resolves.toBe('Missing ticket nonce'); + expect(env.STREAM_TICKET_NONCE_DO.idFromName).not.toHaveBeenCalled(); + expect(doFetch).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 0d35349e7c0..6260d6f32d0 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -123,6 +123,10 @@ type MockEnv = { idFromName: ReturnType; get: ReturnType; }; + STREAM_TICKET_NONCE_DO: { + idFromName: ReturnType; + get: ReturnType; + }; }; function createEnv(): MockEnv { @@ -139,6 +143,10 @@ function createEnv(): MockEnv { idFromName: vi.fn(), get: vi.fn(), }, + STREAM_TICKET_NONCE_DO: { + idFromName: vi.fn(), + get: vi.fn(), + }, }; } @@ -219,7 +227,7 @@ describe('server /stream', () => { cloudAgentSessionId: 'session-1', }, secret, - { algorithm: 'HS256', expiresIn: 60 } + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-stream' } ); const env = createEnv(); requireCurrentSessionAccessMock.mockRejectedValue( @@ -294,11 +302,16 @@ describe('server /terminal', () => { userId: 'user-1', cloudAgentSessionId: 'session-1', ptyId: 'pty_123', + nonce: 'nonce-1', }, secret, - { algorithm: 'HS256', expiresIn: 60 } + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } ); const env = createEnv(); + env.STREAM_TICKET_NONCE_DO.idFromName.mockReturnValue('nonce-do-id'); + env.STREAM_TICKET_NONCE_DO.get.mockReturnValue({ + consume: vi.fn().mockResolvedValue(true), + }); const metadata = { metadataSchemaVersion: 2, identity: { @@ -359,7 +372,7 @@ describe('server /terminal', () => { ptyId: 'pty_123', }, secret, - { algorithm: 'HS256', expiresIn: 60 } + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } ); requireCurrentSessionAccessMock.mockRejectedValue( Object.assign(new Error('Session access denied'), { code: 'FORBIDDEN' }) @@ -387,7 +400,7 @@ describe('server /terminal', () => { cloudAgentSessionId: 'session-1', }, secret, - { algorithm: 'HS256', expiresIn: 60 } + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } ); const env = createEnv(); const request = new Request( @@ -414,7 +427,7 @@ describe('server /terminal', () => { ptyId: 'pty_other', }, secret, - { algorithm: 'HS256', expiresIn: 60 } + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } ); const env = createEnv(); const request = new Request( diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 533a3128100..50c6bd5815c 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -10,8 +10,11 @@ import { validateStreamTicket, validateKiloToken, validateWrapperDispatchTicket, + STREAM_TICKET_AUDIENCE, + TERMINAL_TICKET_AUDIENCE, type WrapperAuthClaims, } from './auth.js'; +import { consumeStreamTicketNonce } from './persistence/StreamTicketNonceDO.js'; import { createErrorHandler, createNotFoundHandler } from '@kilocode/worker-utils'; import { createCallbackQueueConsumer } from './callbacks/index.js'; import type { CallbackJob } from './callbacks/index.js'; @@ -75,7 +78,7 @@ async function handleTerminalWebSocket(request: Request, env: Env): Promise) => { } const nextAuthSecret = await resolveSecret(c.env.NEXTAUTH_SECRET); - const ticketResult = validateStreamTicket(ticket, nextAuthSecret); + const ticketResult = validateStreamTicket(ticket, nextAuthSecret, STREAM_TICKET_AUDIENCE); if (!ticketResult.success) { logger .withFields({ cloudAgentSessionId, error: ticketResult.error }) @@ -314,6 +328,17 @@ app.get('/stream', async (c: Context) => { return projectSessionAccessHttpError(error); } + const nonce = ticketResult.payload.nonce; + if (!nonce) { + logger.withFields({ cloudAgentSessionId, userId }).warn('/stream: Missing ticket nonce'); + return c.text('Missing ticket nonce', 401); + } + const consumed = await consumeStreamTicketNonce(c.env, nonce); + if (!consumed) { + logger.withFields({ cloudAgentSessionId, userId }).warn('/stream: Ticket nonce already used'); + return c.text('Ticket nonce already used', 401); + } + logger.withFields({ cloudAgentSessionId, userId }).info('/stream: WebSocket upgrade authorized'); const doId = c.env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${cloudAgentSessionId}`); diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index fc74fbc80f4..7f0f65329fb 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -3,6 +3,7 @@ import type { CloudAgentSession } from './persistence/CloudAgentSession.js'; import type { CloudAgentQueueReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; import type { AccessibleCloudAgentSession } from '@kilocode/worker-utils/cloud-agent-session-access'; import type { UserKiloFacade } from './kilo-facade/user-kilo-facade.js'; +import type { StreamTicketNonceDO } from './persistence/StreamTicketNonceDO.js'; import type { CallbackJob } from './callbacks/index.js'; import type { NotificationsBinding } from './notifications-binding.js'; import type { SessionIngestBinding } from './session-ingest-binding.js'; @@ -513,6 +514,8 @@ export type Env = { CLOUD_AGENT_SESSION: DurableObjectNamespace; /** Durable Object namespace for per-user Kilo SDK facade coordination */ USER_KILO_FACADE: DurableObjectNamespace; + /** Durable Object namespace for one-time stream/terminal ticket nonce consumption */ + STREAM_TICKET_NONCE_DO: DurableObjectNamespace; /** One-way shared sandbox failover overrides keyed by shared identity */ SHARED_SANDBOX_OVERRIDES: KVNamespace; /** Service binding for the session ingest worker */ diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index 4b28db0973a..cb4dd29e5da 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -46,6 +46,7 @@ interface __BaseEnv_Env { SandboxSmallContainment: DurableObjectNamespace; SandboxCodeReviewContainment: DurableObjectNamespace; CLOUD_AGENT_SESSION: DurableObjectNamespace; + STREAM_TICKET_NONCE_DO: DurableObjectNamespace; USER_KILO_FACADE: DurableObjectNamespace; SESSION_INGEST: Service /* entrypoint SessionIngestRPC from session-ingest */; GIT_TOKEN_SERVICE: Service /* entrypoint GitTokenRPCEntrypoint from git-token-service-dev */ | Service /* entrypoint GitTokenRPCEntrypoint from git-token-service */; @@ -105,6 +106,7 @@ declare namespace Cloudflare { SandboxSmallContainment: DurableObjectNamespace; SandboxCodeReviewContainment: DurableObjectNamespace; CLOUD_AGENT_SESSION: DurableObjectNamespace; + STREAM_TICKET_NONCE_DO: DurableObjectNamespace; USER_KILO_FACADE: DurableObjectNamespace; SESSION_INGEST: Service /* entrypoint SessionIngestRPC from session-ingest */; GIT_TOKEN_SERVICE: Service /* entrypoint GitTokenRPCEntrypoint from git-token-service-dev */; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index 5d412d4c6e9..d8e49a67250 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -290,6 +290,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "StreamTicketNonceDO", + "name": "STREAM_TICKET_NONCE_DO", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", @@ -329,6 +333,10 @@ ], "tag": "v7", }, + { + "new_sqlite_classes": ["StreamTicketNonceDO"], + "tag": "v8", + }, ], /** * Hyperdrive Bindings @@ -575,6 +583,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "StreamTicketNonceDO", + "name": "STREAM_TICKET_NONCE_DO", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", diff --git a/services/cloud-agent-next/wrangler.test.jsonc b/services/cloud-agent-next/wrangler.test.jsonc index b5e59d06f4b..bb9c4f14535 100644 --- a/services/cloud-agent-next/wrangler.test.jsonc +++ b/services/cloud-agent-next/wrangler.test.jsonc @@ -11,6 +11,10 @@ "class_name": "CloudAgentSession", "name": "CLOUD_AGENT_SESSION", }, + { + "class_name": "StreamTicketNonceDO", + "name": "STREAM_TICKET_NONCE_DO", + }, { "class_name": "UserKiloFacade", "name": "USER_KILO_FACADE", @@ -26,6 +30,10 @@ "new_sqlite_classes": ["UserKiloFacade"], "tag": "v5", }, + { + "new_sqlite_classes": ["StreamTicketNonceDO"], + "tag": "v8", + }, ], "kv_namespaces": [{ "binding": "SHARED_SANDBOX_OVERRIDES" }], "queues": { From 77d0813b1f19cf009520f223fc9db4c3d91410b7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 11:32:39 +0200 Subject: [PATCH 06/23] fix(organizations): read primary db for group policy admission P2-D-24b: getOrganizationGroupPolicyContext without tx opened readDb. The openrouter gateway admission path consumes that membership read. Read the primary (db) instead, keeping repeatable-read read-only. --- .../organization-group-policy-context.server.test.ts | 8 ++++---- .../organization-group-policy-context.server.ts | 4 ++-- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/apps/web/src/lib/organizations/organization-group-policy-context.server.test.ts b/apps/web/src/lib/organizations/organization-group-policy-context.server.test.ts index 5217b41b80c..20baf72349a 100644 --- a/apps/web/src/lib/organizations/organization-group-policy-context.server.test.ts +++ b/apps/web/src/lib/organizations/organization-group-policy-context.server.test.ts @@ -11,9 +11,9 @@ const mockedPrimaryTransaction = jest.mocked(db.transaction); const mockedReplicaTransaction = jest.mocked(readDb.transaction); describe('getOrganizationGroupPolicyContext', () => { - it('opens its repeatable-read snapshot on the read replica', async () => { + it('opens its repeatable-read snapshot on the primary database', async () => { const transactionError = new Error('stop after selecting the transaction client'); - mockedReplicaTransaction.mockRejectedValueOnce(transactionError); + mockedPrimaryTransaction.mockRejectedValueOnce(transactionError); await expect( getOrganizationGroupPolicyContext({ @@ -22,10 +22,10 @@ describe('getOrganizationGroupPolicyContext', () => { }) ).rejects.toBe(transactionError); - expect(mockedReplicaTransaction).toHaveBeenCalledWith(expect.any(Function), { + expect(mockedPrimaryTransaction).toHaveBeenCalledWith(expect.any(Function), { isolationLevel: 'repeatable read', accessMode: 'read only', }); - expect(mockedPrimaryTransaction).not.toHaveBeenCalled(); + expect(mockedReplicaTransaction).not.toHaveBeenCalled(); }); }); diff --git a/apps/web/src/lib/organizations/organization-group-policy-context.server.ts b/apps/web/src/lib/organizations/organization-group-policy-context.server.ts index b33670919ce..f5123323eab 100644 --- a/apps/web/src/lib/organizations/organization-group-policy-context.server.ts +++ b/apps/web/src/lib/organizations/organization-group-policy-context.server.ts @@ -15,7 +15,7 @@ import { import { and, eq, isNull } from 'drizzle-orm'; import { TRPCError } from '@trpc/server'; import { captureException } from '@sentry/nextjs'; -import { readDb, type DrizzleTransaction } from '@/lib/drizzle'; +import { db, type DrizzleTransaction } from '@/lib/drizzle'; export type OrganizationPolicySubject = | { @@ -106,7 +106,7 @@ export async function getOrganizationGroupPolicyContext(params: { tx?: DrizzleTransaction; }): Promise { if (!params.tx) { - return await readDb.transaction( + return await db.transaction( async tx => await getOrganizationGroupPolicyContext({ ...params, tx }), { isolationLevel: 'repeatable read', accessMode: 'read only' } ); From af3f65a76af773133d6048fcf0ff9de3736a9c54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 11:32:43 +0200 Subject: [PATCH 07/23] fix(mobile): store iOS bearer tokens with WHEN_UNLOCKED_THIS_DEVICE_ONLY P2-D-23a: bearer writes and sign-out deletes used no keychain class. Use SecureStore.WHEN_UNLOCKED_THIS_DEVICE_ONLY so the tokens are not in iCloud or iTunes backup and do not migrate to a new device. --- apps/mobile/src/lib/auth/auth-context.test.ts | 7 +- .../mobile/src/lib/auth/auth-context.test.tsx | 26 ++- apps/mobile/src/lib/auth/auth-context.tsx | 10 +- apps/mobile/src/lib/auth/credentials.test.ts | 204 ++++++++++++++++++ apps/mobile/src/lib/auth/credentials.ts | 17 +- apps/mobile/src/lib/trpc.test.ts | 1 + 6 files changed, 251 insertions(+), 14 deletions(-) create mode 100644 apps/mobile/src/lib/auth/credentials.test.ts diff --git a/apps/mobile/src/lib/auth/auth-context.test.ts b/apps/mobile/src/lib/auth/auth-context.test.ts index 7487f5edce0..5768763cc59 100644 --- a/apps/mobile/src/lib/auth/auth-context.test.ts +++ b/apps/mobile/src/lib/auth/auth-context.test.ts @@ -23,6 +23,7 @@ function releaseCredentialWrites(): void { /* eslint-disable import/first */ // vi.mock is hoisted by Vitest before the real import resolves. vi.mock('expo-secure-store', () => ({ + WHEN_UNLOCKED_THIS_DEVICE_ONLY: 'WHEN_UNLOCKED_THIS_DEVICE_ONLY', getItemAsync: vi.fn(async (key: string) => { await Promise.resolve(); return store.get(key) ?? null; @@ -599,7 +600,11 @@ describe('exchangeLegacyToken', () => { // Let the exchange pass its pre-persist epoch check and block at the // fenced AUTH_TOKEN write, then supersede it. await vi.waitFor(() => { - expect(vi.mocked(SecureStore.setItemAsync)).toHaveBeenCalledWith(AUTH_TOKEN_KEY, 'exchanged'); + expect(vi.mocked(SecureStore.setItemAsync)).toHaveBeenCalledWith( + AUTH_TOKEN_KEY, + 'exchanged', + expect.anything() + ); }); bumpAuthEpoch(); releaseCredentialWrites(); diff --git a/apps/mobile/src/lib/auth/auth-context.test.tsx b/apps/mobile/src/lib/auth/auth-context.test.tsx index 37e3afba700..4b6edb993df 100644 --- a/apps/mobile/src/lib/auth/auth-context.test.tsx +++ b/apps/mobile/src/lib/auth/auth-context.test.tsx @@ -94,6 +94,7 @@ vi.mock('expo-secure-store', () => ({ setItem: hoisted.secureStore.setItem, setItemAsync: hoisted.secureStore.setItemAsync, deleteItemAsync: hoisted.secureStore.deleteItemAsync, + WHEN_UNLOCKED_THIS_DEVICE_ONLY: 'WHEN_UNLOCKED_THIS_DEVICE_ONLY', })); vi.mock('@sentry/react-native', () => ({ @@ -322,7 +323,10 @@ describe('sign-out teardown ordering', () => { await ctx.signOut(); }); - expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('auth-token'); + expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith( + 'auth-token', + expect.anything() + ); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('organization'); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('session-filters'); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('notification-prompt-seen'); @@ -1070,7 +1074,10 @@ describe('auth-transition queue and sign-out failure matrix', () => { expect(hoisted.posthog.discardPostHog).toHaveBeenCalledTimes(1); const { queryClient: queryClientMock } = await import('@/lib/query-client'); expect(vi.mocked(queryClientMock.clear)).toHaveBeenCalledTimes(1); - expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('auth-token'); + expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith( + 'auth-token', + expect.anything() + ); expect(getCtx().token).toBeUndefined(); expect(getCtx().sessionEnded).toBe(false); @@ -1087,7 +1094,10 @@ describe('auth-transition queue and sign-out failure matrix', () => { // Cleanup still ran, the deletion batch still ran, and auth state reset. expect(logoutCleanupMock.runLogoutCleanup).toHaveBeenCalledTimes(1); - expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('auth-token'); + expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith( + 'auth-token', + expect.anything() + ); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('organization'); const { queryClient: queryClientMock } = await import('@/lib/query-client'); expect(vi.mocked(queryClientMock.clear)).toHaveBeenCalledTimes(1); @@ -1116,7 +1126,10 @@ describe('auth-transition queue and sign-out failure matrix', () => { expect(currentAuthEpoch()).toBeGreaterThan(before); const tokenOwner = await import('@/lib/auth/token-owner'); expect(tokenOwner.getActiveToken()).toBeNull(); - expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('auth-token'); + expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith( + 'auth-token', + expect.anything() + ); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('active-user-id'); const { queryClient: queryClientMock } = await import('@/lib/query-client'); expect(vi.mocked(queryClientMock.clear)).toHaveBeenCalledTimes(1); @@ -1136,7 +1149,10 @@ describe('auth-transition queue and sign-out failure matrix', () => { }); // All independent batch members still ran despite the one rejection. - expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('auth-token'); + expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith( + 'auth-token', + expect.anything() + ); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('organization'); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('session-filters'); expect(hoisted.secureStore.deleteItemAsync).toHaveBeenCalledWith('active-user-id'); diff --git a/apps/mobile/src/lib/auth/auth-context.tsx b/apps/mobile/src/lib/auth/auth-context.tsx index 5b76dd51ab3..112357da0ae 100644 --- a/apps/mobile/src/lib/auth/auth-context.tsx +++ b/apps/mobile/src/lib/auth/auth-context.tsx @@ -21,6 +21,7 @@ import { setTrpcUnauthorizedHandler } from '@/lib/auth/trpc-unauthorized'; import { exchangeLegacyToken } from '@/lib/auth/exchange-legacy-token'; import { bumpAuthEpoch, currentAuthEpoch, isCurrentAuthEpoch } from '@/lib/auth/auth-epoch'; import { + IOS_BEARER_SECURE_STORE_OPTIONS, performRefresh, persistSignInCredentialsAtEpoch, REFRESH_MARGIN_MS, @@ -256,9 +257,12 @@ export function AuthProvider({ children }: { readonly children: ReactNode }) { // deletion are members of the same always-attempted batch. await Promise.allSettled([ writeCredentials(async () => { - await SecureStore.deleteItemAsync(AUTH_TOKEN_KEY); - await SecureStore.deleteItemAsync(REFRESH_TOKEN_KEY); - await SecureStore.deleteItemAsync(TOKEN_EXPIRES_AT_KEY); + await SecureStore.deleteItemAsync(AUTH_TOKEN_KEY, IOS_BEARER_SECURE_STORE_OPTIONS); + await SecureStore.deleteItemAsync(REFRESH_TOKEN_KEY, IOS_BEARER_SECURE_STORE_OPTIONS); + await SecureStore.deleteItemAsync( + TOKEN_EXPIRES_AT_KEY, + IOS_BEARER_SECURE_STORE_OPTIONS + ); await SecureStore.deleteItemAsync(LEGACY_EXCHANGE_DONE_KEY); }), deleteAccountMetadata(ACTIVE_USER_ID_KEY), diff --git a/apps/mobile/src/lib/auth/credentials.test.ts b/apps/mobile/src/lib/auth/credentials.test.ts new file mode 100644 index 00000000000..3bf60154ba7 --- /dev/null +++ b/apps/mobile/src/lib/auth/credentials.test.ts @@ -0,0 +1,204 @@ +/* oxlint-disable typescript-eslint/no-deprecated -- react-test-renderer is the DOM-free renderer for RN trees under vitest (node env, no jsdom) */ +/* oxlint-disable @typescript-eslint/no-unsafe-call @typescript-eslint/no-unsafe-member-access */ +import { createElement } from 'react'; +import TestRenderer, { act } from 'react-test-renderer'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +const store = new Map(); + +/* eslint-disable import/first */ +// vi.mock is hoisted by Vitest before the real import resolves. +vi.mock('expo-secure-store', () => ({ + WHEN_UNLOCKED_THIS_DEVICE_ONLY: 'WHEN_UNLOCKED_THIS_DEVICE_ONLY', + getItem: vi.fn((key: string) => store.get(key) ?? null), + setItem: vi.fn((key: string, value: string) => { + store.set(key, value); + }), + getItemAsync: vi.fn(async (key: string) => { + await Promise.resolve(); + return store.get(key) ?? null; + }), + setItemAsync: vi.fn(async (key: string, value: string) => { + await Promise.resolve(); + store.set(key, value); + }), + deleteItemAsync: vi.fn(async (key: string) => { + await Promise.resolve(); + store.delete(key); + }), +})); + +vi.mock('@/lib/config', () => ({ API_BASE_URL: 'https://api.example.com' })); + +// The sign-out deletes live in auth-context.tsx; mounting it pulls in the full +// teardown graph, so stub every side-effecting collaborator. +vi.mock('@sentry/react-native', () => ({ setUser: vi.fn() })); +vi.mock('@/lib/analytics/posthog', () => ({ + discardPostHog: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/appsflyer', () => ({ resetAppsFlyerState: vi.fn(), trackEvent: vi.fn() })); +vi.mock('@/lib/auth/account-metadata-write', () => ({ + deleteAccountMetadata: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/auth/logout-cleanup', () => ({ + runLogoutCleanup: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/query-client', () => ({ + queryClient: { clear: vi.fn(), invalidateQueries: vi.fn() }, +})); +vi.mock('@/lib/auth/trpc-unauthorized', () => ({ setTrpcUnauthorizedHandler: vi.fn() })); +vi.mock('@/lib/hooks/use-persisted-agent-model', () => ({ clearAgentModelPreference: vi.fn() })); +vi.mock('@/lib/hooks/use-keep-screen-on-preference', () => ({ + clearKeepScreenOnPreference: vi.fn(), +})); +vi.mock('@/lib/hooks/use-reasoning-preference', () => ({ clearReasoningPreference: vi.fn() })); +vi.mock('@/lib/kiloclaw-tab-ownership', () => ({ + gateKiloClawOwned: vi.fn(), + clearKiloClawOwned: vi.fn(), +})); +vi.mock('@/lib/last-active-instance', () => ({ + clearLastActiveInstance: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/kilo-pass/use-store-kilo-pass-purchase', () => ({ + resetPurchaseErrorToastDedup: vi.fn(), +})); +vi.mock('@/lib/persist/read-cache', () => ({ + clearCacheScopeForSignOut: vi.fn().mockResolvedValue(undefined), + readCachedUserId: vi.fn().mockReturnValue(null), +})); +vi.mock('@/lib/pr-review/recent-prs', () => ({ + clearRecentPrs: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/pr-review/viewed-files', () => ({ + clearViewedFiles: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('@/lib/telemetry/controller', () => ({ clearTelemetryDecision: vi.fn() })); +vi.mock('@/lib/telemetry/posthog-storage', () => ({ + purgePostHogPersistence: vi.fn().mockResolvedValue(undefined), +})); +vi.mock('react-native', () => ({ + AppState: { addEventListener: vi.fn(() => ({ remove: vi.fn() })) }, +})); + +import * as SecureStore from 'expo-secure-store'; +import { persistSignInCredentialsAtEpoch } from '@/lib/auth/credentials'; +import { bumpAuthEpoch } from '@/lib/auth/auth-epoch'; +import { clearActiveToken, setSignOutTeardownActive } from '@/lib/auth/token-owner'; +import { + AUTH_TOKEN_KEY, + LEGACY_EXCHANGE_DONE_KEY, + REFRESH_TOKEN_KEY, + TOKEN_EXPIRES_AT_KEY, +} from '@/lib/storage-keys'; +/* eslint-enable import/first */ + +// Apple `kSecAttrAccessibleWhenUnlockedThisDeviceOnly` is not in iCloud or +// iTunes backup and does not migrate to a new device. The assertions below +// prove every bearer write and delete pins the keychain to that class. +const expectedOptions = { keychainAccessible: SecureStore.WHEN_UNLOCKED_THIS_DEVICE_ONLY }; + +describe('bearer credential writes', () => { + beforeEach(() => { + vi.clearAllMocks(); + store.clear(); + clearActiveToken(); + setSignOutTeardownActive(false); + }); + + it('writes every bearer key with WHEN_UNLOCKED_THIS_DEVICE_ONLY', async () => { + const published = await persistSignInCredentialsAtEpoch('token', 'refresh', { + expiresIn: 3600, + }); + + expect(published).toBe(true); + expect(SecureStore.setItemAsync).toHaveBeenCalledWith(AUTH_TOKEN_KEY, 'token', expectedOptions); + expect(SecureStore.setItemAsync).toHaveBeenCalledWith( + REFRESH_TOKEN_KEY, + 'refresh', + expectedOptions + ); + expect(SecureStore.setItemAsync).toHaveBeenCalledWith( + TOKEN_EXPIRES_AT_KEY, + expect.any(String), + expectedOptions + ); + }); + + it('deletes the prior refresh pair with the option on a token-only sign-in', async () => { + store.set(REFRESH_TOKEN_KEY, 'old-refresh'); + store.set(TOKEN_EXPIRES_AT_KEY, '999'); + + await persistSignInCredentialsAtEpoch('token-only', undefined, {}); + + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(REFRESH_TOKEN_KEY, expectedOptions); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(TOKEN_EXPIRES_AT_KEY, expectedOptions); + }); + + it('clears a partial pair with the option when the epoch moves mid-write', async () => { + vi.mocked(SecureStore.setItemAsync).mockImplementationOnce(async (key, value) => { + store.set(key, value); + bumpAuthEpoch(); + await Promise.resolve(); + }); + + const published = await persistSignInCredentialsAtEpoch('token', 'refresh', { + expiresIn: 3600, + }); + + expect(published).toBe(false); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(AUTH_TOKEN_KEY, expectedOptions); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(REFRESH_TOKEN_KEY, expectedOptions); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(TOKEN_EXPIRES_AT_KEY, expectedOptions); + }); +}); + +describe('sign-out deletes', () => { + beforeEach(() => { + vi.clearAllMocks(); + store.clear(); + }); + + it('deletes the three bearer keys with WHEN_UNLOCKED_THIS_DEVICE_ONLY', async () => { + await mountAndSignOut(); + + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(AUTH_TOKEN_KEY, expectedOptions); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(REFRESH_TOKEN_KEY, expectedOptions); + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(TOKEN_EXPIRES_AT_KEY, expectedOptions); + // The legacy-exchange marker is not a bearer key: it keeps the default class. + expect(SecureStore.deleteItemAsync).toHaveBeenCalledWith(LEGACY_EXCHANGE_DONE_KEY); + }); +}); + +/** Mount the AuthProvider and run sign-out so the credential deletes execute. */ +async function mountAndSignOut(): Promise { + vi.resetModules(); + const mod = await import('./auth-context'); + + let captured: { signOut: (ended?: boolean) => Promise } | undefined = undefined; + function Consumer(): null { + captured = mod.useAuth(); + return null; + } + + let renderer: TestRenderer.ReactTestRenderer | undefined = undefined; + await act(async () => { + renderer = TestRenderer.create(createElement(mod.AuthProvider, null, createElement(Consumer))); + await Promise.resolve(); + }); + await act(async () => { + await new Promise(resolve => { + void setTimeout(resolve, 0); + }); + }); + + // oxlint-disable-next-line @typescript-eslint/no-unnecessary-condition -- safety net for test failures + if (!captured) { + throw new Error('auth context not captured'); + } + + await act(async () => { + await captured.signOut(); + }); + + renderer?.unmount(); +} diff --git a/apps/mobile/src/lib/auth/credentials.ts b/apps/mobile/src/lib/auth/credentials.ts index e7b608bfcd2..ef6af237b8e 100644 --- a/apps/mobile/src/lib/auth/credentials.ts +++ b/apps/mobile/src/lib/auth/credentials.ts @@ -8,6 +8,13 @@ import { chainSave } from '@/lib/hooks/save-chain'; import { AUTH_TOKEN_KEY, REFRESH_TOKEN_KEY, TOKEN_EXPIRES_AT_KEY } from '@/lib/storage-keys'; import { CONTROL_PLANE_DEADLINE_MS, withDeadline } from '@kilocode/event-service'; +// Apple `kSecAttrAccessibleWhenUnlockedThisDeviceOnly` is not in iCloud or +// iTunes backup and does not migrate to a new device. Pin every bearer-token +// write and delete to it so credentials never follow a device restore. +export const IOS_BEARER_SECURE_STORE_OPTIONS = { + keychainAccessible: SecureStore.WHEN_UNLOCKED_THIS_DEVICE_ONLY, +}; + /** * Credential persistence and refresh rotation, with no React dependency. * @@ -71,9 +78,9 @@ export async function persistSignInCredentialsAtEpoch( // a newer sign-in or sign-out: their own credential write is queued // strictly behind this one. const clearPartialCredentials = async (): Promise => { - await SecureStore.deleteItemAsync(AUTH_TOKEN_KEY); - await SecureStore.deleteItemAsync(REFRESH_TOKEN_KEY); - await SecureStore.deleteItemAsync(TOKEN_EXPIRES_AT_KEY); + await SecureStore.deleteItemAsync(AUTH_TOKEN_KEY, IOS_BEARER_SECURE_STORE_OPTIONS); + await SecureStore.deleteItemAsync(REFRESH_TOKEN_KEY, IOS_BEARER_SECURE_STORE_OPTIONS); + await SecureStore.deleteItemAsync(TOKEN_EXPIRES_AT_KEY, IOS_BEARER_SECURE_STORE_OPTIONS); }; // Fence one credential operation: skip it when the epoch moved before the @@ -83,8 +90,8 @@ export async function persistSignInCredentialsAtEpoch( return false; } await (value === undefined - ? SecureStore.deleteItemAsync(key) - : SecureStore.setItemAsync(key, value)); + ? SecureStore.deleteItemAsync(key, IOS_BEARER_SECURE_STORE_OPTIONS) + : SecureStore.setItemAsync(key, value, IOS_BEARER_SECURE_STORE_OPTIONS)); if (!isCurrentAuthEpoch(epoch)) { await clearPartialCredentials(); return false; diff --git a/apps/mobile/src/lib/trpc.test.ts b/apps/mobile/src/lib/trpc.test.ts index eea3a3153f3..1b2ebb30ca6 100644 --- a/apps/mobile/src/lib/trpc.test.ts +++ b/apps/mobile/src/lib/trpc.test.ts @@ -61,6 +61,7 @@ vi.mock('expo-secure-store', () => ({ getItemAsync: secureStoreMock.getItemAsync, setItemAsync: secureStoreMock.setItemAsync, deleteItemAsync: secureStoreMock.deleteItemAsync, + WHEN_UNLOCKED_THIS_DEVICE_ONLY: 'WHEN_UNLOCKED_THIS_DEVICE_ONLY', })); vi.mock('@/lib/config', () => ({ From e4c971ddd6c5885a5cda9a6278a7ce83099998cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 11:57:04 +0200 Subject: [PATCH 08/23] feat(cloud-agent): return ticket and expiresAt from every getTicket client P2-GH-54c: mobile and extension getTicket parsed only { ticket } and returned a string, so the SDK set expiresAt to undefined. Return { ticket, expiresAt } and fail when expiresAt is missing. CodeReviewStreamView keeps a string handoff to websocket-manager. --- .../extension-agent-session-manager.test.ts | 27 ++++- .../shared/extension-agent-session-manager.ts | 15 ++- .../agents/mobile-session-manager.test.ts | 105 +++++++++++++++++- .../agents/mobile-session-manager.ts | 19 +++- .../code-reviews/CodeReviewStreamView.tsx | 30 ++--- .../code-reviews/fetch-stream-ticket.test.ts | 62 +++++++++++ .../code-reviews/fetch-stream-ticket.ts | 28 +++++ 7 files changed, 247 insertions(+), 39 deletions(-) create mode 100644 apps/web/src/components/code-reviews/fetch-stream-ticket.test.ts create mode 100644 apps/web/src/components/code-reviews/fetch-stream-ticket.ts diff --git a/apps/extension/src/shared/extension-agent-session-manager.test.ts b/apps/extension/src/shared/extension-agent-session-manager.test.ts index 3f852ba2827..25ab9811522 100644 --- a/apps/extension/src/shared/extension-agent-session-manager.test.ts +++ b/apps/extension/src/shared/extension-agent-session-manager.test.ts @@ -418,14 +418,14 @@ describe('createExtensionAgentSessionManager', () => { describe('getTicket', () => { it('builds correct ticket URL and returns ticket', async () => { const { mock, restore } = withFakeFetch( - Response.json({ ticket: 'ticket-123' }, { status: 200 }) + Response.json({ expiresAt: 1_700_000_000, ticket: 'ticket-123' }, { status: 200 }) ); try { const opts = makeDefaultOptions(); vi.spyOn(opts, 'getToken').mockReturnValue('bearer-token'); createExtensionAgentSessionManager(opts); - const ticket = await capturedConfig!.getTicket(CLOUD_AGENT_ID as never); - expect(ticket).toBe('ticket-123'); + const result = await capturedConfig!.getTicket(CLOUD_AGENT_ID as never); + expect(result).toStrictEqual({ expiresAt: 1_700_000_000, ticket: 'ticket-123' }); expect(mock).toHaveBeenCalledTimes(1); const [callUrl] = mock.mock.calls[0] as unknown as [string]; expect(callUrl).toBe('https://api.test/api/cloud-agent-next/sessions/stream-ticket'); @@ -447,7 +447,7 @@ describe('createExtensionAgentSessionManager', () => { it('includes organizationId in body when org is set', async () => { const { mock, restore } = withFakeFetch( - Response.json({ ticket: 'ticket-org' }, { status: 200 }) + Response.json({ expiresAt: 1_700_000_000, ticket: 'ticket-org' }, { status: 200 }) ); try { const opts = { @@ -469,7 +469,7 @@ describe('createExtensionAgentSessionManager', () => { it('omits Authorization header when getToken returns undefined', async () => { const { mock, restore } = withFakeFetch( - Response.json({ ticket: 'ticket-noauth' }, { status: 200 }) + Response.json({ expiresAt: 1_700_000_000, ticket: 'ticket-noauth' }, { status: 200 }) ); try { const opts = makeDefaultOptions(); @@ -488,7 +488,7 @@ describe('createExtensionAgentSessionManager', () => { it('omits Authorization header when getToken returns empty string', async () => { const { mock, restore } = withFakeFetch( - Response.json({ ticket: 'ticket-empty' }, { status: 200 }) + Response.json({ expiresAt: 1_700_000_000, ticket: 'ticket-empty' }, { status: 200 }) ); try { const opts = makeDefaultOptions(); @@ -527,6 +527,21 @@ describe('createExtensionAgentSessionManager', () => { restore(); } }); + + it('throws when expiresAt is missing from response', async () => { + const { restore } = withFakeFetch( + Response.json({ ticket: 'ticket-noexpiry' }, { status: 200 }) + ); + try { + const opts = makeDefaultOptions(); + createExtensionAgentSessionManager(opts); + await expect(capturedConfig!.getTicket(CLOUD_AGENT_ID as never)).rejects.toThrow( + 'Missing expiresAt' + ); + } finally { + restore(); + } + }); }); // ---- fetchSnapshot ---- diff --git a/apps/extension/src/shared/extension-agent-session-manager.ts b/apps/extension/src/shared/extension-agent-session-manager.ts index 0804615d889..f2a52869450 100644 --- a/apps/extension/src/shared/extension-agent-session-manager.ts +++ b/apps/extension/src/shared/extension-agent-session-manager.ts @@ -649,7 +649,9 @@ export function createExtensionAgentSessionManager({ }), // ---- getTicket ---- - getTicket: async (sessionId: CloudAgentSessionId): Promise => { + getTicket: async ( + sessionId: CloudAgentSessionId + ): Promise<{ ticket: string; expiresAt: number }> => { const token = getToken(); const body: Record = { cloudAgentSessionId: sessionId }; if (organizationId !== null) { @@ -664,14 +666,21 @@ export function createExtensionAgentSessionManager({ method: 'POST', }); // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion -- fetch json is untyped - const data = (await response.json()) as { ticket?: string; error?: string }; + const data = (await response.json()) as { + ticket?: string; + expiresAt?: number; + error?: string; + }; if (!response.ok) { throw new Error(data.error ?? 'Failed to get stream ticket'); } if (data.ticket === undefined) { throw new Error('Missing ticket in stream-ticket response'); } - return data.ticket; + if (data.expiresAt === undefined) { + throw new Error('Missing expiresAt in stream-ticket response'); + } + return { expiresAt: data.expiresAt, ticket: data.ticket }; }, // ---- initiate ---- diff --git a/apps/mobile/src/components/agents/mobile-session-manager.test.ts b/apps/mobile/src/components/agents/mobile-session-manager.test.ts index 2be65c829f3..a14cde4b520 100644 --- a/apps/mobile/src/components/agents/mobile-session-manager.test.ts +++ b/apps/mobile/src/components/agents/mobile-session-manager.test.ts @@ -4,7 +4,11 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import { type AgentAttachmentSubmissionPayload } from '@/lib/agent-attachments/agent-attachment-types'; import { SPAWNED_NOT_FOUND_MAX_ATTEMPTS } from '@/lib/spawned-not-found-retry'; -import { type KiloSessionId } from '@kilocode/cloud-agent-sdk'; +import { + type KiloSessionId, + type SessionManager, + type SessionManagerConfig, +} from '@kilocode/cloud-agent-sdk'; // Keep this suite on the pure vitest project: mock every RN / Expo / SDK // side-effect import that `mobile-session-manager.ts` pulls transitively @@ -15,8 +19,21 @@ vi.mock('expo-secure-store', () => ({ vi.mock('sonner-native', () => ({ toast: { error: vi.fn(), success: vi.fn() }, })); +// Capture the config passed to createSessionManager so the getTicket closure +// can be asserted directly, mirroring the extension suite. +const { configHolder, mockCreateSessionManager } = vi.hoisted(() => { + const holder: { current: SessionManagerConfig | null } = { current: null }; + const createSessionManagerMock = vi.fn((config: SessionManagerConfig): SessionManager => { + holder.current = config; + return { atoms: {} } as unknown as SessionManager; + }); + return { configHolder: holder, mockCreateSessionManager: createSessionManagerMock }; +}); vi.mock('@kilocode/cloud-agent-sdk', () => ({ - createSessionManager: vi.fn(), + createSessionManager: mockCreateSessionManager, +})); +vi.mock('@/lib/auth/token-owner', () => ({ + getAuthTokenForRequest: vi.fn(async () => 'test-token'), })); vi.mock('@/components/agents/mobile-session-transport-payload', () => ({ normalizeTransportPayload: vi.fn((x: unknown) => x), @@ -57,8 +74,12 @@ vi.mock('@/lib/trpc', () => ({ const { buildRemoteAttachmentParts } = await import('@/components/agents/mobile-session-manager-helpers'); -const { fetchSessionWithNotFoundRetry, isCloudPrepareRetryableError, readFetchSessionErrorCode } = - await import('@/components/agents/mobile-session-manager'); +const { + createMobileAgentSessionManager, + fetchSessionWithNotFoundRetry, + isCloudPrepareRetryableError, + readFetchSessionErrorCode, +} = await import('@/components/agents/mobile-session-manager'); const SESSION_ID = 'ses_test_session_id_0000000001' as KiloSessionId; @@ -72,6 +93,18 @@ function withCode(code: string, message: string): Error { return Object.assign(new Error(message), { data: { code } }); } +function withFakeFetch(response: Response) { + const originalFetch = globalThis.fetch; + const mock = vi.fn(async () => response); + globalThis.fetch = mock as never; + return { + mock, + restore: () => { + globalThis.fetch = originalFetch; + }, + }; +} + describe('buildRemoteAttachmentParts', () => { beforeEach(() => { mutate.mockReset(); @@ -330,3 +363,67 @@ describe('fetchSessionWithNotFoundRetry', () => { expect(sleep).not.toHaveBeenCalled(); }); }); + +describe('getTicket', () => { + const CLOUD_AGENT_ID = 'agent_12345678-1234-1234-1234-123456789abc'; + + beforeEach(() => { + configHolder.current = null; + mockCreateSessionManager.mockClear(); + }); + + function setup(): SessionManagerConfig { + const options = { + store: {}, + userWebConnection: {}, + }; + createMobileAgentSessionManager(options as never); + const config = configHolder.current; + if (config === null) { + throw new Error('createSessionManager did not capture a config'); + } + return config; + } + + it('returns the ticket and expiresAt from the stream-ticket response', async () => { + const { mock, restore } = withFakeFetch( + Response.json({ ticket: 'ticket-123', expiresAt: 1_700_000_000 }, { status: 200 }) + ); + try { + const config = setup(); + await expect(config.getTicket(CLOUD_AGENT_ID as never)).resolves.toEqual({ + ticket: 'ticket-123', + expiresAt: 1_700_000_000, + }); + expect(mock).toHaveBeenCalledTimes(1); + } finally { + restore(); + } + }); + + it('throws when expiresAt is missing from the response', async () => { + const { restore } = withFakeFetch( + Response.json({ ticket: 'ticket-noexpiry' }, { status: 200 }) + ); + try { + const config = setup(); + await expect(config.getTicket(CLOUD_AGENT_ID as never)).rejects.toThrow( + 'Missing expiresAt in stream-ticket response' + ); + } finally { + restore(); + } + }); + + it('throws when ticket is missing from the response', async () => { + const { restore } = withFakeFetch(Response.json({ expiresAt: 1_700_000_000 }, { status: 200 })); + try { + const config = setup(); + await expect(config.getTicket(CLOUD_AGENT_ID as never)).rejects.toThrow( + 'Missing ticket in stream-ticket response' + ); + } finally { + restore(); + } + }); +}); diff --git a/apps/mobile/src/components/agents/mobile-session-manager.ts b/apps/mobile/src/components/agents/mobile-session-manager.ts index 55880e624dd..8cde2507d7d 100644 --- a/apps/mobile/src/components/agents/mobile-session-manager.ts +++ b/apps/mobile/src/components/agents/mobile-session-manager.ts @@ -207,8 +207,10 @@ export function createMobileAgentSessionManager({ ...(activeSession.capabilities ? { capabilities: activeSession.capabilities } : {}), }; }, - getTicket: async (sessionId: CloudAgentSessionId): Promise => { - const ticket = await withCloudAgentDiagnostics('getTicket', organizationId, async () => { + getTicket: async ( + sessionId: CloudAgentSessionId + ): Promise<{ ticket: string; expiresAt: number }> => { + const result = await withCloudAgentDiagnostics('getTicket', organizationId, async () => { const token = await getAuthTokenForRequest(); const body: Record = { cloudAgentSessionId: sessionId }; if (organizationId) { @@ -225,16 +227,23 @@ export function createMobileAgentSessionManager({ body: JSON.stringify(body), } ); - const data = (await response.json()) as { ticket?: string; error?: string }; + const data = (await response.json()) as { + ticket?: string; + expiresAt?: number; + error?: string; + }; if (!response.ok) { throw new Error(data.error ?? 'Failed to get stream ticket'); } if (!data.ticket) { throw new Error('Missing ticket in stream-ticket response'); } - return data.ticket; + if (data.expiresAt === undefined) { + throw new Error('Missing expiresAt in stream-ticket response'); + } + return { ticket: data.ticket, expiresAt: data.expiresAt }; }); - return ticket; + return result; }, fetchSnapshot: async (id: KiloSessionId) => { const [sessionData, messagesResult] = await Promise.all([ diff --git a/apps/web/src/components/code-reviews/CodeReviewStreamView.tsx b/apps/web/src/components/code-reviews/CodeReviewStreamView.tsx index fe70744d06e..0b39f74a95d 100644 --- a/apps/web/src/components/code-reviews/CodeReviewStreamView.tsx +++ b/apps/web/src/components/code-reviews/CodeReviewStreamView.tsx @@ -21,6 +21,7 @@ import { import type { CloudAgentEvent, StreamError } from '@/lib/cloud-agent-next/event-types'; import { CLOUD_AGENT_NEXT_WS_URL } from '@/lib/constants'; import { getCodeReviewDisplayBehavior } from './code-review-stream-behavior'; +import { fetchStreamTicket } from './fetch-stream-ticket'; type CodeReviewStreamViewProps = { reviewId: string; @@ -303,24 +304,8 @@ export function CodeReviewStreamView({ // --------------------------------------------------------------------------- const getTicket = useCallback( - async (sessionId: string): Promise => { - const body: { cloudAgentSessionId: string; organizationId?: string } = { - cloudAgentSessionId: sessionId, - }; - if (organizationId) { - body.organizationId = organizationId; - } - const response = await fetch('/api/cloud-agent-next/sessions/stream-ticket', { - method: 'POST', - headers: { 'Content-Type': 'application/json' }, - body: JSON.stringify(body), - }); - if (!response.ok) { - const errorData = (await response.json()) as { error?: string }; - throw new Error(errorData.error ?? 'Failed to get stream ticket'); - } - const result = (await response.json()) as { ticket: string }; - return result.ticket; + async (sessionId: string): Promise<{ ticket: string; expiresAt: number }> => { + return fetchStreamTicket(sessionId, organizationId); }, [organizationId] ); @@ -360,7 +345,7 @@ export function CodeReviewStreamView({ async function connect() { if (cancelled || !cloudAgentSessionId) return; try { - const ticket = await getTicket(cloudAgentSessionId); + const result = await getTicket(cloudAgentSessionId); if (cancelled) return; const url = new URL('/stream', CLOUD_AGENT_NEXT_WS_URL); @@ -368,11 +353,14 @@ export function CodeReviewStreamView({ const manager = createWebSocketManager({ url: url.toString(), - ticket, + ticket: result.ticket, onEvent: handleEvent, onError: handleWsError, onStateChange: setConnectionState, - onRefreshTicket: async () => getTicket(cloudAgentSessionId), + onRefreshTicket: async () => { + const refreshed = await getTicket(cloudAgentSessionId); + return refreshed.ticket; + }, }); wsManagerRef.current = manager; diff --git a/apps/web/src/components/code-reviews/fetch-stream-ticket.test.ts b/apps/web/src/components/code-reviews/fetch-stream-ticket.test.ts new file mode 100644 index 00000000000..e91f6d01c36 --- /dev/null +++ b/apps/web/src/components/code-reviews/fetch-stream-ticket.test.ts @@ -0,0 +1,62 @@ +import { fetchStreamTicket } from './fetch-stream-ticket'; + +const mockFetch = jest.fn(); +global.fetch = mockFetch; + +describe('fetchStreamTicket', () => { + beforeEach(() => { + mockFetch.mockReset(); + }); + + it('returns the ticket and expiresAt from a 200 response', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ ticket: 't', expiresAt: 123 }), + }); + + await expect(fetchStreamTicket('session-1', 'org-1')).resolves.toEqual({ + ticket: 't', + expiresAt: 123, + }); + + expect(mockFetch).toHaveBeenCalledWith( + '/api/cloud-agent-next/sessions/stream-ticket', + expect.objectContaining({ + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ cloudAgentSessionId: 'session-1', organizationId: 'org-1' }), + }) + ); + }); + + it('throws when the response is missing expiresAt', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ ticket: 't' }), + }); + + await expect(fetchStreamTicket('session-1')).rejects.toThrow( + 'Missing expiresAt in stream-ticket response' + ); + }); + + it('throws when the response is missing ticket', async () => { + mockFetch.mockResolvedValue({ + ok: true, + json: () => Promise.resolve({ expiresAt: 123 }), + }); + + await expect(fetchStreamTicket('session-1')).rejects.toThrow( + 'Missing ticket in stream-ticket response' + ); + }); + + it('throws the server error from a non-OK response', async () => { + mockFetch.mockResolvedValue({ + ok: false, + json: () => Promise.resolve({ error: 'boom' }), + }); + + await expect(fetchStreamTicket('session-1')).rejects.toThrow('boom'); + }); +}); diff --git a/apps/web/src/components/code-reviews/fetch-stream-ticket.ts b/apps/web/src/components/code-reviews/fetch-stream-ticket.ts new file mode 100644 index 00000000000..1f4f76ee5ee --- /dev/null +++ b/apps/web/src/components/code-reviews/fetch-stream-ticket.ts @@ -0,0 +1,28 @@ +export async function fetchStreamTicket( + sessionId: string, + organizationId?: string +): Promise<{ ticket: string; expiresAt: number }> { + const body: { cloudAgentSessionId: string; organizationId?: string } = { + cloudAgentSessionId: sessionId, + }; + if (organizationId) { + body.organizationId = organizationId; + } + const response = await fetch('/api/cloud-agent-next/sessions/stream-ticket', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + }); + if (!response.ok) { + const errorData = (await response.json()) as { error?: string }; + throw new Error(errorData.error ?? 'Failed to get stream ticket'); + } + const result = (await response.json()) as { ticket?: string; expiresAt?: number }; + if (result.ticket === undefined) { + throw new Error('Missing ticket in stream-ticket response'); + } + if (result.expiresAt === undefined) { + throw new Error('Missing expiresAt in stream-ticket response'); + } + return { ticket: result.ticket, expiresAt: result.expiresAt }; +} From a44f1cd6b494446cb33e93f01ca436ca4a4225c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 11:57:08 +0200 Subject: [PATCH 09/23] feat(session-ingest): recheck membership on command and close sockets on removal P2-D-22: handleWebCommand rechecks current org membership before forwarding. Member removal closes the member's live UserConnection web sockets and Cloud Agent stream sockets via closeViewerSockets and closeOrgStreams. --- .../cloud-agent-client.test.ts | 44 ++++ .../cloud-agent-next/cloud-agent-client.ts | 27 +++ .../lib/organizations/organizations.test.ts | 43 ++++ .../src/lib/organizations/organizations.ts | 13 ++ .../CloudAgentSession-stream-ticket.test.ts | 65 ++++++ .../src/persistence/CloudAgentSession.ts | 20 ++ services/cloud-agent-next/src/server.test.ts | 130 +++++++++++ services/cloud-agent-next/src/server.ts | 50 +++++ services/session-ingest/src/app.ts | 3 + .../src/dos/UserConnectionDO.test.ts | 203 ++++++++++++++---- .../src/dos/UserConnectionDO.ts | 88 +++++++- services/session-ingest/src/index.test.ts | 46 ++++ services/session-ingest/src/routes/api.ts | 3 + 13 files changed, 686 insertions(+), 49 deletions(-) diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index 6e046ac84f1..3e0f1632712 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -56,6 +56,10 @@ beforeEach(() => { mockCreateAppBuilderCloudAgentNextClient.mockClear(); }); +// Load the real `closeCloudAgentOrgStreams` (the module mock above does not +// expose it) so the test exercises the actual fetch call, not a stub. +const { closeCloudAgentOrgStreams } = jest.requireActual('./cloud-agent-client'); + describe('createCloudAgentNextClientForModel', () => { it('returns the default client when the model is paid and has no BYOK', () => { const result = createCloudAgentNextClientForModel('token', { @@ -87,3 +91,43 @@ describe('createCloudAgentNextClientForModel', () => { expect(mockCreateCloudAgentNextClient).not.toHaveBeenCalled(); }); }); + +describe('closeCloudAgentOrgStreams', () => { + const originalFetch = global.fetch; + + afterEach(() => { + global.fetch = originalFetch; + }); + + it('POSTs to the close endpoint with the internal key header and body', async () => { + const fetchMock = jest.fn().mockResolvedValue(new Response(null, { status: 204 })); + global.fetch = fetchMock as unknown as typeof fetch; + + await closeCloudAgentOrgStreams('usr_1', 'org_1'); + + expect(fetchMock).toHaveBeenCalledWith( + 'http://cloud-agent-next/internal/streams/close', + expect.objectContaining({ + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': 'test-secret', + }, + body: JSON.stringify({ userId: 'usr_1', organizationId: 'org_1' }), + }) + ); + }); + + it('throws when the close endpoint returns a non-OK response', async () => { + const fetchMock = jest + .fn() + .mockResolvedValue( + new Response('boom', { status: 500, statusText: 'Internal Server Error' }) + ); + global.fetch = fetchMock as unknown as typeof fetch; + + await expect(closeCloudAgentOrgStreams('usr_1', 'org_1')).rejects.toThrow( + 'Cloud Agent stream close failed: 500 Internal Server Error - boom' + ); + }); +}); diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts index 06ef96766f4..6475870789c 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.ts @@ -866,3 +866,30 @@ export function createCloudAgentNextClientForModel( } return createCloudAgentNextClient(authToken); } + +/** + * Close live Cloud Agent stream sockets for a removed organization member. + * Best-effort: callers log and continue on failure so member removal never + * fails because stream close failed. + */ +export async function closeCloudAgentOrgStreams( + userId: string, + organizationId: string +): Promise { + const response = await fetch(`${CLOUD_AGENT_NEXT_API_URL}/internal/streams/close`, { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': INTERNAL_API_SECRET, + }, + body: JSON.stringify({ userId, organizationId }), + signal: AbortSignal.timeout(30_000), + }); + + if (!response.ok) { + const errorText = await response.text().catch(() => ''); + throw new Error( + `Cloud Agent stream close failed: ${response.status} ${response.statusText}${errorText ? ` - ${errorText}` : ''}` + ); + } +} diff --git a/apps/web/src/lib/organizations/organizations.test.ts b/apps/web/src/lib/organizations/organizations.test.ts index 5988074dd50..746786439aa 100644 --- a/apps/web/src/lib/organizations/organizations.test.ts +++ b/apps/web/src/lib/organizations/organizations.test.ts @@ -26,11 +26,16 @@ import { import { fromMicrodollars } from '@/lib/utils'; import { DEFAULT_MEMBER_DAILY_LIMIT_USD } from '@/lib/organizations/constants'; import { invalidateOrganizationSessionAccess } from '@/lib/session-ingest-client'; +import { closeCloudAgentOrgStreams } from '@/lib/cloud-agent-next/cloud-agent-client'; jest.mock('@/lib/session-ingest-client', () => ({ invalidateOrganizationSessionAccess: jest.fn().mockResolvedValue(undefined), })); +jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ + closeCloudAgentOrgStreams: jest.fn().mockResolvedValue(undefined), +})); + describe('Organizations', () => { afterEach(async () => { // eslint-disable-next-line drizzle/enforce-delete-with-where @@ -407,6 +412,7 @@ describe('Organizations', () => { describe('removeUserFromOrganization', () => { beforeEach(() => { jest.mocked(invalidateOrganizationSessionAccess).mockClear(); + jest.mocked(closeCloudAgentOrgStreams).mockClear(); }); test('should remove user from organization', async () => { @@ -459,6 +465,43 @@ describe('Organizations', () => { expect(invalidateOrganizationSessionAccess).toHaveBeenCalledWith(member.id, organization.id); }); + test('closes Cloud Agent org streams when a member is removed', async () => { + const owner = await insertTestUser(); + const member = await insertTestUser(); + const organization = await createOrganization('Test Org', owner.id); + await addUserToOrganization(organization.id, member.id, 'member'); + + await removeUserFromOrganization(organization.id, member.id); + + expect(closeCloudAgentOrgStreams).toHaveBeenCalledWith(member.id, organization.id); + }); + + test('does not close Cloud Agent org streams when no membership was removed', async () => { + const owner = await insertTestUser(); + const nonMember = await insertTestUser(); + const organization = await createOrganization('Test Org', owner.id); + + await removeUserFromOrganization(organization.id, nonMember.id); + + expect(closeCloudAgentOrgStreams).not.toHaveBeenCalled(); + }); + + test('completes removal when Cloud Agent stream close fails', async () => { + const owner = await insertTestUser(); + const member = await insertTestUser(); + const organization = await createOrganization('Test Org', owner.id); + await addUserToOrganization(organization.id, member.id, 'member'); + jest + .mocked(closeCloudAgentOrgStreams) + .mockRejectedValueOnce(new Error('stream close unavailable')); + + await expect(removeUserFromOrganization(organization.id, member.id)).resolves.toBeDefined(); + + const memberOrgs = await getUserOrganizationsWithSeats(member.id); + expect(memberOrgs).toHaveLength(0); + expect(closeCloudAgentOrgStreams).toHaveBeenCalledWith(member.id, organization.id); + }); + test('should remove specific user without affecting others', async () => { const owner = await insertTestUser(); const member1 = await insertTestUser(); diff --git a/apps/web/src/lib/organizations/organizations.ts b/apps/web/src/lib/organizations/organizations.ts index 1a688a13d55..8bf8a4405e5 100644 --- a/apps/web/src/lib/organizations/organizations.ts +++ b/apps/web/src/lib/organizations/organizations.ts @@ -32,6 +32,7 @@ import { resolveEffectiveOrganizationSsoPolicy } from './organization-sso-policy import { classifyOrganizationEntitlement } from './trial-utils'; import { errorExceptInTest, logExceptInTest } from '@/lib/utils.server'; import { invalidateOrganizationSessionAccess } from '@/lib/session-ingest-client'; +import { closeCloudAgentOrgStreams } from '@/lib/cloud-agent-next/cloud-agent-client'; import { APP_URL } from '@/lib/constants'; import { createAuditLog } from '@/lib/organizations/organization-audit-logs'; import { failureResult, successResult } from '@/lib/maybe-result'; @@ -494,6 +495,18 @@ async function invalidateRemovedMemberSessionAccess( } ); } + + // Best-effort: close the removed member's live Cloud Agent stream sockets. + // A failure here must not fail member removal. + try { + await closeCloudAgentOrgStreams(userId, organizationId); + } catch (error) { + errorExceptInTest('Failed to close Cloud Agent streams for removed organization member', { + organizationId, + userId, + error: error instanceof Error ? error.message : String(error), + }); + } } export async function updateUserRoleInOrganization( diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts index c6eb348c925..76ec7ee6651 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession-stream-ticket.test.ts @@ -103,6 +103,7 @@ function createCtx() { id: { name: 'user-1:session-1' }, storage, blockConcurrencyWhile: vi.fn().mockResolvedValue(undefined), + getWebSockets: vi.fn(), }; return ctx; } @@ -179,3 +180,67 @@ describe('CloudAgentSession /stream ticket audience', () => { expect(createStreamHandlerMock).not.toHaveBeenCalled(); }); }); + +function makeStreamSocket() { + return { close: vi.fn() }; +} + +function makeMetadata(orgId?: string) { + return { + metadataSchemaVersion: 2, + identity: { + sessionId: 'session-1', + userId: 'user-1', + ...(orgId ? { orgId } : {}), + }, + auth: {}, + lifecycle: { version: 1, timestamp: 1700000000000 }, + }; +} + +describe('closeOrgStreams', () => { + it('closes matching stream sockets and returns the count', async () => { + const ctx = createCtx(); + const doInstance = new CloudAgentSession(ctx as never, createEnv() as never); + + ctx.storage.get.mockResolvedValue(makeMetadata('org_1')); + const stream1 = makeStreamSocket(); + const stream2 = makeStreamSocket(); + const terminal1 = makeStreamSocket(); + ctx.getWebSockets.mockImplementation((tag: string) => + tag === 'stream' ? [stream1, stream2] : [terminal1] + ); + + const closed = await doInstance.closeOrgStreams('org_1'); + + expect(closed).toBe(2); + expect(ctx.getWebSockets).toHaveBeenCalledWith('stream'); + expect(stream1.close).toHaveBeenCalledWith(1000, 'session access revoked'); + expect(stream2.close).toHaveBeenCalledWith(1000, 'session access revoked'); + expect(terminal1.close).not.toHaveBeenCalled(); + }); + + it('closes nothing when identity.orgId is missing', async () => { + const ctx = createCtx(); + const doInstance = new CloudAgentSession(ctx as never, createEnv() as never); + + ctx.storage.get.mockResolvedValue(makeMetadata()); + + const closed = await doInstance.closeOrgStreams('org_1'); + + expect(closed).toBe(0); + expect(ctx.getWebSockets).not.toHaveBeenCalled(); + }); + + it('closes nothing when identity.orgId does not match', async () => { + const ctx = createCtx(); + const doInstance = new CloudAgentSession(ctx as never, createEnv() as never); + + ctx.storage.get.mockResolvedValue(makeMetadata('org_other')); + + const closed = await doInstance.closeOrgStreams('org_1'); + + expect(closed).toBe(0); + expect(ctx.getWebSockets).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 99bb7d1b52f..1589a78f95f 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -1303,6 +1303,26 @@ export class CloudAgentSession extends DurableObject { return getConnectedStreamClientCount(this.ctx); } + /** + * Close every connected /stream client whose session belongs to the given + * organization. Called on member removal so a removed member's live stream + * sockets are torn down immediately. Returns the number of sockets closed. + * Sockets live on this DO, so the HTTP handler must delegate here rather + * than close anything itself. + */ + async closeOrgStreams(organizationId: string): Promise { + const metadata = await this.getStoredMetadata(); + const orgId = metadata?.identity.orgId; + if (!orgId || orgId !== organizationId) return 0; + + let closed = 0; + for (const ws of this.ctx.getWebSockets('stream')) { + ws.close(1000, 'session access revoked'); + closed++; + } + return closed; + } + // --------------------------------------------------------------------------- // Metadata RPC Methods // --------------------------------------------------------------------------- diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 6260d6f32d0..668c7951890 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -8,11 +8,13 @@ const { consumeCloudAgentReportBatchMock, removeExpiredCloudAgentReportDataMock, requireCurrentSessionAccessMock, + getPgDbMock, } = vi.hoisted(() => ({ getRunningTerminalClientMock: vi.fn(), consumeCloudAgentReportBatchMock: vi.fn().mockResolvedValue(undefined), removeExpiredCloudAgentReportDataMock: vi.fn().mockResolvedValue(undefined), requireCurrentSessionAccessMock: vi.fn(), + getPgDbMock: vi.fn(), })); vi.mock('./logger.js', () => { @@ -93,6 +95,10 @@ vi.mock('./persistence/CloudAgentSession.js', () => ({ CloudAgentSession: class CloudAgentSession {}, })); +vi.mock('./db/pg.js', () => ({ + getPgDb: getPgDbMock, +})); + vi.mock('@kilocode/db/client', () => ({ getWorkerDb: () => ({ select: () => ({ @@ -115,6 +121,7 @@ type MockEnv = { SandboxSmall: unknown; WS_ALLOWED_ORIGINS?: string; HYPERDRIVE: { connectionString: string }; + INTERNAL_API_SECRET?: string; CLOUD_AGENT_SESSION: { idFromName: ReturnType; get: ReturnType; @@ -135,6 +142,7 @@ function createEnv(): MockEnv { Sandbox: {}, SandboxSmall: {}, HYPERDRIVE: { connectionString: 'postgres://test' }, + INTERNAL_API_SECRET: 'test-internal-secret', CLOUD_AGENT_SESSION: { idFromName: vi.fn(), get: vi.fn(), @@ -185,6 +193,7 @@ beforeEach(() => { getRunningTerminalClientMock.mockReset(); consumeCloudAgentReportBatchMock.mockClear(); removeExpiredCloudAgentReportDataMock.mockClear(); + getPgDbMock.mockReset(); requireCurrentSessionAccessMock.mockReset().mockResolvedValue({ kiloSessionId: 'ses_12345678901234567890123456', organizationId: null, @@ -920,3 +929,124 @@ describe('server wrapper log upload route', () => { expect(env.R2_BUCKET.put).not.toHaveBeenCalled(); }); }); + +describe('server /internal/streams/close', () => { + it('rejects without the internal API key', async () => { + const env = createEnv(); + const response = await fetchWorker( + new Request('http://worker.test/internal/streams/close', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ userId: 'usr_removed', organizationId: 'org_1' }), + }), + env + ); + + expect(response.status).toBe(401); + expect(getPgDbMock).not.toHaveBeenCalled(); + }); + + it('rejects with an incorrect internal API key', async () => { + const env = createEnv(); + const response = await fetchWorker( + new Request('http://worker.test/internal/streams/close', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': 'wrong-key', + }, + body: JSON.stringify({ userId: 'usr_removed', organizationId: 'org_1' }), + }), + env + ); + + expect(response.status).toBe(401); + expect(getPgDbMock).not.toHaveBeenCalled(); + }); + + it('rejects a malformed body', async () => { + const env = createEnv(); + const response = await fetchWorker( + new Request('http://worker.test/internal/streams/close', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': 'test-internal-secret', + }, + body: JSON.stringify({ userId: 'usr_removed' }), + }), + env + ); + + expect(response.status).toBe(400); + expect(getPgDbMock).not.toHaveBeenCalled(); + }); + + it('closes org stream sockets for every matching session', async () => { + const env = createEnv(); + const closeOrgStreams = vi.fn().mockResolvedValue(1); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ closeOrgStreams }); + + getPgDbMock.mockReturnValue({ + select: vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(async () => [ + { cloudAgentSessionId: 'agent_1' }, + { cloudAgentSessionId: 'agent_2' }, + ]), + })), + })), + }); + + const response = await fetchWorker( + new Request('http://worker.test/internal/streams/close', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': 'test-internal-secret', + }, + body: JSON.stringify({ userId: 'usr_removed', organizationId: 'org_1' }), + }), + env + ); + + expect(response.status).toBe(204); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_removed:agent_1'); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_removed:agent_2'); + expect(env.CLOUD_AGENT_SESSION.get).toHaveBeenCalledTimes(2); + expect(closeOrgStreams).toHaveBeenCalledTimes(2); + expect(closeOrgStreams).toHaveBeenCalledWith('org_1'); + }); + + it('skips rows without a cloud agent session id', async () => { + const env = createEnv(); + const closeOrgStreams = vi.fn().mockResolvedValue(0); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ closeOrgStreams }); + + getPgDbMock.mockReturnValue({ + select: vi.fn(() => ({ + from: vi.fn(() => ({ + where: vi.fn(async () => [{ cloudAgentSessionId: null }]), + })), + })), + }); + + const response = await fetchWorker( + new Request('http://worker.test/internal/streams/close', { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'x-internal-api-key': 'test-internal-secret', + }, + body: JSON.stringify({ userId: 'usr_removed', organizationId: 'org_1' }), + }), + env + ); + + expect(response.status).toBe(204); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + expect(closeOrgStreams).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 50c6bd5815c..a6c994934f6 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -29,6 +29,10 @@ import { resolveTerminalWrapperClient } from './terminal/access.js'; import { requestMethodAllowsBody } from './shared/http-proxy.js'; import { hasDuplicateQueryParameters } from './shared/http-query.js'; import { projectSessionAccessHttpError, requireCurrentSessionAccess } from './session-access.js'; +import { timingSafeEqual } from '@kilocode/encryption'; +import { and, eq, isNotNull } from 'drizzle-orm'; +import { cli_sessions_v2 } from '@kilocode/db/schema'; +import { getPgDb } from './db/pg.js'; import { KILO_FACADE_AUTH_TOKEN_HEADER, KILO_FACADE_GLOBAL_FEED_PATH, @@ -617,6 +621,52 @@ app.put( } ); +app.post('/internal/streams/close', async (c: Context) => { + const internalApiKey = c.req.header('x-internal-api-key'); + if (!c.env.INTERNAL_API_SECRET) { + return c.text('Internal API secret not configured', 500); + } + if (!internalApiKey || !timingSafeEqual(internalApiKey, c.env.INTERNAL_API_SECRET)) { + return c.text('Invalid or missing internal API key', 401); + } + + const body = (await c.req.json().catch(() => null)) as { + userId?: unknown; + organizationId?: unknown; + } | null; + const userId = body?.userId; + const organizationId = body?.organizationId; + if ( + typeof userId !== 'string' || + userId.length === 0 || + typeof organizationId !== 'string' || + organizationId.length === 0 + ) { + return c.text('Invalid request', 400); + } + + const db = getPgDb(c.env); + const rows = await db + .select({ cloudAgentSessionId: cli_sessions_v2.cloud_agent_session_id }) + .from(cli_sessions_v2) + .where( + and( + eq(cli_sessions_v2.kilo_user_id, userId), + eq(cli_sessions_v2.organization_id, organizationId), + isNotNull(cli_sessions_v2.cloud_agent_session_id) + ) + ); + + for (const row of rows) { + if (!row.cloudAgentSessionId) continue; + const doId = c.env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${row.cloudAgentSessionId}`); + const stub = c.env.CLOUD_AGENT_SESSION.get(doId); + await stub.closeOrgStreams(organizationId); + } + + return c.body(null, 204); +}); + app.use('/trpc/*', authMiddleware); app.use('/trpc/*', balanceMiddleware); diff --git a/services/session-ingest/src/app.ts b/services/session-ingest/src/app.ts index 069b377324e..ee05bce6458 100644 --- a/services/session-ingest/src/app.ts +++ b/services/session-ingest/src/app.ts @@ -12,6 +12,7 @@ import { api } from './routes/api'; import { cloudAgentSessionScopeApi } from './routes/cloud-agent-session-scope'; import { getSessionIngestDO } from './dos/SessionIngestDO'; import { getSessionAccessCacheDO } from './dos/SessionAccessCacheDO'; +import { getUserConnectionDO } from './dos/UserConnectionDO'; import { getSessionExport } from './services/session-export'; import { withDORetry } from '@kilocode/worker-utils'; @@ -132,6 +133,8 @@ app.post('/internal/session-access/invalidate', async c => { 'SessionAccessCacheDO.invalidateOrganization' ); + await getUserConnectionDO(c.env, { kiloUserId: parsed.data.kiloUserId }).closeViewerSockets(); + return c.body(null, 204); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index ae1d4251b98..675001f70f9 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -22,6 +22,14 @@ vi.mock('./SessionIngestDO', () => ({ getSessionIngestDO: sessionIngestMocks.getSessionIngestDO, })); +const sessionAccessMocks = vi.hoisted(() => ({ + resolveAccessibleKiloSession: vi.fn(), +})); + +vi.mock('../services/session-access', () => ({ + resolveAccessibleKiloSession: sessionAccessMocks.resolveAccessibleKiloSession, +})); + import { MAX_CATALOG_RESULT_BYTES, MAX_DURABLE_RESULT_BYTES, @@ -208,7 +216,7 @@ function connectWebSocket(doInstance: UserConnectionDO, connectionId: string): M ); doInstance.fetch( - new Request(`http://local/web?connectionId=${connectionId}`, { + new Request(`http://local/web?connectionId=${connectionId}&kiloUserId=usr_1`, { headers: { Upgrade: 'websocket' }, }) ); @@ -271,9 +279,10 @@ function addCliSocket( function addWebSocket( mockCtx: ReturnType, connectionId = 'web-1', - subscribedSessions: string[] = [] + subscribedSessions: string[] = [], + kiloUserId = 'usr_1' ): MockWS { - const attachment = { role: 'web' as const, connectionId, subscribedSessions }; + const attachment = { role: 'web' as const, connectionId, subscribedSessions, kiloUserId }; const ws = createMockWs(['web'], attachment); mockCtx.addSocket(ws); return ws; @@ -307,9 +316,10 @@ function sendHeartbeat( } /** Send a subscribe from a web ws */ -function sendSubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { +async function sendSubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { const msg = JSON.stringify({ type: 'subscribe', sessionId }); - doInstance.webSocketMessage(webWs as never, msg); + await doInstance.webSocketMessage(webWs as never, msg); + await flushAsync(); } /** Send an unsubscribe from a web ws */ @@ -339,7 +349,7 @@ async function sendCommand( } ): Promise { const msg = JSON.stringify({ type: 'command', ...opts }); - doInstance.webSocketMessage(webWs as never, msg); + await doInstance.webSocketMessage(webWs as never, msg); await flushAsync(); } @@ -406,6 +416,12 @@ describe('UserConnectionDO', () => { resetAttentionStatusOnCliDisconnect: sessionIngestMocks.resetAttentionStatusOnCliDisconnect, claimSessionReadyPush: sessionIngestMocks.claimSessionReadyPush, }); + sessionAccessMocks.resolveAccessibleKiloSession.mockReset(); + sessionAccessMocks.resolveAccessibleKiloSession.mockResolvedValue({ + kiloSessionId: 'ses_12345678901234567890123456', + organizationId: null, + cloudAgentSessionScopeId: null, + }); }); afterEach(() => { @@ -587,7 +603,7 @@ describe('UserConnectionDO', () => { // web subscribes to s1 — subscribe sent to cli1 (the current owner) const webWs = mockCtx.sockets.find(s => s._tags.includes('web'))!; - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); cli1.send.mockClear(); cli2.send.mockClear(); @@ -613,8 +629,8 @@ describe('UserConnectionDO', () => { // subWeb subscribes to s1, otherWeb subscribes to s2 (subscriptions are // irrelevant to broadcast delivery — both should still receive cli1's heartbeat) - sendSubscribe(doInstance, subWeb, 's1'); - sendSubscribe(doInstance, otherWeb, 's2'); + await sendSubscribe(doInstance, subWeb, 's1'); + await sendSubscribe(doInstance, otherWeb, 's2'); subWeb.send.mockClear(); otherWeb.send.mockClear(); @@ -675,7 +691,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { protocolVersion: '1', }); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { @@ -695,7 +711,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); @@ -713,7 +729,7 @@ describe('UserConnectionDO', () => { // cli1 owns s1 sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, subWeb, 's1'); + await sendSubscribe(doInstance, subWeb, 's1'); subWeb.send.mockClear(); otherWeb.send.mockClear(); @@ -832,7 +848,7 @@ describe('UserConnectionDO', () => { // Establish ownership and subscription first so the heartbeat broadcast // is targeted at this web socket. sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')], { @@ -895,7 +911,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')], { capabilities: { attachments: true }, }); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); // Latest heartbeat omits capabilities. @@ -1079,7 +1095,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); // CLI should receive subscribe expect(cliWs.send).toHaveBeenCalledTimes(1); @@ -1093,7 +1109,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1', 'busy', 'Fix bug')]); webWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); expect(parseSent(webWs)).toEqual({ type: 'system', @@ -1119,7 +1135,7 @@ describe('UserConnectionDO', () => { // No heartbeat sent, so no owner for 's1' // Trigger ensureState via a harmless message first - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); // Both CLIs should receive subscribe expect(cli1.send).toHaveBeenCalled(); @@ -1133,8 +1149,8 @@ describe('UserConnectionDO', () => { addCliSocket(mockCtx, 'cli-1'); const webWs = addWebSocket(mockCtx, 'web-1'); - sendSubscribe(doInstance, webWs, 's1'); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); const att = webWs.deserializeAttachment() as { subscribedSessions: string[]; @@ -1150,7 +1166,7 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); cliWs.send.mockClear(); sendUnsubscribe(doInstance, webWs, 's1'); @@ -1171,8 +1187,8 @@ describe('UserConnectionDO', () => { sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); cliWs.send.mockClear(); - sendSubscribe(doInstance, web1, 's1'); - sendSubscribe(doInstance, web2, 's1'); + await sendSubscribe(doInstance, web1, 's1'); + await sendSubscribe(doInstance, web2, 's1'); cliWs.send.mockClear(); // Unsubscribe first — CLI should NOT get unsubscribe @@ -1217,6 +1233,7 @@ describe('UserConnectionDO', () => { role: 'web', connectionId: 'viewer-1', subscribedSessions: [], + kiloUserId: 'usr_1', }); }); }); @@ -1258,7 +1275,7 @@ describe('UserConnectionDO', () => { cliWs.send.mockClear(); const oldWeb = connectWebSocket(doInstance, 'viewer-1'); - sendSubscribe(doInstance, oldWeb, 's1'); + await sendSubscribe(doInstance, oldWeb, 's1'); cliWs.send.mockClear(); const newWeb = connectWebSocket(doInstance, 'viewer-1'); @@ -1272,6 +1289,7 @@ describe('UserConnectionDO', () => { role: 'web', connectionId: 'viewer-1', subscribedSessions: [], + kiloUserId: 'usr_1', }); }); @@ -1535,7 +1553,7 @@ describe('UserConnectionDO', () => { // State should NOT be cleaned up — cli2 is live // Verify by routing a command to s1 — should reach cli2 cli2.send.mockClear(); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); expect(cli2.send).toHaveBeenCalled(); expect(parseSent(cli2)).toEqual({ type: 'subscribe', sessionId: 's1' }); }); @@ -1801,8 +1819,8 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1'), makeSession('s2')]); - sendSubscribe(doInstance, webWs, 's1'); - sendSubscribe(doInstance, webWs, 's2'); + await sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's2'); mockCtx.removeSocket(webWs); disconnectWeb(doInstance, webWs); @@ -1824,7 +1842,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); cliWs.send.mockClear(); mockCtx.removeSocket(webWs); @@ -3501,7 +3519,7 @@ describe('UserConnectionDO', () => { const otherWeb = addWebSocket(mockCtx, 'web-other'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, subWeb, 's1'); + await sendSubscribe(doInstance, subWeb, 's1'); subWeb.send.mockClear(); otherWeb.send.mockClear(); @@ -3531,8 +3549,8 @@ describe('UserConnectionDO', () => { const childWeb = addWebSocket(mockCtx, 'web-child'); sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, parentWeb, 'parent-session'); - sendSubscribe(doInstance, childWeb, 'child-session-1'); + await sendSubscribe(doInstance, parentWeb, 'parent-session'); + await sendSubscribe(doInstance, childWeb, 'child-session-1'); parentWeb.send.mockClear(); childWeb.send.mockClear(); @@ -3564,8 +3582,8 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, webWs, 'parent-session'); - sendSubscribe(doInstance, webWs, 'child-session-1'); + await sendSubscribe(doInstance, webWs, 'parent-session'); + await sendSubscribe(doInstance, webWs, 'child-session-1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ @@ -3587,7 +3605,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('parent-session')]); - sendSubscribe(doInstance, webWs, 'parent-session'); + await sendSubscribe(doInstance, webWs, 'parent-session'); webWs.send.mockClear(); // CLI sends event for a child session with parentSessionId @@ -3616,7 +3634,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('other-session')]); - sendSubscribe(doInstance, webWs, 'other-session'); + await sendSubscribe(doInstance, webWs, 'other-session'); webWs.send.mockClear(); // Child event with parent that nobody subscribes to @@ -3638,7 +3656,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ @@ -3664,7 +3682,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); webWs.send.mockClear(); const eventMsg = JSON.stringify({ @@ -3725,7 +3743,7 @@ describe('UserConnectionDO', () => { const webWs = addWebSocket(mockCtx, 'web-1'); // Trigger ensureState by calling any method (e.g., webSocketMessage with subscribe) - sendSubscribe(doInstance, webWs, 's1'); + await sendSubscribe(doInstance, webWs, 's1'); // Verify state was reconstructed by routing a command const web2 = addWebSocket(mockCtx, 'web-2'); @@ -6612,7 +6630,7 @@ describe('UserConnectionDO', () => { }) ); - doInstance.webSocketMessage( + await doInstance.webSocketMessage( webWs as never, JSON.stringify({ type: 'command', @@ -6678,7 +6696,7 @@ describe('UserConnectionDO', () => { }) ); - doInstance.webSocketMessage( + await doInstance.webSocketMessage( webWs as never, JSON.stringify({ type: 'command', @@ -7656,4 +7674,111 @@ describe('UserConnectionDO', () => { }); }); }); + + // ------------------------------------------------------------------------- + // Command and subscribe access recheck + // ------------------------------------------------------------------------- + + describe('command and subscribe access recheck', () => { + it('rejects a command on an inaccessible org session without forwarding', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + webWs.send.mockClear(); + sessionAccessMocks.resolveAccessibleKiloSession.mockResolvedValueOnce(null); + + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + + expect(sessionAccessMocks.resolveAccessibleKiloSession).toHaveBeenCalledWith( + expect.anything(), + { kiloUserId: 'usr_1', kiloSessionId: 's1' } + ); + expect(parseSent(webWs)).toEqual({ + type: 'response', + id: 'cmd-1', + error: { + source: 'relay', + code: 'SESSION_ACCESS_DENIED', + message: 'You no longer have access to this session', + }, + }); + expect(cliWs.send).not.toHaveBeenCalled(); + }); + + it('forwards a command on an accessible personal session', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + sessionAccessMocks.resolveAccessibleKiloSession.mockResolvedValueOnce({ + kiloSessionId: 's1', + organizationId: null, + cloudAgentSessionScopeId: null, + }); + + await sendCommand(doInstance, webWs, { + id: 'cmd-1', + command: 'send_message', + sessionId: 's1', + }); + + expect(cliWs.send).toHaveBeenCalledTimes(1); + expect(parseSent(cliWs)).toMatchObject({ + type: 'command', + command: 'send_message', + sessionId: 's1', + }); + }); + + it('does not tell the CLI to forward events when a subscribe loses org access', async () => { + const { doInstance, mockCtx } = setup(); + const cliWs = addCliSocket(mockCtx, 'cli-1'); + const webWs = addWebSocket(mockCtx, 'web-1'); + + sendHeartbeat(doInstance, cliWs, [makeSession('s1')]); + cliWs.send.mockClear(); + sessionAccessMocks.resolveAccessibleKiloSession.mockResolvedValueOnce(null); + + await sendSubscribe(doInstance, webWs, 's1'); + + expect(cliWs.send).not.toHaveBeenCalled(); + }); + }); +}); + +describe('closeViewerSockets', () => { + it('closes every web socket and returns the count', () => { + const { doInstance, mockCtx } = setup(); + const web1 = addWebSocket(mockCtx, 'web-1'); + const web2 = addWebSocket(mockCtx, 'web-2'); + const cli = addCliSocket(mockCtx, 'cli-1'); + + const closed = doInstance.closeViewerSockets(); + + expect(closed).toBe(2); + expect(web1.close).toHaveBeenCalledWith(1000, 'session access revoked'); + expect(web2.close).toHaveBeenCalledWith(1000, 'session access revoked'); + expect(cli.close).not.toHaveBeenCalled(); + }); + + it('does not close non-web sockets', () => { + const { doInstance, mockCtx } = setup(); + const cli1 = addCliSocket(mockCtx, 'cli-1'); + const cli2 = addCliSocket(mockCtx, 'cli-2'); + + const closed = doInstance.closeViewerSockets(); + + expect(closed).toBe(0); + expect(cli1.close).not.toHaveBeenCalled(); + expect(cli2.close).not.toHaveBeenCalled(); + }); }); diff --git a/services/session-ingest/src/dos/UserConnectionDO.ts b/services/session-ingest/src/dos/UserConnectionDO.ts index 9b4f67d6960..0f090233a6c 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.ts @@ -2,6 +2,7 @@ import { DurableObject } from 'cloudflare:workers'; import type { Env } from '../env'; import { getSessionIngestDO } from './SessionIngestDO'; +import { resolveAccessibleKiloSession } from '../services/session-access'; import { CLIOutboundMessageSchema, type CLIInboundMessage, @@ -65,6 +66,10 @@ type WSAttachment = connectionId: string; subscribedSessions: string[]; replaced?: true; + // Set from the authenticated /user/web route; undefined on sockets + // accepted before this field existed. Needed for command/subscribe + // access rechecks against current session membership. + kiloUserId?: string; }; // Type re-export so test files and other internal callers can reference the @@ -128,6 +133,12 @@ const SESSION_OWNER_CHANGED_ERROR = { message: 'Session owner changed', }; +const SESSION_ACCESS_DENIED_ERROR = { + source: 'relay', + code: 'SESSION_ACCESS_DENIED', + message: 'You no longer have access to this session', +}; + const CATALOG_TOO_LARGE_ERROR = { source: 'relay', code: 'CATALOG_TOO_LARGE', @@ -377,10 +388,12 @@ export class UserConnectionDO extends DurableObject { } else { this.replaceWebSocket(connectionId); + const kiloUserId = url.searchParams.get('kiloUserId') ?? undefined; const attachment: WSAttachment = { role: 'web', connectionId, subscribedSessions: [], + kiloUserId, }; this.ctx.acceptWebSocket(server, ['web']); server.serializeAttachment(attachment); @@ -403,7 +416,7 @@ export class UserConnectionDO extends DurableObject { return new Response(null, { status: 101, webSocket: client }); } - webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): void { + async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer): Promise { this.ensureState(); const attachment = ws.deserializeAttachment() as WSAttachment | null; @@ -429,7 +442,7 @@ export class UserConnectionDO extends DurableObject { if (attachment.role === 'cli') { this.handleCliMessage(ws, attachment, parsed, raw, binaryByteCount); } else if (!attachment.replaced) { - this.handleWebMessage(ws, attachment, parsed); + await this.handleWebMessage(ws, attachment, parsed); } } @@ -1195,11 +1208,11 @@ export class UserConnectionDO extends DurableObject { // Web message handling // --------------------------------------------------------------------------- - private handleWebMessage( + private async handleWebMessage( ws: WebSocket, attachment: WSAttachment & { role: 'web' }, parsed: unknown - ): void { + ): Promise { const result = WebOutboundMessageSchema.safeParse(parsed); if (!result.success) { console.warn('Invalid web message', { @@ -1212,13 +1225,13 @@ export class UserConnectionDO extends DurableObject { switch (msg.type) { case 'subscribe': - this.handleWebSubscribe(ws, attachment, msg.sessionId); + await this.handleWebSubscribe(ws, attachment, msg.sessionId); break; case 'unsubscribe': this.handleWebUnsubscribe(ws, attachment, msg.sessionId); break; case 'command': - this.handleWebCommand(ws, msg); + await this.handleWebCommand(ws, attachment, msg); break; case 'ping': this.sendToWeb(ws, { type: 'pong', nonce: msg.nonce }); @@ -1226,11 +1239,25 @@ export class UserConnectionDO extends DurableObject { } } - private handleWebSubscribe( + private async handleWebSubscribe( ws: WebSocket, attachment: WSAttachment & { role: 'web' }, sessionId: string - ): void { + ): Promise { + // Recheck current membership before subscribing: a removed member must not + // receive events from an org session they no longer have access to. + const kiloUserId = attachment.kiloUserId; + if (!kiloUserId) { + return; + } + const accessible = await resolveAccessibleKiloSession(this.env, { + kiloUserId, + kiloSessionId: sessionId, + }); + if (!accessible) { + return; + } + let subs = this.webSubscriptions.get(sessionId); if (!subs) { subs = new Set(); @@ -1291,8 +1318,9 @@ export class UserConnectionDO extends DurableObject { } } - private handleWebCommand( + private async handleWebCommand( ws: WebSocket, + attachment: WSAttachment & { role: 'web' }, msg: { id: string; command: string; @@ -1301,7 +1329,7 @@ export class UserConnectionDO extends DurableObject { data?: unknown; mutationId?: string; } - ): void { + ): Promise { const now = Date.now(); this.expirePendingCommands(now); @@ -1334,6 +1362,33 @@ export class UserConnectionDO extends DurableObject { return; } + // Command-time access recheck: a removed member must not forward a command + // to a session they no longer have access to. Catalog commands without a + // sessionId skip this check. + if (msg.sessionId) { + const kiloUserId = attachment.kiloUserId; + if (!kiloUserId) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: SESSION_ACCESS_DENIED_ERROR, + }); + return; + } + const accessible = await resolveAccessibleKiloSession(this.env, { + kiloUserId, + kiloSessionId: msg.sessionId, + }); + if (!accessible) { + this.sendToWeb(ws, { + type: 'response', + id: msg.id, + error: SESSION_ACCESS_DENIED_ERROR, + }); + return; + } + } + // Find target CLI let targetCli: WebSocket | undefined; @@ -1847,6 +1902,19 @@ export class UserConnectionDO extends DurableObject { return { instances }; } + /** + * Close every live web viewer socket. Called on member removal so a removed + * member's open viewer connections are torn down immediately. Returns the + * number of sockets closed. + */ + closeViewerSockets(): number { + const sockets = this.ctx.getWebSockets('web'); + for (const ws of sockets) { + ws.close(1000, 'session access revoked'); + } + return sockets.length; + } + async notifySessionEvent(event: SessionEventPayload): Promise<{ delivered: number }> { this.ensureState(); const parsed = SessionEventPayloadSchema.parse(event); diff --git a/services/session-ingest/src/index.test.ts b/services/session-ingest/src/index.test.ts index 724191e521d..473ee19910c 100644 --- a/services/session-ingest/src/index.test.ts +++ b/services/session-ingest/src/index.test.ts @@ -30,9 +30,14 @@ vi.mock('./dos/SessionAccessCacheDO', () => ({ getSessionAccessCacheDO: vi.fn(), })); +vi.mock('./dos/UserConnectionDO', () => ({ + getUserConnectionDO: vi.fn(), +})); + import { app } from './index'; import { getWorkerDb } from '@kilocode/db/client'; import { getSessionAccessCacheDO } from './dos/SessionAccessCacheDO'; +import { getUserConnectionDO } from './dos/UserConnectionDO'; type TestBindings = { HYPERDRIVE: { connectionString: string }; @@ -94,6 +99,7 @@ describe('session access invalidation route', () => { expect(res.status).toBe(401); expect(cache.invalidateOrganization).not.toHaveBeenCalled(); + expect(getUserConnectionDO).not.toHaveBeenCalled(); }); it.each(['wrong-secretxxx', 'wrong'])( @@ -123,6 +129,7 @@ describe('session access invalidation route', () => { expect(res.status).toBe(401); expect(getSessionAccessCacheDO).not.toHaveBeenCalled(); expect(cache.invalidateOrganization).not.toHaveBeenCalled(); + expect(getUserConnectionDO).not.toHaveBeenCalled(); } ); @@ -131,6 +138,9 @@ describe('session access invalidation route', () => { vi.mocked(getSessionAccessCacheDO).mockReturnValue( cache as unknown as ReturnType ); + vi.mocked(getUserConnectionDO).mockReturnValue({ + closeViewerSockets: vi.fn(async () => 0), + } as unknown as ReturnType); const res = await app.request( '/internal/session-access/invalidate', @@ -156,6 +166,42 @@ describe('session access invalidation route', () => { '11111111-1111-4111-8111-111111111111' ); }); + + it('closes the removed member viewer sockets', async () => { + const cache = { invalidateOrganization: vi.fn(async () => undefined) }; + vi.mocked(getSessionAccessCacheDO).mockReturnValue( + cache as unknown as ReturnType + ); + const closeViewerSockets = vi.fn(async () => 1); + vi.mocked(getUserConnectionDO).mockReturnValue({ + closeViewerSockets, + } as unknown as ReturnType); + + const res = await app.request( + '/internal/session-access/invalidate', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'X-Internal-Secret': 'internal-secret', + }, + body: JSON.stringify({ + kiloUserId: 'usr_removed', + organizationId: '11111111-1111-4111-8111-111111111111', + }), + }, + defaultEnv + ); + + expect(res.status).toBe(204); + expect(getUserConnectionDO).toHaveBeenCalledWith(defaultEnv, { + kiloUserId: 'usr_removed', + }); + expect(closeViewerSockets).toHaveBeenCalled(); + expect(cache.invalidateOrganization).toHaveBeenCalledWith( + '11111111-1111-4111-8111-111111111111' + ); + }); }); describe('public session route', () => { diff --git a/services/session-ingest/src/routes/api.ts b/services/session-ingest/src/routes/api.ts index e96931bde1a..e84c3ef1519 100644 --- a/services/session-ingest/src/routes/api.ts +++ b/services/session-ingest/src/routes/api.ts @@ -728,6 +728,9 @@ api.get('/user/web', async c => { const stub = getUserConnectionDO(c.env, { kiloUserId }); const wsUrl = new URL(c.req.url); wsUrl.pathname = '/web'; + // The DO can't recover the user from its idFromName-derived id; the web + // attachment needs it for command/subscribe access rechecks. + wsUrl.searchParams.set('kiloUserId', kiloUserId); return stub.fetch(new Request(wsUrl.toString(), c.req.raw)); }); From 1f88e88934ecef4a95e8aa235541005c845fb72d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 11:57:12 +0200 Subject: [PATCH 10/23] feat(user): reauthenticate and perform GDPR removal on account deletion P1-D-34a/b: requestAccountDeletionChallenge emails a sign-in code and returns devCode in non-production. requestAccountDeletion reauthenticates with the code, checks preconditions, then calls performGdprRemoval. Mobile keeps the challenge on precondition and invalid-code errors so the user can resubmit. --- .../profile-screen.mounted.test.tsx | 234 ++++++++++++++++++ apps/mobile/src/components/profile-screen.tsx | 53 ++-- .../src/components/use-delete-account.ts | 112 +++++++++ apps/web/src/routers/user-router.test.ts | 153 ++++++++++++ apps/web/src/routers/user-router.ts | 119 +++++++-- 5 files changed, 629 insertions(+), 42 deletions(-) create mode 100644 apps/mobile/src/components/profile-screen.mounted.test.tsx create mode 100644 apps/mobile/src/components/use-delete-account.ts diff --git a/apps/mobile/src/components/profile-screen.mounted.test.tsx b/apps/mobile/src/components/profile-screen.mounted.test.tsx new file mode 100644 index 00000000000..692eb299ec2 --- /dev/null +++ b/apps/mobile/src/components/profile-screen.mounted.test.tsx @@ -0,0 +1,234 @@ +/* eslint-disable typescript-eslint/no-deprecated -- react-test-renderer is the DOM-free renderer used to mount React/RN trees under vitest (same pattern as src/test/render-with-providers.tsx) */ +import { createElement } from 'react'; +import { act } from 'react-test-renderer'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { renderWithProviders, waitFor } from '@/test/render-with-providers'; +import { type DeleteAccountPhase, useDeleteAccount } from './use-delete-account'; + +// ── Hoisted mocks ────────────────────────────────────────────────────────── + +const challengeFn = vi.hoisted(() => vi.fn()); +const executeFn = vi.hoisted(() => vi.fn()); +const signOutFn = vi.hoisted(() => vi.fn()); +const toastSuccess = vi.hoisted(() => vi.fn()); +const toastError = vi.hoisted(() => vi.fn()); + +vi.mock('@/lib/trpc', () => ({ + useTRPC: () => ({ + user: { + requestAccountDeletionChallenge: { + mutationOptions: (opts: object) => ({ + ...opts, + mutationFn: challengeFn, + mutationKey: ['delete-account-challenge'], + }), + }, + requestAccountDeletion: { + mutationOptions: (opts: object) => ({ + ...opts, + mutationFn: executeFn, + mutationKey: ['delete-account-execute'], + }), + }, + }, + }), +})); + +vi.mock('@/lib/auth/auth-context', () => ({ + useAuth: () => ({ signOut: signOutFn }), +})); + +vi.mock('sonner-native', () => ({ + toast: { success: toastSuccess, error: toastError }, +})); + +// ── Helpers ──────────────────────────────────────────────────────────────── + +type DeleteAccountResult = { + phase: DeleteAccountPhase; + isPending: boolean; + devCode: string | null; + beginDelete: () => void; + submitCode: () => void; + setCode: (code: string) => void; +}; + +function Probe({ holder }: { holder: { current: DeleteAccountResult | null } }) { + holder.current = useDeleteAccount(); + return null; +} + +async function mountProbe(): Promise<{ current: DeleteAccountResult | null }> { + const holder: { current: DeleteAccountResult | null } = { current: null }; + await renderWithProviders(createElement(Probe, { holder })); + return holder; +} + +function current(holder: { current: DeleteAccountResult | null }): DeleteAccountResult { + const result = holder.current; + if (!result) { + throw new Error('probe did not render'); + } + return result; +} + +function createDeferred(): { promise: Promise; resolve: (value: T) => void } { + const state: { resolve: ((value: T) => void) | undefined } = { resolve: undefined }; + const promise = new Promise(resolve => { + state.resolve = resolve; + }); + return { + promise, + resolve: (value: T) => { + state.resolve?.(value); + }, + }; +} + +// ── Tests ────────────────────────────────────────────────────────────────── + +describe('useDeleteAccount', () => { + beforeEach(() => { + vi.stubGlobal('__DEV__', false); + challengeFn.mockReset(); + executeFn.mockReset(); + signOutFn.mockReset(); + toastSuccess.mockReset(); + toastError.mockReset(); + }); + + it('happy: challenge then valid code deletes and signs out', async () => { + challengeFn.mockResolvedValue({ challengeId: 'challenge-1', devCode: '123456' }); + executeFn.mockResolvedValue({ status: 'deleted' }); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + await waitFor(() => current(holder).phase === 'awaiting-code'); + + act(() => { + current(holder).setCode('123456'); + current(holder).submitCode(); + }); + await waitFor(() => current(holder).phase === 'deleted'); + + expect(executeFn).toHaveBeenCalledTimes(1); + expect(executeFn.mock.calls[0]?.[0]).toEqual({ challengeId: 'challenge-1', code: '123456' }); + expect(toastSuccess).toHaveBeenCalledWith('Your account has been deleted.'); + expect(signOutFn).toHaveBeenCalledWith(true); + }); + + it('retryable: a failed challenge toasts and stays retryable', async () => { + challengeFn.mockRejectedValue(new Error('Network error')); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + await waitFor(() => toastError.mock.calls.length > 0); + + expect(toastError).toHaveBeenCalledWith('Network error'); + expect(signOutFn).not.toHaveBeenCalled(); + // Back to idle: the caller can re-request a code (retry / resend). + expect(current(holder).phase).toBe('idle'); + }); + + it('retryable: an execute network failure keeps the challenge', async () => { + challengeFn.mockResolvedValue({ challengeId: 'challenge-1' }); + executeFn.mockRejectedValue(new Error('Network error')); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + await waitFor(() => current(holder).phase === 'awaiting-code'); + + act(() => { + current(holder).setCode('123456'); + current(holder).submitCode(); + }); + await waitFor(() => toastError.mock.calls.length > 0); + + expect(toastError).toHaveBeenCalledWith('Network error'); + expect(signOutFn).not.toHaveBeenCalled(); + // Challenge kept: the user can resubmit the same code. + expect(current(holder).phase).toBe('awaiting-code'); + }); + + it('non-retryable: a precondition failure keeps the challenge so the user can resubmit', async () => { + challengeFn.mockResolvedValue({ challengeId: 'challenge-1' }); + executeFn.mockRejectedValue({ + data: { code: 'PRECONDITION_FAILED' }, + message: 'Cancel the subscription before deleting the account.', + }); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + await waitFor(() => current(holder).phase === 'awaiting-code'); + + act(() => { + current(holder).setCode('654321'); + current(holder).submitCode(); + }); + await waitFor(() => toastError.mock.calls.length > 0); + + expect(toastError).toHaveBeenCalledWith('Cancel the subscription before deleting the account.'); + expect(signOutFn).not.toHaveBeenCalled(); + // Challenge kept: the user can fix the blocker and resubmit the same code. + expect(current(holder).phase).toBe('awaiting-code'); + }); + + it('non-retryable: an invalid code keeps the challenge so the user can re-enter', async () => { + challengeFn.mockResolvedValue({ challengeId: 'challenge-1' }); + executeFn.mockRejectedValue({ + data: { code: 'UNAUTHORIZED' }, + message: 'Invalid confirmation code', + }); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + await waitFor(() => current(holder).phase === 'awaiting-code'); + + act(() => { + current(holder).setCode('000000'); + current(holder).submitCode(); + }); + await waitFor(() => toastError.mock.calls.length > 0); + + expect(toastError).toHaveBeenCalledWith('Invalid confirmation code'); + expect(signOutFn).not.toHaveBeenCalled(); + // Challenge kept: the user can re-enter a correct code. + expect(current(holder).phase).toBe('awaiting-code'); + }); + + it('pending: isPending is true while the challenge is in flight', async () => { + const deferredChallenge = createDeferred<{ challengeId: string }>(); + challengeFn.mockReturnValue(deferredChallenge.promise); + + const holder = await mountProbe(); + + act(() => { + current(holder).beginDelete(); + }); + + await waitFor(() => current(holder).isPending); + expect(current(holder).phase).toBe('requesting'); + + act(() => { + deferredChallenge.resolve({ challengeId: 'challenge-1' }); + }); + await waitFor(() => current(holder).phase === 'awaiting-code'); + expect(current(holder).isPending).toBe(false); + }); +}); diff --git a/apps/mobile/src/components/profile-screen.tsx b/apps/mobile/src/components/profile-screen.tsx index 5b6baaef70f..182055b5b1b 100644 --- a/apps/mobile/src/components/profile-screen.tsx +++ b/apps/mobile/src/components/profile-screen.tsx @@ -1,5 +1,5 @@ /* eslint-disable max-lines -- The profile screen composes Credits, Agents, Reviews, Organization, Linked accounts, App, Restore Purchases, and Actions; each section is a small rendered surface that mirrors the shared ConfigureRow/Text-header pattern. Splitting would re-encode the same hooks. */ -import { useMutation, useQuery } from '@tanstack/react-query'; +import { useQuery } from '@tanstack/react-query'; import * as Application from 'expo-application'; import { type Href, useRouter } from 'expo-router'; import { @@ -16,7 +16,6 @@ import { Trash2, } from '@/components/ui/icons'; import { Alert, Platform, View } from 'react-native'; -import { toast } from 'sonner-native'; import Animated, { FadeIn, FadeOut } from 'react-native-reanimated'; import { RestorePurchasesButton } from '@/components/kilo-pass/restore-purchases-button'; @@ -25,9 +24,12 @@ import { CreditsCard } from '@/components/profile-credits-card'; import { QueryError } from '@/components/query-error'; import { ScreenHeader } from '@/components/screen-header'; import { TabScreenScrollView } from '@/components/tab-screen'; +import { Button } from '@/components/ui/button'; import { ConfigureRow } from '@/components/ui/configure-row'; +import { FormField } from '@/components/ui/form-field'; import { Skeleton } from '@/components/ui/skeleton'; import { Text } from '@/components/ui/text'; +import { useDeleteAccount } from '@/components/use-delete-account'; import { FEATURE_FLAG_PR_REVIEW, useFeatureFlag } from '@/lib/analytics/posthog'; import { useAuth } from '@/lib/auth/auth-context'; import { showFeedbackPrompt } from '@/lib/feedback'; @@ -82,29 +84,25 @@ export function ProfileScreen() { const { userId } = useCurrentUserId({ enabled: isAuthenticated }); - const deleteAccount = useMutation( - trpc.user.requestAccountDeletion.mutationOptions({ - onSuccess: () => { - toast.success('Account deletion request sent. Check your email for confirmation.'); - }, - onError: error => { - toast.error(error.message); - }, - }) - ); + const { + phase: deletePhase, + isPending: deletePending, + devCode, + beginDelete, + submitCode, + setCode, + } = useDeleteAccount(); const confirmDeleteAccount = () => { Alert.alert( 'Delete account?', - 'This will send a request to permanently delete your account and all associated data. This action cannot be undone.', + 'This permanently deletes your account and signs you out. We will email a confirmation code.', [ { text: 'Cancel', style: 'cancel' }, { text: 'Delete account', style: 'destructive', - onPress: () => { - deleteAccount.mutate(); - }, + onPress: beginDelete, }, ] ); @@ -328,10 +326,31 @@ export function ProfileScreen() { label="Delete Account" color={colors.destructive} destructive - disabled={deleteAccount.isPending} + disabled={deletePending} onPress={confirmDeleteAccount} /> + {(deletePhase === 'awaiting-code' || deletePhase === 'executing') && ( + + + + + )} + v{Application.nativeApplicationVersion} ({Application.nativeBuildVersion}) diff --git a/apps/mobile/src/components/use-delete-account.ts b/apps/mobile/src/components/use-delete-account.ts new file mode 100644 index 00000000000..6e004c1be42 --- /dev/null +++ b/apps/mobile/src/components/use-delete-account.ts @@ -0,0 +1,112 @@ +import { useRef, useState } from 'react'; +import { useMutation } from '@tanstack/react-query'; +import { toast } from 'sonner-native'; + +import { useAuth } from '@/lib/auth/auth-context'; +import { useTRPC } from '@/lib/trpc'; + +export type DeleteAccountPhase = 'idle' | 'requesting' | 'awaiting-code' | 'executing' | 'deleted'; + +type DeleteAccountChallenge = { challengeId: string; devCode?: string }; + +/** + * Derive the deletion flow phase from the mutation states. Order matters: + * a completed delete wins over an in-flight one, and a held challenge wins + * over a still-pending challenge request. + */ +function derivePhase(state: { + isSuccess: boolean; + isExecuting: boolean; + hasChallenge: boolean; + isRequesting: boolean; +}): DeleteAccountPhase { + if (state.isSuccess) { + return 'deleted'; + } + if (state.isExecuting) { + return 'executing'; + } + if (state.hasChallenge) { + return 'awaiting-code'; + } + if (state.isRequesting) { + return 'requesting'; + } + return 'idle'; +} + +/** + * Self-service account deletion, reauthenticated with an emailed sign-in code. + * + * `beginDelete` requests a confirmation-code challenge. On success the hook + * moves to `awaiting-code`; the caller renders a code field and calls + * `submitCode`. `submitCode` executes the deletion and, on success, signs the + * user out. Any execute failure keeps the challenge so the caller can fix the + * blocker (or correct the code) and resubmit the same `challengeId`. + */ +export function useDeleteAccount() { + const trpc = useTRPC(); + const { signOut } = useAuth(); + + const [challenge, setChallenge] = useState(null); + const codeRef = useRef(''); + + const requestChallenge = useMutation( + trpc.user.requestAccountDeletionChallenge.mutationOptions({ + onSuccess: data => { + setChallenge(data); + if (__DEV__ && data.devCode) { + codeRef.current = data.devCode; + } + }, + onError: error => { + toast.error(error.message); + }, + }) + ); + + const executeDeletion = useMutation( + trpc.user.requestAccountDeletion.mutationOptions({ + onSuccess: () => { + toast.success('Your account has been deleted.'); + void signOut(true); + }, + onError: error => { + toast.error(error.message); + }, + }) + ); + + const phase = derivePhase({ + isSuccess: executeDeletion.isSuccess, + isExecuting: executeDeletion.isPending, + hasChallenge: challenge !== null, + isRequesting: requestChallenge.isPending, + }); + + const isPending = requestChallenge.isPending || executeDeletion.isPending; + + const beginDelete = () => { + requestChallenge.mutate(); + }; + + const submitCode = () => { + if (!challenge) { + return; + } + executeDeletion.mutate({ challengeId: challenge.challengeId, code: codeRef.current }); + }; + + const setCode = (code: string) => { + codeRef.current = code; + }; + + return { + phase, + isPending, + devCode: challenge?.devCode ?? null, + beginDelete, + submitCode, + setCode, + }; +} diff --git a/apps/web/src/routers/user-router.test.ts b/apps/web/src/routers/user-router.test.ts index a85269a85c4..7a23e0c6dc2 100644 --- a/apps/web/src/routers/user-router.test.ts +++ b/apps/web/src/routers/user-router.test.ts @@ -4,11 +4,39 @@ import { credit_transactions, device_sessions, kilocode_users, + magic_link_tokens, user_notification_preferences, } from '@kilocode/db/schema'; import { eq, inArray } from 'drizzle-orm'; import { insertTestUser } from '@/tests/helpers/user.helper'; import type { User } from '@kilocode/db/schema'; +import { sendSignInCodeEmail } from '@/lib/email'; +import { performGdprRemoval } from '@/lib/user/gdpr-removal'; +import { assertUserCanBeSoftDeleted, SoftDeletePreconditionError } from '@/lib/user'; + +jest.mock('@/lib/email', () => { + const actual = jest.requireActual('@/lib/email'); + return { + ...actual, + sendSignInCodeEmail: jest.fn(), + }; +}); + +jest.mock('@/lib/user/gdpr-removal', () => ({ + performGdprRemoval: jest.fn(), +})); + +jest.mock('@/lib/user', () => { + const actual = jest.requireActual('@/lib/user'); + return { + ...actual, + assertUserCanBeSoftDeleted: jest.fn(), + }; +}); + +const mockSendSignInCodeEmail = jest.mocked(sendSignInCodeEmail); +const mockPerformGdprRemoval = jest.mocked(performGdprRemoval); +const mockAssertUserCanBeSoftDeleted = jest.mocked(assertUserCanBeSoftDeleted); let testUser: User; let surveyTestUser: User; @@ -996,3 +1024,128 @@ describe('user router - device sessions', () => { expect(row?.revoked_at).toBeNull(); }); }); + +describe('user router - account deletion', () => { + let deletionUser: User; + + beforeEach(async () => { + deletionUser = await insertTestUser({ + google_user_email: `deletion-${crypto.randomUUID()}@example.com`, + google_user_name: 'Deletion Test User', + }); + + mockSendSignInCodeEmail.mockReset(); + mockPerformGdprRemoval.mockReset(); + mockAssertUserCanBeSoftDeleted.mockReset(); + + mockAssertUserCanBeSoftDeleted.mockResolvedValue(undefined); + mockPerformGdprRemoval.mockResolvedValue({ warnings: [] }); + mockSendSignInCodeEmail.mockResolvedValue({ sent: false, reason: 'provider_not_configured' }); + }); + + afterEach(async () => { + await db + .delete(magic_link_tokens) + .where(eq(magic_link_tokens.email, deletionUser.google_user_email)); + await db.delete(kilocode_users).where(eq(kilocode_users.id, deletionUser.id)); + }); + + it('challenge sends the sign-in code email to the authenticated email', async () => { + const caller = await createCallerForUser(deletionUser.id); + + const result = await caller.user.requestAccountDeletionChallenge(); + + expect(mockSendSignInCodeEmail).toHaveBeenCalledTimes(1); + expect(mockSendSignInCodeEmail).toHaveBeenCalledWith( + deletionUser.google_user_email, + expect.any(String) + ); + expect(result.challengeId).toEqual(expect.any(String)); + expect(result.devCode).toEqual(expect.any(String)); + }); + + it('challenge does not accept a client-supplied email', async () => { + const caller = await createCallerForUser(deletionUser.id); + + await caller.user.requestAccountDeletionChallenge(); + + // The mutation takes no input; the only email used is the authenticated + // user's own address, never a client-provided value. + expect(mockSendSignInCodeEmail).toHaveBeenCalledWith( + deletionUser.google_user_email, + expect.any(String) + ); + }); + + it('challenge returns devCode when the email sends in non-production', async () => { + mockSendSignInCodeEmail.mockResolvedValue({ sent: true }); + const caller = await createCallerForUser(deletionUser.id); + + const result = await caller.user.requestAccountDeletionChallenge(); + + expect(result.challengeId).toEqual(expect.any(String)); + expect(result.devCode).toEqual(expect.any(String)); + }); + + it('challenge omits devCode in production', async () => { + mockSendSignInCodeEmail.mockResolvedValue({ sent: true }); + const caller = await createCallerForUser(deletionUser.id); + const originalNodeEnv = process.env.NODE_ENV; + process.env.NODE_ENV = 'production'; + try { + const result = await caller.user.requestAccountDeletionChallenge(); + + expect(result.challengeId).toEqual(expect.any(String)); + expect(result.devCode).toBeUndefined(); + } finally { + process.env.NODE_ENV = originalNodeEnv; + } + }); + + it('wrong code does not delete the account', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + const wrongCode = devCode === '000000' ? '111111' : '000000'; + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: wrongCode }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + + expect(mockPerformGdprRemoval).not.toHaveBeenCalled(); + }); + + it('valid code calls performGdprRemoval once', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + const result = await caller.user.requestAccountDeletion({ + challengeId, + code: devCode as string, + }); + + expect(result).toEqual({ status: 'deleted' }); + expect(mockPerformGdprRemoval).toHaveBeenCalledTimes(1); + expect(mockPerformGdprRemoval).toHaveBeenCalledWith(deletionUser.id, { + destroyReason: 'admin_request', + actor: { + id: deletionUser.id, + email: deletionUser.google_user_email, + name: deletionUser.google_user_name, + }, + }); + }); + + it('precondition error maps to PRECONDITION_FAILED without deleting', async () => { + mockAssertUserCanBeSoftDeleted.mockRejectedValue( + new SoftDeletePreconditionError('active subscription') + ); + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).rejects.toMatchObject({ code: 'PRECONDITION_FAILED' }); + + expect(mockPerformGdprRemoval).not.toHaveBeenCalled(); + }); +}); diff --git a/apps/web/src/routers/user-router.ts b/apps/web/src/routers/user-router.ts index 096677f012e..1c26a2a4ec5 100644 --- a/apps/web/src/routers/user-router.ts +++ b/apps/web/src/routers/user-router.ts @@ -1,9 +1,17 @@ import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; -import { getUserAuthProviders, unlinkAuthProviderFromUser } from '@/lib/user'; import { - sendAccountDeletionConfirmationEmail, - sendAccountDeletionSupportNotification, -} from '@/lib/email'; + assertUserCanBeSoftDeleted, + getUserAuthProviders, + SoftDeletePreconditionError, + unlinkAuthProviderFromUser, +} from '@/lib/user'; +import { sendSignInCodeEmail } from '@/lib/email'; +import { + consumeSignInCode, + createSignInCode, + deleteSignInCode, +} from '@/lib/auth/magic-link-tokens'; +import { performGdprRemoval } from '@/lib/user/gdpr-removal'; import { createAccountLinkingSession } from '@/lib/account-linking-session'; import { TRPCError } from '@trpc/server'; import { captureException } from '@sentry/nextjs'; @@ -877,33 +885,94 @@ export const userRouter = createTRPCRouter({ return successResult({ is_member: isMember }); }), - requestAccountDeletion: baseProcedure.mutation(async ({ ctx }) => { - const userEmail = ctx.user.google_user_email; - const userId = ctx.user.id; + requestAccountDeletionChallenge: baseProcedure + .output(z.object({ challengeId: z.uuid(), devCode: z.string().optional() })) + .mutation(async ({ ctx }) => { + const userEmail = ctx.user.google_user_email; + const userId = ctx.user.id; + + const lastRequested = ctx.user.account_deletion_requested_at; + if ( + lastRequested && + Date.now() - new Date(lastRequested).getTime() < ACCOUNT_DELETION_COOLDOWN_MS + ) { + throw new TRPCError({ + code: 'TOO_MANY_REQUESTS', + message: 'Account deletion already requested. Please wait before trying again.', + }); + } + + const { code, challengeId } = await createSignInCode(userEmail); + const sendResult = await sendSignInCodeEmail(userEmail, code); + const includeDevCode = process.env.NODE_ENV !== 'production'; - const lastRequested = ctx.user.account_deletion_requested_at; - if ( - lastRequested && - Date.now() - new Date(lastRequested).getTime() < ACCOUNT_DELETION_COOLDOWN_MS - ) { + if (sendResult.sent) { + await db + .update(kilocode_users) + .set({ account_deletion_requested_at: new Date().toISOString() }) + .where(eq(kilocode_users.id, userId)); + return includeDevCode ? { challengeId, devCode: code } : { challengeId }; + } + + if (sendResult.reason === 'provider_not_configured') { + // Local/dev has no mail provider. Keep the code row so the caller can + // still complete the flow with the returned dev code. + await db + .update(kilocode_users) + .set({ account_deletion_requested_at: new Date().toISOString() }) + .where(eq(kilocode_users.id, userId)); + return includeDevCode ? { challengeId, devCode: code } : { challengeId }; + } + + // neverbounce_rejected: discard the code and leave the stamp untouched so + // a retry / resend stays available. + await deleteSignInCode(userEmail, code); throw new TRPCError({ - code: 'TOO_MANY_REQUESTS', - message: 'Account deletion already requested. Please wait before trying again.', + code: 'INTERNAL_SERVER_ERROR', + message: 'Failed to send confirmation code. Please try again.', }); - } + }), - await Promise.all([ - sendAccountDeletionConfirmationEmail(userEmail), - sendAccountDeletionSupportNotification(userEmail, userId), - ]); + requestAccountDeletion: baseProcedure + .input(z.object({ challengeId: z.uuid(), code: z.string().min(1) })) + .mutation(async ({ ctx, input }) => { + const userEmail = ctx.user.google_user_email; + const userId = ctx.user.id; - await db - .update(kilocode_users) - .set({ account_deletion_requested_at: new Date().toISOString() }) - .where(eq(kilocode_users.id, userId)); + try { + await assertUserCanBeSoftDeleted(userId); + } catch (error) { + if (error instanceof SoftDeletePreconditionError) { + throw new TRPCError({ code: 'PRECONDITION_FAILED', message: error.message }); + } + throw error; + } - return successResult(); - }), + const consumed = await consumeSignInCode(userEmail, input.code, input.challengeId); + if (!consumed) { + throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid confirmation code' }); + } + + try { + await performGdprRemoval(userId, { + destroyReason: 'admin_request', + actor: { + id: userId, + email: userEmail, + name: ctx.user.google_user_name, + }, + }); + } catch (error) { + // A precondition race inside performGdprRemoval (re-asserted there) maps + // to the same PRECONDITION_FAILED surface. + if (error instanceof SoftDeletePreconditionError) { + throw new TRPCError({ code: 'PRECONDITION_FAILED', message: error.message }); + } + throw error; + } + + return { status: 'deleted' as const }; + }), // ─── Push Notification Tokens ────────────────────────────────────── From 075a45f0873959a2a58e93bad597e78da7c39bfd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 12:11:57 +0200 Subject: [PATCH 11/23] test: satisfy typecheck and lint in changed test files Fix tsgo errors (body cast, NODE_ENV replaceProperty, nonce-DO binding, holder pattern, requireActual cast, fetch mock typing) and oxlint no-floating-promises / require-await in the slice test files. --- ...r-web-connection-provider.mounted.test.tsx | 45 +++++++------- apps/mobile/src/lib/auth/credentials.test.ts | 20 ++++--- .../cloud-agent-client.test.ts | 10 +++- apps/web/src/routers/user-router.test.ts | 5 +- services/cloud-agent-next/src/router.test.ts | 6 ++ .../src/dos/UserConnectionDO.test.ts | 59 ++++++++++--------- .../session-ingest/src/routes/api.test.ts | 2 +- 7 files changed, 84 insertions(+), 63 deletions(-) diff --git a/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx b/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx index 2e8c1e36021..129a4039959 100644 --- a/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx +++ b/apps/mobile/src/components/agents/user-web-connection-provider.mounted.test.tsx @@ -8,7 +8,7 @@ import { UserWebConnectionProvider } from './user-web-connection-provider'; type AuthConfig = { getAuthToken: () => Promise }; const mocks = vi.hoisted(() => ({ - mutate: vi.fn(async () => ({ token: 'ticket-1', expiresAt: 1_700_000_060 })), + mutate: vi.fn(() => ({ token: 'ticket-1', expiresAt: 1_700_000_060 })), query: vi.fn(), createUserWebConnection: vi.fn(), capturedConfig: null as AuthConfig | null, @@ -19,19 +19,19 @@ vi.mock('@kilocode/cloud-agent-sdk/user-web-connection', () => ({ mocks.capturedConfig = config; mocks.createUserWebConnection(config); return { - retain: () => () => {}, - connect: () => {}, - disconnect: () => {}, - destroy: () => {}, - isConnected: () => false, - onConnectionChange: () => () => {}, - subscribeToCliSession: () => () => {}, - sendCommand: async () => ({}), - sendCommandToConnection: async () => ({}), - onCliEvent: () => () => {}, - onSystemEvent: () => () => {}, - onReconnect: () => () => {}, - onSessionEvent: () => () => {}, + retain: vi.fn(() => vi.fn()), + connect: vi.fn(), + disconnect: vi.fn(), + destroy: vi.fn(), + isConnected: vi.fn(() => false), + onConnectionChange: vi.fn(() => vi.fn()), + subscribeToCliSession: vi.fn(() => vi.fn()), + sendCommand: vi.fn(), + sendCommandToConnection: vi.fn(), + onCliEvent: vi.fn(() => vi.fn()), + onSystemEvent: vi.fn(() => vi.fn()), + onReconnect: vi.fn(() => vi.fn()), + onSessionEvent: vi.fn(() => vi.fn()), }; }, })); @@ -64,22 +64,25 @@ describe('UserWebConnectionProvider', () => { }); it('mints the ingest ticket via the getToken mutation', async () => { - let renderer: TestRenderer.ReactTestRenderer | undefined; - await act(async () => { - renderer = TestRenderer.create(createElement(UserWebConnectionProvider, { children: null })); + const holder: { renderer?: TestRenderer.ReactTestRenderer } = {}; + await act(() => { + holder.renderer = TestRenderer.create(createElement(UserWebConnectionProvider, null)); }); expect(mocks.createUserWebConnection).toHaveBeenCalledTimes(1); - expect(mocks.capturedConfig).not.toBeNull(); + const config = mocks.capturedConfig; + if (!config) { + throw new Error('user web connection config not captured'); + } - const token = await mocks.capturedConfig!.getAuthToken(); + const token = await config.getAuthToken(); expect(token).toBe('ticket-1'); expect(mocks.mutate).toHaveBeenCalledTimes(1); expect(mocks.query).not.toHaveBeenCalled(); - await act(async () => { - renderer!.unmount(); + await act(() => { + holder.renderer?.unmount(); }); }); }); diff --git a/apps/mobile/src/lib/auth/credentials.test.ts b/apps/mobile/src/lib/auth/credentials.test.ts index 3bf60154ba7..35ddadc6c91 100644 --- a/apps/mobile/src/lib/auth/credentials.test.ts +++ b/apps/mobile/src/lib/auth/credentials.test.ts @@ -174,15 +174,20 @@ async function mountAndSignOut(): Promise { vi.resetModules(); const mod = await import('./auth-context'); - let captured: { signOut: (ended?: boolean) => Promise } | undefined = undefined; + const holder: { + captured?: { signOut: (ended?: boolean) => Promise }; + renderer?: TestRenderer.ReactTestRenderer; + } = {}; + function Consumer(): null { - captured = mod.useAuth(); + holder.captured = mod.useAuth(); return null; } - let renderer: TestRenderer.ReactTestRenderer | undefined = undefined; await act(async () => { - renderer = TestRenderer.create(createElement(mod.AuthProvider, null, createElement(Consumer))); + holder.renderer = TestRenderer.create( + createElement(mod.AuthProvider, null, createElement(Consumer)) + ); await Promise.resolve(); }); await act(async () => { @@ -192,13 +197,14 @@ async function mountAndSignOut(): Promise { }); // oxlint-disable-next-line @typescript-eslint/no-unnecessary-condition -- safety net for test failures - if (!captured) { + if (!holder.captured) { throw new Error('auth context not captured'); } + const signOut = holder.captured.signOut; await act(async () => { - await captured.signOut(); + await signOut(); }); - renderer?.unmount(); + holder.renderer?.unmount(); } diff --git a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts index 3e0f1632712..7a48ea84472 100644 --- a/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts +++ b/apps/web/src/lib/cloud-agent-next/cloud-agent-client.test.ts @@ -58,7 +58,9 @@ beforeEach(() => { // Load the real `closeCloudAgentOrgStreams` (the module mock above does not // expose it) so the test exercises the actual fetch call, not a stub. -const { closeCloudAgentOrgStreams } = jest.requireActual('./cloud-agent-client'); +const { closeCloudAgentOrgStreams } = jest.requireActual('./cloud-agent-client') as { + closeCloudAgentOrgStreams: (userId: string, organizationId: string) => Promise; +}; describe('createCloudAgentNextClientForModel', () => { it('returns the default client when the model is paid and has no BYOK', () => { @@ -100,7 +102,9 @@ describe('closeCloudAgentOrgStreams', () => { }); it('POSTs to the close endpoint with the internal key header and body', async () => { - const fetchMock = jest.fn().mockResolvedValue(new Response(null, { status: 204 })); + const fetchMock = jest + .fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>() + .mockResolvedValue(new Response(null, { status: 204 })); global.fetch = fetchMock as unknown as typeof fetch; await closeCloudAgentOrgStreams('usr_1', 'org_1'); @@ -120,7 +124,7 @@ describe('closeCloudAgentOrgStreams', () => { it('throws when the close endpoint returns a non-OK response', async () => { const fetchMock = jest - .fn() + .fn<(input: RequestInfo | URL, init?: RequestInit) => Promise>() .mockResolvedValue( new Response('boom', { status: 500, statusText: 'Internal Server Error' }) ); diff --git a/apps/web/src/routers/user-router.test.ts b/apps/web/src/routers/user-router.test.ts index 7a23e0c6dc2..dfd4b820a8f 100644 --- a/apps/web/src/routers/user-router.test.ts +++ b/apps/web/src/routers/user-router.test.ts @@ -1090,15 +1090,14 @@ describe('user router - account deletion', () => { it('challenge omits devCode in production', async () => { mockSendSignInCodeEmail.mockResolvedValue({ sent: true }); const caller = await createCallerForUser(deletionUser.id); - const originalNodeEnv = process.env.NODE_ENV; - process.env.NODE_ENV = 'production'; + const restoreNodeEnv = jest.replaceProperty(process.env, 'NODE_ENV', 'production'); try { const result = await caller.user.requestAccountDeletionChallenge(); expect(result.challengeId).toEqual(expect.any(String)); expect(result.devCode).toBeUndefined(); } finally { - process.env.NODE_ENV = originalNodeEnv; + restoreNodeEnv.restore(); } }); diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 66135d15215..a0787b9e39c 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -364,6 +364,7 @@ describe('router sessionId validation', () => { })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -797,6 +798,7 @@ describe('router sessionId validation', () => { get: vi.fn(() => mockSessionStub), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -917,6 +919,7 @@ describe('router sessionId validation', () => { })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -1221,6 +1224,7 @@ describe('router sessionId validation', () => { })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn() } as unknown as TRPCContext['env']['SESSION_INGEST'], CONTAINER_USAGE_METER: {} as TRPCContext['env']['CONTAINER_USAGE_METER'], @@ -1337,6 +1341,7 @@ describe('router sessionId validation', () => { })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), @@ -1625,6 +1630,7 @@ describe('router sessionId validation', () => { })), } as unknown as TRPCContext['env']['CLOUD_AGENT_SESSION'], USER_KILO_FACADE: {} as TRPCContext['env']['USER_KILO_FACADE'], + STREAM_TICKET_NONCE_DO: {} as TRPCContext['env']['STREAM_TICKET_NONCE_DO'], SHARED_SANDBOX_OVERRIDES: {} as TRPCContext['env']['SHARED_SANDBOX_OVERRIDES'], SESSION_INGEST: { fetch: vi.fn(), diff --git a/services/session-ingest/src/dos/UserConnectionDO.test.ts b/services/session-ingest/src/dos/UserConnectionDO.test.ts index 675001f70f9..eeb0571fd09 100644 --- a/services/session-ingest/src/dos/UserConnectionDO.test.ts +++ b/services/session-ingest/src/dos/UserConnectionDO.test.ts @@ -312,7 +312,7 @@ function sendHeartbeat( ...(options.capabilities ? { capabilities: options.capabilities } : {}), ...(options.instance ? { instance: options.instance } : {}), }); - doInstance.webSocketMessage(cliWs as never, msg); + void doInstance.webSocketMessage(cliWs as never, msg); } /** Send a subscribe from a web ws */ @@ -325,13 +325,13 @@ async function sendSubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessio /** Send an unsubscribe from a web ws */ function sendUnsubscribe(doInstance: UserConnectionDO, webWs: MockWS, sessionId: string) { const msg = JSON.stringify({ type: 'unsubscribe', sessionId }); - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); } /** Send a viewer ping from a web ws */ function sendPing(doInstance: UserConnectionDO, webWs: MockWS, nonce: string) { const msg = JSON.stringify({ type: 'ping', nonce }); - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); } /** Send a command from a web ws. Auto-flushes so durable-before-send @@ -362,7 +362,7 @@ async function sendCliResponse( opts: { id: string; result?: unknown; error?: unknown } ): Promise { const msg = JSON.stringify({ type: 'response', ...opts }); - doInstance.webSocketMessage(cliWs as never, msg); + void doInstance.webSocketMessage(cliWs as never, msg); await flushAsync(); } @@ -1832,7 +1832,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: {}, }); - doInstance.webSocketMessage(cliWs as never, cliEventMsg); + void doInstance.webSocketMessage(cliWs as never, cliEventMsg); // No web sockets to receive the event — no crash = success }); @@ -3530,7 +3530,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); expect(subWeb.send).toHaveBeenCalledTimes(1); expect(parseSent(subWeb)).toEqual({ @@ -3561,7 +3561,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); expect(parentWeb.send).toHaveBeenCalledTimes(1); expect(childWeb.send).toHaveBeenCalledTimes(1); @@ -3593,7 +3593,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); // Should only receive once despite subscribing to both expect(webWs.send).toHaveBeenCalledTimes(1); @@ -3616,7 +3616,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-child-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ @@ -3645,7 +3645,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-child-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).not.toHaveBeenCalled(); }); @@ -3665,7 +3665,7 @@ describe('UserConnectionDO', () => { event: 'message.updated', data: { id: 'msg-1' }, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); expect(webWs.send).toHaveBeenCalledTimes(1); expect(parseSent(webWs)).toEqual({ @@ -3691,7 +3691,7 @@ describe('UserConnectionDO', () => { event: 'session.status', data: {}, }); - doInstance.webSocketMessage(cliWs as never, eventMsg); + void doInstance.webSocketMessage(cliWs as never, eventMsg); const sent = parseSent(webWs); expect(sent).not.toHaveProperty('parentSessionId'); @@ -3778,7 +3778,7 @@ describe('UserConnectionDO', () => { event: 'test', data: {}, }); - doInstance.webSocketMessage(cliWs as never, triggerMsg); + void doInstance.webSocketMessage(cliWs as never, triggerMsg); // webWs should have received the event because it was subscribed via attachment expect(webWs.send).toHaveBeenCalledTimes(1); @@ -3799,7 +3799,7 @@ describe('UserConnectionDO', () => { replaced: true, }); - doInstance.webSocketMessage( + void doInstance.webSocketMessage( cliWs as never, JSON.stringify({ type: 'event', @@ -4216,7 +4216,7 @@ describe('UserConnectionDO', () => { const cliWs = addCliSocket(mockCtx, 'cli-1'); // Should not throw - doInstance.webSocketMessage(cliWs as never, 'not-json'); + void doInstance.webSocketMessage(cliWs as never, 'not-json'); }); it('logs invalid CLI JSON metadata without raw payload content', async () => { @@ -4225,7 +4225,7 @@ describe('UserConnectionDO', () => { const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); const malformed = '{"secret":"raw-secret-must-not-be-logged"'; - doInstance.webSocketMessage(cliWs as never, malformed); + void doInstance.webSocketMessage(cliWs as never, malformed); expect(warn).toHaveBeenCalledWith('Failed to parse WebSocket message as JSON', { role: 'cli', @@ -4241,7 +4241,10 @@ describe('UserConnectionDO', () => { mockCtx.addSocket(ws); // Trigger ensureState first - doInstance.webSocketMessage(ws as never, JSON.stringify({ type: 'heartbeat', sessions: [] })); + void doInstance.webSocketMessage( + ws as never, + JSON.stringify({ type: 'heartbeat', sessions: [] }) + ); // Should not throw }); @@ -4252,12 +4255,12 @@ describe('UserConnectionDO', () => { // Invalid CLI message const badMsg = JSON.stringify({ type: 'invalid_type' }); - doInstance.webSocketMessage(cliWs as never, badMsg); + void doInstance.webSocketMessage(cliWs as never, badMsg); // Should not throw // Invalid web message const webWs = addWebSocket(mockCtx, 'web-1'); - doInstance.webSocketMessage(webWs as never, badMsg); + void doInstance.webSocketMessage(webWs as never, badMsg); // Should not throw }); @@ -4272,7 +4275,7 @@ describe('UserConnectionDO', () => { result: { secret }, }); - doInstance.webSocketMessage(cliWs as never, malformed); + void doInstance.webSocketMessage(cliWs as never, malformed); expect(warn).toHaveBeenCalledWith('CLI message parse failed', { role: 'cli', @@ -5963,7 +5966,7 @@ describe('UserConnectionDO', () => { }); // Send raw — schema validation must reject it. - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); // No response forwarded to CLI. const cliCommands = allSent(mockCtx.sockets.find(s => s._tags.includes('cli'))!).filter( @@ -6095,7 +6098,7 @@ describe('UserConnectionDO', () => { // Start the first CLI reply — the handler rehydrates the durable // entry, builds the in-memory entry, then blocks on the deferred // storage.put before sending the live response. - doInstance.webSocketMessage( + void doInstance.webSocketMessage( cliWs as never, JSON.stringify({ type: 'response', id: correlationId, result: { first: true } }) ); @@ -6104,7 +6107,7 @@ describe('UserConnectionDO', () => { // Start the second CLI reply while the first is held on storage.put. // The completedCorrelationIds reservation is still set, so the // second reply returns early without delivering. - doInstance.webSocketMessage( + void doInstance.webSocketMessage( cliWs as never, JSON.stringify({ type: 'response', id: correlationId, result: { second: true } }) ); @@ -6652,7 +6655,7 @@ describe('UserConnectionDO', () => { await Promise.resolve(); const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; expect(markers.has(correlationId!)).toBe(true); - doInstance.webSocketMessage( + void doInstance.webSocketMessage( cliWs as never, JSON.stringify({ type: 'response', id: correlationId, result: { wrong: true } }) ); @@ -6717,7 +6720,7 @@ describe('UserConnectionDO', () => { await Promise.resolve(); const markers = Reflect.get(doInstance, 'completedCorrelationIds') as Set; expect(markers.has(correlationId!)).toBe(true); - doInstance.webSocketMessage( + void doInstance.webSocketMessage( firstOwner as never, JSON.stringify({ type: 'response', id: correlationId, result: { wrong: true } }) ); @@ -7448,7 +7451,7 @@ describe('UserConnectionDO', () => { command: 'send_message', sessionId: 's1', }); - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); // Flush enough for the initial write to fail and the .catch + .finally // to execute. The command was never forwarded to the CLI. @@ -7496,7 +7499,7 @@ describe('UserConnectionDO', () => { command: 'send_message', sessionId: 's1', }); - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); // Drain microtasks so the waitUntil launches and the initial write // promise is awaited (but not yet resolved). @@ -7563,7 +7566,7 @@ describe('UserConnectionDO', () => { command: 'send_message', sessionId: 's1', }); - doInstance.webSocketMessage(webWs as never, msg); + void doInstance.webSocketMessage(webWs as never, msg); await flushAsync(); // Get the correlationId that was assigned. diff --git a/services/session-ingest/src/routes/api.test.ts b/services/session-ingest/src/routes/api.test.ts index 303527f2ece..22488a3f600 100644 --- a/services/session-ingest/src/routes/api.test.ts +++ b/services/session-ingest/src/routes/api.test.ts @@ -1892,7 +1892,7 @@ describe('api routes', () => { now.mockRestore(); expect(res.status).toBe(200); - const body = await res.json(); + const body = (await res.json()) as { ticket: string; expiresAt: number }; expect(body).toEqual({ ticket: expect.any(String), expiresAt: 1_700_000_060, From 6068c1087edd002b84cf61e590ba7f713abc7312 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 12:38:37 +0200 Subject: [PATCH 12/23] test(extension): return expiresAt from stream-ticket e2e mock The extension getTicket now requires expiresAt; the e2e fixture mock omitted it and would fail Cloud Agent connect. --- apps/extension/tests/e2e/agents-fixture.ts | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/apps/extension/tests/e2e/agents-fixture.ts b/apps/extension/tests/e2e/agents-fixture.ts index 062de2ea46d..3460ccf6c5c 100644 --- a/apps/extension/tests/e2e/agents-fixture.ts +++ b/apps/extension/tests/e2e/agents-fixture.ts @@ -281,7 +281,10 @@ export const mockAgentsApi = async ( // ---- /api/cloud-agent-next/sessions/stream-ticket ---- await context.route('https://app.kilo.ai/api/cloud-agent-next/sessions/stream-ticket', route => - route.fulfill({ json: { ticket: 'mock-stream-ticket' }, status: 200 }) + route.fulfill({ + json: { expiresAt: Math.floor(Date.now() / 1000) + 3600, ticket: 'mock-stream-ticket' }, + status: 200, + }) ); // ---- tRPC batch dispatcher ---- @@ -520,7 +523,7 @@ export const mockAgentsApi = async ( if (proc === 'activeSessions.getToken') { return { - result: { data: { token: 'mock-ingest-token', expiresAt: 1_700_000_000 } }, + result: { data: { expiresAt: 1_700_000_000, token: 'mock-ingest-token' } }, }; } From 460a397ef86cfdbcda6095f6aa7cf8a563eb454c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 12:52:36 +0200 Subject: [PATCH 13/23] test(cloud-agent): prove /terminal nonce replay is rejected P3-D-08ba requires stream and terminal replay rejection; /stream had a test but /terminal did not. --- .../src/server-stream-ticket.test.ts | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/services/cloud-agent-next/src/server-stream-ticket.test.ts b/services/cloud-agent-next/src/server-stream-ticket.test.ts index 0d5863c53aa..b39a9ca6e34 100644 --- a/services/cloud-agent-next/src/server-stream-ticket.test.ts +++ b/services/cloud-agent-next/src/server-stream-ticket.test.ts @@ -184,6 +184,29 @@ function streamRequest(ticket: string): Request { ); } +function signTerminalTicket(overrides: Record = {}): string { + return jwt.sign( + { + type: 'stream_ticket', + purpose: 'terminal', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + ptyId: 'pty-1', + nonce: 'nonce-1', + ...overrides, + }, + secret, + { algorithm: 'HS256', expiresIn: 60, audience: 'cloud-agent-terminal' } + ); +} + +function terminalRequest(ticket: string): Request { + return new Request( + `http://worker.test/terminal?cloudAgentSessionId=session-1&ptyId=pty-1&ticket=${encodeURIComponent(ticket)}`, + { headers: { Upgrade: 'websocket' } } + ); +} + beforeEach(() => { getRunningTerminalClientMock.mockReset(); consumeCloudAgentReportBatchMock.mockClear(); @@ -267,3 +290,42 @@ describe('server /stream ticket nonce consume', () => { expect(doFetch).not.toHaveBeenCalled(); }); }); + +describe('server /terminal ticket nonce consume', () => { + it('rejects an immediate replay of the same nonce on the terminal path', async () => { + const env = createEnv(); + installNonceStore(env); + + const connectTerminal = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); + getRunningTerminalClientMock.mockResolvedValue({ + status: 'running', + client: { connectTerminal }, + }); + + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + getMetadata: vi.fn().mockResolvedValue({ + metadataSchemaVersion: 2, + identity: { + sessionId: 'session-1', + userId: 'user-1', + createdOnPlatform: 'cloud-agent', + }, + auth: {}, + lifecycle: { version: 1, timestamp: Date.now(), preparedAt: Date.now() }, + workspace: { workspacePath: '/workspace' }, + }), + }); + + const ticket = signTerminalTicket(); + + const first = await fetchWorker(terminalRequest(ticket), env); + expect(first.status).toBe(200); + expect(connectTerminal).toHaveBeenCalledOnce(); + + const replay = await fetchWorker(terminalRequest(ticket), env); + expect(replay.status).toBe(401); + await expect(replay.text()).resolves.toBe('Ticket nonce already used'); + expect(connectTerminal).toHaveBeenCalledTimes(1); + }); +}); From 326131603f1a2a5b7a460bd07ee732e33d4c6599 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 13:21:07 +0200 Subject: [PATCH 14/23] test(seed): add w4c-org-pair fixture for member-revocation E2E Creates an organization with an owner and a member for the S4 socket-close scenario. --- dev/seed/app/w4c-org-pair.ts | 101 +++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) create mode 100644 dev/seed/app/w4c-org-pair.ts diff --git a/dev/seed/app/w4c-org-pair.ts b/dev/seed/app/w4c-org-pair.ts new file mode 100644 index 00000000000..51b829aa8bb --- /dev/null +++ b/dev/seed/app/w4c-org-pair.ts @@ -0,0 +1,101 @@ +/** + * Seed fixture for the W4C member-revocation E2E scenario (S4). + * + * Creates one organization with an owner and a member, both of whom must + * already exist (created via `app:create-user`). Used to test that removing a + * member closes their socket. + * + * Usage: pnpm dev:seed app:w4c-org-pair + */ + +import { randomUUID } from 'node:crypto'; + +import { kilocode_users, organizations, organization_memberships } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; + +import { getSeedDb } from '../lib/db'; +import { normalizeSeedEmail } from '../lib/email'; +import type { SeedResult } from '../index'; + +export const usage = ' '; + +function printUsage(): void { + console.log(`Usage: pnpm dev:seed app:w4c-org-pair ${usage}`); + console.log(''); + console.log('Creates one organization with an owner and a member. Both users must'); + console.log('already exist (create them first with app:create-user).'); + console.log(''); + console.log('Examples:'); + console.log(' pnpm dev:seed app:w4c-org-pair owner@example.com member@example.com'); +} + +async function lookupUserId(email: string): Promise { + const db = getSeedDb(); + const normalizedEmail = normalizeSeedEmail(email); + + const rows = await db + .select({ id: kilocode_users.id }) + .from(kilocode_users) + .where(eq(kilocode_users.normalized_email, normalizedEmail)) + .limit(1); + + if (rows.length === 0) { + throw new Error(`No user with email ${email} exists. Create it first with app:create-user.`); + } + + return rows[0].id; +} + +export async function run(...args: string[]): Promise { + if (args.includes('--help') || args.includes('-h')) { + printUsage(); + return; + } + + const [ownerEmail, memberEmail, ...rest] = args; + if (!ownerEmail || !memberEmail) { + printUsage(); + throw new Error('owner-email and member-email are required'); + } + if (rest.length > 0) { + printUsage(); + throw new Error(`Unexpected extra arguments: ${rest.join(' ')}`); + } + + const db = getSeedDb(); + const trimmedOwnerEmail = ownerEmail.trim(); + const trimmedMemberEmail = memberEmail.trim(); + + const ownerUserId = await lookupUserId(trimmedOwnerEmail); + const memberUserId = await lookupUserId(trimmedMemberEmail); + + if (ownerUserId === memberUserId) { + throw new Error('owner-email and member-email must refer to different users'); + } + + const organizationId = randomUUID(); + + await db.insert(organizations).values({ + id: organizationId, + name: `[seed:w4c-org-pair] ${trimmedOwnerEmail}`, + }); + + await db.insert(organization_memberships).values([ + { + organization_id: organizationId, + kilo_user_id: ownerUserId, + role: 'owner', + }, + { + organization_id: organizationId, + kilo_user_id: memberUserId, + role: 'member', + }, + ]); + + return { + organizationId, + ownerUserId, + memberUserId, + }; +} From 32ec70d017b2b81144f500e5b839b3575d7fa1f5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 14:41:48 +0200 Subject: [PATCH 15/23] refactor(session-ingest): drop unused ticket consume response schema connectionTicketConsumeResponseSchema is only used via z.infer to derive a type, never safeParse'd. Replace it with a plain type. --- services/session-ingest/src/dos/connection-ticket-do.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/services/session-ingest/src/dos/connection-ticket-do.ts b/services/session-ingest/src/dos/connection-ticket-do.ts index de3c1440884..c9908028861 100644 --- a/services/session-ingest/src/dos/connection-ticket-do.ts +++ b/services/session-ingest/src/dos/connection-ticket-do.ts @@ -8,13 +8,9 @@ const ticketStateSchema = z.object({ expiresAt: z.number().int(), }); -export const connectionTicketConsumeResponseSchema = z.object({ - userId: z.string().min(1), -}); - type TicketState = z.infer; export type TicketMintRequest = TicketState; -export type ConnectionTicketConsumeResponse = z.infer; +export type ConnectionTicketConsumeResponse = { userId: string }; export class ConnectionTicketDO extends DurableObject { async mint(input: TicketMintRequest): Promise { From e0b441c73a967665da9d1a2f392ff0f3bbcf4d9b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 14:56:59 +0200 Subject: [PATCH 16/23] fix(cloud-agent): isolate pg-dependent kilo token validation validateKiloToken imports verifyKiloBearerAgainstCurrentPepper, which transitively imports pg. auth.ts is imported by CloudAgentSession, which the integration-test worker imports, so pg leaked into the Workers runtime and failed to load. Move validateKiloToken into its own module so auth.ts stays pg-free. --- services/cloud-agent-next/src/auth.test.ts | 103 +----------------- services/cloud-agent-next/src/auth.ts | 36 ------ .../src/middleware/auth.test.ts | 8 +- .../cloud-agent-next/src/middleware/auth.ts | 3 +- services/cloud-agent-next/src/server.ts | 2 +- .../src/validate-kilo-token.test.ts | 103 ++++++++++++++++++ .../src/validate-kilo-token.ts | 37 +++++++ 7 files changed, 148 insertions(+), 144 deletions(-) create mode 100644 services/cloud-agent-next/src/validate-kilo-token.test.ts create mode 100644 services/cloud-agent-next/src/validate-kilo-token.ts diff --git a/services/cloud-agent-next/src/auth.test.ts b/services/cloud-agent-next/src/auth.test.ts index 4c07fd7d1b6..9f8f9c319ce 100644 --- a/services/cloud-agent-next/src/auth.test.ts +++ b/services/cloud-agent-next/src/auth.test.ts @@ -1,36 +1,13 @@ -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { describe, expect, it } from 'vitest'; import jwt from 'jsonwebtoken'; -import { signKiloToken } from '@kilocode/worker-utils'; import { mintWrapperDispatchTicket, - validateKiloToken, validateStreamTicket, validateWrapperDispatchTicket, type WrapperDispatchTicketClaims, } from './auth.js'; const secret = 'test-secret'; -const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; - -const userRowByUserId = vi.hoisted( - () => new Map() -); - -vi.mock('@kilocode/db/client', () => ({ - getWorkerDb: () => ({ - select: () => ({ - from: () => ({ - where: () => ({ - limit: async () => { - const row = userRowByUserId.get('user-1'); - if (!row) return []; - return [{ api_token_pepper: row.pepper, blocked_reason: row.blockedReason }]; - }, - }), - }), - }), - }), -})); describe('validateStreamTicket', () => { it('returns a configuration error when NEXTAUTH_SECRET is missing', () => { @@ -241,81 +218,3 @@ describe('validateWrapperDispatchTicket', () => { }); }); }); - -describe('validateKiloToken', () => { - beforeEach(() => { - userRowByUserId.clear(); - }); - - it('returns a configuration error when NEXTAUTH_SECRET is missing', async () => { - await expect( - validateKiloToken('Bearer token', { secret: null, connectionString: 'postgres://test' }) - ).resolves.toEqual({ - success: false, - error: 'NEXTAUTH_SECRET is not configured on the worker', - }); - }); - - it('returns an error when the Authorization header is missing or malformed', async () => { - await expect( - validateKiloToken(null, { secret: TEST_JWT_SECRET, connectionString: 'postgres://test' }) - ).resolves.toEqual({ - success: false, - error: 'Missing or malformed Authorization header', - }); - }); - - it('accepts a token for an active account with the current pepper', async () => { - userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); - const { token } = await signKiloToken({ - userId: 'user-1', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - }); - - await expect( - validateKiloToken(`Bearer ${token}`, { - secret: TEST_JWT_SECRET, - connectionString: 'postgres://test', - }) - ).resolves.toMatchObject({ success: true, userId: 'user-1', token }); - }); - - it('rejects a token with a stale pepper', async () => { - userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); - const { token } = await signKiloToken({ - userId: 'user-1', - pepper: 'pepper-stale', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - }); - - await expect( - validateKiloToken(`Bearer ${token}`, { - secret: TEST_JWT_SECRET, - connectionString: 'postgres://test', - }) - ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); - }); - - it('rejects a token for a blocked account', async () => { - userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: 'manual block' }); - const { token } = await signKiloToken({ - userId: 'user-1', - pepper: 'pepper-current', - secret: TEST_JWT_SECRET, - expiresInSeconds: 3600, - env: 'production', - }); - - await expect( - validateKiloToken(`Bearer ${token}`, { - secret: TEST_JWT_SECRET, - connectionString: 'postgres://test', - }) - ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); - }); -}); diff --git a/services/cloud-agent-next/src/auth.ts b/services/cloud-agent-next/src/auth.ts index 261174f56f2..5c3ed50a34a 100644 --- a/services/cloud-agent-next/src/auth.ts +++ b/services/cloud-agent-next/src/auth.ts @@ -1,6 +1,5 @@ import jwt from 'jsonwebtoken'; import { verifyKiloToken, extractBearerToken } from '@kilocode/worker-utils'; -import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; type StreamTicketPayload = { type: 'stream_ticket'; @@ -63,41 +62,6 @@ export async function resolveSecret( } } -export async function validateKiloToken( - authHeader: string | null, - params: { - secret: string | null | undefined; - connectionString: string; - workerEnv?: string; - } -): Promise< - | { success: true; userId: string; token: string; botId?: string } - | { success: false; error: string } -> { - const { secret, connectionString, workerEnv } = params; - if (!secret) { - return { success: false, error: 'NEXTAUTH_SECRET is not configured on the worker' }; - } - - const token = extractBearerToken(authHeader); - if (!token) { - return { success: false, error: 'Missing or malformed Authorization header' }; - } - - const auth = await verifyKiloBearerAgainstCurrentPepper({ - token, - nextAuthSecret: secret, - connectionString, - ...(workerEnv ? { workerEnv } : {}), - }); - - if (!auth) { - return { success: false, error: 'Invalid or expired token' }; - } - - return { success: true, userId: auth.userId, token, botId: auth.botId }; -} - export function validateStreamTicket( ticket: string | null, secret: string | null | undefined, diff --git a/services/cloud-agent-next/src/middleware/auth.test.ts b/services/cloud-agent-next/src/middleware/auth.test.ts index d7ee71075f5..c5640c94a93 100644 --- a/services/cloud-agent-next/src/middleware/auth.test.ts +++ b/services/cloud-agent-next/src/middleware/auth.test.ts @@ -1,11 +1,11 @@ import { Hono } from 'hono'; import { beforeEach, describe, expect, it, vi } from 'vitest'; -import type * as AuthModule from '../auth.js'; +import type * as ValidateKiloTokenModule from '../validate-kilo-token.js'; import type { HonoContext } from '../hono-context.js'; import type { Env } from '../types.js'; -vi.mock('../auth.js', async () => { - const actual = await vi.importActual('../auth.js'); +vi.mock('../validate-kilo-token.js', async () => { + const actual = await vi.importActual('../validate-kilo-token.js'); return { ...actual, validateKiloToken: vi.fn(), @@ -36,7 +36,7 @@ vi.mock('../logger.js', () => { }); const { authMiddleware } = await import('./auth.js'); -const { validateKiloToken } = await import('../auth.js'); +const { validateKiloToken } = await import('../validate-kilo-token.js'); describe('authMiddleware', () => { beforeEach(() => { diff --git a/services/cloud-agent-next/src/middleware/auth.ts b/services/cloud-agent-next/src/middleware/auth.ts index 6277a6aa227..4a0b0ccd6da 100644 --- a/services/cloud-agent-next/src/middleware/auth.ts +++ b/services/cloud-agent-next/src/middleware/auth.ts @@ -1,7 +1,8 @@ import { createMiddleware } from 'hono/factory'; import type { Context, Next } from 'hono'; import type { HonoContext } from '../hono-context.js'; -import { resolveSecret, validateKiloToken } from '../auth.js'; +import { resolveSecret } from '../auth.js'; +import { validateKiloToken } from '../validate-kilo-token.js'; import { logger } from '../logger.js'; import { buildTrpcErrorResponse } from '../trpc-error.js'; import { extractProcedureName } from '../balance-validation.js'; diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index a6c994934f6..3581b4d9dd8 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -8,12 +8,12 @@ import { logger, withLogTags } from './logger.js'; import { resolveSecret, validateStreamTicket, - validateKiloToken, validateWrapperDispatchTicket, STREAM_TICKET_AUDIENCE, TERMINAL_TICKET_AUDIENCE, type WrapperAuthClaims, } from './auth.js'; +import { validateKiloToken } from './validate-kilo-token.js'; import { consumeStreamTicketNonce } from './persistence/StreamTicketNonceDO.js'; import { createErrorHandler, createNotFoundHandler } from '@kilocode/worker-utils'; import { createCallbackQueueConsumer } from './callbacks/index.js'; diff --git a/services/cloud-agent-next/src/validate-kilo-token.test.ts b/services/cloud-agent-next/src/validate-kilo-token.test.ts new file mode 100644 index 00000000000..943ad70af6a --- /dev/null +++ b/services/cloud-agent-next/src/validate-kilo-token.test.ts @@ -0,0 +1,103 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { signKiloToken } from '@kilocode/worker-utils'; +import { validateKiloToken } from './validate-kilo-token.js'; + +const TEST_JWT_SECRET = 'test-secret-that-is-long-enough-for-hs256'; + +const userRowByUserId = vi.hoisted( + () => new Map() +); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => { + const row = userRowByUserId.get('user-1'); + if (!row) return []; + return [{ api_token_pepper: row.pepper, blocked_reason: row.blockedReason }]; + }, + }), + }), + }), + }), +})); + +describe('validateKiloToken', () => { + beforeEach(() => { + userRowByUserId.clear(); + }); + + it('returns a configuration error when NEXTAUTH_SECRET is missing', async () => { + await expect( + validateKiloToken('Bearer token', { secret: null, connectionString: 'postgres://test' }) + ).resolves.toEqual({ + success: false, + error: 'NEXTAUTH_SECRET is not configured on the worker', + }); + }); + + it('returns an error when the Authorization header is missing or malformed', async () => { + await expect( + validateKiloToken(null, { secret: TEST_JWT_SECRET, connectionString: 'postgres://test' }) + ).resolves.toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); + }); + + it('accepts a token for an active account with the current pepper', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toMatchObject({ success: true, userId: 'user-1', token }); + }); + + it('rejects a token with a stale pepper', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: null }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-stale', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); + }); + + it('rejects a token for a blocked account', async () => { + userRowByUserId.set('user-1', { pepper: 'pepper-current', blockedReason: 'manual block' }); + const { token } = await signKiloToken({ + userId: 'user-1', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + env: 'production', + }); + + await expect( + validateKiloToken(`Bearer ${token}`, { + secret: TEST_JWT_SECRET, + connectionString: 'postgres://test', + }) + ).resolves.toEqual({ success: false, error: 'Invalid or expired token' }); + }); +}); diff --git a/services/cloud-agent-next/src/validate-kilo-token.ts b/services/cloud-agent-next/src/validate-kilo-token.ts new file mode 100644 index 00000000000..82a05ef7edc --- /dev/null +++ b/services/cloud-agent-next/src/validate-kilo-token.ts @@ -0,0 +1,37 @@ +import { extractBearerToken } from '@kilocode/worker-utils'; +import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; + +export async function validateKiloToken( + authHeader: string | null, + params: { + secret: string | null | undefined; + connectionString: string; + workerEnv?: string; + } +): Promise< + | { success: true; userId: string; token: string; botId?: string } + | { success: false; error: string } +> { + const { secret, connectionString, workerEnv } = params; + if (!secret) { + return { success: false, error: 'NEXTAUTH_SECRET is not configured on the worker' }; + } + + const token = extractBearerToken(authHeader); + if (!token) { + return { success: false, error: 'Missing or malformed Authorization header' }; + } + + const auth = await verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: secret, + connectionString, + ...(workerEnv ? { workerEnv } : {}), + }); + + if (!auth) { + return { success: false, error: 'Invalid or expired token' }; + } + + return { success: true, userId: auth.userId, token, botId: auth.botId }; +} From a747102ff07f669e6a52870682b751aba4e3757d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 16:09:09 +0200 Subject: [PATCH 17/23] fix(user): enforce code expiry and attempt budget on account deletion Reserve the sign-in code before consuming it so an expired (>10 min) code and the 5-attempt budget are enforced, matching the normal sign-in flow. Treat production provider_not_configured as a send failure instead of stamping a cooldown and returning a codeless success. --- apps/web/src/routers/user-router.test.ts | 73 ++++++++++++++++++++++++ apps/web/src/routers/user-router.ts | 31 +++++++++- 2 files changed, 102 insertions(+), 2 deletions(-) diff --git a/apps/web/src/routers/user-router.test.ts b/apps/web/src/routers/user-router.test.ts index dfd4b820a8f..c8058abea47 100644 --- a/apps/web/src/routers/user-router.test.ts +++ b/apps/web/src/routers/user-router.test.ts @@ -1134,6 +1134,79 @@ describe('user router - account deletion', () => { }); }); + it('rejects an expired code without deleting', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + await db + .update(magic_link_tokens) + .set({ + // Shift both timestamps into the past so the code is expired while + // still satisfying the `expires_at > created_at` check constraint. + created_at: new Date(Date.now() - 2 * 60_000).toISOString(), + expires_at: new Date(Date.now() - 60_000).toISOString(), + }) + .where(eq(magic_link_tokens.challenge_id, challengeId)); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + + expect(mockPerformGdprRemoval).not.toHaveBeenCalled(); + }); + + it('wrong code increments attempts and does not delete', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + const wrongCode = devCode === '000000' ? '111111' : '000000'; + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: wrongCode }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + + expect(mockPerformGdprRemoval).not.toHaveBeenCalled(); + + const [row] = await db + .select() + .from(magic_link_tokens) + .where(eq(magic_link_tokens.challenge_id, challengeId)); + expect(row?.attempts).toBe(1); + }); + + it('maps too_many_attempts to TOO_MANY_REQUESTS', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + await db + .update(magic_link_tokens) + .set({ attempts: 5 }) + .where(eq(magic_link_tokens.challenge_id, challengeId)); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).rejects.toMatchObject({ code: 'TOO_MANY_REQUESTS' }); + + expect(mockPerformGdprRemoval).not.toHaveBeenCalled(); + }); + + it('production provider_not_configured does not stamp and throws INTERNAL_SERVER_ERROR', async () => { + const caller = await createCallerForUser(deletionUser.id); + const restoreNodeEnv = jest.replaceProperty(process.env, 'NODE_ENV', 'production'); + try { + await expect(caller.user.requestAccountDeletionChallenge()).rejects.toMatchObject({ + code: 'INTERNAL_SERVER_ERROR', + }); + + const [user] = await db + .select() + .from(kilocode_users) + .where(eq(kilocode_users.id, deletionUser.id)); + expect(user?.account_deletion_requested_at).toBeNull(); + } finally { + restoreNodeEnv.restore(); + } + }); + it('precondition error maps to PRECONDITION_FAILED without deleting', async () => { mockAssertUserCanBeSoftDeleted.mockRejectedValue( new SoftDeletePreconditionError('active subscription') diff --git a/apps/web/src/routers/user-router.ts b/apps/web/src/routers/user-router.ts index 1c26a2a4ec5..9b160ac2415 100644 --- a/apps/web/src/routers/user-router.ts +++ b/apps/web/src/routers/user-router.ts @@ -10,6 +10,7 @@ import { consumeSignInCode, createSignInCode, deleteSignInCode, + reserveSignInCode, } from '@/lib/auth/magic-link-tokens'; import { performGdprRemoval } from '@/lib/user/gdpr-removal'; import { createAccountLinkingSession } from '@/lib/account-linking-session'; @@ -916,12 +917,21 @@ export const userRouter = createTRPCRouter({ if (sendResult.reason === 'provider_not_configured') { // Local/dev has no mail provider. Keep the code row so the caller can - // still complete the flow with the returned dev code. + // still complete the flow with the returned dev code. In production, + // treat this like any other send failure: discard the code and fail so + // the user is not left stuck with a cooldown and no code. + if (!includeDevCode) { + await deleteSignInCode(userEmail, code); + throw new TRPCError({ + code: 'INTERNAL_SERVER_ERROR', + message: 'Failed to send confirmation code. Please try again.', + }); + } await db .update(kilocode_users) .set({ account_deletion_requested_at: new Date().toISOString() }) .where(eq(kilocode_users.id, userId)); - return includeDevCode ? { challengeId, devCode: code } : { challengeId }; + return { challengeId, devCode: code }; } // neverbounce_rejected: discard the code and leave the stamp untouched so @@ -948,6 +958,23 @@ export const userRouter = createTRPCRouter({ throw error; } + const reserveResult = await reserveSignInCode(userEmail, input.code, input.challengeId); + if (reserveResult === 'invalid') { + throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid confirmation code' }); + } + if (reserveResult === 'too_many_attempts') { + throw new TRPCError({ + code: 'TOO_MANY_REQUESTS', + message: 'Too many attempts. Please request a new confirmation code.', + }); + } + if (reserveResult === 'in_progress') { + throw new TRPCError({ + code: 'TOO_MANY_REQUESTS', + message: 'Confirmation code is already being verified. Please wait.', + }); + } + const consumed = await consumeSignInCode(userEmail, input.code, input.challengeId); if (!consumed) { throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid confirmation code' }); From 00fb97a825b64bbcaa31aa5604438f162485ffa6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 16:09:13 +0200 Subject: [PATCH 18/23] fix(sdk): refresh one-use ingest ticket on every reconnect The ingest ticket is consumed once at the /api/user/web upgrade. Refresh it before a reconnect so a non-auth-failure close (1006) does not retry with an already-consumed ticket and loop forever. hasEverOpened scopes the refresh to reconnects so the initial connect does not double-mint. --- .../cloud-agent-sdk/src/base-connection.ts | 7 ++-- .../src/user-web-connection.test.ts | 33 ++++++++++++++++++- .../src/user-web-connection.ts | 12 +++++++ 3 files changed, 48 insertions(+), 4 deletions(-) diff --git a/packages/cloud-agent-sdk/src/base-connection.ts b/packages/cloud-agent-sdk/src/base-connection.ts index b19832be7cd..9254df6f7fc 100644 --- a/packages/cloud-agent-sdk/src/base-connection.ts +++ b/packages/cloud-agent-sdk/src/base-connection.ts @@ -116,7 +116,7 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec return; } authRefreshAttempted = true; - connectInternal(0, expectedGeneration); + connectInternal(0, expectedGeneration, true); } catch (err) { console.error('[Connection] Failed to refresh auth:', err); if (destroyed || intentionalDisconnect || expectedGeneration !== generation) return; @@ -141,7 +141,7 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec } if (destroyed || intentionalDisconnect || expectedGeneration !== generation) return; } - connectInternal(0, expectedGeneration); + connectInternal(0, expectedGeneration, true); } finally { if (expectedGeneration === generation) { preconnectAuthRefreshAttempted = false; @@ -167,12 +167,13 @@ export function createBaseConnection(config: BaseConnectionConfig): Connec }, delay); } - function connectInternal(attempt = 0, expectedGeneration = generation) { + function connectInternal(attempt = 0, expectedGeneration = generation, skipAuthRefresh = false) { if (destroyed || intentionalDisconnect || expectedGeneration !== generation) return; if ( config.refreshAuth && (config.shouldRefreshAuthBeforeConnect?.() ?? false) && + !skipAuthRefresh && !preconnectAuthRefreshAttempted ) { void refreshAndConnect(expectedGeneration); diff --git a/packages/cloud-agent-sdk/src/user-web-connection.test.ts b/packages/cloud-agent-sdk/src/user-web-connection.test.ts index 3b4fed409bb..2f07ce7f745 100644 --- a/packages/cloud-agent-sdk/src/user-web-connection.test.ts +++ b/packages/cloud-agent-sdk/src/user-web-connection.test.ts @@ -944,7 +944,7 @@ describe('createUserWebConnection', () => { client.destroy(); }); - it('resubscribes retained sessions and calls reconnect listeners', () => { + it('resubscribes retained sessions and calls reconnect listeners', async () => { jest.useFakeTimers(); try { const onReconnect = jest.fn(); @@ -962,6 +962,9 @@ describe('createUserWebConnection', () => { }); sockets[0].onclose?.({ code: 1006 } as CloseEvent); jest.advanceTimersByTime(60_000); + // The reconnect now refreshes auth before opening the new socket. + await Promise.resolve(); + await Promise.resolve(); open(sockets[1]); inbound( { type: 'system', event: 'sessions.list', data: { connectionId: 'c2', sessions: [] } }, @@ -978,6 +981,34 @@ describe('createUserWebConnection', () => { } }); + it('mints a fresh ticket on a non-auth-failure (1006) reconnect', async () => { + jest.useFakeTimers(); + try { + const getAuthToken = jest.fn().mockReturnValueOnce('old-token').mockReturnValue('new-token'); + const client = createUserWebConnection({ websocketUrl: WS_URL, getAuthToken }); + client.subscribeToCliSession('ses-1'); + open(); + inbound({ type: 'system', event: 'sessions.list', data: { sessions: [] } }); + expect(getAuthToken).toHaveBeenCalledTimes(1); + expect(webSocketConstructor).toHaveBeenLastCalledWith( + `${WS_URL}?ticket=old-token&connectionId=uuid-1` + ); + + sockets[0].onclose?.({ code: 1006 } as CloseEvent); + jest.advanceTimersByTime(60_000); + await Promise.resolve(); + await Promise.resolve(); + + expect(getAuthToken).toHaveBeenCalledTimes(2); + expect(webSocketConstructor).toHaveBeenLastCalledWith( + `${WS_URL}?ticket=new-token&connectionId=uuid-1` + ); + client.destroy(); + } finally { + jest.useRealTimers(); + } + }); + it('reopens cleanly after disconnect', () => { const client = createUserWebConnection({ websocketUrl: WS_URL, getAuthToken: () => 'token' }); client.connect(); diff --git a/packages/cloud-agent-sdk/src/user-web-connection.ts b/packages/cloud-agent-sdk/src/user-web-connection.ts index 98c28e82607..c0a429afd96 100644 --- a/packages/cloud-agent-sdk/src/user-web-connection.ts +++ b/packages/cloud-agent-sdk/src/user-web-connection.ts @@ -152,6 +152,10 @@ function createUserWebConnection( let retainCount = 0; let commandRetainCount = 0; let legacyRetained = false; + // True once the current base connection's socket has opened (i.e. the one-use + // ingest ticket was consumed at the upgrade). Drives the pre-connect auth + // refresh so only reconnects mint a fresh ticket, not the initial connect. + let hasEverOpened = false; let pingInterval: ReturnType | null = null; let pongTimeout: ReturnType | null = null; let initialAuthRetryTimeout: ReturnType | null = null; @@ -424,6 +428,7 @@ function createUserWebConnection( function ensureBaseConnection(): void { if (baseConnection) return; removePreSocketLifecycleListeners(); + hasEverOpened = false; baseConnection = createBaseConnection({ lifecycleHooks: createLifecycleHooks(), buildUrl, @@ -440,6 +445,7 @@ function createUserWebConnection( }, onEvent: handleInboundMessage, onOpen: ws => { + hasEverOpened = true; if (!hasLifetime()) return; currentWs = ws; resolveOpenWaiters(ws); @@ -468,6 +474,12 @@ function createUserWebConnection( }, onError: config.onError, isAuthFailure: event => event.code === 4001 || event.code === 1008, + // The ingest ticket is one-use (consumed at the `/api/user/web` upgrade). + // Refresh it before a reconnect so a non-auth-failure close (e.g. 1006) + // does not retry with an already-consumed ticket and loop forever. The + // initial connect already mints a fresh ticket in `startConnection`, so + // `hasEverOpened` (set on first `onOpen`) scopes the refresh to reconnects. + shouldRefreshAuthBeforeConnect: () => hasEverOpened, refreshAuth: async () => { token = await config.getAuthToken(); }, From 412d87868bf0b788cd4d48f6958384fc1be7e754 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 16:09:17 +0200 Subject: [PATCH 19/23] fix(cloud-agent): expire stream-ticket nonce storage after ticket expiry Set an alarm on nonce consume and delete the used flag in alarm() so nonce DOs do not accumulate storage forever. The flag only needs to persist until the ticket JWT expires, after which validateStreamTicket rejects the replay. --- .../src/persistence/StreamTicketNonceDO.ts | 21 +++- .../src/server-stream-ticket.test.ts | 97 +++++++++++++++++++ services/cloud-agent-next/src/server.ts | 12 ++- 3 files changed, 124 insertions(+), 6 deletions(-) diff --git a/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts b/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts index b94064f8dda..42ddeef76d6 100644 --- a/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts +++ b/services/cloud-agent-next/src/persistence/StreamTicketNonceDO.ts @@ -9,8 +9,8 @@ import type { Env } from '../types.js'; * marks the nonce used and returns true, or returns false on any replay. */ export class StreamTicketNonceDO extends DurableObject { - async consume(): Promise { - return this.ctx.storage.transaction(async txn => { + async consume(expiresAtMs: number): Promise { + const consumed = await this.ctx.storage.transaction(async txn => { const used = await txn.get('used'); if (used !== undefined) { return false; @@ -18,6 +18,15 @@ export class StreamTicketNonceDO extends DurableObject { await txn.put('used', { usedAt: Date.now() }); return true; }); + + if (consumed) { + await this.ctx.storage.setAlarm(expiresAtMs); + } + return consumed; + } + + async alarm(): Promise { + await this.ctx.storage.delete('used'); } } @@ -26,8 +35,12 @@ export class StreamTicketNonceDO extends DurableObject { * from `validateStreamTicket` so signature/type/expiry/audience validation * never consumes the nonce. */ -export async function consumeStreamTicketNonce(env: Env, nonce: string): Promise { +export async function consumeStreamTicketNonce( + env: Env, + nonce: string, + expiresAtMs: number +): Promise { const doId = env.STREAM_TICKET_NONCE_DO.idFromName(nonce); const stub = env.STREAM_TICKET_NONCE_DO.get(doId); - return stub.consume(); + return stub.consume(expiresAtMs); } diff --git a/services/cloud-agent-next/src/server-stream-ticket.test.ts b/services/cloud-agent-next/src/server-stream-ticket.test.ts index b39a9ca6e34..19d83c7526f 100644 --- a/services/cloud-agent-next/src/server-stream-ticket.test.ts +++ b/services/cloud-agent-next/src/server-stream-ticket.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import jwt from 'jsonwebtoken'; import type { Env } from './types.js'; +import { StreamTicketNonceDO } from './persistence/StreamTicketNonceDO.js'; const { getRunningTerminalClientMock, @@ -163,6 +164,51 @@ function fetchWorker(request: Request, env: MockEnv): Promise | Respon return worker.fetch(request, env as unknown as Env, {} as ExecutionContext); } +/** + * Instantiates StreamTicketNonceDO against an in-memory storage mock so the + * alarm lifecycle can be asserted directly without a Workers runtime. + */ +function createNonceDO(): { + instance: StreamTicketNonceDO; + map: Map; + alarms: number[]; +} { + const map = new Map(); + const alarms: number[] = []; + const storage = { + transaction: async ( + fn: (txn: { + get: (key: string) => Promise; + put: (key: string, value: unknown) => Promise; + delete: (key: string) => Promise; + }) => Promise + ) => + fn({ + get: async (key: string) => map.get(key), + put: async (key: string, value: unknown) => { + map.set(key, value); + }, + delete: async (key: string) => { + map.delete(key); + }, + }), + get: async (key: string) => map.get(key), + put: async (key: string, value: unknown) => { + map.set(key, value); + }, + delete: async (key: string) => { + map.delete(key); + }, + setAlarm: async (at: number) => { + alarms.push(at); + }, + deleteAlarm: async () => {}, + }; + const instance = new StreamTicketNonceDO({} as never, {} as never); + (instance as unknown as { ctx: { storage: typeof storage } }).ctx = { storage }; + return { instance, map, alarms }; +} + function signStreamTicket(overrides: Record = {}): string { return jwt.sign( { @@ -329,3 +375,54 @@ describe('server /terminal ticket nonce consume', () => { expect(connectTerminal).toHaveBeenCalledTimes(1); }); }); + +describe('server stream ticket nonce expiry pass-through', () => { + it('passes the ticket expiry to the nonce DO on consume', async () => { + const env = createEnv(); + const consume = vi.fn().mockResolvedValue(true); + env.STREAM_TICKET_NONCE_DO.idFromName.mockReturnValue('nonce-do-id'); + env.STREAM_TICKET_NONCE_DO.get.mockReturnValue({ consume }); + const doFetch = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ fetch: doFetch }); + const ticket = signStreamTicket(); + + const response = await fetchWorker(streamRequest(ticket), env); + + expect(response.status).toBe(200); + const decoded = jwt.decode(ticket) as { exp: number }; + expect(consume).toHaveBeenCalledWith(decoded.exp * 1000); + }); +}); + +describe('StreamTicketNonceDO alarm lifecycle', () => { + it('sets an alarm at the ticket expiry on first consume', async () => { + const { instance, alarms } = createNonceDO(); + + const consumed = await instance.consume(1_700_000_000_000); + + expect(consumed).toBe(true); + expect(alarms).toEqual([1_700_000_000_000]); + }); + + it('does not set a second alarm on replay', async () => { + const { instance, alarms } = createNonceDO(); + + await instance.consume(1_700_000_000_000); + const replay = await instance.consume(1_700_000_000_000); + + expect(replay).toBe(false); + expect(alarms).toHaveLength(1); + }); + + it('alarm() deletes the used flag', async () => { + const { instance, map } = createNonceDO(); + + await instance.consume(1_700_000_000_000); + expect(map.has('used')).toBe(true); + + await instance.alarm(); + + expect(map.has('used')).toBe(false); + }); +}); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 3581b4d9dd8..98e354384d0 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -132,7 +132,11 @@ async function handleTerminalWebSocket(request: Request, env: Env): Promise) => { logger.withFields({ cloudAgentSessionId, userId }).warn('/stream: Missing ticket nonce'); return c.text('Missing ticket nonce', 401); } - const consumed = await consumeStreamTicketNonce(c.env, nonce); + const consumed = await consumeStreamTicketNonce( + c.env, + nonce, + (ticketResult.payload as unknown as { exp: number }).exp * 1000 + ); if (!consumed) { logger.withFields({ cloudAgentSessionId, userId }).warn('/stream: Ticket nonce already used'); return c.text('Ticket nonce already used', 401); From 54c4c1fdd4a4e6b1ae54ee962194a0ac7bfb7e6a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 16:25:53 +0200 Subject: [PATCH 20/23] chore: retrigger review From d1c62c4e02a115699616e2c677663eff5e38af2d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 17:31:27 +0200 Subject: [PATCH 21/23] fix(auth): consume the deletion code only after GDPR removal succeeds A failed performGdprRemoval burned the confirmation code, and the 1 h request cooldown then blocked a new one. Release the reservation on failure and consume the code after removal succeeds, so the client retry with the same challenge works. --- apps/web/src/routers/user-router.test.ts | 42 ++++++++++++++++++++++++ apps/web/src/routers/user-router.ts | 12 ++++--- 2 files changed, 49 insertions(+), 5 deletions(-) diff --git a/apps/web/src/routers/user-router.test.ts b/apps/web/src/routers/user-router.test.ts index c8058abea47..8eb7034709f 100644 --- a/apps/web/src/routers/user-router.test.ts +++ b/apps/web/src/routers/user-router.test.ts @@ -1134,6 +1134,48 @@ describe('user router - account deletion', () => { }); }); + it('keeps the code usable when performGdprRemoval fails, so a retry succeeds', async () => { + mockPerformGdprRemoval.mockRejectedValueOnce(new Error('kiloclaw destroy failed')); + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).rejects.toThrow(); + + const [afterFailure] = await db + .select() + .from(magic_link_tokens) + .where(eq(magic_link_tokens.challenge_id, challengeId)); + expect(afterFailure?.consumed_at).toBeNull(); + expect(afterFailure?.reserved_until).toBeNull(); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).resolves.toEqual({ status: 'deleted' }); + + expect(mockPerformGdprRemoval).toHaveBeenCalledTimes(2); + + const [afterRetry] = await db + .select() + .from(magic_link_tokens) + .where(eq(magic_link_tokens.challenge_id, challengeId)); + expect(afterRetry?.consumed_at).not.toBeNull(); + }); + + it('rejects a replay of the code after a successful deletion', async () => { + const caller = await createCallerForUser(deletionUser.id); + const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); + + await caller.user.requestAccountDeletion({ challengeId, code: devCode as string }); + + await expect( + caller.user.requestAccountDeletion({ challengeId, code: devCode as string }) + ).rejects.toMatchObject({ code: 'UNAUTHORIZED' }); + + expect(mockPerformGdprRemoval).toHaveBeenCalledTimes(1); + }); + it('rejects an expired code without deleting', async () => { const caller = await createCallerForUser(deletionUser.id); const { challengeId, devCode } = await caller.user.requestAccountDeletionChallenge(); diff --git a/apps/web/src/routers/user-router.ts b/apps/web/src/routers/user-router.ts index 9b160ac2415..52f8c55a905 100644 --- a/apps/web/src/routers/user-router.ts +++ b/apps/web/src/routers/user-router.ts @@ -10,6 +10,7 @@ import { consumeSignInCode, createSignInCode, deleteSignInCode, + releaseSignInCode, reserveSignInCode, } from '@/lib/auth/magic-link-tokens'; import { performGdprRemoval } from '@/lib/user/gdpr-removal'; @@ -975,11 +976,9 @@ export const userRouter = createTRPCRouter({ }); } - const consumed = await consumeSignInCode(userEmail, input.code, input.challengeId); - if (!consumed) { - throw new TRPCError({ code: 'UNAUTHORIZED', message: 'Invalid confirmation code' }); - } - + // The code is consumed only after removal succeeds. Burning it first + // stranded the caller: a failed removal left no usable code and the + // 1 h request cooldown blocked a new one. try { await performGdprRemoval(userId, { destroyReason: 'admin_request', @@ -990,6 +989,7 @@ export const userRouter = createTRPCRouter({ }, }); } catch (error) { + await releaseSignInCode(userEmail, input.code, input.challengeId); // A precondition race inside performGdprRemoval (re-asserted there) maps // to the same PRECONDITION_FAILED surface. if (error instanceof SoftDeletePreconditionError) { @@ -998,6 +998,8 @@ export const userRouter = createTRPCRouter({ throw error; } + await consumeSignInCode(userEmail, input.code, input.challengeId); + return { status: 'deleted' as const }; }), From 151820c6c3c164e56caba2dbcaf58673b15a7f77 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 17:31:27 +0200 Subject: [PATCH 22/23] fix(session-ingest): report auth dependency outages as 503, not 401 verifyKiloBearerAgainstCurrentPepper swallowed secret-store and database failures into null, so Session Ingest answered an outage with 'Invalid or expired token'. Let dependency errors propagate and map them to a retryable 503; null now means an invalid or revoked credential only. --- .../worker-utils/src/kilo-token-auth.test.ts | 34 ++++++++++ packages/worker-utils/src/kilo-token-auth.ts | 68 +++++++++++-------- .../src/middleware/kilo-jwt-auth.test.ts | 31 +++++++++ .../src/middleware/kilo-jwt-auth.ts | 22 ++++-- 4 files changed, 121 insertions(+), 34 deletions(-) diff --git a/packages/worker-utils/src/kilo-token-auth.test.ts b/packages/worker-utils/src/kilo-token-auth.test.ts index 989ad111272..539450e6d68 100644 --- a/packages/worker-utils/src/kilo-token-auth.test.ts +++ b/packages/worker-utils/src/kilo-token-auth.test.ts @@ -59,6 +59,40 @@ describe('verifyKiloBearerAgainstCurrentPepper', () => { await expect(verifyToken(token)).resolves.toEqual({ userId: 'user-xyz-789' }); }); + it('propagates a pepper-lookup failure instead of reporting an invalid token', async () => { + const { token } = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { get: async () => TEST_JWT_SECRET }, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper: async () => { + throw new Error('connection refused'); + }, + }) + ).rejects.toThrow('connection refused'); + }); + + it('propagates a secret-store failure instead of reporting an invalid token', async () => { + const { token } = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); + + await expect( + verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: { + get: async () => { + throw new Error('secrets store unavailable'); + }, + }, + workerEnv: 'production', + connectionString: 'postgres://test', + getUserPepper, + }) + ).rejects.toThrow('secrets store unavailable'); + }); + it('rejects tokens for missing users', async () => { userResultByUserId.clear(); const { token } = await signToken({ pepper: 'pepper-current', tokenSource: 'kilo-chat' }); diff --git a/packages/worker-utils/src/kilo-token-auth.ts b/packages/worker-utils/src/kilo-token-auth.ts index 1c52b51c676..d2f91bd1d44 100644 --- a/packages/worker-utils/src/kilo-token-auth.ts +++ b/packages/worker-utils/src/kilo-token-auth.ts @@ -39,6 +39,14 @@ export async function findKiloUserPepper( return { pepper: row.api_token_pepper ?? null, blockedReason: row.blocked_reason }; } +/** + * Verify a Kilo bearer against the account's current pepper and active state. + * + * Returns null only for a credential the caller must not retry: malformed, + * expired, wrong env, unknown user, blocked user, or a stale pepper. + * A dependency failure (secret store, database) throws, so a caller can map an + * outage to a retryable 503 instead of reporting it as an invalid token. + */ export async function verifyKiloBearerAgainstCurrentPepper(params: { token: string | null; nextAuthSecret: KiloSecretBinding | string; @@ -50,40 +58,42 @@ export async function verifyKiloBearerAgainstCurrentPepper(params: { const getUserPepper = params.getUserPepper ?? findKiloUserPepper; - try { - const secret = - typeof params.nextAuthSecret === 'string' - ? params.nextAuthSecret - : await getCachedSecret(params.nextAuthSecret, 'NEXTAUTH_SECRET'); - const payload = await verifyKiloToken(params.token, secret); + const secret = + typeof params.nextAuthSecret === 'string' + ? params.nextAuthSecret + : await getCachedSecret(params.nextAuthSecret, 'NEXTAUTH_SECRET'); - // Env check is skipped only when the caller does not pass a workerEnv. - // When workerEnv is set, a token without env (or with a mismatched env) fails. - if (params.workerEnv && payload.env !== params.workerEnv) { - return null; - } + let payload: Awaited>; + try { + payload = await verifyKiloToken(params.token, secret); + } catch { + return null; + } - const result = await getUserPepper(params.connectionString, payload.kiloUserId); - if (!result) { - return null; - } + // Env check is skipped only when the caller does not pass a workerEnv. + // When workerEnv is set, a token without env (or with a mismatched env) fails. + if (params.workerEnv && payload.env !== params.workerEnv) { + return null; + } - if (result.blockedReason !== null) { - return null; - } + const result = await getUserPepper(params.connectionString, payload.kiloUserId); + if (!result) { + return null; + } - // Pepper equality is skipped only when the claim is absent (internal - // service tokens). A present claim — string or null — is always compared. - if (payload.apiTokenPepper !== undefined && result.pepper !== payload.apiTokenPepper) { - return null; - } + if (result.blockedReason !== null) { + return null; + } - const authResult: KiloBearerAuthResult = { userId: payload.kiloUserId }; - if (payload.botId) { - authResult.botId = payload.botId; - } - return authResult; - } catch { + // Pepper equality is skipped only when the claim is absent (internal + // service tokens). A present claim — string or null — is always compared. + if (payload.apiTokenPepper !== undefined && result.pepper !== payload.apiTokenPepper) { return null; } + + const authResult: KiloBearerAuthResult = { userId: payload.kiloUserId }; + if (payload.botId) { + authResult.botId = payload.botId; + } + return authResult; } diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts index 4ff8257cdd3..9fd985b617d 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts @@ -11,12 +11,15 @@ const userRowByUserId = vi.hoisted( () => new Map() ); +const dbState = vi.hoisted(() => ({ fails: false })); + vi.mock('@kilocode/db/client', () => ({ getWorkerDb: () => ({ select: () => ({ from: () => ({ where: () => ({ limit: async () => { + if (dbState.fails) throw new Error('connection refused'); const row = userRowByUserId.get('usr_123'); if (!row) return []; return [{ api_token_pepper: row.pepper, blocked_reason: row.blockedReason }]; @@ -116,6 +119,34 @@ describe('kiloJwtAuthMiddleware', () => { beforeEach(() => { clearSecretCacheForTest(); userRowByUserId.clear(); + dbState.fails = false; + }); + + it('returns a retryable 503 when the secret store fails', async () => { + const token = await signUserToken('pepper-current'); + const env = makeEnv(TEST_JWT_SECRET); + env.NEXTAUTH_SECRET_PROD.get = async () => { + throw new Error('secrets store unavailable'); + }; + + const res = await makeApp().fetch( + new Request('http://local/api/me', { headers: { Authorization: `Bearer ${token}` } }), + env + ); + + expect(res.status).toBe(503); + }); + + it('returns a retryable 503 when the database fails', async () => { + dbState.fails = true; + const token = await signUserToken('pepper-current'); + + const res = await makeApp().fetch( + new Request('http://local/api/me', { headers: { Authorization: `Bearer ${token}` } }), + makeEnv(TEST_JWT_SECRET) + ); + + expect(res.status).toBe(503); }); it('rejects missing Authorization header', async () => { diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.ts index d4841ce542e..a19ebe16077 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.ts @@ -39,11 +39,23 @@ export const kiloJwtAuthMiddleware = createMiddleware<{ return c.json({ success: false, error: 'Missing or malformed Authorization header' }, 401); } - const auth = await verifyKiloBearerAgainstCurrentPepper({ - token, - nextAuthSecret: c.env.NEXTAUTH_SECRET_PROD, - connectionString: c.env.HYPERDRIVE.connectionString, - }); + let auth; + try { + auth = await verifyKiloBearerAgainstCurrentPepper({ + token, + nextAuthSecret: c.env.NEXTAUTH_SECRET_PROD, + connectionString: c.env.HYPERDRIVE.connectionString, + }); + } catch (error) { + // Secret-store or database failure. Stay retryable: a 401 here would tell + // the client its credential is bad and stop it from retrying. + console.error('Auth infrastructure failure', { + operation: 'kilo-bearer-verify', + errorClass: error instanceof Error ? error.name : typeof error, + errorMessage: error instanceof Error ? error.message : String(error), + }); + return c.json({ success: false, error: 'Service temporarily unavailable' }, 503); + } if (!auth) { return c.json({ success: false, error: 'Invalid or expired token' }, 401); From 581fc812b25b626d0fc3a43205dfa665f89eb720 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Igor=20=C5=A0=C4=87eki=C4=87?= Date: Tue, 18 Aug 2026 20:32:26 +0200 Subject: [PATCH 23/23] test(mobile): cast api.send fixture to branded SDK types --- .../src/components/agents/mobile-session-manager.test.ts | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apps/mobile/src/components/agents/mobile-session-manager.test.ts b/apps/mobile/src/components/agents/mobile-session-manager.test.ts index 4b60ffde8b4..742095ce497 100644 --- a/apps/mobile/src/components/agents/mobile-session-manager.test.ts +++ b/apps/mobile/src/components/agents/mobile-session-manager.test.ts @@ -459,7 +459,8 @@ describe('createMobileAgentSessionManager api.send', () => { sendMessageMutate.mockReset(); sendMessageMutate.mockResolvedValue(undefined); const config = setup(); - await config.api.send({ sessionId: 'c-1', payload: 'hi', messageId: 'm-1' }); + const input = { sessionId: 'c-1', payload: 'hi', messageId: 'm-1' }; + await config.api.send(input as never); expect(sendMessageMutate).toHaveBeenCalledTimes(1); const payload = sendMessageMutate.mock.calls[0]?.[0] as Record;