diff --git a/apps/web/src/routers/active-sessions-router.test.ts b/apps/web/src/routers/active-sessions-router.test.ts index 382ab7a79c..f5a9cd910d 100644 --- a/apps/web/src/routers/active-sessions-router.test.ts +++ b/apps/web/src/routers/active-sessions-router.test.ts @@ -1,5 +1,6 @@ import { describe, expect, it, jest, beforeAll, afterEach } from '@jest/globals'; import { TRPCError } from '@trpc/server'; +import jwt from 'jsonwebtoken'; import { insertTestUser } from '@/tests/helpers/user.helper'; import { db } from '@/lib/drizzle'; import { organizations, organization_memberships } from '@kilocode/db/schema'; @@ -30,6 +31,8 @@ import type { createCallerForUser as CreateCallerForUser } from '@/routers/test- // `require()`. process.env.SESSION_INGEST_WORKER_URL = 'https://test-ingest.example.com'; let createCallerForUser: typeof CreateCallerForUser; +let JWT_TOKEN_VERSION: number; +let TOKEN_EXPIRY: { oneHour: number }; let regularUser: User; let testOrganization: Organization; @@ -37,6 +40,7 @@ let testOrganization: Organization; describe('active-sessions-router', () => { beforeAll(async () => { ({ createCallerForUser } = await import('@/routers/test-utils')); + ({ JWT_TOKEN_VERSION, TOKEN_EXPIRY } = await import('@/lib/tokens')); regularUser = await insertTestUser({ google_user_email: 'active-sessions-user@example.com', @@ -72,6 +76,32 @@ describe('active-sessions-router', () => { jest.restoreAllMocks(); }); + describe('getToken', () => { + it('returns a one-hour pepper-bearing API token for the caller', async () => { + const knownPepper = 'active-sessions-get-token-pepper'; + const tokenUser = await insertTestUser({ + google_user_email: `active-sessions-token-${crypto.randomUUID()}@example.com`, + google_user_name: 'Active Sessions Token User', + api_token_pepper: knownPepper, + }); + + const caller = await createCallerForUser(tokenUser.id); + const result = await caller.activeSessions.getToken(); + const payload = jwt.decode(result.token) as jwt.JwtPayload & { + kiloUserId: string; + apiTokenPepper: string; + version: number; + }; + + expect(payload.kiloUserId).toBe(tokenUser.id); + expect(payload.apiTokenPepper).toBe(knownPepper); + expect(payload.version).toBe(JWT_TOKEN_VERSION); + expect(payload.exp).toBeDefined(); + expect(payload.iat).toBeDefined(); + expect((payload.exp ?? 0) - (payload.iat ?? 0)).toBe(TOKEN_EXPIRY.oneHour); + }); + }); + 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 376e7ae5ac..60c5c85592 100644 --- a/apps/web/src/routers/active-sessions-router.ts +++ b/apps/web/src/routers/active-sessions-router.ts @@ -3,7 +3,7 @@ import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; import { TRPCError } from '@trpc/server'; import { z } from 'zod'; import { SESSION_INGEST_WORKER_URL } from '@/lib/config.server'; -import { generateInternalServiceToken } from '@/lib/tokens'; +import { generateApiToken, generateInternalServiceToken, TOKEN_EXPIRY } from '@/lib/tokens'; import { db } from '@/lib/drizzle'; import { cli_sessions_v2, cloud_agent_session_runs } from '@kilocode/db/schema'; import { and, desc, eq, gt, inArray, isNotNull, isNull, or, sql } from 'drizzle-orm'; @@ -221,7 +221,9 @@ function throwOrgContextFailure(error: unknown): never { export const activeSessionsRouter = createTRPCRouter({ getToken: baseProcedure.query(async ({ ctx }) => { - const token = generateInternalServiceToken(ctx.user.id); + const token = generateApiToken(ctx.user, undefined, { + expiresIn: TOKEN_EXPIRY.oneHour, + }); return { token }; }), diff --git a/packages/worker-utils/src/kilo-token-auth.ts b/packages/worker-utils/src/kilo-token-auth.ts index 6f4f8b9cd4..0b9978ec88 100644 --- a/packages/worker-utils/src/kilo-token-auth.ts +++ b/packages/worker-utils/src/kilo-token-auth.ts @@ -61,6 +61,9 @@ export async function verifyKiloBearerAgainstCurrentPepper(params: { return null; } + // This shared helper intentionally normalizes legacy tokens without a + // pepper to null. Session-ingest has a separate internal-token class and + // must not reuse this normalization for its admission boundary. const tokenPepper = payload.apiTokenPepper ?? null; if (result.pepper !== tokenPepper) { return null; diff --git a/services/session-ingest/src/app.ts b/services/session-ingest/src/app.ts index 46005acaad..88ea315c64 100644 --- a/services/session-ingest/src/app.ts +++ b/services/session-ingest/src/app.ts @@ -4,7 +4,7 @@ import { createMiddleware } from 'hono/factory'; import type { Env } from './env'; import { z } from 'zod'; -import { kiloJwtAuthMiddleware } from './middleware/kilo-jwt-auth'; +import { kiloJwtAuthMiddleware, USER_AUTH_CACHE_KEY_PREFIX } from './middleware/kilo-jwt-auth'; import { api } from './routes/api'; import { cloudAgentSessionScopeApi } from './routes/cloud-agent-session-scope'; import { getSessionIngestDO } from './dos/SessionIngestDO'; @@ -18,6 +18,9 @@ const invalidateSessionAccessSchema = z.object({ kiloUserId: z.string().min(1), organizationId: z.uuid(), }); +const invalidateUserAuthSchema = z.object({ + kiloUserId: z.string().min(1), +}); async function hasValidInternalSecret(c: { req: { header(name: string): string | undefined }; @@ -44,7 +47,18 @@ const requireValidInternalSecret = createMiddleware<{ Bindings: Env; Variables: { user_id: string }; }>(async (c, next) => { - if (!(await hasValidInternalSecret(c))) { + let isValid: boolean; + try { + isValid = await hasValidInternalSecret(c); + } catch (error) { + console.error('Auth infrastructure failure', { + operation: 'internal-api-secret-get', + 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 (!isValid) { return c.json({ success: false, error: 'Unauthorized' }, 401); } return next(); @@ -112,11 +126,7 @@ app.get('/session/:shareToken/metadata', async c => { ); }); -app.post('/internal/session-access/invalidate', async c => { - if (!(await hasValidInternalSecret(c))) { - return c.json({ success: false, error: 'Unauthorized' }, 401); - } - +app.post('/internal/session-access/invalidate', requireValidInternalSecret, async c => { const parsed = invalidateSessionAccessSchema.safeParse(await c.req.json().catch(() => null)); if (!parsed.success) { return c.json({ success: false, error: 'Invalid request', issues: parsed.error.issues }, 400); @@ -131,12 +141,18 @@ app.post('/internal/session-access/invalidate', async c => { return c.body(null, 204); }); -// Internal route for service-binding HTTP fetch (secret-protected) -app.get('/internal/session/:sessionId/export', async c => { - if (!(await hasValidInternalSecret(c))) { - return c.json({ success: false, error: 'Unauthorized' }, 401); +app.post('/internal/user-auth/invalidate', requireValidInternalSecret, async c => { + const parsed = invalidateUserAuthSchema.safeParse(await c.req.json().catch(() => null)); + if (!parsed.success) { + return c.json({ success: false, error: 'Invalid request', issues: parsed.error.issues }, 400); } + await c.env.USER_EXISTS_CACHE.delete(`${USER_AUTH_CACHE_KEY_PREFIX}${parsed.data.kiloUserId}`); + return c.body(null, 204); +}); + +// Internal route for service-binding HTTP fetch (secret-protected) +app.get('/internal/session/:sessionId/export', requireValidInternalSecret, async c => { const kiloUserId = c.req.header('X-Kilo-User-Id'); if (!kiloUserId) return c.json({ success: false, error: 'Missing X-Kilo-User-Id' }, 400); diff --git a/services/session-ingest/src/index.test.ts b/services/session-ingest/src/index.test.ts index 5948eaa58f..eafebb943a 100644 --- a/services/session-ingest/src/index.test.ts +++ b/services/session-ingest/src/index.test.ts @@ -45,6 +45,7 @@ type TestBindings = { INTERNAL_API_SECRET_PROD: { get(): Promise }; SESSION_SHARE_JWT_SECRET_PROD: { get(): Promise }; SESSION_SHARE_TOKEN_MIN_IAT: string; + USER_EXISTS_CACHE?: { delete: ReturnType Promise>> }; }; function makeDbFakes() { @@ -167,6 +168,142 @@ describe('session access invalidation route', () => { }); }); +describe('user auth invalidation route', () => { + beforeEach(() => { + vi.resetAllMocks(); + }); + + it('rejects invalidation without the internal secret', async () => { + const cache = { delete: vi.fn(async () => undefined) }; + + const res = await app.request( + '/internal/user-auth/invalidate', + { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ kiloUserId: 'usr_blocked' }), + }, + { ...defaultEnv, USER_EXISTS_CACHE: cache } + ); + + expect(res.status).toBe(401); + expect(cache.delete).not.toHaveBeenCalled(); + }); + + it('rejects invalidation with an incorrect internal secret', async () => { + const cache = { delete: vi.fn(async () => undefined) }; + + const res = await app.request( + '/internal/user-auth/invalidate', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'X-Internal-Secret': 'wrong-secret', + }, + body: JSON.stringify({ kiloUserId: 'usr_blocked' }), + }, + { ...defaultEnv, USER_EXISTS_CACHE: cache } + ); + + expect(res.status).toBe(401); + expect(cache.delete).not.toHaveBeenCalled(); + }); + + it('rejects invalidation with a malformed body', async () => { + const cache = { delete: vi.fn(async () => undefined) }; + + const res = await app.request( + '/internal/user-auth/invalidate', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'X-Internal-Secret': 'internal-secret', + }, + body: JSON.stringify({ kiloUserId: 123 }), + }, + { ...defaultEnv, USER_EXISTS_CACHE: cache } + ); + + expect(res.status).toBe(400); + expect(cache.delete).not.toHaveBeenCalled(); + }); + + it('deletes the versioned user-auth cache key', async () => { + const cache = { delete: vi.fn(async () => undefined) }; + + const res = await app.request( + '/internal/user-auth/invalidate', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'X-Internal-Secret': 'internal-secret', + }, + body: JSON.stringify({ kiloUserId: 'usr_blocked' }), + }, + { ...defaultEnv, USER_EXISTS_CACHE: cache } + ); + + expect(res.status).toBe(204); + expect(cache.delete).toHaveBeenCalledWith('user-auth:v1:usr_blocked'); + }); + + it('protects the session export route with the shared internal-secret middleware', async () => { + const res = await app.request( + '/internal/session/ses_12345678901234567890123456/export', + { method: 'GET' }, + defaultEnv + ); + + expect(res.status).toBe(401); + }); + + it('returns 503 when the Secrets Store cannot resolve the internal secret', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + const cache = { delete: vi.fn(async () => undefined) }; + const suppliedSecret = 'caller-supplied-secret'; + + const res = await app.request( + '/internal/user-auth/invalidate', + { + method: 'POST', + headers: { + 'content-type': 'application/json', + 'X-Internal-Secret': suppliedSecret, + }, + body: JSON.stringify({ kiloUserId: 'usr_blocked' }), + }, + { + ...defaultEnv, + USER_EXISTS_CACHE: cache, + INTERNAL_API_SECRET_PROD: { + get: async () => { + throw new Error('secret store unavailable'); + }, + }, + } + ); + + const body = await res.json(); + expect(res.status).toBe(503); + expect(body).toEqual({ success: false, error: 'Service temporarily unavailable' }); + expect(JSON.stringify(body)).not.toContain('secret store unavailable'); + expect(JSON.stringify(body)).not.toContain(suppliedSecret); + expect(cache.delete).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith( + 'Auth infrastructure failure', + expect.objectContaining({ + operation: 'internal-api-secret-get', + errorClass: 'Error', + }) + ); + expect(JSON.stringify(error.mock.calls)).not.toContain(suppliedSecret); + error.mockRestore(); + }); +}); + describe('public session route', () => { beforeEach(() => { vi.resetAllMocks(); 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 52c440012e..8902330515 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,446 @@ import { Hono } from 'hono'; import { SignJWT } from 'jose'; +import { vi } from 'vitest'; +vi.mock('@kilocode/worker-utils/kilo-token-auth', () => ({ + findKiloUserPepper: vi.fn(), +})); + +import { findKiloUserPepper } from '@kilocode/worker-utils/kilo-token-auth'; import { kiloJwtAuthMiddleware } from './kilo-jwt-auth'; +type CachedUserAuthV1 = + | { v: 1; exists: false } + | { v: 1; exists: true; pepper: string | null; blockedReason: string | null }; + type TestEnv = { NEXTAUTH_SECRET_PROD: { get: () => Promise; }; USER_EXISTS_CACHE: { - get: (key: string) => Promise; - put: (key: string, value: string, options?: { expirationTtl: number }) => Promise; + get: ReturnType Promise>>; + put: ReturnType< + typeof vi.fn< + (key: string, value: string, options?: { expirationTtl: number }) => Promise + > + >; }; HYPERDRIVE: { connectionString: string; }; }; -function makeEnv(secret: string, opts?: { cachedUserState?: '1' | '0' | null }): TestEnv { +const SECRET = 'test-secret'; +const USER_ID = 'usr_123'; +const PEPPER = 'pepper-current'; +const GENERIC_403 = 'User account not found'; + +const unblockedUser: CachedUserAuthV1 = { + v: 1, + exists: true, + pepper: PEPPER, + blockedReason: null, +}; + +function makeEnv(opts?: { cached?: string | null }): TestEnv { return { NEXTAUTH_SECRET_PROD: { - get: async () => secret, + get: async () => SECRET, }, USER_EXISTS_CACHE: { - get: async () => opts?.cachedUserState ?? null, - put: async () => {}, + get: vi.fn(async () => opts?.cached ?? null), + put: vi.fn(async () => undefined), }, - // 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') })); + return app; +} + +async function sign( + payload: Record, + opts?: { audience?: string } +): Promise { + let jwt = new SignJWT(payload) .setProtectedHeader({ alg: 'HS256' }) .setIssuedAt() - .setExpirationTime('1h') - .sign(new TextEncoder().encode(secret)); + .setExpirationTime('1h'); + if (opts?.audience) jwt = jwt.setAudience(opts.audience); + return jwt.sign(new TextEncoder().encode(SECRET)); +} + +function userToken(pepper: string | null = PEPPER) { + return sign({ kiloUserId: USER_ID, version: 3, apiTokenPepper: pepper }); +} + +function internalToken(userId = USER_ID) { + return sign({ kiloUserId: userId, version: 3 }); +} + +function authRequest(token: string) { + return new Request('http://local/api/me', { + headers: { Authorization: `Bearer ${token}` }, + }); } describe('kiloJwtAuthMiddleware', () => { + beforeEach(() => { + vi.mocked(findKiloUserPepper).mockReset(); + }); + 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()); + expect(res.status).toBe(401); + }); + + it('rejects a token with an audience', async () => { + const token = await sign( + { kiloUserId: USER_ID, version: 3, apiTokenPepper: PEPPER }, + { audience: 'session-ingest' } + ); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); - const res = await app.fetch(new Request('http://local/api/me'), makeEnv('secret')); expect(res.status).toBe(401); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + }); + + it('authorizes a user token with matching pepper when unblocked', async () => { + const token = await userToken(); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: USER_ID }); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + }); + + it('rejects a user token with a stale pepper', async () => { + const token = await userToken('pepper-stale'); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(403); + expect(await res.json()).toEqual({ success: false, error: GENERIC_403 }); + }); + + it('rejects a user token with matching pepper when blocked', async () => { + const token = await userToken(); + const env = makeEnv({ + cached: JSON.stringify({ + v: 1, + exists: true, + pepper: PEPPER, + blockedReason: 'tos', + } satisfies CachedUserAuthV1), + }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(403); + expect(await res.json()).toEqual({ success: false, error: GENERIC_403 }); + }); + + it('authorizes a user token with null pepper when cached pepper is null and unblocked', async () => { + const token = await userToken(null); + const env = makeEnv({ + cached: JSON.stringify({ + v: 1, + exists: true, + pepper: null, + blockedReason: null, + } satisfies CachedUserAuthV1), + }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: USER_ID }); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + }); + + it('rejects a user token with null pepper when cached pepper is a string', async () => { + const token = await userToken(null); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(403); + expect(await res.json()).toEqual({ success: false, error: GENERIC_403 }); + }); + + it('rejects a missing user', async () => { + const token = await userToken(); + const env = makeEnv({ + cached: JSON.stringify({ v: 1, exists: false } satisfies CachedUserAuthV1), + }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(403); + expect(await res.json()).toEqual({ success: false, error: GENERIC_403 }); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + }); + + it('treats malformed KV JSON as a miss, then reads Postgres', async () => { + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const token = await userToken(); + const env = makeEnv({ cached: '{"not":"auth-state"' }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(findKiloUserPepper).toHaveBeenCalledWith('postgres://test', USER_ID); + expect(env.USER_EXISTS_CACHE.put).toHaveBeenCalledWith( + `user-auth:v1:${USER_ID}`, + JSON.stringify(unblockedUser), + { expirationTtl: 60 } + ); + }); + + it('treats an otherwise valid cache state with extra fields as a miss', async () => { + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const token = await userToken(); + const env = makeEnv({ cached: JSON.stringify({ ...unblockedUser, unexpected: true }) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(findKiloUserPepper).toHaveBeenCalledWith('postgres://test', USER_ID); + }); + + it.each(['1', '0'] as const)( + 'treats legacy %j cache values as a miss, then reads Postgres', + async cached => { + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const token = await userToken(); + const env = makeEnv({ cached }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(findKiloUserPepper).toHaveBeenCalledWith('postgres://test', USER_ID); + expect(env.USER_EXISTS_CACHE.put).toHaveBeenCalledWith( + `user-auth:v1:${USER_ID}`, + JSON.stringify(unblockedUser), + { expirationTtl: 60 } + ); + } + ); + + it('does not read Postgres on a warm cache hit', async () => { + const token = await userToken(); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + expect(env.USER_EXISTS_CACHE.put).not.toHaveBeenCalled(); + }); + + it('reads Postgres on a cache miss and puts user-auth:v1 with TTL 60', async () => { + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const token = await userToken(); + const env = makeEnv(); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(env.USER_EXISTS_CACHE.put).toHaveBeenCalledWith( + `user-auth:v1:${USER_ID}`, + JSON.stringify(unblockedUser), + { expirationTtl: 60 } + ); + expect(res.status).toBe(200); + expect(findKiloUserPepper).toHaveBeenCalledTimes(1); }); - 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('puts a missing-user cache entry with TTL 300', async () => { + vi.mocked(findKiloUserPepper).mockResolvedValue(undefined); + const token = await userToken(); + const env = makeEnv(); - 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(authRequest(token), env); - const res = await app.fetch( - new Request('http://local/api/me', { - headers: { - Authorization: `Bearer ${token}`, - }, - }), - makeEnv(secret, { cachedUserState: '1' }) + expect(env.USER_EXISTS_CACHE.put).toHaveBeenCalledWith( + `user-auth:v1:${USER_ID}`, + JSON.stringify({ v: 1, exists: false }), + { expirationTtl: 300 } ); + expect(res.status).toBe(403); + }); + + it('awaits KV.put before returning 200', async () => { + let resolvePut!: () => void; + const putGate = new Promise(resolve => { + resolvePut = resolve; + }); + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const token = await userToken(); + const env = makeEnv(); + env.USER_EXISTS_CACHE.put.mockImplementation(() => putGate); + + const responsePromise = Promise.resolve(makeApp().fetch(authRequest(token), env)); + await vi.waitFor(() => { + expect(env.USER_EXISTS_CACHE.put).toHaveBeenCalled(); + }); + let settled = false; + void responsePromise.then(() => { + settled = true; + }); + await Promise.resolve(); + expect(settled).toBe(false); + + resolvePut(); + const res = await responsePromise; 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); + it('returns 503 when Postgres throws and does not authorize', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + vi.mocked(findKiloUserPepper).mockRejectedValue(new Error('hyperdrive down')); + const token = await userToken(); + const env = makeEnv(); - 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(authRequest(token), env); - const res = await app.fetch( - new Request('http://local/api/me', { - headers: { - Authorization: `Bearer ${token}`, - }, - }), - makeEnv(secret, { cachedUserState: '0' }) + expect(res.status).toBe(503); + expect(env.USER_EXISTS_CACHE.put).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith( + 'Auth infrastructure failure', + expect.objectContaining({ + operation: 'user-auth-load', + kiloUserId: USER_ID, + errorMessage: 'hyperdrive down', + }) ); + error.mockRestore(); + }); + + it('returns 503 when the cache read throws and does not authorize', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + const token = await userToken(); + const env = makeEnv(); + env.USER_EXISTS_CACHE.get.mockRejectedValueOnce(new Error('kv unavailable')); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(503); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith( + 'Auth infrastructure failure', + expect.objectContaining({ + operation: 'user-auth-load', + kiloUserId: USER_ID, + errorMessage: 'kv unavailable', + }) + ); + error.mockRestore(); + }); + + it('returns the authoritative result when caching the state fails', async () => { + vi.mocked(findKiloUserPepper).mockResolvedValue({ + pepper: PEPPER, + blockedReason: null, + }); + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + const token = await userToken(); + const env = makeEnv(); + env.USER_EXISTS_CACHE.put.mockRejectedValueOnce(new Error('kv unavailable')); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(warn).toHaveBeenCalledWith( + 'Failed to cache user auth state', + expect.objectContaining({ operation: 'user-auth-cache-put', kiloUserId: USER_ID }) + ); + warn.mockRestore(); + }); + + it('returns 503 when the secret store cannot resolve the JWT secret', async () => { + const error = vi.spyOn(console, 'error').mockImplementation(() => undefined); + const token = await userToken(); + const env = makeEnv(); + env.NEXTAUTH_SECRET_PROD.get = vi.fn(async () => { + throw new Error('secret store unavailable'); + }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(503); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + expect(error).toHaveBeenCalledWith( + 'Auth infrastructure failure', + expect.objectContaining({ + operation: 'nextauth-secret-get', + errorMessage: 'secret store unavailable', + }) + ); + error.mockRestore(); + }); + + it('authorizes an internal token when the user is unblocked', async () => { + const token = await internalToken(); + const env = makeEnv({ cached: JSON.stringify(unblockedUser) }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: USER_ID }); + expect(findKiloUserPepper).not.toHaveBeenCalled(); + }); + + it('authorizes an internal token when the user is blocked', async () => { + const token = await internalToken(); + const env = makeEnv({ + cached: JSON.stringify({ + v: 1, + exists: true, + pepper: PEPPER, + blockedReason: 'gdpr', + } satisfies CachedUserAuthV1), + }); + + const res = await makeApp().fetch(authRequest(token), env); + + expect(res.status).toBe(200); + expect(await res.json()).toEqual({ user_id: USER_ID }); + }); + + it('rejects an internal token when the user is missing', async () => { + const token = await internalToken('deleted_user'); + const env = makeEnv({ + cached: JSON.stringify({ v: 1, exists: false } satisfies CachedUserAuthV1), + }); + + const res = await makeApp().fetch(authRequest(token), env); expect(res.status).toBe(403); - expect(await res.json()).toEqual({ success: false, error: 'User account not found' }); + expect(await res.json()).toEqual({ success: false, error: GENERIC_403 }); }); }); diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.ts index b893e7e695..d0e7681cd8 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.ts @@ -1,46 +1,66 @@ 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 { findKiloUserPepper } from '@kilocode/worker-utils/kilo-token-auth'; +import { z } from 'zod'; import type { Env } from '../env'; -const USER_EXISTS_TTL_SECONDS = 24 * 60 * 60; // 24h -const USER_NOT_FOUND_TTL_SECONDS = 5 * 60; // 5m +export const USER_AUTH_CACHE_KEY_PREFIX = 'user-auth:v1:'; +const USER_AUTH_TTL_SECONDS = 60; +const USER_MISSING_TTL_SECONDS = 5 * 60; -/** - * 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}`; +type CachedUserAuthV1 = + | { v: 1; exists: false } + | { v: 1; exists: true; pepper: string | null; blockedReason: string | null }; - const cached = await env.USER_EXISTS_CACHE.get(cacheKey); - if (cached === '1') { - return true; - } - if (cached === '0') { - return false; +const cachedUserAuthV1Schema = z.union([ + z.object({ v: z.literal(1), exists: z.literal(false) }).strict(), + z + .object({ + v: z.literal(1), + exists: z.literal(true), + pepper: z.string().nullable(), + blockedReason: z.string().nullable(), + }) + .strict(), +]); + +const USER_AUTH_DENIED = 'User account not found'; + +function parseCachedUserAuth(raw: string | null): CachedUserAuthV1 | null { + if (raw === null) return null; + try { + const parsed = cachedUserAuthV1Schema.safeParse(JSON.parse(raw)); + return parsed.success ? parsed.data : null; + } catch { + return null; } +} - 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); +async function loadUserAuth(env: Env, userId: string): Promise { + const cacheKey = `${USER_AUTH_CACHE_KEY_PREFIX}${userId}`; + const cached = parseCachedUserAuth(await env.USER_EXISTS_CACHE.get(cacheKey)); + if (cached) return cached; - const row = rows[0]; + const row = await findKiloUserPepper(env.HYPERDRIVE.connectionString, userId); + const state: CachedUserAuthV1 = + row === undefined || row === null + ? { v: 1, exists: false } + : { v: 1, exists: true, pepper: row.pepper, blockedReason: row.blockedReason }; - if (!row) { - void env.USER_EXISTS_CACHE.put(cacheKey, '0', { expirationTtl: USER_NOT_FOUND_TTL_SECONDS }); - return false; + try { + await env.USER_EXISTS_CACHE.put(cacheKey, JSON.stringify(state), { + expirationTtl: state.exists ? USER_AUTH_TTL_SECONDS : USER_MISSING_TTL_SECONDS, + }); + } catch (error) { + console.warn('Failed to cache user auth state', { + operation: 'user-auth-cache-put', + kiloUserId: userId, + errorClass: error instanceof Error ? error.name : typeof error, + errorMessage: error instanceof Error ? error.message : String(error), + }); } - - void env.USER_EXISTS_CACHE.put(cacheKey, '1', { expirationTtl: USER_EXISTS_TTL_SECONDS }); - return true; + return state; } export const kiloJwtAuthMiddleware = createMiddleware<{ @@ -58,19 +78,60 @@ 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(); + let secret: string; + try { + const configuredSecret = await c.env.NEXTAUTH_SECRET_PROD.get(); + if (!configuredSecret) { + console.error('Auth infrastructure failure', { operation: 'nextauth-secret-missing' }); + return c.json({ success: false, error: 'Service temporarily unavailable' }, 503); + } + secret = configuredSecret; + } catch (error) { + console.error('Auth infrastructure failure', { + operation: 'nextauth-secret-get', + 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); + } let kiloUserId: string; + let apiTokenPepper: string | null | undefined; try { const payload = await verifyKiloToken(token, secret); kiloUserId = payload.kiloUserId; + apiTokenPepper = payload.apiTokenPepper; } catch { 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); + let state: CachedUserAuthV1; + try { + state = await loadUserAuth(c.env, kiloUserId); + } catch (error) { + console.error('Auth infrastructure failure', { + operation: 'user-auth-load', + kiloUserId, + 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 (!state.exists) { + return c.json({ success: false, error: USER_AUTH_DENIED }, 403); + } + + // A missing apiTokenPepper is the legacy internal service-token class. It + // intentionally requires only an existing user; ordinary tokens carry a + // pepper (including null) and must pass the blocked/pepper checks below. + if (apiTokenPepper === undefined) { + c.set('user_id', kiloUserId); + return next(); + } + + if (state.blockedReason !== null || state.pepper !== apiTokenPepper) { + return c.json({ success: false, error: USER_AUTH_DENIED }, 403); } c.set('user_id', kiloUserId); diff --git a/services/session-ingest/src/routes/api.test.ts b/services/session-ingest/src/routes/api.test.ts index 3f7eef6a72..f17786b19d 100644 --- a/services/session-ingest/src/routes/api.test.ts +++ b/services/session-ingest/src/routes/api.test.ts @@ -290,6 +290,7 @@ describe('api routes', () => { it('POST /session emits created only for newly inserted rows', async () => { const { db, fns } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([{ blocked_reason: null }]); fns.insertResult.mockResolvedValueOnce([ { session_id: 'ses_12345678901234567890123456', @@ -327,7 +328,6 @@ describe('api routes', () => { ); expect(res.status).toBe(200); - expect(fns.select).not.toHaveBeenCalled(); expect(notifyUserSessionEvent).toHaveBeenCalledWith( expect.anything(), 'usr_test', @@ -341,6 +341,7 @@ describe('api routes', () => { it('POST /session does not emit created when row already exists', async () => { const { db, fns } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([{ blocked_reason: null }]); fns.insertResult.mockResolvedValueOnce([]); const sessionCache = { @@ -364,7 +365,6 @@ describe('api routes', () => { ); expect(res.status).toBe(200); - expect(fns.select).not.toHaveBeenCalled(); expect(notifyUserSessionEvent).not.toHaveBeenCalled(); expect(env.NOTIFICATIONS.sendSessionReadyNotification).not.toHaveBeenCalled(); }); @@ -372,6 +372,7 @@ describe('api routes', () => { it('POST /session caches a newly created personal session', async () => { const { db, fns } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([{ blocked_reason: null }]); fns.insertResult.mockResolvedValueOnce([ { session_id: 'ses_12345678901234567890123456', @@ -423,6 +424,7 @@ describe('api routes', () => { it('POST /session succeeds when cache warming is unavailable during rollout', async () => { const { db, fns } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([{ blocked_reason: null }]); fns.insertResult.mockResolvedValueOnce([ { session_id: 'ses_12345678901234567890123456', @@ -460,6 +462,46 @@ describe('api routes', () => { consoleError.mockRestore(); }); + it('POST /session returns 403 and does not insert when blocked_reason is set', async () => { + const { db, fns } = makeDbFakes(); + vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([{ blocked_reason: 'tos' }]); + + const res = await makeApiApp().fetch( + new Request('http://local/session', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ sessionId: 'ses_12345678901234567890123456' }), + }), + makeTestEnv() + ); + + expect(res.status).toBe(403); + expect(fns.insert).not.toHaveBeenCalled(); + expect(notifyUserSessionEvent).not.toHaveBeenCalled(); + expect(getSessionAccessCacheDO).not.toHaveBeenCalled(); + }); + + it('POST /session returns 403 and does not insert when the user is missing', async () => { + const { db, fns } = makeDbFakes(); + vi.mocked(getWorkerDb).mockReturnValue(db); + fns.selectResult.mockResolvedValueOnce([]); + + const res = await makeApiApp().fetch( + new Request('http://local/session', { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ sessionId: 'ses_12345678901234567890123456' }), + }), + makeTestEnv() + ); + + expect(res.status).toBe(403); + expect(fns.insert).not.toHaveBeenCalled(); + expect(notifyUserSessionEvent).not.toHaveBeenCalled(); + expect(getSessionAccessCacheDO).not.toHaveBeenCalled(); + }); + it('POST /session/:sessionId/ingest streams to R2 after access is resolved', async () => { const { db } = makeDbFakes(); vi.mocked(getWorkerDb).mockReturnValue(db); diff --git a/services/session-ingest/src/routes/api.ts b/services/session-ingest/src/routes/api.ts index 04057d0742..74f8e10dc7 100644 --- a/services/session-ingest/src/routes/api.ts +++ b/services/session-ingest/src/routes/api.ts @@ -21,6 +21,7 @@ import { handleDirectIngestRequest } from '../ingest/direct-ingest'; import { isDefaultSessionTitle } from '../ingest/default-session-title'; import { resolveAccessibleKiloSession } from '../services/session-access'; import { signSessionShareToken } from '../services/session-share-token'; +import { canCreateCliSessionForUser } from '../services/user-session-admission'; export type ApiContext = { Bindings: Env; @@ -86,16 +87,30 @@ api.post('/session', zodJsonValidator(createSessionSchema), async c => { const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); const kiloUserId = c.get('user_id'); - const [createdRow] = await db - .insert(cli_sessions_v2) - .values({ - session_id: body.sessionId, - kilo_user_id: kiloUserId, - }) - .onConflictDoNothing({ - target: [cli_sessions_v2.session_id, cli_sessions_v2.kilo_user_id], - }) - .returning(); + const result = await db.transaction(async tx => { + if (!(await canCreateCliSessionForUser(tx, kiloUserId))) { + return { status: 'not_admitted' } as const; + } + + const [createdRow] = await tx + .insert(cli_sessions_v2) + .values({ + session_id: body.sessionId, + kilo_user_id: kiloUserId, + }) + .onConflictDoNothing({ + target: [cli_sessions_v2.session_id, cli_sessions_v2.kilo_user_id], + }) + .returning(); + + return { status: 'admitted', createdRow } as const; + }); + + if (result.status === 'not_admitted') { + return c.json({ success: false, error: 'User account not found' }, 403); + } + + const { createdRow } = result; if (createdRow) { const session = mapSessionEventRow(createdRow); diff --git a/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts b/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts index 614a065af5..f7f96b2508 100644 --- a/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts +++ b/services/session-ingest/src/routes/cloud-agent-session-scope.test.ts @@ -10,6 +10,9 @@ vi.mock('@kilocode/db/client', () => ({ getWorkerDb: vi.fn() })); vi.mock('../dos/SessionAccessCacheDO', () => ({ getSessionAccessCacheDO: vi.fn() })); vi.mock('../ingest/direct-ingest', () => ({ handleDirectIngestRequest: vi.fn() })); vi.mock('../services/session-access', () => ({ resolveAccessibleKiloSession: vi.fn() })); +vi.mock('../services/user-session-admission', () => ({ + canCreateCliSessionForUser: vi.fn(), +})); vi.mock('../session-events', () => ({ mapSessionEventRow: vi.fn(row => ({ id: row.session_id, updatedAt: row.updated_at })), notifyUserSessionEvent: vi.fn(), @@ -26,6 +29,7 @@ import { cloudAgentSessionScopeHeaders } from '@kilocode/session-ingest-contract import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; import { handleDirectIngestRequest } from '../ingest/direct-ingest'; import { resolveAccessibleKiloSession } from '../services/session-access'; +import { canCreateCliSessionForUser } from '../services/user-session-admission'; import { cloudAgentSessionScopeApi } from './cloud-agent-session-scope'; const rootSessionId = 'ses_12345678901234567890123456'; @@ -113,6 +117,27 @@ describe('Cloud Agent session scope routes', () => { beforeEach(() => { vi.resetAllMocks(); workerUtils.hasOrganizationAccess.mockResolvedValue(true); + vi.mocked(canCreateCliSessionForUser).mockResolvedValue(true); + }); + + it('rejects a blocked or missing user before locking the root or inserting a child', async () => { + const { db, insertedValues } = makeDb([], []); + vi.mocked(canCreateCliSessionForUser).mockResolvedValueOnce(false); + vi.mocked(getWorkerDb).mockReturnValue(db as never); + + const response = await makeApp().fetch( + new Request('http://local/session', { + method: 'POST', + headers: assertionHeaders(), + body: JSON.stringify({ sessionId: childSessionId }), + }), + env + ); + + expect(response.status).toBe(403); + expect(insertedValues).toHaveLength(0); + expect(getSessionAccessCacheDO).not.toHaveBeenCalled(); + expect(workerUtils.hasOrganizationAccess).not.toHaveBeenCalled(); }); it('rejects an incomplete or invalid session scope assertion before database access', async () => { diff --git a/services/session-ingest/src/routes/cloud-agent-session-scope.ts b/services/session-ingest/src/routes/cloud-agent-session-scope.ts index a14bbf4b3c..2c8bafe4b6 100644 --- a/services/session-ingest/src/routes/cloud-agent-session-scope.ts +++ b/services/session-ingest/src/routes/cloud-agent-session-scope.ts @@ -15,6 +15,7 @@ import { getSessionAccessCacheDO } from '../dos/SessionAccessCacheDO'; import { handleDirectIngestRequest } from '../ingest/direct-ingest'; import { mapSessionEventRow, notifyUserSessionEvent } from '../session-events'; import { resolveAccessibleKiloSession } from '../services/session-access'; +import { canCreateCliSessionForUser } from '../services/user-session-admission'; import type { ApiContext } from './api'; const createScopedSessionSchema = z.object({ @@ -68,6 +69,10 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS const kiloUserId = c.get('user_id'); const db = getWorkerDb(c.env.HYPERDRIVE.connectionString); const result = await db.transaction(async tx => { + if (!(await canCreateCliSessionForUser(tx, kiloUserId))) { + return { status: 'user_not_admitted' } as const; + } + const [root] = await tx .select({ sessionId: cli_sessions_v2.session_id, @@ -163,6 +168,9 @@ cloudAgentSessionScopeApi.post('/session', zodJsonValidator(createScopedSessionS return { status: 'existing', row: existing } as const; }); + if (result.status === 'user_not_admitted') { + return c.json({ success: false, error: 'User account not found' }, 403); + } if (result.status === 'root_not_found') { return c.json({ success: false, error: 'session_not_found' }, 404); } diff --git a/services/session-ingest/src/services/user-session-admission.test.ts b/services/session-ingest/src/services/user-session-admission.test.ts new file mode 100644 index 0000000000..8ff1fc37df --- /dev/null +++ b/services/session-ingest/src/services/user-session-admission.test.ts @@ -0,0 +1,23 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { canCreateCliSessionForUser } from './user-session-admission'; + +describe('canCreateCliSessionForUser', () => { + it.each([ + { name: 'unblocked user', rows: [{ blocked_reason: null }], expected: true }, + { name: 'blocked user', rows: [{ blocked_reason: 'tos' }], expected: false }, + { name: 'missing user', rows: [], expected: false }, + ])('returns $expected for a $name', async ({ rows, expected }) => { + const forUpdate = vi.fn(async () => rows); + const query = { + from: vi.fn(() => query), + where: vi.fn(() => query), + limit: vi.fn(() => query), + for: forUpdate, + }; + const db = { select: vi.fn(() => query) }; + + await expect(canCreateCliSessionForUser(db as never, 'usr_test')).resolves.toBe(expected); + expect(forUpdate).toHaveBeenCalledWith('update'); + }); +}); diff --git a/services/session-ingest/src/services/user-session-admission.ts b/services/session-ingest/src/services/user-session-admission.ts new file mode 100644 index 0000000000..9f7fda3461 --- /dev/null +++ b/services/session-ingest/src/services/user-session-admission.ts @@ -0,0 +1,25 @@ +import type { WorkerDb } from '@kilocode/db/client'; +import { kilocode_users } from '@kilocode/db/schema'; +import { eq } from 'drizzle-orm'; + +type UserSessionAdmissionDb = Pick; + +/** + * Check the user row under the caller's transaction lock immediately before + * creating a cli_sessions_v2 row. + */ +export async function canCreateCliSessionForUser( + db: UserSessionAdmissionDb, + kiloUserId: string +): Promise { + const [user] = await db + .select({ blocked_reason: kilocode_users.blocked_reason }) + .from(kilocode_users) + .where(eq(kilocode_users.id, kiloUserId)) + .limit(1) + .for('update'); + + return user?.blocked_reason === null; +} + +export const USER_SESSION_ADMISSION_ERROR = 'User session creation is not allowed'; diff --git a/services/session-ingest/src/session-ingest-rpc.test.ts b/services/session-ingest/src/session-ingest-rpc.test.ts index 47642d4cc7..db218ad920 100644 --- a/services/session-ingest/src/session-ingest-rpc.test.ts +++ b/services/session-ingest/src/session-ingest-rpc.test.ts @@ -43,6 +43,11 @@ vi.mock('./session-events', () => ({ notifyUserSessionEvent: vi.fn(), })); +vi.mock('./services/user-session-admission', () => ({ + canCreateCliSessionForUser: vi.fn(), + USER_SESSION_ADMISSION_ERROR: 'User session creation is not allowed', +})); + import { getWorkerDb } from '@kilocode/db/client'; import { cli_sessions_v2, organization_memberships } from '@kilocode/db/schema'; import { @@ -57,6 +62,7 @@ import { import { desc, gte, isNotNull, or } from 'drizzle-orm'; import { getSessionIngestDO } from './dos/SessionIngestDO'; import { SessionIngestRPC } from './session-ingest-rpc'; +import { canCreateCliSessionForUser } from './services/user-session-admission'; const sdkSessionInfoFixture = { id: 'ses_12345678901234567890123456', @@ -164,6 +170,10 @@ describe('createSessionForCloudAgent', () => { createdOnPlatform: 'cloud-agent', }; + beforeEach(() => { + vi.mocked(canCreateCliSessionForUser).mockReset().mockResolvedValue(true); + }); + it('creates a root with both Cloud Agent identity columns', async () => { const row = { session_id: params.sessionId, @@ -204,6 +214,17 @@ describe('createSessionForCloudAgent', () => { 'Cloud Agent root session identity conflict' ); }); + + it('rejects a blocked or missing user before any session write', async () => { + const fake = makeRootWriteDb({ created: { session_id: params.sessionId } }); + vi.mocked(canCreateCliSessionForUser).mockResolvedValueOnce(false); + const rpc = makeRpc(fake.db as never); + + await expect(rpc.createSessionForCloudAgent(params)).rejects.toThrow( + 'User session creation is not allowed' + ); + expect(fake.values).not.toHaveBeenCalled(); + }); }); describe('Kilo SDK persisted identity schemas', () => { diff --git a/services/session-ingest/src/session-ingest-rpc.ts b/services/session-ingest/src/session-ingest-rpc.ts index 7d7a9fbfa7..7ce7ed9b07 100644 --- a/services/session-ingest/src/session-ingest-rpc.ts +++ b/services/session-ingest/src/session-ingest-rpc.ts @@ -33,6 +33,10 @@ import { getSessionAccessCacheDO } from './dos/SessionAccessCacheDO'; import { withDORetry } from '@kilocode/worker-utils'; import { app } from './app'; import { mapSessionEventRow, notifyUserSessionEvent } from './session-events'; +import { + canCreateCliSessionForUser, + USER_SESSION_ADMISSION_ERROR, +} from './services/user-session-admission'; const MAX_CLOUD_AGENT_ROOT_SESSION_TITLE_CHARACTERS = 512; @@ -75,6 +79,10 @@ export class SessionIngestRPC extends WorkerEntrypoint implements SessionIn const db = getWorkerDb(this.env.HYPERDRIVE.connectionString); const { existingRow, persistedRow } = await db.transaction(async tx => { + if (!(await canCreateCliSessionForUser(tx, parsed.kiloUserId))) { + throw new Error(USER_SESSION_ADMISSION_ERROR); + } + const [created] = await tx .insert(cli_sessions_v2) .values({